def find_m44(accelerator, indices=None, energy_offset=0.0, closed_orbit=None):
"""Calculate 4D transfer matrices of elements in an accelerator.
Keyword arguments:
accelerator
indices
energy_offset
closed_orbit
Return values:
m44
cumul_trans_matrices -- values at the start of each lattice element
"""
if indices is None:
indices = list(range(len(accelerator)))
if closed_orbit is None:
# calcs closed orbit if it was not passed.
fixed_point_guess = _trackcpp.CppDoublePos()
fixed_point_guess.de = energy_offset
_closed_orbit = _trackcpp.CppDoublePosVector()
r = _trackcpp.track_findorbit4(accelerator._accelerator, _closed_orbit,
fixed_point_guess)
if r > 0:
raise TrackingException(_trackcpp.string_error_messages[r])
else:
_closed_orbit = _4Numpy2CppDoublePosVector(closed_orbit,
de=energy_offset)
_cumul_trans_matrices = _trackcpp.CppMatrixVector()
_m44 = _trackcpp.Matrix()
_v0 = _trackcpp.CppDoublePos()
r = _trackcpp.track_findm66(accelerator._accelerator, _closed_orbit,
_cumul_trans_matrices, _m44, _v0)
if r > 0:
raise TrackingException(_trackcpp.string_error_messages[r])
m44 = _CppMatrix24Numpy(_m44)
if indices == 'm44':
return m44
cumul_trans_matrices = []
for i in range(len(_cumul_trans_matrices)):
cumul_trans_matrices.append(_CppMatrix24Numpy(
_cumul_trans_matrices[i]))
return m44, cumul_trans_matrices
def _Numpy2CppDoublePosVector(orbit):
orbit_out = _trackcpp.CppDoublePosVector()
for i in range(orbit.shape[1]):
orbit_out.push_back(_trackcpp.CppDoublePos(
orbit[0,i], orbit[1,i],
orbit[2,i], orbit[3,i],
orbit[4,i], orbit[5,i]))
return orbit_out
def _4Numpy2CppDoublePosVector(orbit, de=0.0):
if isinstance(orbit, _trackcpp.CppDoublePosVector):
return orbit
if isinstance(orbit, _numpy.ndarray):
orbit_out = _trackcpp.CppDoublePosVector()
for i in range(orbit.shape[1]):
orbit_out.push_back(
_trackcpp.CppDoublePos(orbit[0, i], orbit[1, i], orbit[2, i],
orbit[3, i], de, 0.0))
elif isinstance(orbit, (list, tuple)):
orbit_out = _trackcpp.CppDoublePosVector()
orbit_out.push_back(
_trackcpp.CppDoublePos(orbit[0], orbit[1], orbit[2], orbit[3], de,
0.0))
else:
raise TrackingException('invalid orbit argument')
return orbit_out
def find_orbit6(accelerator, indices=None, fixed_point_guess=None):
"""Calculate 6D closed orbit of accelerator and return it.
Accepts an optional list of indices of ring elements where closed orbit
coordinates are to be returned. If this argument is not passed, closed orbit
positions are returned at the start of the first element. In addition a guess
fixed point at the entrance of the ring may be provided.
Keyword arguments:
accelerator : Accelerator object
indices : may be a (list,tuple, numpy.ndarray) of element indices where
closed orbit data is to be returned or a string:
'open' : return the closed orbit at the entrance of all elements.
'closed': besides all the points of 'open' also return the orbit
at the exit of the last element.
If indices is None the closed orbit is returned only at the
entrance of the first element.
fixed_point_guess -- A 6D position where to start the search of the closed
orbit at the entrance of the first element. If not
provided the algorithm will start with zero orbit.
Returns:
orbit : 6D closed orbit at the entrance of the selected elements as a 2D
numpy array with the 6 phase space variables in the first dimension and
the indices of the elements in the second dimension.
Raises TrackingException
"""
if fixed_point_guess is None:
fixed_point_guess = _trackcpp.CppDoublePos()
else:
fixed_point_guess = _Numpy2CppDoublePos(fixed_point_guess)
_closed_orbit = _trackcpp.CppDoublePosVector()
r = _trackcpp.track_findorbit6(accelerator._accelerator, _closed_orbit,
fixed_point_guess)
if r > 0:
raise TrackingException(_trackcpp.string_error_messages[r])
if indices is None:
closed_orbit = _CppDoublePos2Numpy(_closed_orbit[0])[None, :]
elif indices == 'open':
closed_orbit = _numpy.zeros((len(accelerator), 6))
for i in range(len(accelerator)):
closed_orbit[i] = _CppDoublePos2Numpy(_closed_orbit[i])
elif indices == 'closed':
closed_orbit = _numpy.zeros((len(accelerator) + 1, 6))
for i in range(len(accelerator)):
closed_orbit[i] = _CppDoublePos2Numpy(_closed_orbit[i])
closed_orbit[-1] = closed_orbit[0]
elif isinstance(indices, (list, tuple, _numpy.ndarray)):
closed_orbit = _numpy.zeros((len(indices), 6))
for i, ind in enumerate(indices):
closed_orbit[i] = _CppDoublePos2Numpy(_closed_orbit[ind])
else:
raise TrackingException("invalid value for 'indices' in findorbit6")
return closed_orbit.T
def find_m66(accelerator, indices=None, closed_orbit=None):
"""Calculate 6D transfer matrices of elements in an accelerator.
Keyword arguments:
accelerator
indices
closed_orbit
Return values:
m66
cumul_trans_matrices -- values at the start of each lattice element
"""
if indices is None:
indices = list(range(len(accelerator)))
if closed_orbit is None:
# Closed orbit is calculated by trackcpp
_closed_orbit = _trackcpp.CppDoublePosVector()
else:
_closed_orbit = _Numpy2CppDoublePosVector(closed_orbit)
_cumul_trans_matrices = _trackcpp.CppMatrixVector()
_m66 = _trackcpp.Matrix()
_v0 = _trackcpp.CppDoublePos()
r = _trackcpp.track_findm66(accelerator._accelerator, _closed_orbit,
_cumul_trans_matrices, _m66, _v0)
if r > 0:
raise TrackingException(_trackcpp.string_error_messages[r])
m66 = _CppMatrix2Numpy(_m66)
if indices == 'm66':
return m66
cumul_trans_matrices = MatrixList(_cumul_trans_matrices)
return m66, cumul_trans_matrices
def _Numpy2CppDoublePos(p_in):
p_out = _trackcpp.CppDoublePos()
p_out.rx, p_out.px = float(p_in[0]), float(p_in[1])
p_out.ry, p_out.py = float(p_in[2]), float(p_in[3])
p_out.de, p_out.dl = float(p_in[4]), float(p_in[5])
return p_out
def calc_twiss(accelerator=None,
init_twiss=None,
fixed_point=None,
indices='open',
energy_offset=None):
"""Return Twiss parameters of uncoupled dynamics.
Keyword arguments:
accelerator -- Accelerator object
init_twiss -- Twiss parameters at the start of first element
fixed_point -- 6D position at the start of first element
indices -- Open or closed
energy_offset -- float denoting the energy deviation (used only for periodic
solutions).
Returns:
tw -- list of Twiss objects (closed orbit data is in the objects vector)
m66 -- one-turn transfer matrix
"""
if indices == 'open':
closed_flag = False
elif indices == 'closed':
closed_flag = True
else:
raise OpticsException("invalid value for 'indices' in calc_twiss")
_m66 = _trackcpp.Matrix()
_twiss = _trackcpp.CppTwissVector()
if init_twiss is not None:
''' as a transport line: uses init_twiss '''
_init_twiss = init_twiss._t
if fixed_point is None:
_fixed_point = _init_twiss.co
else:
raise OpticsException(
'arguments init_twiss and fixed_point are mutually exclusive')
r = _trackcpp.calc_twiss(accelerator._accelerator, _fixed_point, _m66,
_twiss, _init_twiss, closed_flag)
else:
''' as a periodic system: try to find periodic solution '''
if accelerator.harmonic_number == 0:
raise OpticsException(
'Either harmonic number was not set or calc_twiss was'
'invoked for transport line without initial twiss')
if fixed_point is None:
_closed_orbit = _trackcpp.CppDoublePosVector()
_fixed_point_guess = _trackcpp.CppDoublePos()
if energy_offset is not None: _fixed_point_guess.de = energy_offset
if not accelerator.cavity_on and not accelerator.radiation_on:
r = _trackcpp.track_findorbit4(accelerator._accelerator,
_closed_orbit,
_fixed_point_guess)
elif not accelerator.cavity_on and accelerator.radiation_on:
raise OpticsException(
'The radiation is on but the cavity is off')
else:
r = _trackcpp.track_findorbit6(accelerator._accelerator,
_closed_orbit,
_fixed_point_guess)
if r > 0:
raise _tracking.TrackingException(
_trackcpp.string_error_messages[r])
_fixed_point = _closed_orbit[0]
else:
_fixed_point = _tracking._Numpy2CppDoublePos(fixed_point)
if energy_offset is not None: _fixed_point.de = energy_offset
r = _trackcpp.calc_twiss(accelerator._accelerator, _fixed_point, _m66,
_twiss)
if r > 0:
raise OpticsException(_trackcpp.string_error_messages[r])
twiss = TwissList(_twiss)
m66 = _tracking._CppMatrix2Numpy(_m66)
return twiss, m66