Change coordinate system

In our meeting today, we talked about changing the default coordinate system in Cheetah. @cr-xu @jp-ga can we assign this to either one of you? I think you have more expertise regarding this than I do.

I'll try to clearly document the changes we want to make.
@jp-ga Please check if there's anything incorrect.

The coordinate system should be explained properly in the documentation.

Transverse

Change the current coordinates $(x, x', y, y')$ to the canonical cooridnates $(x,p_x, y, p_y)$.

The new transverse momentum should be the scaled one w.r.t. the reference momentum $p_x = \frac{P_x}{P_0}$ (as in bmad).

This should be merely a change in the documentation at this point as the linear transfer matrices are still correct.

Longitudinal

For the longitudinal coordinate we will stay with $(\tau, \delta)$, where $\tau = ct - \frac{s}{\beta_0} = c (t - t_0(s)) = c \Delta t$, and $\delta=\frac{E}{p_0 c} - \frac{1}{\beta_0} = \frac{E-E_0}{p_0 c} = \frac{E - E_0}{\beta_0 E_0}$. For particle arriving at a earlier time than the reference particle (bunch head), $\tau <0$.
The longitudinal coordinates are derived from the $F_2 = (\frac{z}{\beta_0} - ct)(\frac{1}{\beta_0} + \delta)$ (A. Wolski Eq. 2.45).

Note:

This coordinate $(\tau, \delta)$ is the same as in OCELOT
In A. Wolski's book, the sign of $\tau$ is inverted
Related to the Mad-x coordinates $(T, PT)$ as $(\tau, \delta) = (- T, PT)$, (according to this paper also in TEAPOT, PyORBIT, )
Related to the Bmad coordinates $(z, p_z)$ (Bmad manual 15.28, 15.32) as $z_\text{Bmad} = -\beta \tau$ (not the reference velocity $\beta_0$) and $p_{z, \text{Bmad}} = \frac{\Delta E}{E_0} = \beta_0 \delta$

In addition, we should also clearly distinguish the reference energy $E_0$ and reference momentum $p_0$ ($p_0 c$), which will be different for non-relativisitic particles.
Right now the attribute is called Beam.energy and should return the energy. We could add another property Beam.p0c as convenient method

If we are already in here, maybe we should also think about the fact that beam's currently don't track their reference s and whether doing so would be make sense (for example in light of the space charge implementation).

@jank324 Now that we're clarifying the meaning of a few variables, does it make sense to rename also the properties $s$ and $p$?

In Cheetah now $s$ actually stands for $\tau$ as in OCELOT (or $-z$ as in other simulation code like mad-x). $s$ is usually used as the position of the reference particle along the beamline, (independent variable). Also having the real $s$ would make sense now since we have space charge integration etc.

$p$ is the dimensionless relative energy offset, which is $\delta$ in most literatures, probably we can rename it to something like $dp$?
So that we can provide a new property $p$ or $momentum$ to denote the actual momentum of a particle.

@jank324 Now that we're clarifying the meaning of a few variables, does it make sense to rename also the properties s and p?

In Cheetah now s actually stands for τ as in OCELOT (or −z as in other simulation code like mad-x). s is usually used as the position of the reference particle along the beamline, (independent variable). Also having the real s would make sense now since we have space charge integration etc.

p is the dimensionless relative energy offset, which is δ in most literatures, probably we can rename it to something like dp? So that we can provide a new property p or momentum to denote the actual momentum of a particle.

Yes, this sounds like it would make a lot of sense to do this now.

I'll try to clearly document the changes we want to make. @jp-ga Please check if there's anything incorrect.

The coordinate system should be explained properly in the documentation.

Transverse

Change the current coordinates (x,x′,y,y′) to the canonical cooridnates (x,px,y,py).

The new transverse momentum should be the scaled one w.r.t. the reference momentum px=PxP0 (as in bmad).

This should be merely a change in the documentation at this point as the linear transfer matrices are still correct.

Longitudinal

For the longitudinal coordinate we will stay with (τ,δ), where τ=ct−sβ0=c(t−t0(s))=cΔt, and δ=Ep0c−1β0=E−E0p0c=E−E0β0E0. For particle arriving at a earlier time than the reference particle (bunch head), τ<0. The longitudinal coordinates are derived from the F2=(zβ0−ct)(1β0+δ) (A. Wolski Eq. 2.45).

Note:

This coordinate (τ,δ) is the same as in OCELOT

In A. Wolski's book, the sign of τ is inverted

Related to the Mad-x coordinates (T,PT) as (τ,δ)=(−T,PT), (according to this paper also in TEAPOT, PyORBIT, )

Related to the Bmad coordinates (z,pz) (Bmad manual 15.28, 15.32) as zBmad=−βτ (not the reference velocity β0) and pz,Bmad=ΔEE0=β0δ

The longitudinal momentum coordinate in Bmad is $p_{z, \mathrm{Bmad}} = \frac{\Delta P}{P_0}$, where $P$ is the total momentum of the particle (Bmad manual Eq. 15.31). The equation you referenced is an approximation for ultra-relativistic particles. Other than that, it looks good to me.

In addition, we should also clearly distinguish the reference energy E0 and reference momentum p0 (p0c), which will be different for non-relativisitic particles. Right now the attribute is called Beam.energy and should return the energy. We could add another property Beam.p0c as convenient method

I think the distinction between $E_0$ and $p_0$ is clear. What is not clear in is that Beam.energy is $E_0$. It could be renamed to Beam.ref_energy or something similar.

As a reference, this is how I implemented the coordinate transform for bmadx:

cheetah/cheetah/bmadx_utils.py

Lines 4 to 62 in 0d702e6

    
           def cheetah_to_bmad_coords( 
        
               cheetah_coords: torch.Tensor, ref_energy: torch.Tensor, mc2: torch.Tensor 
        
           ) -> torch.Tensor: 
        
               """Transforms Cheetah coordinates to Bmad coordinates. 
        
               :param cheetah_coords: 7-dimensional particle vectors in Cheetah coordinates. 
        
               :param ref_energy: reference energy in eV. 
        
               """ 
        
               # initialize Bmad coordinates: 
        
               bmad_coords = cheetah_coords[..., :6].clone() 
        
               # Cheetah longitudinal coords: 
        
               tau = cheetah_coords[..., 4] 
        
               delta = cheetah_coords[..., 5] 
        
               # compute p0c and bmad z, pz 
        
               p0c = torch.sqrt(ref_energy**2 - mc2**2) 
        
               energy = ref_energy + delta * p0c 
        
               p = torch.sqrt(energy**2 - mc2**2) 
        
               beta = p / energy 
        
               z = -beta * tau 
        
               pz = (p - p0c) / p0c 
        
               # Bmad coords: 
        
               bmad_coords[..., 4] = z 
        
               bmad_coords[..., 5] = pz 
        
               return bmad_coords, p0c 
        
           def bmad_to_cheetah_coords( 
        
               bmad_coords: torch.Tensor, p0c: torch.Tensor, mc2: torch.Tensor 
        
           ) -> torch.Tensor: 
        
               """Transforms Bmad coordinates to Cheetah coordinates. 
        
               :param bmad_coords: 6-dimensional particle vectors in Bmad coordinates. 
        
               :param p0c: reference momentum in eV/c. 
        
               """ 
        
               # initialize Cheetah coordinates: 
        
               cheetah_coords = torch.ones( 
        
                   (*bmad_coords.shape[:-1], 7), dtype=bmad_coords.dtype, device=bmad_coords.device 
        
               ) 
        
               cheetah_coords[..., :6] = bmad_coords.clone() 
        
               # Bmad longitudinal coords: 
        
               z = bmad_coords[..., 4] 
        
               pz = bmad_coords[..., 5] 
        
               # compute ref_energy and Cheetah tau, delta 
        
               ref_energy = torch.sqrt(p0c**2 + mc2**2) 
        
               p = (1 + pz) * p0c 
        
               energy = torch.sqrt(p**2 + mc2**2) 
        
               beta = p / energy 
        
               tau = -z / beta 
        
               delta = (energy - ref_energy) / p0c 
        
               # Cheetah coords: 
        
               cheetah_coords[..., 4] = tau 
        
               cheetah_coords[..., 5] = delta 
        
               return cheetah_coords, ref_energy

	def cheetah_to_bmad_coords(
	cheetah_coords: torch.Tensor, ref_energy: torch.Tensor, mc2: torch.Tensor
	) -> torch.Tensor:
	"""Transforms Cheetah coordinates to Bmad coordinates.
	:param cheetah_coords: 7-dimensional particle vectors in Cheetah coordinates.
	:param ref_energy: reference energy in eV.
	"""
	# initialize Bmad coordinates:
	bmad_coords = cheetah_coords[..., :6].clone()

	# Cheetah longitudinal coords:
	tau = cheetah_coords[..., 4]
	delta = cheetah_coords[..., 5]

	# compute p0c and bmad z, pz
	p0c = torch.sqrt(ref_energy2 - mc22)
	energy = ref_energy + delta * p0c
	p = torch.sqrt(energy2 - mc22)
	beta = p / energy
	z = -beta * tau
	pz = (p - p0c) / p0c

	# Bmad coords:
	bmad_coords[..., 4] = z
	bmad_coords[..., 5] = pz

	return bmad_coords, p0c


	def bmad_to_cheetah_coords(
	bmad_coords: torch.Tensor, p0c: torch.Tensor, mc2: torch.Tensor
	) -> torch.Tensor:
	"""Transforms Bmad coordinates to Cheetah coordinates.
	:param bmad_coords: 6-dimensional particle vectors in Bmad coordinates.
	:param p0c: reference momentum in eV/c.
	"""
	# initialize Cheetah coordinates:
	cheetah_coords = torch.ones(
	(*bmad_coords.shape[:-1], 7), dtype=bmad_coords.dtype, device=bmad_coords.device
	)
	cheetah_coords[..., :6] = bmad_coords.clone()

	# Bmad longitudinal coords:
	z = bmad_coords[..., 4]
	pz = bmad_coords[..., 5]

	# compute ref_energy and Cheetah tau, delta
	ref_energy = torch.sqrt(p0c2 + mc22)
	p = (1 + pz) * p0c
	energy = torch.sqrt(p2 + mc22)
	beta = p / energy
	tau = -z / beta
	delta = (energy - ref_energy) / p0c

	# Cheetah coords:
	cheetah_coords[..., 4] = tau
	cheetah_coords[..., 5] = delta

	return cheetah_coords, ref_energy