"""
Closed orbit related functions
"""
import numpy
from at.constants import clight
from at.lattice import AtError, AtWarning, check_6d, DConstant, Orbit
from at.lattice import Lattice, get_s_pos, Refpts, frequency_control
from at.tracking import internal_lpass
from .energy_loss import ELossMethod, get_timelag_fromU0
import warnings
__all__ = ['find_orbit4', 'find_sync_orbit', 'find_orbit6', 'find_orbit']
@check_6d(False)
def _orbit_dp(ring: Lattice, dp: float = None, guess: Orbit = None, **kwargs):
"""Solver for fixed energy deviation"""
# We seek
# - f(x) = x
# - g(x) = f(x) - x = 0
# - g'(x) = f'(x) - 1
# Use a Newton-Raphson-type algorithm:
# - r_n+1 = r_n - g(r_n) / g'(r_n)
# - r_n+1 = r_n - (f(r_n) - r_n) / (f'(r_n) - 1)
#
# (f(r_n) - r_n) / (f'(r_n) - 1) can be seen as x = b/a where we use least
# squares fitting to determine x when ax = b
# f(r_n) - r_n is denoted b
# f'(r_n) is the 4x4 jacobian, denoted j4
keep_lattice = kwargs.pop('keep_lattice', False)
convergence = kwargs.pop('convergence', DConstant.OrbConvergence)
max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter)
xy_step = kwargs.pop('XYStep', DConstant.XYStep)
kwargs.pop('DPStep', DConstant.DPStep)
rem = kwargs.keys()
if len(rem) > 0:
raise AtError(f'Unexpected keywords for orbit_dp: {", ".join(rem)}')
ref_in = numpy.zeros((6,)) if guess is None else numpy.copy(guess)
ref_in[4] = 0.0 if dp is None else dp
scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0])
delta_matrix = numpy.zeros((6, 5), order='F')
for i in range(4):
delta_matrix[i, i] = scaling[i]
id4 = numpy.asfortranarray(numpy.identity(4))
change = 1
itercount = 0
while (change > convergence) and itercount < max_iterations:
in_mat = ref_in.reshape((6, 1)) + delta_matrix
_ = internal_lpass(ring, in_mat, refpts=[], keep_lattice=keep_lattice)
# the reference particle after one turn
ref_out = in_mat[:, 4]
# 4x4 jacobian matrix from numerical differentiation:
# f(x+d) - f(x) / d
j4 = (in_mat[:4, :4] - in_mat[:4, 4:]) / scaling
a = j4 - id4 # f'(r_n) - 1
b = ref_out[:4] - ref_in[:4]
b_over_a = numpy.linalg.solve(a, b)
r_next = ref_in - numpy.append(b_over_a, numpy.zeros((2,)))
# determine if we are close enough
change = numpy.linalg.norm(r_next - ref_in)
itercount += 1
ref_in = r_next
keep_lattice = True
if itercount == max_iterations:
warnings.warn(AtWarning('Maximum number of iterations reached. '
'Possible non-convergence'))
return ref_in
@check_6d(False)
def _orbit_dct(ring: Lattice, dct: float = None, guess: Orbit = None, **kwargs):
"""Solver for fixed path lengthening"""
keep_lattice = kwargs.pop('keep_lattice', False)
convergence = kwargs.pop('convergence', DConstant.OrbConvergence)
max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter)
xy_step = kwargs.pop('XYStep', DConstant.XYStep)
kwargs.pop('DPStep', DConstant.DPStep)
rem = kwargs.keys()
if len(rem) > 0:
raise AtError(f'Unexpected keywords for orbit_dct: {", ".join(rem)}')
ref_in = numpy.zeros((6,)) if guess is None else numpy.copy(guess)
scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 1.0])
delta_matrix = numpy.zeros((6, 6), order='F')
for i in range(5):
delta_matrix[i, i] = scaling[i]
theta5 = numpy.zeros((5,), order='F')
theta5[4] = 0.0 if dct is None else dct
id5 = numpy.zeros((5, 5), order='F')
for i in range(4):
id5[i, i] = 1.0
idx = numpy.array([0, 1, 2, 3, 5])
change = 1
itercount = 0
while (change > convergence) and itercount < max_iterations:
in_mat = ref_in.reshape((6, 1)) + delta_matrix
_ = internal_lpass(ring, in_mat, refpts=[], keep_lattice=keep_lattice)
# the reference particle after one turn
ref_out = in_mat[:, -1]
# 5x5 jacobian matrix from numerical differentiation:
# f(x+d) - f(x) / d
j5 = (in_mat[idx, :5] - in_mat[idx, 5:]) / scaling
a = j5 - id5 # f'(r_n) - 1
b = ref_out[idx] - numpy.append(ref_in[:4], 0.0) - theta5
b_over_a = numpy.linalg.solve(a, b)
r_next = ref_in - numpy.append(b_over_a, 0.0)
# determine if we are close enough
change = numpy.linalg.norm(r_next - ref_in)
itercount += 1
ref_in = r_next
keep_lattice = True
if itercount == max_iterations:
warnings.warn(AtWarning('Maximum number of iterations reached. '
'Possible non-convergence'))
return ref_in
[docs]
def find_orbit4(ring: Lattice, dp: float = None, refpts: Refpts = None, *,
dct: float = None,
df: float = None,
orbit: Orbit = None,
keep_lattice: bool = False, **kwargs):
r"""Gets the 4D closed orbit for a given dp
Finds the closed orbit in the 4-d transverse phase space by numerically
solving for a fixed point of the one turn map M calculated with
:py:func:`.internal_lpass`.
.. math:: \begin{pmatrix}x \\ p_x \\ y \\ p_y \\ dp \\ c\tau_2\end{pmatrix}
=\mathbf{M} \cdot
\begin{pmatrix}x \\ p_x \\ y \\ p_y \\ dp \\ c\tau_1\end{pmatrix}
under the **CONSTANT MOMENTUM** constraint **dp** and with **NO**
constraint on the 6-th coordinate :math:`c\tau`
Important:
:py:func:`find_orbit4` imposes a constraint on *dp* and relaxes the
constraint on the revolution frequency. A physical storage ring does
exactly the opposite: the momentum deviation of a particle on the
closed orbit settles at the value such that the revolution is
synchronous with the RF cavity: :math:`f_{RF} = h*f_{rev}`
To impose this artificial constraint in :py:func:`find_orbit4`,
the ``PassMethod`` used for any element **SHOULD NOT**:
1. Change the longitudinal momentum *dp* (cavities ,
magnets with radiation)
2. Have any time dependence (localized impedance, fast kickers etc)
Parameters:
ring: Lattice description (:py:attr:`~.Lattice.is_6d` must be
:py:obj:`False`)
dp: Momentum deviation. Defaults to 0
refpts: Observation points.
See ":ref:`Selecting elements in a lattice <refpts>`"
dct: Path lengthening. If specified, *dp* is ignored and
the off-momentum is deduced from the path lengthening.
df: Deviation from the nominal RF frequency. If specified,
*dp* is ignored and the off-momentum is deduced from the frequency
deviation.
orbit: Avoids looking for initial the closed orbit if it is
already known ((6,) array). :py:func:`find_orbit4` propagates it to
the specified *refpts*.
keep_lattice: Assume no lattice change since the previous tracking.
Default: False
Keyword Args:
guess (Orbit): (6,) initial value for the
closed orbit. It may help convergence. Default: (0, 0, 0, 0, 0, 0)
convergence (float): Convergence criterion.
Default: :py:data:`DConstant.OrbConvergence <.DConstant>`
max_iterations (int): Maximum number of iterations.
Default: :py:data:`DConstant.OrbMaxIter <.DConstant>`
XYStep (float): Step size.
Default: :py:data:`DConstant.XYStep <.DConstant>`
Returns:
orbit0: (6,) closed orbit vector at the entrance of the
1-st element (x,px,y,py,dp,0)
orbit: (Nrefs, 6) closed orbit vector at each location
specified in *refpts*
See also:
:py:func:`find_sync_orbit`, :py:func:`find_orbit6`
"""
if len([v for v in (dp, dct, df) if v is not None]) > 1:
raise AtError("For off-momentum specification, only one of "
"dp, dct and df may be specified")
if orbit is None:
if df is not None:
frf = ring.cell_revolution_frequency * ring.cell_harmnumber
dct = -ring.cell_length * df / (frf+df)
orbit = _orbit_dct(ring, dct, keep_lattice=keep_lattice, **kwargs)
elif dct is not None:
orbit = _orbit_dct(ring, dct, keep_lattice=keep_lattice, **kwargs)
else:
orbit = _orbit_dp(ring, dp, keep_lattice=keep_lattice, **kwargs)
keep_lattice = True
# bug in numpy < 1.13
if ring.refcount(refpts) == 0:
all_points = numpy.empty((0, 6), dtype=float)
else:
all_points = internal_lpass(ring, orbit.copy(order='K'), refpts=refpts,
keep_lattice=keep_lattice)
all_points = numpy.squeeze(all_points, axis=(1, 3)).T
return orbit, all_points
[docs]
def find_sync_orbit(ring: Lattice, dct: float = None, refpts: Refpts = None, *,
dp: float = None,
df: float = None,
orbit: Orbit = None,
keep_lattice: bool = False, **kwargs):
r"""Gets the 4D closed orbit for a given dct
Finds the closed orbit, synchronous with the RF cavity (first 5
components of the phase space vector) by numerically solving for a fixed
point of the one turn map M calculated with :py:func:`.internal_lpass`
.. math:: \begin{pmatrix}x \\ p_x \\ y \\ p_y \\ dp \\ c\tau_1+
dc\tau\end{pmatrix} =\mathbf{M} \cdot
\begin{pmatrix}x \\ p_x \\ y \\ p_y \\ dp \\ c\tau_1\end{pmatrix}
under the constraint :math:`dc\tau = c\tau_2 - c\tau_1 = c/f_{rev} -
c/f_{rev_0}`, where :math:`f_{rev_0} = f_{RF_0}/h` is the design revolution
frequency and :math:`f_{rev} = (f_{RF_0} + df_{RF})/h` is the imposed
revolution frequency.
Important:
:py:func:`find_sync_orbit` imposes a constraint :math:`c\tau_2 -
c\tau_1` and :math:`dp_2 = dp_1` but no constraint on the value of
:math:`dp_2` or :math:`dp_1`. The algorithm assumes time-independent
fixed-momentum ring to reduce the dimensionality of the problem.
To impose this artificial constraint in :py:func:`find_sync_orbit`,
the ``PassMethod`` used for any element **SHOULD NOT**:
1. Change the longitudinal momentum *dp* (cavities ,
magnets with radiation)
2. Have any time dependence (localized impedance, fast kickers etc)
Parameters:
ring: Lattice description (:py:attr:`~.Lattice.is_6d` must be
:py:obj:`False`)
dct: Path lengthening.
refpts: Observation points.
See ":ref:`Selecting elements in a lattice <refpts>`"
dp: Momentum deviation. Defaults to :py:obj:`None`
df: Deviation from the nominal RF frequency. If specified,
*dct* is ignored and the off-momentum is deduced from the frequency
deviation.
orbit: Avoids looking for initial the closed orbit if it is
already known ((6,) array). :py:func:`find_sync_orbit` propagates it
to the specified *refpts*.
keep_lattice: Assume no lattice change since the previous tracking.
Default: False
Keyword Args:
guess (Orbit): (6,) initial value for the
closed orbit. It may help convergence. Default: (0, 0, 0, 0, 0, 0)
convergence (float): Convergence criterion.
Default: :py:data:`DConstant.OrbConvergence <.DConstant>`
max_iterations (int): Maximum number of iterations.
Default: :py:data:`DConstant.OrbMaxIter <.DConstant>`
XYStep (float): Step size.
Default: :py:data:`DConstant.XYStep <.DConstant>`
Returns:
orbit0: (6,) closed orbit vector at the entrance of the
1-st element (x,px,y,py,dp,0)
orbit: (Nrefs, 6) closed orbit vector at each location
specified in *refpts*
See also:
:py:func:`find_orbit4`, :py:func:`find_orbit6`
"""
if len([v for v in (dp, dct, df) if v is not None]) > 1:
raise AtError("For off-momentum specification, only one of "
"dp, dct and df may be specified")
if orbit is None:
if df is not None:
frf = ring.cell_revolution_frequency * ring.cell_harmnumber
dct = -ring.cell_length * df / (frf+df)
orbit = _orbit_dct(ring, dct, keep_lattice=keep_lattice, **kwargs)
elif dp is not None:
orbit = _orbit_dp(ring, dp, keep_lattice=keep_lattice, **kwargs)
else:
orbit = _orbit_dct(ring, dct, keep_lattice=keep_lattice, **kwargs)
keep_lattice = True
# bug in numpy < 1.13
if ring.refcount(refpts) == 0:
all_points = numpy.empty((0, 6), dtype=float)
else:
all_points = internal_lpass(ring, orbit.copy(order='K'), refpts=refpts,
keep_lattice=keep_lattice)
all_points = numpy.squeeze(all_points, axis=(1, 3)).T
return orbit, all_points
def _orbit6(ring: Lattice, cavpts=None, guess=None, keep_lattice=False,
**kwargs):
"""Solver for 6D motion"""
convergence = kwargs.pop('convergence', DConstant.OrbConvergence)
max_iterations = kwargs.pop('max_iterations', DConstant.OrbMaxIter)
xy_step = kwargs.pop('XYStep', DConstant.XYStep)
dp_step = kwargs.pop('DPStep', DConstant.DPStep)
method = kwargs.pop('method', ELossMethod.TRACKING)
rem = kwargs.keys()
if len(rem) > 0:
raise AtError(f'Unexpected keywords for orbit6: {", ".join(rem)}')
l0 = get_s_pos(ring, len(ring))[0]
f_rf = ring.get_rf_frequency()
harm_number = round(f_rf*l0/ring.beta/clight)
if guess is None:
_, dt = get_timelag_fromU0(ring, method=method, cavpts=cavpts)
# Getting timelag by tracking uses a different lattice,
# so we cannot now use the same one again.
if method is ELossMethod.TRACKING:
keep_lattice = False
ref_in = numpy.zeros((6,), order='F')
ref_in[5] = -dt
else:
ref_in = numpy.copy(guess)
theta = numpy.zeros((6,))
theta[5] = ring.beta * clight * harm_number / f_rf - l0
scaling = xy_step * numpy.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0]) + \
dp_step * numpy.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0])
delta_matrix = numpy.asfortranarray(
numpy.concatenate((numpy.diag(scaling), numpy.zeros((6, 1))), axis=1))
id6 = numpy.asfortranarray(numpy.identity(6))
change = 1
itercount = 0
while (change > convergence) and itercount < max_iterations:
in_mat = ref_in.reshape((6, 1)) + delta_matrix
_ = internal_lpass(ring, in_mat, refpts=[], keep_lattice=keep_lattice)
# the reference particle after one turn
ref_out = in_mat[:, 6]
# 6x6 jacobian matrix from numerical differentiation:
# f(x+d) - f(x) / d
j6 = (in_mat[:, :6] - in_mat[:, 6:]) / scaling
a = j6 - id6 # f'(r_n) - 1
b = ref_out[:] - ref_in[:] - theta
# b_over_a, _, _, _ = numpy.linalg.lstsq(a, b, rcond=-1)
b_over_a = numpy.linalg.solve(a, b)
r_next = ref_in - b_over_a
# determine if we are close enough
change = numpy.linalg.norm(r_next - ref_in)
itercount += 1
ref_in = r_next
keep_lattice = True
if itercount == max_iterations:
warnings.warn(AtWarning('Maximum number of iterations reached. '
'Possible non-convergence'))
return ref_in
# noinspection PyIncorrectDocstring
[docs]
@frequency_control
def find_orbit6(ring: Lattice, refpts: Refpts = None, *,
orbit: Orbit = None, keep_lattice: bool = False, **kwargs):
r"""Gets the closed orbit in the full 6-D phase space
Finds the closed orbit in the full 6-D phase space
by numerically solving for a fixed point of the one turn
map M calculated with :py:func:`.internal_lpass`
.. math:: \begin{pmatrix}x \\ p_x \\ y \\ p_y \\ dp \\ c\tau_2\end{pmatrix}
=\mathbf{M} \cdot
\begin{pmatrix}x \\ p_x \\ y \\ p_y \\ dp \\ c\tau_1\end{pmatrix}
with constraint :math:`c\tau_2 - c\tau_1 = c.h (1/f_{RF} - 1/f_{RF_0})`
Important:
:py:func:`find_orbit6` is a realistic simulation:
1. The requency in the RF cavities :math:`f_{RF}` (may be different
from :math:`f_{RF_0}`) imposes the synchronous condition
:math:`c\tau_2 - c\tau_1 = c.h (1/f_{RF} - 1/f_{RF_0})`,
2. The algorithm numerically calculates the
6-by-6 Jacobian matrix J6. In order for (J-E) matrix
to be non-singular it is **NECESSARY** to use a realistic
``PassMethod`` for cavities with non-zero momentum kick
(such as ``RFCavityPass``).
3. :py:func:`find_orbit6` can find orbits with radiation.
In order for the solution to exist the cavity must supply an
adequate energy compensation.
In the simplest case of a single cavity, it must have
``Voltage`` set so that :math:`V > E_{loss}`, the energy loss
per turn
4. There is a family of solutions that correspond to different RF
buckets They differ in the 6-th coordinate by
:math:`c*Nb/f_{RF}`, Nb = 0:h-1
5. The value of the 6-th coordinate found at the cavity gives
the equilibrium RF phase. If there is no radiation it is 0.
6. ``dp``, ``dct`` and ``df`` arguments are applied with respect
to the **NOMINAL** on-momentum frequency. They overwrite
exisiting frequency offsets
Parameters:
ring: Lattice description
refpts: Observation points
orbit: Avoids looking for initial the closed orbit if it is
already known ((6,) array). :py:func:`find_sync_orbit` propagates it
to the specified *refpts*.
keep_lattice: Assume no lattice change since the previous tracking.
Default: False
Keyword Args:
guess (Orbit): (6,) initial value for the
closed orbit. It may help convergence. Default: (0, 0, 0, 0, 0, 0)
convergence (float): Convergence criterion.
Default: :py:data:`DConstant.OrbConvergence <.DConstant>`
max_iterations (int): Maximum number of iterations.
Default: :py:data:`DConstant.OrbMaxIter <.DConstant>`
XYStep (float): Step size.
Default: :py:data:`DConstant.XYStep <.DConstant>`
DPStep (float): Momentum step size.
Default: :py:data:`DConstant.DPStep <.DConstant>`
method (ELossMethod): Method for energy loss computation.
See :py:class:`.ELossMethod`.
cavpts (Refpts): Cavity location. If :py:obj:`None`, use all
cavities. This is used to compute the initial synchronous phase.
Returns:
orbit0: (6,) closed orbit vector at the entrance of the
1-st element (x,px,y,py,dp,0)
orbit: (Nrefs, 6) closed orbit vector at each location
specified in *refpts*
See also:
:py:func:`find_orbit4`, :py:func:`find_sync_orbit`
"""
if orbit is None:
orbit = _orbit6(ring, keep_lattice=keep_lattice, **kwargs)
keep_lattice = True
# bug in numpy < 1.13
if ring.refcount(refpts) == 0:
all_points = numpy.empty((0, 6), dtype=float)
else:
all_points = internal_lpass(ring, orbit.copy(order='K'), refpts=refpts,
keep_lattice=keep_lattice)
all_points = numpy.squeeze(all_points, axis=(1, 3)).T
return orbit, all_points
[docs]
def find_orbit(ring, refpts: Refpts = None, **kwargs):
r"""Gets the closed orbit in the general case
Depending on the lattice, :py:func:`find_orbit` will:
* use :py:func:`find_orbit6` if :py:attr:`~.Lattice.is_6d` is
:py:obj:`True`,
* use :py:func:`find_sync_orbit` if :py:attr:`~.Lattice.is_6d` is
:py:obj:`False` and
*dct* or *df* is specified,
* use :py:func:`find_orbit4` otherwise.
Parameters:
ring: Lattice description
refpts: Observation points.
See ":ref:`Selecting elements in a lattice <refpts>`"
Keyword Args:
orbit (Orbit): Avoids looking for initial the closed
orbit if it is already known. :py:func:`find_orbit` propagates it
to the specified *refpts*.
dp (float): Momentum deviation. Defaults to :py:obj:`None`
dct (float): Path lengthening. Defaults to :py:obj:`None`
df (float): Deviation from the nominal RF frequency.
Defaults to :py:obj:`None`
guess (Orbit): (6,) initial value for the closed orbit.
It may help convergence. Default: (0, 0, 0, 0, 0, 0)
convergence (float): Convergence criterion.
Default: :py:data:`DConstant.OrbConvergence <.DConstant>`
max_iterations (int): Maximum number of iterations.
Default: :py:data:`DConstant.OrbMaxIter <.DConstant>`
XYStep (float): Step size.
Default: :py:data:`DConstant.XYStep <.DConstant>`
DPStep (float): Momentum step size.
Default: :py:data:`DConstant.DPStep <.DConstant>`
For other keywords, refer to the underlying methods
Returns:
orbit0: (6,) closed orbit vector at the entrance of the
1-st element (x,px,y,py,dp,0)
orbit: (Nrefs, 6) closed orbit vector at each location
specified in *refpts*
See also:
:py:func:`find_orbit4`, :py:func:`find_sync_orbit`,
:py:func:`find_orbit6`
"""
if ring.is_6d:
return find_orbit6(ring, refpts=refpts, **kwargs)
else:
return find_orbit4(ring, refpts=refpts, **kwargs)
Lattice.find_orbit4 = find_orbit4
Lattice.find_sync_orbit = find_sync_orbit
Lattice.find_orbit6 = find_orbit6
Lattice.find_orbit = find_orbit
