"""Lattice geometry"""
from __future__ import annotations
__all__ = ["get_geometry"]
from collections.abc import Sequence
from math import sin, cos, atan2, sqrt
import numpy as np
from .elements import Element, Dipole
from .utils import Refpts, All, refpts_count, get_bool_index
_GEOMETRY_EPSIL = 1.0e-3
_GEOMETRY_DTYPE = [
("x", np.float64),
("y", np.float64),
("z", np.float64),
("angle", np.float64),
("v_angle", np.float64),
]
def _get_geometry2(
ring: list[Element],
refpts: Refpts = All,
start_coordinates: tuple[float, float, float] = (0, 0, 0),
centered: bool = False,
regex: bool = False,
):
# noinspection PyShadowingNames
r"""Compute the 2D ring geometry in cartesian coordinates
Parameters:
ring: Lattice description.
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
start_coordinates: *x*, *y*, *angle* at starting point. *x*
and *y* are ignored if *centered* is :py:obj:`True`.
centered: if :py:obj:`True` the coordinates origin is the
centre of the ring.
regex: Use regular expression for *refpts* string matching instead of
Unix shell-style wildcards.
Returns:
geomdata: recarray containing, x, y, angle.
radius: machine radius at the beginning of the lattice.
.. attention::
This radius is different from the radius usually defined as
:math:`C/2\pi`
Example:
>>> geomdata, radius = get_geometry(ring)
"""
geom_dtype = [("x", np.float64), ("y", np.float64), ("angle", np.float64)]
boolrefs = get_bool_index(ring, refpts, endpoint=True, regex=regex)
nrefs = refpts_count(boolrefs, len(ring))
geomdata = np.recarray((nrefs,), dtype=geom_dtype)
xx = np.zeros(len(ring) + 1)
yy = np.zeros(len(ring) + 1)
angle = np.zeros(len(ring) + 1)
x0, y0, t0 = start_coordinates
x, y = 0.0, 0.0
t = t0
xx[0] = x
yy[0] = y
angle[0] = t
for ind, el in enumerate(ring):
ll = el.Length
if isinstance(el, Dipole) and el.BendingAngle != 0:
ang = 0.5 * el.BendingAngle
ll *= np.sin(ang) / ang
else:
ang = 0.0
t -= ang
x += ll * np.cos(t)
y += ll * np.sin(t)
t -= ang
xx[ind + 1] = x
yy[ind + 1] = y
angle[ind + 1] = t
dff = (t + _GEOMETRY_EPSIL) % (2.0 * np.pi) - _GEOMETRY_EPSIL
if abs(dff) < _GEOMETRY_EPSIL:
xcenter = np.mean(xx)
ycenter = np.mean(yy)
elif abs(dff - np.pi) < _GEOMETRY_EPSIL:
xcenter = 0.5 * x
ycenter = 0.5 * y
else:
num = np.cos(t) * x + np.sin(t) * y
den = np.sin(t - t0)
xcenter = -num * np.sin(t0) / den
ycenter = num * np.cos(t0) / den
radius = np.sqrt(xcenter * xcenter + ycenter * ycenter)
if centered:
xx -= xcenter
yy -= ycenter
else:
xx += x0
yy += y0
geomdata["x"] = xx[boolrefs]
geomdata["y"] = yy[boolrefs]
geomdata["angle"] = angle[boolrefs]
return geomdata, radius
[docs]
def get_geometry(
ring: Sequence[Element],
refpts: Refpts = All,
start_coordinates: tuple[float, float, float] = (0.0, 0.0, 0.0),
h_angle: float = 0.0,
v_angle: float = 0.0,
centred: bool = False,
**kwargs,
):
# noinspection PyShadowingNames
r"""Compute the 3D ring geometry in cartesian coordinates
Parameters:
ring: Lattice description.
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
start_coordinates: *x*, *y*, *z* at starting point. *x* and *y* are ignored
if *centred* is :py:obj:`True`.
h_angle: initial horizontal angle.
v_angle: initial vertical angle.
centred: if :py:obj:`True` the coordinates origin is the
centre of the ring.
Returns:
geomdata: :py:class:`~numpy.recarray` containing the fields *x*, *y*,
*z*, *angle*, *v_angle*.
radius: machine radius at the beginning of the lattice.
.. attention::
This radius is different from the radius usually defined as
:math:`C/2\pi`
Example:
>>> geomdata, _ = get_geometry(ring)
>>> xcoord = geomdata.x
"""
def get_centre(xx, yy, tt):
dff = (tt[-1] - tt[0] + _GEOMETRY_EPSIL) % (2.0 * np.pi) - _GEOMETRY_EPSIL
if abs(dff) < _GEOMETRY_EPSIL: # Total angle is 2*pi: full ring
xcentre = np.mean(xx)
ycentre = np.mean(yy)
elif abs(dff - np.pi) < _GEOMETRY_EPSIL: # Total angle is pi: half ring
xcentre = 0.5 * xx[-1]
ycentre = 0.5 * yy[-1]
else:
c1 = np.cos(tt[0])
s1 = np.sin(tt[0])
c2 = np.cos(tt[-1])
s2 = np.sin(tt[-1])
den = s2 * c1 - s1 * c2
a1 = xx[0] * c1 + yy[0] * s1
a2 = xx[-1] * c2 + yy[-1] * s2
xcentre = (a1 * s2 - a2 * s1) / den
ycentre = (a2 * c1 - a1 * c2) / den
return xcentre, ycentre
def rots(rotmat):
cns = rotmat[0, 0]
sns = rotmat[0, 2]
return np.array([[1.0, 0.0, 0.0], [0, cns, -sns], [0.0, sns, cns]])
def hkick(ang: float):
"""positive: to the left"""
cns = cos(ang)
sns = sin(ang)
return np.array([[cns, -sns, 0.0], [sns, cns, 0.0], [0.0, 0.0, 1.0]])
def vkick(ang: float):
"""positive: upward"""
cns = cos(ang)
sns = sin(ang)
return np.array([[cns, 0.0, -sns], [0.0, 1.0, 0.0], [sns, 0.0, cns]])
def increment(xyz, conv, elem):
"""Propagation across one element"""
length = elem.Length
if hasattr(elem, "R1"):
# inplace matrix multiplication requires numpy >= 1.25
# conv @= rots(elem.R1)
ctemp = conv.copy()
np.matmul(ctemp, rots(elem.R1), out=conv)
if hasattr(elem, "BendingAngle") and elem.BendingAngle != 0.0:
ang = 0.5 * elem.BendingAngle
rm = hkick(-ang)
# conv @= rm
ctemp = conv @ rm
xyz += ctemp @ np.array([length / ang * sin(ang), 0.0, 0.0])
# conv @= rm
np.matmul(ctemp, rm, out=conv)
else:
xyz += conv @ np.array([length, 0.0, 0.0])
if hasattr(elem, "R2"):
# conv @= rots(elem.R2)
ctemp = conv.copy()
np.matmul(ctemp, rots(elem.R2), out=conv)
hangle = atan2(conv[1, 0], conv[0, 0])
vangle = atan2(conv[2, 0], sqrt(conv[0, 0] ** 2 + conv[1, 0] ** 2))
return tuple(xyz) + (hangle, vangle)
# Accept the US wording "centered" for compatibility
centred = kwargs.pop("centered", centred)
boolrefs = get_bool_index(ring, refpts, endpoint=True)
x0, y0, z0 = start_coordinates
xyzc = np.array(start_coordinates, dtype=np.float64)
coord0 = start_coordinates + (h_angle, v_angle)
convi = hkick(h_angle) @ vkick(v_angle)
coords = [coord0] + [increment(xyzc, convi, el) for el in ring]
geomdata = np.rec.array(coords, dtype=_GEOMETRY_DTYPE)
xc, yc = get_centre(geomdata.x, geomdata.y, geomdata.angle)
radius = np.sqrt((xc - x0) ** 2 + (yc - y0) ** 2)
if centred:
geomdata.x -= xc
geomdata.y -= yc
return geomdata[boolrefs], radius
