Read Bmad lattice files

[jank324]

To complete this, the conversion of sbend needs to be properly implemented and there should also be some tests that the conversion runs without error and produces the expected results.

[jank324]

sband is now implemented. Now this needs to be checked for correctness.

Also, before merging this PR we should delete the test Notebook in the root directory.

[jank324]

So, as it stands I am seeing two potential and potentially related issues.

Firstly, there are NaNs appearing when tracking through the LCLS Cu-HXR lattice. I have no idea why, but it might be caused by the next issue.

Secondly, the Twiss parameters in the LCLS lattice start deviating from the end of the l0a element, which is a cavity. Here is the MRE for that cavity in Cheetah

incoming_beam = cheetah.ParticleBeam.from_twiss(
    num_particles=1_000_000,
    beta_x=0.44865302,
    beta_y=0.44865302,
    alpha_x=0.18863809,
    alpha_y=0.18863809,
    emittance_x=3.494768647122823e-09,
    emittance_y=3.497810737006068e-09,
    energy=6e6,
)
cavity_element = cheetah.Cavity(
    length=3.0952440,
    voltage=5.8010691e7,
    phase=np.degrees(-3.0555556e-3 * 2 * np.pi),
    frequency=2.8560000e9,
)

outgoing_beam = cavity_element.track(incoming_beam)

print(f"{outgoing_beam.beta_x = }")
print(f"{outgoing_beam.beta_y = }")
print(f"{outgoing_beam.alpha_x = }")
print(f"{outgoing_beam.alpha_y = }")
print(f"{outgoing_beam.energy = }")

which outputs

outgoing_beam.beta_x = tensor(12.8575)
outgoing_beam.beta_y = tensor(12.8490)
outgoing_beam.alpha_x = tensor(-3.5258)
outgoing_beam.alpha_y = tensor(-3.5226)
outgoing_beam.energy = tensor(64000000.3325, dtype=torch.float64)

and here is the same in Ocelot

tws = ocelot.Twiss()
tws.beta_x = 0.44865302
tws.beta_y = 0.44865302
tws.alpha_x = 0.18863809
tws.alpha_y = 0.18863809
tws.emit_x = 3.494768647122823e-09
tws.emit_y = 3.497810737006068e-09
tws.gamma_x = (1 + tws.alpha_x**2) / tws.beta_x
tws.gamma_y = (1 + tws.alpha_y**2) / tws.beta_y
tws.E = 6e-3

p_array = ocelot.generate_parray(tws=tws)

cell = [
    ocelot.Cavity(
        l=3.0952440,
        v=5.8010691e7 * 1e-9,
        freq=2.8560000e9,
        phi=np.degrees(-3.0555556e-3 * 2 * np.pi),
    )
]
lattice = ocelot.MagneticLattice(cell)
navigator = ocelot.Navigator(lattice=lattice)

_, outgoing_parray = ocelot.track(lattice, deepcopy(p_array), navigator)
derived_twiss = ocelot.cpbd.beam.get_envelope(outgoing_parray)

print(derived_twiss)

which outputs

emit_x  = 3.2688336805880204e-10
emit_y  = 3.284914343857117e-10
beta_x  = 12.856815106445865
beta_y  = 12.856815106442994
alpha_x = -3.524868604372471
alpha_y = -3.5248686043537316
gamma_x = 1.0441698482044868
gamma_y = 1.0441698481944446
Dx      = 0.0
Dy      = 0.0
Dxp     = 0.0
Dyp     = 0.0
mux     = 0.0
muy     = 0.0
nu_x    = 0.0
nu_y    = 0.0
E       = 0.06400000033252268
s        = 0.0

Apparently they arrive at roughly the same result. However, their result differs from Bmad/Tao, which outputs

Tao> show element l0a
Element # 3179
Element Name: L0A
Element Type:  "DUALFEED"
Key: Lcavity
S_start, S:      1.459000,      4.554244
Ref_time_start, Ref_time:  4.884447E-09,  1.521259E-08

Attribute values [Only non-zero values shown]:
    1  L                           =  3.0952440E+00 m        31  L_ACTIVE                    =  3.0952440E+00 m
    7  GRADIENT_ERR                =  0.0000000E+00 eV/m
    8  VOLTAGE                     =  5.8010691E+07 Volt      6  GRADIENT                    =  1.8741880E+07 eV/m
    9  VOLTAGE_ERR                 =  0.0000000E+00 Volt
   10  FRINGE_TYPE                 =  Full (4)               11  FRINGE_AT                   =  Both_Ends (3)
   13  SPIN_FRINGE_ON              =  T (1)
   15  RF_FREQUENCY                =  2.8560000E+09 Hz       16  RF_WAVELENGTH               =  1.0496935E-01 m
   17  STATIC_LINEAR_MAP           =  F (0)
   18  AUTOSCALE_AMPLITUDE         =  T (1)                  19  AUTOSCALE_PHASE             =  T (1)
   20  LONGITUDINAL_MODE           = 0
   21  E_LOSS                      =  0.0000000E+00 eV
   24  PHI0                        = -3.0555556E-03 rad/2pi  26  PHI0_MULTIPASS              =  0.0000000E+00 rad/2pi
   25  PHI0_ERR                    =  0.0000000E+00 rad/2pi
   29  FIELD_AUTOSCALE             =  1.0000000E+00          27  PHI0_AUTOSCALE              =  0.0000000E+00 rad/2pi
   30  VOLTAGE_TOT                 =  5.8010691E+07 Volt      5  GRADIENT_TOT                =  1.8741880E+07 eV/m
   33  N_CELL                      = 1                       22  CAVITY_TYPE                 =  Traveling_Wave (2)
   44  COUPLER_PHASE               =  0.0000000E+00 rad/2pi
   47  COUPLER_AT                  =  Exit_End (2)           46  COUPLER_STRENGTH            =  0.0000000E+00
   51  P0C_START                   =  5.9782004E+06 eV           BETA_START                  =  9.9636673E-01
   52  E_TOT_START                 =  6.0000000E+06 eV           DELTA_E                     =  5.8000000E+07 eV
   53  P0C                         =  6.3997960E+07 eV           BETA                        =  9.9996812E-01
   54  E_TOT                       =  6.4000000E+07 eV           GAMMA                       =  1.2524488E+02
   64  REF_TIME_START              =  4.8844466E-09          50  DELTA_REF_TIME              =  1.0328140E-08 sec
   65  INTEGRATOR_ORDER            = 0
   67  DS_STEP                     =  3.0952440E-01 m        66  NUM_STEPS                   = 10
   68  CSR_DS_STEP                 =  0.0000000E+00 m

       TRACKING_METHOD              =  Bmad_Standard             APERTURE_AT                =  Exit_End
       MAT6_CALC_METHOD             =  Auto                      APERTURE_TYPE              =  Rectangular
       SPIN_TRACKING_METHOD         =  Tracking                  OFFSET_MOVES_APERTURE      =  F
       PTC_INTEGRATION_TYPE         =  Matrix_Kick               SYMPLECTIFY                =  F
       CSR_METHOD                   =  Off                       FIELD_MASTER               =  F
       SPACE_CHARGE_METHOD          =  Off                       LONGITUDINAL ORIENTATION   =       1
       FIELD_CALC                   =  Bmad_Standard

Slave_status: Free

Lord_status:  Super_Lord
Slaves:
   Index   Name        Type                     S
      44   L0A#1       Lcavity                  1.517646
      46   L0A#2       Lcavity                  1.717000
      48   L0A#3       Lcavity                  2.366320
      51   L0A#4       Lcavity                  3.006622
      53   L0A#5       Lcavity                  4.158468
      56   L0A#6       Lcavity                  4.493325
      58   L0A#7       Lcavity                  4.554244

Twiss at end of element:
                          A              B            Cbar                        C_mat
  Beta (m)        16.62625295    16.62625295  |   0.00000000   0.00000000      0.00000000   0.00000000
  Alpha           -4.71572777    -4.71572777  |   0.00000000   0.00000000      0.00000000   0.00000000
  Gamma (1/m)      1.39767442     1.39767442  |   Gamma_c =   1.00000000       Mode_Flip = F
  Phi (rad)        2.99653305     2.99653305            X              Y              Z
  Eta (m)          0.00000000     0.00000000     0.00000000     0.00000000     0.14515452
  Etap             0.00000000     0.00000000     0.00000000     0.00000000     0.00000000

Short-Range Wake:
  scale_with_length = T
  amp_scale         = 1
  z_scale           = 1
  z_max             = 0.01

  Short-Range Longitudinal Pseudo Modes:
   #        Amp        Damp           K         Phi   Transverse_Dependence
   1  1.0215E+14  2.0074E+02  0.0000E+00  2.5000E-01   none
   2  2.0265E+13  1.6475E+05  0.0000E+00  2.5000E-01   none
   3  9.1115E+13  1.2068E+03  0.0000E+00  2.5000E-01   none
   4  4.9582E+13  9.0508E+03  0.0000E+00  2.5000E-01   none

Orbit:  Electron   State: Alive
         Position[mm] Momentum[mrad]        Spin   |
  X:       0.00000000     0.00000000               | t_particle [sec]:        1.52125867E-08  E_tot: 6.40000E+07
  Y:       0.00000000     0.00000000               | t_part-t_ref [sec]:      0.00000000E+00  PC:    6.39980E+07
  Z:      -0.00000000     0.00000000               | (t_ref-t_part)*Vel [m]:  0.00000000E+00  Beta:  0.999968125
  Particle Phase relative to RF Phase (rad/2pi): -0.00305555556

@cr-xu maybe you have an idea? I'm worried/hoping that some of the parameters are simply not correctly converted. For the other cavity tested in test/test_compare_ocelot.py all three simulation codes agreed. Notably, for that example, phi=0.

[jank324]

The Bmad results for that last test commit:

[jank324]

I fixed some things going through a cavity with a ParameterBeam. However, sigma_p still doesn't match with the one computed for a ParticleBeam. This can be seen easiest when looking at just the first 100 elements of the LCLS CU-HXR lattice.

[jank324]

Removed data classes again because the elements are not exactly pieces of data, and as a result using data classes didn't simplify things, but often actually complicated them. This may have lead to elements having attributes that their user wouldn't be aware of.

[jank324]

Looks better, but definitely deviates towards the end.

[cr-xu]

Is it still only the cavity that's causing the problem? Can you check first if the longitudinal components of the beam/ particle of cheetah and bmad agree?

One thing that is different here:

• In cheetah's Cavity transfer map, we only consider longitudinal parts, but not the transverse focusing part
• In Bmad, c.f. S24.12 Bmad manual, they used the also a transverse part, Eq.9 in Rosenzweig and Serafini 94

Still, I'm not sure if this helps the result. What method are you using in bmad tracking?

[jank324]

Is it still only the cavity that's causing the problem?

I don't think so looking at the section from ca. 550-1550m with the remaining cavities' voltage = 0 in the last screenshot I posted, there is still some deviation starting from ca. 875m in the screenshot (not 875m of the total beam line). I'm not sure if this is simply a gradual compounding error or if any one element is to blame. I will have to go through and check, but there is a couple hundred elements ...

Can you check first if the longitudinal components of the beam/ particle of cheetah and bmad agree?

Which properties exactly? I do have this sigma_s and sigma_p comparison for both beam types in Cheetah ... but this is obviously only within Cheetah.

What method are you using in bmad tracking?

For l0a, the first cavity and the first element where the Twiss parameters deviate, TRACKING_METHOD = Bmad_Standard.

Why would you thing focussing might not be the issue? The main difference is that the beta in Bmad is larger (12 vs 16). Alpha is smaller (-3.5 vs -4.7).

[cr-xu]

Can you check first if the longitudinal components of the beam/ particle of cheetah and bmad agree?

Which properties exactly? I do have this sigma_s and sigma_p comparison for both beam types in Cheetah ... but this is obviously only within Cheetah.

Yes, also the reference energy p etc. just to make sure. Bmad can not give out the sigma_p using single-particle tracking?

What method are you using in bmad tracking?

For l0a, the first cavity and the first element where the Twiss parameters deviate, TRACKING_METHOD = Bmad_Standard.

Why would you thing focussing might not be the issue? The main difference is that the beta in Bmad is larger (12 vs 16). Alpha is smaller (-3.5 vs -4.7).

Oh sorry, I totally missed that the transverse focusing is already included in the _cavity_rmatrix, Another 2 thoughts:

1. the phi0 in bmad is the phase of the reference particle with respect to the RF; The phi in cheetah and Ocelot Cavity seems to be the phase of the RF cavity. So probably ${\phi0}^\mathrm{bmad} = - \phi{0}^\mathrm{OCELOT}$, need to verify this though...

2. I think we already talked about it: we can add an option to enable double precision tracking. Should be fairly straightforward by replacing the torch.float32 everywhere with a class argument like self.precision

[jank324]

the phi0 in bmad is the phase of the reference particle with respect to the RF; The phi in cheetah and Ocelot Cavity seems to be the phase of the RF cavity. So probably, need to verify this though...

So, I've just tested what happens if I negate phi in the converse from Bmad to Cheetah. We go from

element.name = 'l0a'
beam_stepped.beta_x = tensor(12.8963)
beam_stepped.beta_y = tensor(12.9667)
beam_stepped.alpha_x = tensor(-3.5450)
beam_stepped.alpha_y = tensor(-3.5449)

to

element.name = 'l0a'
beam_stepped.beta_x = tensor(12.9823)
beam_stepped.beta_y = tensor(12.9714)
beam_stepped.alpha_x = tensor(-3.5528)
beam_stepped.alpha_y = tensor(-3.5604)

Interestingly, I also noticed that the beam in Bmad is perfectly round, suspiciously so actually:

Twiss at end of element:
                          A              B            Cbar                        C_mat
  Beta (m)        16.62625295    16.62625295  |   0.00000000   0.00000000      0.00000000   0.00000000
  Alpha           -4.71572777    -4.71572777  |   0.00000000   0.00000000      0.00000000   0.00000000
  Gamma (1/m)      1.39767442     1.39767442  |   Gamma_c =   1.00000000       Mode_Flip = F
  Phi (rad)        2.99653305     2.99653305            X              Y              Z
  Eta (m)          0.00000000     0.00000000     0.00000000     0.00000000     0.14515452
  Etap             0.00000000     0.00000000     0.00000000     0.00000000     0.00000000

I think we already talked about it: we can add an option to enable double precision tracking. Should be fairly straightforward by replacing the torch.float32 everywhere with a class argument like self.precision

I like the idea, but I think that, unless we realise we need it for this pull request, we should do that separately later.

Yes, also the reference energy p etc. just to make sure. Bmad can not give out the sigma_p using single-particle tracking?

Excerpt from Bmad output:

51  P0C_START                   =  5.9782004E+06 eV           BETA_START                  =  9.9636673E-01
52  E_TOT_START                 =  6.0000000E+06 eV           DELTA_E                     =  5.8000000E+07 eV
53  P0C                         =  6.3997960E+07 eV           BETA                        =  9.9996812E-01
54  E_TOT                       =  6.4000000E+07 eV           GAMMA                       =  1.2524488E+02

And Cheetah output:

element.name = 'l0a'
beam_stepped.beta_x = tensor(12.5934)
beam_stepped.beta_y = tensor(13.1527)
beam_stepped.alpha_x = tensor(-3.4415)
beam_stepped.alpha_y = tensor(-3.6063)
beam_stepped.energy = tensor(64000000., dtype=torch.float64)

Also here energy over the entire beam line as per Cheetah:

[jank324]

I think this is the corresponding Bmad code (ref. bmad/low_level/track_a_lcavity.f90):

!+
! Subroutine track_a_lcavity (orbit, ele, param, mat6, make_matrix)
!
! Bmad_standard tracking through a lcavity element.
!
! Modified version of:
!       J. Rosenzweig and L. Serafini
!       Phys Rev E, Vol. 49, p. 1599, (1994)
! with b_0 = b_-1 = 1. See the Bmad manual for more details.
!
! One must keep in mind that we are NOT using good canonical coordinates since
!   the energy of the reference particle is changing.
! This means that the resulting matrix will NOT be symplectic.
!
! Input:
!   orbit       -- Coord_struct: Starting position.
!   ele         -- ele_struct: Thick multipole element.
!   param       -- lat_param_struct: Lattice parameters.
!   mat6(6,6)   -- Real(rp), optional: Transfer matrix before the element.
!   make_matrix -- logical, optional: Propagate the transfer matrix? Default is false.
!
! Output:
!   orbit      -- coord_struct: End position.
!   mat6(6,6)  -- real(rp), optional: Transfer matrix through the element.
!-

subroutine track_a_lcavity (orbit, ele, param, mat6, make_matrix)

use bmad_interface, except_dummy => track_a_lcavity
use super_recipes_mod, only: super_zbrent

implicit none

type (coord_struct) :: orbit
type (ele_struct), target :: ele
type (lat_param_struct) :: param
type (em_field_struct) field

real(rp), optional :: mat6(6,6)
real(rp) length, pc_start, pc_end, gradient_ref, gradient_max, dz_factor, rel_p, coef, k2
real(rp) alpha, sin_a, cos_a, r_mat(2,2), dph
real(rp) phase, cos_phi, sin_phi, gradient_net, e_start, e_end, e_ratio, voltage_max, dp_dg, sqrt_8, f, k1
real(rp) dE_start, dE_end, dE, beta_start, beta_end, sqrt_beta12, dsqrt_beta12(6), f_ave, pc_start_ref
real(rp) pxy2, xp1, xp2, yp1, yp2, mc2, om, om_g, m2(2,2), kmat(6,6), ds, r_step, step_len
real(rp) dbeta1_dE1, dbeta2_dE2, dalpha_dt1, dalpha_dE1, dcoef_dt1, dcoef_dE1, z21, z22
real(rp) c_min, c_plu, dc_min, dc_plu, cos_term, dcos_phi, drp1_dr0, drp1_drp0, drp2_dr0, drp2_drp0
real(rp) an(0:n_pole_maxx), bn(0:n_pole_maxx), an_elec(0:n_pole_maxx), bn_elec(0:n_pole_maxx)
real(rp) E_tot_start, E_tot, p0c_start, p0c, phase0, phase1, phase2, dphase, ph_err(2), rf_phase
real(rp), parameter :: phase_abs_tol = 1e-4_rp

integer ix_mag_max, ix_elec_max, n_step, status

logical, optional :: make_matrix

character(*), parameter :: r_name = 'track_a_lcavity'

! 

if (ele%value(rf_frequency$) == 0  .and. (ele%value(voltage$) /= 0 .or. ele%value(voltage_err$) /= 0)) then
  call out_io (s_error$, r_name, 'LCAVITY ELEMENT HAS ZERO RF_FREQUENCY: ' // ele%name)
  orbit%state = lost$
  return
endif

length = orbit%time_dir * ele%value(l$)
if (length == 0) return
n_step = 1
r_step = rp8(orbit%time_dir) / n_step
step_len = length / n_step

call multipole_ele_to_ab (ele, .false., ix_mag_max, an,      bn,      magnetic$, include_kicks$)
call multipole_ele_to_ab (ele, .false., ix_elec_max, an_elec, bn_elec, electric$)

call offset_particle (ele, set$, orbit, mat6 = mat6, make_matrix = make_matrix)

if (ix_mag_max > -1)  call ab_multipole_kicks (an,      bn,      ix_mag_max,  ele, orbit, magnetic$, 0.5_rp*r_step,   mat6, make_matrix)
if (ix_elec_max > -1) call ab_multipole_kicks (an_elec, bn_elec, ix_elec_max, ele, orbit, electric$, 0.5_rp*step_len, mat6, make_matrix)

! The RF phase is defined with respect to the time at the beginning of the element.
! So if dealing with a slave element and absolute time tracking then need to correct.
! Note: phi0_autoscale is not used here since bmad_standard tracking by design gives the correct tracking.
! In fact, using phi0_autoscale would be a mistake if, say, tracking_method = runge_kutta, mat6_calc_method = bmad_standard.

phase = twopi * (ele%value(phi0_err$) + ele%value(phi0$) + ele%value(phi0_multipass$) + &
           (particle_rf_time (orbit, ele, .false.) - rf_ref_time_offset(ele)) * ele%value(rf_frequency$))
if (bmad_com%absolute_time_tracking .and. ele%orientation*orbit%time_dir*orbit%direction == -1) then
  phase = phase - twopi * ele%value(rf_frequency$) * ele%value(delta_ref_time$)
endif
phase = modulo2(phase, pi)

call rf_coupler_kick (ele, param, first_track_edge$, phase, orbit, mat6, make_matrix)

!

if (orbit%time_dir * orbit%direction == 1) then
  E_tot_start = ele%value(E_tot_start$)
  E_tot       = ele%value(E_tot$)
  p0c_start   = ele%value(p0c_start$)
  p0c         = ele%value(p0c$)
else
  E_tot_start = ele%value(E_tot$)
  E_tot       = ele%value(E_tot_start$)
  p0c_start   = ele%value(p0c$)
  p0c         = ele%value(p0c_start$)
endif

rel_p = 1 + orbit%vec(6)
gradient_ref = (E_tot - E_tot_start) / length
mc2 = mass_of(orbit%species)

pc_start = p0c_start * rel_p
beta_start = orbit%beta
E_start = pc_start / beta_start 

gradient_max = e_accel_field(ele, gradient$, .true.)
phase0 = phase
phase1 = phase

! The idea is to find the "average" phase so that forwards tracking followed by backwards time 
! tracking comes back to the original position.
! If the phase change from start to finish is more than pi then this whole calculation is garbage.
! In this case, the particle is considered lost

ph_err(1) = phase_func(phase1, status)
if (abs(dphase) > pi) then
  orbit%state = lost_pz_aperture$
  return

elseif (abs(dphase) < phase_abs_tol) then
  rf_phase = phase1 + 0.5_rp * dphase
  call track_this_lcavity(rf_phase, orbit, make_matrix)

else
  phase2 = phase0 + dphase
  ph_err(2) = phase_func(phase2, status)
  ! At very low energies can happen that phase0 and phase0+dphase do not bracket the solution.
  ! In this case try extrapolating. Factor of 2 to try to make sure root is bracketed.
  do
    if (ph_err(1)*ph_err(2) <= 0) exit
    dph = phase2 - phase1
    phase2 = phase1 + 2 * sign_of(dph) * max(abs(dph), 0.1)
    if (abs(phase2 - phase1) > pi) then
      orbit%state = lost_pz_aperture$
      return
    endif
    ph_err(2) = phase_func(phase2, status)
  enddo

  rf_phase = super_zbrent (phase_func, phase1, phase2, 0.0_rp, phase_abs_tol, status, ph_err)
  call track_this_lcavity(rf_phase, orbit, make_matrix)
endif

! And end

if (nint(ele%value(cavity_type$)) == traveling_wave$ .and. fringe_here(ele, orbit, second_track_edge$)) then
  ds = bmad_com%significant_length / 10  ! Make sure inside field region
  call em_field_calc (ele, param, length - ds, orbit, .true., field, logic_option(.false., make_matrix))
  f = -charge_of(orbit%species) / (2 * p0c)

  if (logic_option(.false., make_matrix)) then
    call mat_make_unit(kmat)
    kmat(2,1) = -f * (field%dE(3,1) * orbit%vec(1) + field%E(3))
    kmat(2,3) = -f * field%dE(3,2) * orbit%vec(1) 
    kmat(2,5) = -f * field%dE(3,3) * orbit%vec(1) * beta_end
    kmat(2,6) =  f * field%E(3) * orbit%vec(1) * f / pc_end
    kmat(4,1) = -f * field%dE(3,1) * orbit%vec(3)
    kmat(4,3) = -f * (field%dE(3,2) * orbit%vec(3) + field%E(3))
    kmat(4,5) = -f * field%dE(3,3) * orbit%vec(3) * beta_end
    kmat(4,6) =  f * field%E(3) * orbit%vec(1) * f / pc_end
    mat6 = matmul(kmat, mat6)
  endif

  orbit%vec(2) = orbit%vec(2) - f * field%E(3) * orbit%vec(1)
  orbit%vec(4) = orbit%vec(4) - f * field%E(3) * orbit%vec(3)
endif

! Coupler kick

call rf_coupler_kick (ele, param, second_track_edge$, phase, orbit, mat6, make_matrix)

if (ix_elec_max > -1) call ab_multipole_kicks (an_elec, bn_elec, ix_elec_max, ele, orbit, electric$, 0.5_rp*step_len, mat6, make_matrix)
if (ix_mag_max > -1)  call ab_multipole_kicks (an,      bn,      ix_mag_max,  ele, orbit, magnetic$, 0.5_rp*r_step,   mat6, make_matrix)

call offset_particle (ele, unset$, orbit, mat6 = mat6, make_matrix = make_matrix)

!---------------------------------------------------------------------------------------------
contains

! Returns rf_phase_at_end - rf_phase_at_start
! Note: rf_phase_at_start = phase0 global variable
!       rf_phase is the phase used to calculate the accelerating gradient.

function dphase_end_minus_start(rf_phase) result (dphase)

type (coord_struct) this_orb
real(rp) rf_phase, dphase

!

this_orb = orbit
call track_this_lcavity (rf_phase, this_orb, .false.)
dphase = twopi * (ele%value(rf_frequency$) / c_light) * (orbit%vec(5) / orbit%beta - this_orb%vec(5) / this_orb%beta)

end function dphase_end_minus_start

!---------------------------------------------------------------------------------------------
! contains

function phase_func (rf_phase, status) result (phase_err)

real(rp), intent(in) :: rf_phase
real(rp) phase_err
integer status

!

dphase = dphase_end_minus_start(rf_phase)
phase_err = rf_phase - (phase0 + 0.5_rp * dphase)
if (abs(phase_err) < phase_abs_tol) phase_err = 0

end function phase_func

!---------------------------------------------------------------------------------------------
! contains

! rf_phase is the phase used to calculate the change in energy.

subroutine track_this_lcavity (rf_phase, orbit, make_matrix)

type (coord_struct) orbit
real(rp) rf_phase
logical, optional :: make_matrix

!

cos_phi = cos(rf_phase)
sin_phi = sin(rf_phase)
gradient_net = gradient_max * cos_phi + gradient_shift_sr_wake(ele, param)

dE = gradient_net * length
E_end = E_start + dE
if (E_end <= mass_of(orbit%species)) then
  orbit%state = lost_pz_aperture$
  orbit%vec(6) = -1.01  ! Something less than -1
  if (present(mat6)) mat6 = 0
  return
endif

call convert_total_energy_to (E_end, orbit%species, pc = pc_end, beta = beta_end)
E_ratio = E_end / E_start
sqrt_beta12 = sqrt(beta_start/beta_end)
mc2 = mass_of(orbit%species)

! Traveling_wave fringe (standing_wave fringe is built-in to the body formulas)

if (nint(ele%value(cavity_type$)) == traveling_wave$ .and. fringe_here(ele, orbit, first_track_edge$)) then
  ds = bmad_com%significant_length / 10  ! Make sure inside field region
  call em_field_calc (ele, param, ds, orbit, .true., field, logic_option(.false., make_matrix))
  f = charge_of(orbit%species) / (2 * p0c_start)

  if (logic_option(.false., make_matrix)) then
    call mat_make_unit(kmat)
    kmat(2,1) = -f * (field%dE(3,1) * orbit%vec(1) + field%E(3))
    kmat(2,3) = -f * field%dE(3,2) * orbit%vec(1) 
    kmat(2,5) = -f * field%dE(3,3) * orbit%vec(1) * beta_start
    kmat(2,6) =  f * field%E(3) * orbit%vec(1) * f / pc_start
    kmat(4,1) = -f * field%dE(3,1) * orbit%vec(3)
    kmat(4,3) = -f * (field%dE(3,2) * orbit%vec(3) + field%E(3))
    kmat(4,5) = -f * field%dE(3,3) * orbit%vec(3) * beta_start
    kmat(4,6) =  f * field%E(3) * orbit%vec(1) * f / pc_start
    mat6 = matmul(kmat, mat6)
  endif

  orbit%vec(2) = orbit%vec(2) - f * field%E(3) * orbit%vec(1)
  orbit%vec(4) = orbit%vec(4) - f * field%E(3) * orbit%vec(3)
endif

! Body tracking longitudinal

dp_dg = length * (2*E_start + dE) / (pc_end + pc_start)   ! = (pc_end - pc_start) / gradient_net

if (logic_option(.false., make_matrix)) then
  om = twopi * ele%value(rf_frequency$) / c_light
  om_g = om * gradient_max * length

  dbeta1_dE1 = mc2**2 / (pc_start * E_start**2)
  dbeta2_dE2 = mc2**2 / (pc_end * E_end**2)

  ! Convert from (x, px, y, py, z, pz) to (x, x', y, y', c(t_ref-t), E) coords 
  mat6(2,:) = mat6(2,:) / rel_p - orbit%vec(2) * mat6(6,:) / rel_p**2
  mat6(4,:) = mat6(4,:) / rel_p - orbit%vec(4) * mat6(6,:) / rel_p**2

  m2(1,:) = [1/orbit%beta, -orbit%vec(5) * mc2**2 * orbit%p0c / (pc_start**2 * E_start)]
  m2(2,:) = [0.0_rp, orbit%p0c * orbit%beta]
  mat6(5:6,:) = matmul(m2, mat6(5:6,:))

  ! longitudinal track
  call mat_make_unit (kmat)
  kmat(6,5) = om_g * sin_phi
endif

! Convert to (x', y', c(t_ref-t), E) coords

orbit%vec(2) = orbit%vec(2) / rel_p    ! Convert to x'
orbit%vec(4) = orbit%vec(4) / rel_p    ! Convert to y'
orbit%vec(5) = orbit%vec(5) / beta_start
orbit%vec(6) = rel_p * orbit%p0c / orbit%beta
orbit%t = orbit%t + dp_dg / c_light

! Body tracking. Note: Transverse kick only happens with standing wave cavities.

!--------------------------------------------------------------------
! Traveling wave
if (nint(ele%value(cavity_type$)) == traveling_wave$) then
  f_ave = (pc_start + pc_end) / (2 * pc_end)
  dz_factor = (orbit%vec(2)**2 + orbit%vec(4)**2) * f_ave**2 * dp_dg / 2

  if (logic_option(.false., make_matrix)) then
    if (abs(dE) <  1d-4*(pc_end+pc_start)) then
      kmat(5,5) = 1 - length * (-mc2**2 * kmat(6,5) / (2 * pc_start**3) + mc2**2 * dE * kmat(6,5) * E_start / pc_start**5)
      kmat(5,6) = -length * (-dbeta1_dE1 / beta_start**2 + 2 * mc2**2 * dE / pc_start**4 + &
                      (mc2 * dE)**2 / (2 * pc_start**5) - 5 * (mc2 * dE)**2 / (2 * pc_start**5))
    else
      kmat(5,5) = 1 - kmat(6,5) / (beta_end * gradient_net) + kmat(6,5) * (pc_end - pc_start) / (gradient_net**2 * length)
      kmat(5,6) = -1 / (beta_end * gradient_net) + 1 / (beta_start * gradient_net)
    endif

    kmat(1,2) = length * f_ave
    kmat(1,5) = -orbit%vec(2) * kmat(5,6) * length * pc_start / (2 * pc_end**2)
    kmat(1,6) =  orbit%vec(2) * (1 - kmat(6,6) * pc_start / pc_end) / (2 * pc_end)

    kmat(2,2) = pc_start / pc_end
    kmat(2,5) = -orbit%vec(2) * kmat(5,6) * pc_start / pc_end**2
    kmat(2,6) =  orbit%vec(2) * (1 - kmat(6,6) * pc_start / pc_end) / pc_end 

    kmat(3,4) = length * f_ave
    kmat(3,5) = -orbit%vec(4) * kmat(5,6) * length * pc_start / (2 * pc_end**2)
    kmat(3,6) =  orbit%vec(4) * (1 - kmat(6,6) * pc_start / pc_end) / (2 * pc_end)

    kmat(4,4) = pc_start / pc_end
    kmat(4,5) = -orbit%vec(4) * kmat(5,6) * pc_start / pc_end**2
    kmat(4,6) =  orbit%vec(4) * (1 - kmat(6,6) * pc_start / pc_end) / pc_end 

    kmat(5,2) = -orbit%vec(2) * f_ave**2 * dp_dg 
    kmat(5,4) = -orbit%vec(4) * f_ave**2 * dp_dg 

    mat6 = matmul(kmat, mat6)
  endif

  orbit%vec(1) = orbit%vec(1) + orbit%vec(2) * length * f_ave
  orbit%vec(2) = orbit%vec(2) * pc_start / pc_end 
  orbit%vec(3) = orbit%vec(3) + orbit%vec(4) * length * f_ave
  orbit%vec(4) = orbit%vec(4) * pc_start / pc_end 

  orbit%vec(5) = orbit%vec(5) - dz_factor
  orbit%t = orbit%t + dz_factor / c_light

!--------------------------------------------------------------------
! Standing wave
else
  sqrt_8 = 2 * sqrt_2
  voltage_max = gradient_max * length

  if (abs(voltage_max * cos_phi) < 1d-5 * E_start) then
    f = voltage_max / E_start
    alpha = f * (1 + f * cos_phi / 2)  / sqrt_8
    coef = length * (1 - voltage_max * cos_phi / (2 * E_start))
  else
    alpha = log(E_ratio) / (sqrt_8 * cos_phi)
    coef = sqrt_8 * E_start * sin(alpha) / gradient_max
  endif

  cos_a = cos(alpha)
  sin_a = sin(alpha)

  if (logic_option(.false., make_matrix)) then
    if (abs(voltage_max * cos_phi) < 1d-5 * E_start) then
      dalpha_dt1 = f * f * om * sin_phi / (2 * sqrt_8) 
      dalpha_dE1 = -(voltage_max / E_start**2 + voltage_max**2 * cos_phi / E_start**3) / sqrt_8
      dcoef_dt1 = -length * sin_phi * om_g / (2 * E_start)
      dcoef_dE1 = length * voltage_max * cos_phi / (2 * E_start**2)
    else
      dalpha_dt1 = kmat(6,5) / (E_end * sqrt_8 * cos_phi) - log(E_ratio) * om * sin_phi / (sqrt_8 * cos_phi**2)
      dalpha_dE1 = 1 / (E_end * sqrt_8 * cos_phi) - 1 / (E_start * sqrt_8 * cos_phi)
      dcoef_dt1 = sqrt_8 * E_start * cos(alpha) * dalpha_dt1 / gradient_max
      dcoef_dE1 = coef / E_start + sqrt_8 * E_start * cos(alpha) * dalpha_dE1 / gradient_max
    endif

    z21 = -gradient_max / (sqrt_8 * E_end)
    z22 = E_start / E_end  

    c_min = cos_a - sqrt_2 * sin_a * cos_phi
    c_plu = cos_a + sqrt_2 * sin_a * cos_phi
    dc_min = -sin_a - sqrt_2 * cos_a * cos_phi 
    dc_plu = -sin_a + sqrt_2 * cos_a * cos_phi 

    cos_term = 1 + 2 * cos_phi**2
    dcos_phi = om * sin_phi

    kmat(1,1) =  c_min
    kmat(1,2) =  coef 
    kmat(2,1) =  z21 * cos_term * sin_a
    kmat(2,2) =  c_plu * z22

    kmat(1,5) = orbit%vec(1) * (dc_min * dalpha_dt1 - sqrt_2 * sin_a * dcos_phi) + & 
                 orbit%vec(2) * (dcoef_dt1)

    kmat(1,6) = orbit%vec(1) * dc_min * dalpha_dE1 + orbit%vec(2) * dcoef_dE1

    kmat(2,5) = orbit%vec(1) * z21 * (4 * cos_phi * dcos_phi * sin_a + cos_term * cos_a * dalpha_dt1) + &
                 orbit%vec(1) * (-kmat(2,1) * kmat(6,5) / (E_end)) + &
                 orbit%vec(2) * z22 * (dc_plu * dalpha_dt1 + sqrt_2 * sin_a * dcos_phi - c_plu * kmat(6,5) / E_end)

    kmat(2,6) = orbit%vec(1) * z21 * (cos_term * cos_a * dalpha_dE1) + &
                 orbit%vec(1) * (-kmat(2,1) / (E_end)) + &
                 orbit%vec(2) * z22 * dc_plu * dalpha_dE1 + &
                 orbit%vec(2) * c_plu * (E_end - E_start) / E_end**2

    kmat(3:4,3:4) = kmat(1:2,1:2)

    kmat(3,5) = orbit%vec(3) * (dc_min * dalpha_dt1 - sqrt_2 * beta_start * sin_a * dcos_phi) + & 
                 orbit%vec(4) * (dcoef_dt1)

    kmat(3,6) = orbit%vec(3) * (dc_min * dalpha_dE1 - sqrt_2 * dbeta1_dE1 * sin_a * cos_phi) + &
                 orbit%vec(4) * (dcoef_dE1)

    kmat(4,5) = orbit%vec(3) * z21 * (4 * cos_phi * dcos_phi * sin_a + cos_term * cos_a * dalpha_dt1) + &
                 orbit%vec(3) * (-kmat(2,1) * kmat(6,5) / (E_end)) + &
                 orbit%vec(4) * z22 * (dc_plu * dalpha_dt1 + sqrt_2 * sin_a * dcos_phi - c_plu * kmat(6,5) / E_end)

    kmat(4,6) = orbit%vec(3) * z21 * (cos_term * cos_a * dalpha_dE1) + &
                 orbit%vec(3) * (-kmat(2,1) / E_end) + &
                 orbit%vec(4) * z22 * dc_plu * dalpha_dE1 + &
                 orbit%vec(4) * c_plu * (E_end - E_start) / E_end**2

    ! Correction to z for finite x', y'
    ! Note: Corrections to kmat(5,5) and kmat(5,6) are ignored since these are small (quadratic in the transvers coords).

    c_plu = sqrt_2 * cos_phi * cos_a + sin_a

    drp1_dr0  = -gradient_net / (2 * E_start)
    drp1_drp0 = 1

    xp1 = drp1_dr0 * orbit%vec(1) + drp1_drp0 * orbit%vec(2)
    yp1 = drp1_dr0 * orbit%vec(3) + drp1_drp0 * orbit%vec(4)

    drp2_dr0  = (c_plu * z21)
    drp2_drp0 = (cos_a * z22)

    xp2 = drp2_dr0 * orbit%vec(1) + drp2_drp0 * orbit%vec(2)
    yp2 = drp2_dr0 * orbit%vec(3) + drp2_drp0 * orbit%vec(4)

    kmat(5,1) = -(orbit%vec(1) * (drp1_dr0**2 + drp1_dr0*drp2_dr0 + drp2_dr0**2) + &
                   orbit%vec(2) * (drp1_dr0 * drp1_drp0 + drp2_dr0 * drp2_drp0 + &
                                (drp1_dr0 * drp2_drp0 + drp1_drp0 * drp2_dr0) / 2)) * dp_dg / 3

    kmat(5,2) = -(orbit%vec(2) * (drp1_drp0**2 + drp1_drp0*drp2_drp0 + drp2_drp0**2) + &
                   orbit%vec(1) * (drp1_dr0 * drp1_drp0 + drp2_dr0 * drp2_drp0 + &
                                (drp1_dr0 * drp2_drp0 + drp1_drp0 * drp2_dr0) / 2)) * dp_dg / 3

    kmat(5,3) = -(orbit%vec(3) * (drp1_dr0**2 + drp1_dr0*drp2_dr0 + drp2_dr0**2) + &
                   orbit%vec(4) * (drp1_dr0 * drp1_drp0 + drp2_dr0 * drp2_drp0 + &
                                (drp1_dr0 * drp2_drp0 + drp1_drp0 * drp2_dr0) / 2)) * dp_dg / 3

    kmat(5,4) = -(orbit%vec(4) * (drp1_drp0**2 + drp1_drp0*drp2_drp0 + drp2_drp0**2) + &
                   orbit%vec(3) * (drp1_dr0 * drp1_drp0 + drp2_dr0 * drp2_drp0 + &
                                (drp1_dr0 * drp2_drp0 + drp1_drp0 * drp2_dr0) / 2)) * dp_dg / 3

    ! beta /= 1 corrections

    dsqrt_beta12 = -0.5_rp * sqrt_beta12 * dbeta2_dE2 * kmat(6,:) / beta_end
    dsqrt_beta12(6) = dsqrt_beta12(6) + 0.5_rp * sqrt_beta12 * dbeta1_dE1 / beta_start 

    kmat(1:4,1:4) = sqrt_beta12 * kmat(1:4,1:4)

    !

    mat6 = matmul(kmat, mat6)
  endif

  k1 = -gradient_net / (2 * E_start)
  orbit%vec(2) = orbit%vec(2) + k1 * orbit%vec(1)    ! Entrance kick
  orbit%vec(4) = orbit%vec(4) + k1 * orbit%vec(3)    ! Entrance kick

  xp1 = orbit%vec(2)
  yp1 = orbit%vec(4)

  r_mat(1,1) =  cos_a
  r_mat(1,2) =  coef 
  r_mat(2,1) = -sin_a * gradient_max / (sqrt_8 * E_end)
  r_mat(2,2) =  cos_a * E_start / E_end

  orbit%vec(1:2) = sqrt_beta12 * matmul(r_mat, orbit%vec(1:2))   ! Modified R&S Eq 9.
  orbit%vec(3:4) = sqrt_beta12 * matmul(r_mat, orbit%vec(3:4))

  xp2 = orbit%vec(2)
  yp2 = orbit%vec(4)

  ! Correction of z for finite transverse velocity assumes a uniform change in slope.

  dz_factor = (xp1**2 + xp2**2 + xp1*xp2 + yp1**2 + yp2**2 + yp1*yp2) * dp_dg / 6
  orbit%vec(5) = orbit%vec(5) - dz_factor
  orbit%t = orbit%t + dz_factor / c_light

  !

  k2 = gradient_net / (2 * E_end) 
  orbit%vec(2) = orbit%vec(2) + k2 * orbit%vec(1)         ! Exit kick
  orbit%vec(4) = orbit%vec(4) + k2 * orbit%vec(3)         ! Exit kick
endif

! Shift ref momentum if reached element body entrance end

orbit%vec(5) = orbit%vec(5) - (dp_dg - length * (E_tot_start + E_tot) / (p0c + p0c_start))

orbit%vec(6) = (pc_end - p0c) / p0c 
orbit%p0c = p0c

! Convert back from (x', y', c(t_ref-t), E) coords

if (logic_option(.false., make_matrix)) then
  rel_p = pc_end / p0c
  mat6(2,:) = rel_p * mat6(2,:) + orbit%vec(2) * mat6(6,:) / (p0c * beta_end)
  mat6(4,:) = rel_p * mat6(4,:) + orbit%vec(4) * mat6(6,:) / (p0c * beta_end)

  m2(1,:) = [beta_end, orbit%vec(5) * mc2**2 / (pc_end * E_end**2)]
  m2(2,:) = [0.0_rp, 1 / (p0c * beta_end)]

  mat6(5:6,:) = matmul(m2, mat6(5:6,:))
endif

orbit%vec(2) = orbit%vec(2) * (1 + orbit%vec(6))  ! Convert back to px
orbit%vec(4) = orbit%vec(4) * (1 + orbit%vec(6))  ! Convert back to py
orbit%vec(5) = orbit%vec(5) * beta_end

orbit%beta = beta_end

end subroutine track_this_lcavity

end subroutine track_a_lcavity

[jank324]

[jank324]

[jank324]

[jank324]

[jank324]

Okay ... it's working now.

The cavity that caused this whole cavity investigation had as output Twiss in Cheetah

beta_x  = 12.856815106445865
beta_y  = 12.856815106442994
alpha_x = -3.524868604372471
alpha_y = -3.5248686043537316

when in Bmad it had

  Beta (m)        16.62625295    16.62625295  |   0.00000000   0.00000000      0.00000000   0.00000000
  Alpha           -4.71572777    -4.71572777  |   0.00000000   0.00000000      0.00000000   0.00000000

After the above fixes and looking at a minimal working example of that cavity

... we find that for the minimal working example where the input beams are identical, the error is significantly smaller, i.e. the large error in the full lattice is an accumulated error, hence it is acceptable.

[jank324]

@cr-xu can you create issues for those todos? I think that's the best way to keep track of them.

Regarding what is not properly converted: I actually don't know. I would have to comb through the Bmad docs to figure it out. I literally just went and fixed all the elements that came up in the LCLS Cu-HXR lattice. "Test-driven development" ... 😄