Add twiss parameter calculations

[jank324]

Adding twiss parameters of properties of the Beam class. This way they can easily be a computed from tracked ParameterBeam and ParticleBeam.

[jank324]

So this seems to work now. Here are a couple points we might need to think about:

• I think right now this doesn't work for ParameterBeam, but it absolutely should. The reason being that the emittance computation uses the actual particles:

  @property
  def emittance_x(self) -> torch.Tensor:
  return torch.sqrt(
      self.sigma_x**2 * self.sigma_xp**2
      - torch.mean((self.xs - self.mu_x) * (self.xps - self.mu_xp)) ** 2
  )

  the second part of this after the - this is inspired by Ocelot's

  tws.xpx = np.mean((x - tws.x) * (px - tws.px))

  which I thought is the same as beam.sigma_x * beam.sigma_xp in Cheetah, but apparently it's not.

• Is the twiss parameter computation from the Astra beam compared to Ocelot enough testing of this?

• We should probably also implement creating either type of beam from twiss parameters.

[cr-xu]

So this seems to work now. Here are a couple points we might need to think about:

• I think right now this doesn't work for ParameterBeam, but it absolutely should. The reason being that the emittance computation uses the actual particles:

@property
def emittance_x(self) -> torch.Tensor:
    return torch.sqrt(
        self.sigma_x**2 * self.sigma_xp**2
        - torch.mean((self.xs - self.mu_x) * (self.xps - self.mu_xp)) ** 2
    )

the second part of this after the - this is inspired by Ocelot's

tws.xpx = np.mean((x - tws.x) * (px - tws.px))

which I thought is the same as beam.sigma_x * beam.sigma_xp in Cheetah, but apparently it's not.

Now this is correct. $\epsilon^2 = \det \begin{pmatrix} \sigma{x}^2 & \sigma{xx'}\ \sigma{xx'} & \sigma{x'}^2\end{pmatrix}$

where $\sigma_{xx'} = \<{} x,x'>{} = \mathrm{E}[(x-\bar{x})(x'-\bar{x'})]$

• Is the twiss parameter computation from the Astra beam compared to Ocelot enough testing of this?

I think it is sufficient for now.

• We should probably also implement creating either type of beam from twiss parameters.

~Well... I'm not sure we are supposed to do that, because twiss is actually a function/property of the lattice, not the particle distribution. On the other hand, we can use it to generate a beam that is matched to the lattice.~ On a second thought, let's do that!

[cr-xu]

So this seems to work now. Here are a couple points we might need to think about:

• I think right now this doesn't work for ParameterBeam, but it absolutely should. The reason being that the emittance computation uses the actual particles:

@property
def emittance_x(self) -> torch.Tensor:
    return torch.sqrt(
        self.sigma_x**2 * self.sigma_xp**2
        - torch.mean((self.xs - self.mu_x) * (self.xps - self.mu_xp)) ** 2
    )

But this term $\sigma_{xx'}$ is actually just the cor_x in the covariance matrix used to initialize a ParameterBeam
One just needs to properly set it to a realistic value c.f. here

cov=torch.tensor(
                [
                    [sigma_x**2, cor_x, 0, 0, 0, 0, 0],
                    [cor_x, sigma_xp**2, 0, 0, 0, 0, 0],
                    [0, 0, sigma_y**2, cor_y, 0, 0, 0],
                    [0, 0, cor_y, sigma_yp**2, 0, 0, 0],
                    [0, 0, 0, 0, sigma_s**2, cor_s, 0],
                    [0, 0, 0, 0, cor_s, sigma_p**2, 0],
                    [0, 0, 0, 0, 0, 0, 0],
                ],
                dtype=torch.float32,
            ),

[cr-xu]

oops somehow the suggestion feature in GitHub messes up the indentation, now it's hopefully fixed

[jank324]

Okay ... so I think all that's missing is to create beams from twiss parameters.

[cr-xu]

Now the from_twiss() method should be fixed, the error was due to missing the cor_x and cory component, i.e. cov[0,1]= $\sigma{}{xx'}$