def test_ohmi_envelope(engine, ml_lattice, py_lattice, refpts):
fields = [('beam66', 'r66'), ('beam44', 'r44'), ('emit66', 'emitXYZ'),
('emit44', 'emitXY')]
nelems = len(py_lattice)
refpts = range(nelems + 1) if refpts is None else refpts
# Python call
py_lattice.radiation_on()
py_emit0, py_beamdata, py_emit = physics.ohmi_envelope(py_lattice, refpts)
# Matlab call
ml_emit = engine.pyproxy('atx', ml_lattice, 0.0, _ml_refs(refpts, nelems))
ml_emit0, ml_params = engine.pyproxy('atx',
ml_lattice,
0.0,
_ml_refs(0, nelems),
nargout=2)
revolution_period = get_s_pos(py_lattice, nelems) / speed_of_light
damping_times = revolution_period / py_beamdata.damping_rates
# Comparison
numpy.testing.assert_almost_equal(damping_times,
_py_data(ml_params['dampingtime']),
decimal=8)
numpy.testing.assert_almost_equal(py_beamdata.mode_emittances,
_py_data(ml_emit0['modemit']),
decimal=8)
_compare_physdata(numpy.expand_dims(py_emit0, 0), ml_emit0, fields)
_compare_physdata(py_emit, ml_emit, fields)
def compute(self, ring, *args, **kwargs):
"""Optics computation before evaluation of all constraints"""
em0, beamdata, em = ohmi_envelope(ring, refpts=self.refpts, **kwargs)
return (em[ref[self.refpts]] for ref in self.refs), (beamdata, )
def sigma_matrix(argin):
"""
Calculate the correlation matrix to be used for particle generation
PARAMETERS
argin Lattice object or list of twiss parameters.
Depending on the length of the list the function
will do different things:
[espread, blength]
computes 2x2 longitudinal sigma matrix
[beta, alpha, emit]
computes 2x2 transverse sigma matrix
[betax, alphax, emitx, betay, alphay, emity]
computes the 4x4 sigma matrix for horizontal
and vertical.
[betax, alphax, emitx, betay, alphay, emity, espread,
blength]
computes the 6x6 sigma matrix.
KEYWORDS
twiss=False Flag which states whether the input is a
lattice object or list of twiss parameters
OUTPUT
sigma_matrix correlation matrix (either 2x2, 4x4 or 6x6)
"""
if isinstance(argin, Lattice):
argin = argin.radiation_on(copy=True)
emit0, beamdata, emit = ohmi_envelope(argin, refpts=[0])
sig_matrix = emit.r66[0]
else:
if len(argin) == 2:
[espread, blength] = argin
sig_matrix = numpy.array([[espread * espread, 0],
[0, blength * blength]])
elif len(argin) == 3:
[bx, ax, epsx] = argin
sig_matrix = epsx * numpy.array([[bx, -ax],
[-ax, (1 + ax * ax) / bx]])
elif len(argin) == 6:
[bx, ax, epsx, by, ay, epsy] = argin
sig_matrix = numpy.block(
[[sigma_matrix([bx, ax, epsx]),
numpy.zeros((2, 2))],
[numpy.zeros((2, 2)),
sigma_matrix([by, ay, epsy])]])
elif len(argin) == 8:
[bx, ax, epsx, by, ay, epsy, espread, blength] = argin
sig_matrix = numpy.block(
[[sigma_matrix([bx, ax, epsx]),
numpy.zeros((2, 4))],
[
numpy.zeros((2, 2)),
sigma_matrix([by, ay, epsy]),
numpy.zeros((2, 2))
], [numpy.zeros((2, 4)),
sigma_matrix([espread, blength])]])
else:
raise AttributeError('Wrong number of inputs provided')
return sig_matrix
def sigma_matrix(ring=None,
twiss_in=None,
emitx=None,
emity=None,
blength=None,
espread=None):
"""
Calculate the correlation matrix to be used for particle generation
PARAMETERS
ring Lattice object or list of twiss parameters.
twiss_in Data structure containing inpu twiss parameters.
emitx Horizontal emittance [m.rad]
emity Vertical emittance [m.rad]
blength One sigma bunch length [m]
espread One sigma energy spread [dp/p]
OUTPUT
sigma_matrix 6x6 correlation matrix
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.
"""
flag = emitx or emity or blength or espread
if flag:
assert emitx is not None, 'emitx must be defined'
assert emity is not None, 'emity must be defined'
assert blength is not None, 'blength must be defined'
assert espread is not None, 'espread must be defined'
if ring:
ld0, bd, ld = ring.get_optics()
if ring.radiation and not flag:
print('Generating correlated sigma matrix using ohmi envelope')
emit0, beamdata, emit = ohmi_envelope(ring, refpts=[0])
sig_matrix = emit.r66[0]
elif flag:
print(
'Generating pseudo-correlated matrix from initial optics conditions'
)
if ring.radiation:
print(
'Ignoring provided blength and calculating it based on espread'
)
blength = _compute_bunch_length_from_espread(ring, espread)
sig_matrix = _generate_sigma_matrix(ld0, emitx, emity, blength,
espread, ring.radiation)
else:
raise AttributeError(
'Radiation is off but no emittances are specified')
elif twiss_in:
if not emitx:
raise AttributeError('Emittances must be specified for twiss_in')
print(
'Generating un-correlated sigma matrix from parameters in twiss_in'
)
sig_matrix = _generate_sigma_matrix(twiss_in, emitx, emity, blength,
espread, False)
else:
raise AttributeError('A lattice or twiss_in must be provided')
return sig_matrix