diff --git a/src/Tao.jl b/src/Tao.jl index b763c61..b905538 100644 --- a/src/Tao.jl +++ b/src/Tao.jl @@ -3,13 +3,14 @@ using DelimitedFiles using LinearAlgebra using Printf using NLsolve +using Interpolations export metadata_path, PolData, BAGELS_1, BAGELS_2, - calc_refill_sensitivity, - calc_refill_time, + calc_P_t, + calc_T, pol_scan, get_pol_data @@ -250,18 +251,59 @@ function BAGELS_2(lat, phi_start, phi_step, N_knobs; suffix="", outf="BAGELS.bma end end -function calc_refill_sensitivity(P_dk, tau_eq; P_0 = 0.85, T = 4.8) + +""" + calc_P_t(P_dk, tau_eq; P_0 = 0.85, T = 4.8) + +Calculates the time-averaged e-polarization one could maintain in a lattice with +equal bunches parallel and antiparallel to the arc dipole fields, assuming an +average bunch replacement rate `T` and initial injected polarization `P_0`. Defaults +are values used in the ESR of the EIC (`P_0` = 85%, `T` = 4.8 min) + +### Input +- `P_dk` -- Array of DK polarizations +- `tau_eq` -- Array of corresponding equilibrium times +- `P_0` -- (Optional) Initial bunch polarization at injection +- `T` -- (Optional) Average bunch replacement rate + +### Output +- `P_t` -- Time-averaged polarization maintainable in ring +- `T_du_t` -- How often the antiparallel bunches need replacement +- `T_dd_t` -- How often the parallel bunches need replacement +""" +function calc_P_t(P_dk, tau_eq; P_0 = 0.85, T = 4.8) P_dd = x -> (-P_0.*tau_eq .- P_dk.*tau_eq .+ P_dk.*x .+ (P_0.*tau_eq .+ P_dk.*tau_eq).*exp.(-x./tau_eq))./x P_du = x -> (P_0.*tau_eq .- P_dk.*tau_eq .+ P_dk.*x .+ (-P_0.*tau_eq .+ P_dk.*tau_eq).*exp.(-x./tau_eq))./x P = T_du -> (P_dd((T_du.*T./(2 .*T_du.-T))) + P_du(T_du)).^2 T_du_t = nlsolve(P, 10. .*ones(length(P_dk))).zero T_dd_t = (T_du_t.*T./(2 .*T_du_t.-T)) - - return P_du(T_du_t), T_du_t, T_dd_t + P_t = P_du(T_du_t) + + return P_t, T_du_t, T_dd_t end -function calc_refill_time(P_dk, tau_eq; P_0 = 0.85, P_min_avg = 0.7) + +""" + calc_T(P_dk, tau_eq; P_0 = 0.85, P_min_avg = 0.7) + +Calculates how often the bunches would need to be replaced in a lattice with +equal bunches parallel and antiparallel to the arc dipole fields to maintain +a minimum time-averaged polarization of `P_min_avg` Defaults are values used +in the ESR of the EIC (P_0 = 85%, P_min_avg = 70%) + +### Input +- `P_dk` -- Array of DK polarizations +- `tau_eq` -- Array of corresponding equilibrium times +- `P_0` -- (Optional) Initial bunch polarization at injection +- `P_min_avg` -- (Optional) Minimum-allowable time-averaged polarization + +### Output +- `T` -- Average bunch replacement time to maintain `P_min_avg` +- `T_du` -- Antiparallel bunch replacement time to maintain `P_min_avg` +- `T_dd` -- Parallel bunch replacement time to maintain `P_min_avg` +""" +function calc_T(P_dk, tau_eq; P_0 = 0.85, P_min_avg = 0.7) P_dd = x -> (-P_0.*tau_eq .- P_dk.*tau_eq .+ P_dk.*x .+ (P_0.*tau_eq .+ P_dk.*tau_eq).*exp.(-x./tau_eq))./x .+ P_min_avg P_du = x -> (P_0.*tau_eq .- P_dk.*tau_eq .+ P_dk.*x .+ (-P_0.*tau_eq .+ P_dk.*tau_eq).*exp.(-x./tau_eq))./x .- P_min_avg T_dd = nlsolve(P_dd, 10. .*ones(length(P_dk))).zero @@ -276,11 +318,12 @@ end pol_scan(lat, agamma0) Performs a polarization scan across the agamma0 range specified, calculates -important quantities and stores them for fast reference in the metadata. +important quantities and stores them for fast reference in the metadata. +The data can be retrieved using the `get_pol_data` method. ### Input - `lat` -- lat file name -- `agamma0` -- range of a*gamma_0 to scan over +- `agamma0` -- range of `agamma0` to scan over ### Output - `pol_data` -- PolData struct containing polarization quantities for lattice @@ -294,7 +337,7 @@ function pol_scan(lat, agamma0) # Generate tao command script tao_cmd = open("$(path)/spin.tao", "w") - a_e = 1.15965218e-3 + a_e = 0.00115965218128 m_e = 0.51099895e6 println(tao_cmd, "set ele * spin_tracking_method = sprint") println(tao_cmd, "set bmad_com spin_tracking_on = T") @@ -333,8 +376,8 @@ function pol_scan(lat, agamma0) tau_dep_2b = 1 ./data[:,17] ./ 60 tau_dep_2c = 1 ./data[:,18] ./ 60 tau_eq = (tau_dep.^-1+tau_bks.^-1).^-1 - T, T_du, T_dd = calc_refill_time(P_dk,tau_eq) - P_t, T_du_t, T_dd_t = calc_refill_sensitivity(P_dk,tau_eq) + T, T_du, T_dd = calc_T(P_dk,tau_eq) + P_t, T_du_t, T_dd_t = calc_P_t(P_dk,tau_eq) pol_data = PolData( agamma0, spin_tune, @@ -427,9 +470,12 @@ function get_pol_data(lat) println("Lattice file $(lat) not found!") return end - + if !isfile("$(path)/pol_data.dlm") + println("First-order polarization data not generated for lattice $(lattice). Please use the pol_scan method to generate the data.") + end pol_data_dlm = readdlm("$(path)/pol_data.dlm", ';')[:,2:end] + return PolData( pol_data_dlm[1,:], pol_data_dlm[2,:], pol_data_dlm[3,:], @@ -457,4 +503,46 @@ function get_pol_data(lat) pol_data_dlm[25,:]) end + +""" + pol_track_scan(lat, n_damp, data_ave="data.ave") + +Reads tracking data for the energies tracked in a `tracking` directory with +the lattice file, calculates important polarization quantities and stores them +for fast reference in the metadata. The data can then be retrieved using the +`get_pol_track_data` method. The `tracking` directory should contain subdirectories +with names equal to the `agamma0` and contain the turn-by-turn averages tracking +output from the `long_term_tracking` Bmad program. + +### Input +- `lat` -- lat file name +- `n_damp` -- Number of turns where equilibrium emittances are reached +- `data_ave` -- (Optional) Averages output file name, default is `data.ave` + +### Output +- `pol_track_data` -- PolData struct containing polarization quantities for all tracked energies +""" +function pol_track_scan(lat, n_damp, data_ave="data.ave") + path = metadata_path(lat) + if path == "" + println("Lattice file $(lat) not found!") + return + end + + # Use first-order P_st and tau_bks to calculate polarization data with nonlinear tau_dep + pol_data = get_pol_data(lat) + if isnothing(pol_data) + return + end + + # Tracking energies not particularly at linear evaluation energies, so interpolate + P_st_interp = linear_interpolation(pol_data.agamma0, pol_data.P_st) + tau_bks_interp = linear_interpolation(pol_data.agamma0, pol_data.tau_bks) + + subdirs = readdir("$(tracking)") + energies = subdirs[isdir.(subdirs, join=true))] + + +end + end