Pythonised the sigma matrix function

[lcarver]

Re-wrote atbeam. Now it takes either a lattice object as input or you can give a twiss_in kwarg and specify emittances. I realised there was an issue before which was that if you gave at.beam a 2x2 or 4x4 matrix, the final particle output from at.beam only had 2 or 4 dimensions, i.e. not suitable for tracking. This issue is now avoided by always returning a 6xN sigma_matrix and therefore a 6xN array of particles.

A small issue that I observed was that if blength and espread are both 0, then the cholesky decomposition of the sigma_matrix (in at.beam) fails due to negative eigenvectors. That is why when neither blength or espread are specified I set them to very small positive values (with a warning).

Example usage below.

import at
import numpy as np

ring = at.load_mat('./S28F.mat', mat_key = 'LOW_EMIT_RING_INJ')
ring.radiation_on()

sig0 = at.sigma_matrix(ring)

ld0, bd, ld = ring.get_optics()
sig1 = at.sigma_matrix(twiss_in=ld0, emitx=135e-12, emity=10e-12)
sig2 = at.sigma_matrix(twiss_in=ld0, emitx=135e-12, emity=10e-12, espread=1e-3, blength=1e-3)

nump = 10
par0 = at.beam(nump, sig0, orbit=ld0.closed_orbit)
par1 = at.beam(nump, sig1)
par2 = at.beam(nump, sig2)

[lcarver]

After offline discussion with @swhite2401, modified the handling of the linalg error. Now if the original 6x6 cholesky fails then it reverts to doing each 2x2 matrix individually, if this fails then it returns 0s.

I also added the ability to specify a vertical emittance when providing a lattice object for the sigma matrix. This takes the uncoupled 2x2 matrix for the vertical plane and inserts it into the 6x6 sigma matrix from ohmi_envelope.

sig0 = at.sigma_matrix(ring, emity=10e-12)

[lcarver]

There are three cases that can occur,

sig0 = at.sigma_matrix(ring)
sig1 = at.sigma_matrix(ring, emitx=135e-12, emity=0, blength=1e-3, espread=1e-3)

ld0, bd, ld = ring.get_optics()
sig2 = at.sigma_matrix(twiss_in=ld0, emitx=135e-12, emity=0, espread=1e-3, blength=1e-3)

Each one is doing something slightly different

If the lattice object is provided with no other arguments, ohmi_envelope is used
to compute the correlated sigma matrix.

If the lattice object and emittances and longitudinal parameters are provided, then
the 2x2 uncorrelated matrices are computed for each plane (x,y,z) using the initial
optics computed from ring.get_optics, and are combined together into the 6x6 matrix.

If the twiss_in is provided alongside the emittances and longitudinal parameters, then
the 2x2 uncorrelated matrices are computed for each plane and combined into the 6x6
matrix.

[lfarv]

@lcarver : when calling get_optics, why not using the R-matrices (either 2x4x4 if no radiation, or 3x6x6 if radiation), that you can multiply by the provided emittance to get a "correlated" sigma matrix:

• if radiation is on, espread and blength are coupled, so:
  sig_matrix = emitx*ld0.R[0] + emity*ld0.R[1] + f(espread)*ld0.R[2]

• if radiation is off, you get a 4x4 sigma matrix with:
  sig_matrix = emitx*ld0.R[0] + emity*ld0.R[1],
  that you can block-assemble with you uncorrelated longitudinal matrix to get the full one.

Better that using alpha and beta! If there is no coupling in the lattice, the result will be uncorrelated anyway.

[lcarver]

@lfarv : Thanks a lot. I think I implemented your suggestions correctly. For me I had the following possibilities in mind:

ring with rad on, no emittances specified -> ohmi_envelope
ring with rad on and emittances ->your sum of 3 6x6 matrices as above with blength calculated from espread.
ring with rad off and emittances -> 4x4 matrices + uncorrelated long. matrix
ring with rad off and no emittances -> AttributeError
twiss_in and emittances -> generates the 4x4 matrices + uncorrelated long.matrix

I hope its clear. I checked the longitudinal distributions of each method and I saw the same distribution each time, but there could be some missing subtleties.

[lfarv]

One cannot expect that calling sigma_matrix(ring,.. modifies the ring. So it should be:

mcf = ring.radiation_off(copy=True).get_mcf()
envelope = ring.radiation_on(copy=True).envelope_parameters()

or

mcf = get_mcf(ring.radiation_off(copy=True)
envelope = envelope_parameters(ring.radiation_on(copy=True))

as you prefer.

In addition, turning radiation off and then on does not ensure that you come back to the initial state (which multipoles did the user activate?). So it should be avoided.

[lfarv]

A detail: when using twiss_in: only linopt6 generates the R matrices. linopt2 and linopt4 don't. So you could try to access R, and if it fails, fall back to alpha and beta (uncoupled ring). This is what _linopt does. However, it's a minor point, so you could leave this for a future upgrade if you prefer.

[swhite2401]

Ok for me to merge after the latest changes proposed by @lfarv are implemented

[lcarver]

Not ready to merge just yet.
EDIT: Now it is ready.

[lfarv]

Ok for me

[swhite2401]

OK!