"""Element translations and rotations.
.. Caution::
The geometrical transformations are not commutative. When combining several
transformations on the same element by using multiple function calls, the
transformations are applied in a fixed order, independent of the order of the
function calls:
.. centered:: translations -> tilt (z-axis) -> yaw (y-axis) -> pitch (x-axis)
"""
from __future__ import annotations
__all__ = [
"set_rotation",
"set_shift",
"set_tilt",
"shift_elem",
"tilt_elem",
"transform_elem",
]
from collections.abc import Sequence
import numpy as np
from .elements import Element, ReferencePoint, transform_options
from .utils import Refpts, All, refpts_iterator, _refcount
from .exceptions import AtError
_x_axis = np.array([1.0, 0.0, 0.0])
_y_axis = np.array([0.0, 1.0, 0.0])
_z_axis = np.array([0.0, 0.0, 1.0])
# noinspection PyPep8Naming
def _rotation(rotations):
"""
The implementation follows the one described in:
https://doi.org/10.1016/j.nima.2022.167487
All the comments featuring 'Eq' points to the paper's equations.
3D rotation matrix (extrinsic rotation) using the Tait–Bryan angles
convention.
For more details, refer to https://en.wikipedia.org/wiki/Euler_angles
Corresponds to Eq. (3)
alpha: Rotation about the X-axis (pitch).
beta: Rotation about the Y-axis (yaw).
gamma: Rotation about the Z-axis (roll/tilt).
"""
alpha, beta, gamma = rotations # ZYX intrinsic rotations (pitch, yaw, tilt)
R_x = np.array(
[
[1, 0, 0],
[0, np.cos(alpha), -np.sin(alpha)],
[0, np.sin(alpha), np.cos(alpha)],
]
)
R_y = np.array(
[[np.cos(beta), 0, np.sin(beta)], [0, 1, 0], [-np.sin(beta), 0, np.cos(beta)]]
)
R_z = np.array(
[
[np.cos(gamma), -np.sin(gamma), 0],
[np.sin(gamma), np.cos(gamma), 0],
[0, 0, 1],
]
)
return R_x @ R_y @ R_z
# noinspection PyPep8Naming
def _translation_vector(ld, r3d, offsets, X_axis, Y_axis):
"""
The implementation follows the one described in:
https://doi.org/10.1016/j.nima.2022.167487
All the comments featuring 'Eq' points to the paper's equations.
Translation vector resulting from the joint effect of a longitudinal
displacement (in the rotated frame), 3D offsets, and the 3D rotation matrix.
Corresponds to Eqs. (8-11)
ld: Longitudinal displacement [m]
r3d: 3D rotation matrix
offsets: 3D offsets [m]
X_axis: X unit axis in rotated frame expressed in the xyz coordinate system
Y_axis: Y unit axis in rotated frame expressed in the xyz coordinate system
"""
tD0 = np.array([-offsets @ X_axis, 0, -offsets @ Y_axis, 0, 0, 0])
T0 = np.array(
[
ld * r3d[2, 0] / r3d[2, 2],
r3d[2, 0],
ld * r3d[2, 1] / r3d[2, 2],
r3d[2, 1],
0,
ld / r3d[2, 2],
]
)
return T0 + tD0
def _r_matrix(ld, r3d):
"""
The implementation follows the one described in:
https://doi.org/10.1016/j.nima.2022.167487
All the comments featuring 'Eq' points to the paper's equations.
Rotation matrix operator (R1, R2).
Can take into account the effect of a longitudinal displacement.
ld: Longitudinal displacement [m]
r3d: 3D rotation matrix
Corresponds to Eq. (9)
"""
return np.array(
[
[
r3d[1, 1] / r3d[2, 2],
ld * r3d[1, 1] / r3d[2, 2] ** 2,
-r3d[0, 1] / r3d[2, 2],
-ld * r3d[0, 1] / r3d[2, 2] ** 2,
0,
0,
],
[0, r3d[0, 0], 0, r3d[1, 0], r3d[2, 0], 0],
[
-r3d[1, 0] / r3d[2, 2],
-ld * r3d[1, 0] / r3d[2, 2] ** 2,
r3d[0, 0] / r3d[2, 2],
ld * r3d[0, 0] / r3d[2, 2] ** 2,
0,
0,
],
[0, r3d[0, 1], 0, r3d[1, 1], r3d[2, 1], 0],
[0, 0, 0, 0, 1, 0],
[
-r3d[0, 2] / r3d[2, 2],
-ld * r3d[0, 2] / r3d[2, 2] ** 2,
-r3d[1, 2] / r3d[2, 2],
-ld * r3d[1, 2] / r3d[2, 2] ** 2,
0,
1,
],
]
)
def _tilt_frame_mat(rots: float) -> np.ndarray:
cs = np.cos(rots)
sn = np.sin(rots)
rm = np.asfortranarray(np.diag([cs, cs, cs, cs, 1.0, 1.0]))
rm[0, 2] = sn
rm[1, 3] = sn
rm[2, 0] = -sn
rm[3, 1] = -sn
return rm
# noinspection PyPep8Naming
[docs]
def tilt_elem(
elem: Element,
rots: float | None = None,
rots_frame: float | None = None,
relative: bool = False,
reference: ReferencePoint | None = None,
) -> None:
r"""Set the tilt angle :math:`\theta` of an :py:class:`.Element`
The rotation matrices are stored in the :pycode:`R1` and :pycode:`R2`
attributes.
:math:`R_1=\begin{pmatrix} cos\theta & sin\theta \\
-sin\theta & cos\theta \end{pmatrix}`,
:math:`R_2=\begin{pmatrix} cos\theta & -sin\theta \\
sin\theta & cos\theta \end{pmatrix}`
Parameters:
elem: Element to be tilted
rots: Tilt angle :math:`\theta` [rd]. *rots* > 0 corresponds
to a corkscrew rotation of the element looking in the direction of
the beam. Use :py:obj:`None` to keep the current value.
rots_frame: Tilt angle :math:`\theta` [rd]. *rots* > 0 corresponds
to a corkscrew rotation of the reference frame looking in the direction of
the beam. Use :py:obj:`None` to keep the current value.
relative: If :py:obj:`True`, the rotation is added to the
existing one
reference: Transformation reference, either
:py:obj:`ReferencePoint.CENTRE` or
:py:obj:`ReferencePoint.ENTRANCE`.
Default: :py:obj:`ReferencePoint.CENTRE` if the reference
point was not previously set, otherwise to set reference
point
See Also:
:py:func:`shift_elem`
:py:func:`.transform_elem`
"""
transform_elem(
elem, tilt=rots, tilt_frame=rots_frame, relative=relative, reference=reference
)
[docs]
def shift_elem(
elem: Element,
dx: float | None = None,
dy: float | None = None,
dz: float | None = None,
*,
reference: ReferencePoint | None = None,
relative: bool = False,
) -> None:
r"""Sets the translations of an :py:class:`.Element`
The translation vectors are stored in the :pycode:`T1` and :pycode:`T2`
attributes.
Parameters:
elem: Element to be shifted
dx: Horizontal translation [m]. Default no change.
dy: Vertical translation [m]. Default no change.
dz: Longitudinal translation [m]. Default no change.
reference: Transformation reference, either
:py:obj:`ReferencePoint.CENTRE` or
:py:obj:`ReferencePoint.ENTRANCE`.
Default: :py:obj:`ReferencePoint.CENTRE` if the reference
point was not previously set, otherwise to set reference
point
relative: If :py:obj:`True`, the translation is added to the
existing one
See Also:
:py:func:`tilt_elem`
:py:func:`.transform_elem`
"""
transform_elem(elem, reference=reference, dx=dx, dy=dy, dz=dz, relative=relative)
[docs]
def set_rotation(
ring: Sequence[Element],
tilts: Sequence[float] | float | None = None,
pitches: Sequence[float] | float | None = None,
yaws: Sequence[float] | float | None = None,
*,
tilts_frame: Sequence[float] | float | None = None,
refpts: Refpts = All,
reference: ReferencePoint | None = None,
relative=False,
) -> None:
r"""Sets the rotations of a list of elements.
Parameters:
ring: Lattice description.
tilts: Scalar or Sequence of tilt values applied to the
selected elements. Use :py:obj:`None` to keep the current values.
tilts_frame: Scalar or Sequence of reference frame tilt values applied to the
selected elements. Use :py:obj:`None` to keep the current values.
pitches: Scalar or Sequence of pitch values applied to the
selected elements. Use :py:obj:`None` to keep the current values.
yaws: Scalar or Sequence of yaw values applied to the
selected elements. Use :py:obj:`None` to keep the current values.
refpts: Element selection key.
See ":ref: `Selecting elements in a lattice <refpts>`"
reference: Transformation reference, either
:py:obj:`ReferencePoint.CENTRE` or
:py:obj:`ReferencePoint.ENTRANCE`.
Default: :py:obj:`ReferencePoint.CENTRE` if the reference
point was not previously set, otherwise to set reference
point
relative: If :py:obj:`True`, the rotations are added to the existing ones.
.. versionadded:: 0.7.0
The *refpts* argument
See Also:
:py:func:`set_tilt`
:py:func:`set_shift`
"""
nb = _refcount(ring, refpts, endpoint=False)
tilts = np.broadcast_to(tilts, (nb,))
pchs = np.broadcast_to(pitches, (nb,))
yaws = np.broadcast_to(yaws, (nb,))
tilts_frame = np.broadcast_to(tilts_frame, (nb,))
for el, tilt, pitch, yaw, tilt_frame in zip(
refpts_iterator(ring, refpts), tilts, pchs, yaws, tilts_frame
):
transform_elem(
el,
reference=reference,
tilt=tilt,
pitch=pitch,
yaw=yaw,
tilt_frame=tilt_frame,
relative=relative,
)
[docs]
def set_tilt(
ring: Sequence[Element],
tilts: Sequence[float] | float | None = None,
*,
tilts_frame: Sequence[float] | float | None = None,
refpts: Refpts = All,
reference: ReferencePoint | None = None,
relative: bool = False,
) -> None:
r"""Sets the tilts of a list of elements.
Parameters:
ring: Lattice description.
tilts: Scalar or Sequence of tilt values applied to the
selected elements. Use :py:obj:`None` to keep the current values.
tilts_frame: Scalar or Sequence of reference frame tilt values applied to the
selected elements. Use :py:obj:`None` to keep the current values.
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
reference: Transformation reference, either
:py:obj:`ReferencePoint.CENTRE` or
:py:obj:`ReferencePoint.ENTRANCE`.
Default: :py:obj:`ReferencePoint.CENTRE` if the reference
point was not previously set, otherwise to set reference
point
relative: If :py:obj:`True`, the rotation is added to the existing one.
.. versionadded:: 0.7.0
The *refpts* argument
See Also:
:py:func:`set_rotation`
:py:func:`set_shift`
"""
nb = _refcount(ring, refpts, endpoint=False)
tilts = np.broadcast_to(tilts, (nb,))
tilts_frame = np.broadcast_to(tilts_frame, (nb,))
for el, tilt, tilt_frame in zip(refpts_iterator(ring, refpts), tilts, tilts_frame):
transform_elem(
el, reference=reference, tilt=tilt, tilt_frame=tilt_frame, relative=relative
)
[docs]
def set_shift(
ring: Sequence[Element],
dxs,
dys,
dzs=None,
*,
refpts: Refpts = All,
reference: ReferencePoint | None = None,
relative: bool = False,
) -> None:
r"""Sets the translations of a list of elements.
Parameters:
ring: Lattice description.
dxs: Scalar or Sequence of horizontal translations values applied
to the selected elements. Use :py:obj:`None` to keep the current values [m].
dys: Scalar or Sequence of vertical translations values applied
to the selected elements. Use :py:obj:`None` to keep the current values [m].
dzs: Scalar or Sequence of longitudinal translations values applied
to the selected elements. Use :py:obj:`None` to keep the current values [m].
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
reference: Transformation reference, either
:py:obj:`ReferencePoint.CENTRE` or
:py:obj:`ReferencePoint.ENTRANCE`.
Default: :py:obj:`ReferencePoint.CENTRE` if the reference
point was not previously set, otherwise to set reference
point
relative: If :py:obj:`True`, the translation is added to the
existing one.
.. versionadded:: 0.7.0
The *refpts* argument
See Also:
:py:func:`set_rotation`
:py:func:`set_tilt`
"""
nb = _refcount(ring, refpts, endpoint=False)
dxs = np.broadcast_to(dxs, (nb,))
dys = np.broadcast_to(dys, (nb,))
dzs = np.broadcast_to(dzs, (nb,))
for el, dx, dy, dz in zip(refpts_iterator(ring, refpts), dxs, dys, dzs):
transform_elem(el, reference=reference, dx=dx, dy=dy, dz=dz, relative=relative)
def _get_referencePoint(elem: Element) -> ReferencePoint:
"""Rotation reference point"""
idx = getattr(elem, "_referencepoint", transform_options.referencepoint.value)
return ReferencePoint(idx)
def _set_referencePoint(elem: Element, value: ReferencePoint) -> None:
elem._referencepoint = value.value
def _get_dx(elem: Element) -> float:
"""Horizontal element shift"""
return getattr(elem, "_dx", 0.0)
def _set_dx(elem: Element, value: float) -> None:
transform_elem(elem, dx=value)
def _get_dy(elem: Element) -> float:
"""Vertical element shift"""
return getattr(elem, "_dy", 0.0)
def _set_dy(elem: Element, value: float) -> None:
transform_elem(elem, dy=value)
def _get_dz(elem: Element) -> float:
"""Longitudinal element shift"""
return getattr(elem, "_dz", 0.0)
def _set_dz(elem: Element, value: float) -> None:
transform_elem(elem, dz=value)
def _get_tilt(elem: Element) -> float:
"""Element tilt"""
return getattr(elem, "_tilt", 0.0)
def _set_tilt(elem: Element, value: float) -> None:
transform_elem(elem, tilt=value)
def _get_pitch(elem: Element) -> float:
"""Element pitch"""
return getattr(elem, "_pitch", 0.0)
def _set_pitch(elem: Element, value: float) -> None:
transform_elem(elem, pitch=value)
def _get_yaw(elem: Element) -> float:
"""Element yaw"""
return getattr(elem, "_yaw", 0.0)
def _set_yaw(elem: Element, value: float) -> None:
transform_elem(elem, yaw=value)
def _get_tilt_frame(elem: Element) -> float:
"""Element tilt frame, different from tilt only for bends"""
return getattr(elem, "_tilt_frame", 0.0)
def _set_tilt_frame(elem: Element, value: float) -> None:
transform_elem(elem, tilt_frame=value)
Element.transform = transform_elem
Element.ReferencePoint = property(_get_referencePoint, _set_referencePoint)
Element.dx = property(_get_dx, _set_dx)
Element.dy = property(_get_dy, _set_dy)
Element.dz = property(_get_dz, _set_dz)
Element.tilt = property(_get_tilt, _set_tilt)
Element.pitch = property(_get_pitch, _set_pitch)
Element.yaw = property(_get_yaw, _set_yaw)
Element.tilt_frame = property(_get_tilt_frame, _set_tilt_frame)
