New element properties: dx, dy, tilt (#902)

* new dx, dy, tilt properties
* Added tests for the new properties

Author: lfarv

pyat/at/lattice/elements.py

@@ -10,20 +10,24 @@
 
 import abc
 import re
+import math
 from abc import ABC
 from collections.abc import Generator, Iterable
 from copy import copy, deepcopy
 from typing import Any, Optional
 
-import numpy
+import numpy as np
 
 # noinspection PyProtectedMember
 from .variables import _nop
 
+_zero6 = np.zeros(6)
+_eye6 = np.eye(6, order="F")
 
-def _array(value, shape=(-1,), dtype=numpy.float64):
+
+def _array(value, shape=(-1,), dtype=np.float64):
     # Ensure proper ordering(F) and alignment(A) for "C" access in integrators
-    return numpy.require(value, dtype=dtype, requirements=["F", "A"]).reshape(
+    return np.require(value, dtype=dtype, requirements=["F", "A"]).reshape(
         shape, order="F"
     )
 
@@ -408,7 +412,7 @@ def definition(self) -> tuple[str, tuple, dict]:
         keywords = {
             k: v
             for k, v in attrs.items()
-            if not numpy.array_equal(v, getattr(defelem, k, None))
+            if not np.array_equal(v, getattr(defelem, k, None))
         }
         return self.__class__.__name__, arguments, keywords
 
@@ -448,6 +452,60 @@ def is_collective(self) -> bool:
         """:py:obj:`True` if the element involves collective effects"""
         return self._get_collective()
 
+    def _getshift(self, idx: int):
+        t1 = getattr(self, "T1", _zero6)
+        t2 = getattr(self, "T2", _zero6)
+        return 0.5 * float(t2[idx] - t1[idx])
+
+    def _setshift(self, value: float, idx: int) -> None:
+        t1 = getattr(self, "T1", _zero6.copy())
+        t2 = getattr(self, "T2", _zero6.copy())
+        sm = 0.5 * (t2[idx] + t1[idx])
+        t2[idx] = sm + value
+        t1[idx] = sm - value
+        self.T1 = t1
+        self.T2 = t2
+
+    @property
+    def dx(self) -> float:
+        """Horizontal element shift"""
+        return self._getshift(0)
+
+    @dx.setter
+    def dx(self, value: float) -> None:
+        self._setshift(value, 0)
+
+    @property
+    def dy(self) -> float:
+        """Vertical element shift"""
+        return self._getshift(2)
+
+    @dy.setter
+    def dy(self, value: float) -> None:
+        self._setshift(value, 2)
+
+    @property
+    def tilt(self) -> float:
+        """Element tilt"""
+        r1 = getattr(self, "R1", _eye6)
+        r2 = getattr(self, "R2", _eye6)
+        c = float(r2[0, 0] + r1[0, 0])
+        s = float(r2[2, 0] - r1[2, 0])
+        return math.atan2(s, c)
+
+    @tilt.setter
+    def tilt(self, value: float) -> None:
+        r1 = getattr(self, "R1", _eye6.copy())
+        r2 = getattr(self, "R2", _eye6.copy())
+        ct, st = math.cos(value), math.sin(value)
+        r44 = np.diag([ct, ct, ct, ct])
+        r44[0, 2] = r44[1, 3] = st
+        r44[2, 0] = r44[3, 1] = -st
+        r1[:4, :4] = r44
+        r2[:4, :4] = r44.T
+        self.R1 = r1
+        self.R2 = r2
+
 
 class LongElement(Element):
     """Base class for long elements"""
@@ -480,15 +538,15 @@ def popattr(element, attr):
             delattr(element, attr)
             return attr, val
 
-        frac = numpy.asarray(frac, dtype=float)
+        frac = np.asarray(frac, dtype=float)
         el = self.copy()
         # Remove entrance and exit attributes
         fin = dict(
             popattr(el, key) for key in vars(self) if key in self._entrance_fields
         )
         fout = dict(popattr(el, key) for key in vars(self) if key in self._exit_fields)
         # Split element
-        element_list = [el._part(f, numpy.sum(frac)) for f in frac]
+        element_list = [el._part(f, np.sum(frac)) for f in frac]
         # Restore entrance and exit attributes
         for key, value in fin.items():
             setattr(element_list[0], key, value)
@@ -505,7 +563,7 @@ def compatible_field(fieldname):
             elif f1 is None or f2 is None:  # only one
                 return False
             else:  # both
-                return numpy.all(f1 == f2)
+                return np.all(f1 == f2)
 
         if not (type(other) is type(self) and self.PassMethod == other.PassMethod):
             return False
@@ -538,13 +596,13 @@ def __init__(self, family_name: str, **kwargs):
         Default PassMethod: ``BeamMomentsPass``
         """
         kwargs.setdefault("PassMethod", "BeamMomentsPass")
-        self._stds = numpy.zeros((6, 1, 1), order="F")
-        self._means = numpy.zeros((6, 1, 1), order="F")
+        self._stds = np.zeros((6, 1, 1), order="F")
+        self._means = np.zeros((6, 1, 1), order="F")
         super().__init__(family_name, **kwargs)
 
     def set_buffers(self, nturns, nbunch):
-        self._stds = numpy.zeros((6, nbunch, nturns), order="F")
-        self._means = numpy.zeros((6, nbunch, nturns), order="F")
+        self._stds = np.zeros((6, nbunch, nturns), order="F")
+        self._means = np.zeros((6, nbunch, nturns), order="F")
 
     @property
     def stds(self):
@@ -583,20 +641,20 @@ def __init__(self, family_name: str, nslice: int, **kwargs):
         self.startturn = self._startturn
         self.endturn = self._endturn
         self._dturns = self.endturn - self.startturn
-        self._stds = numpy.zeros((3, nslice, self._dturns), order="F")
-        self._means = numpy.zeros((3, nslice, self._dturns), order="F")
-        self._spos = numpy.zeros((nslice, self._dturns), order="F")
-        self._weights = numpy.zeros((nslice, self._dturns), order="F")
+        self._stds = np.zeros((3, nslice, self._dturns), order="F")
+        self._means = np.zeros((3, nslice, self._dturns), order="F")
+        self._spos = np.zeros((nslice, self._dturns), order="F")
+        self._weights = np.zeros((nslice, self._dturns), order="F")
         self.set_buffers(self._endturn, 1)
 
     def set_buffers(self, nturns, nbunch):
         self.endturn = min(self.endturn, nturns)
         self._dturns = self.endturn - self.startturn
         self._nbunch = nbunch
-        self._stds = numpy.zeros((3, nbunch * self.nslice, self._dturns), order="F")
-        self._means = numpy.zeros((3, nbunch * self.nslice, self._dturns), order="F")
-        self._spos = numpy.zeros((nbunch * self.nslice, self._dturns), order="F")
-        self._weights = numpy.zeros((nbunch * self.nslice, self._dturns), order="F")
+        self._stds = np.zeros((3, nbunch * self.nslice, self._dturns), order="F")
+        self._means = np.zeros((3, nbunch * self.nslice, self._dturns), order="F")
+        self._spos = np.zeros((nbunch * self.nslice, self._dturns), order="F")
+        self._weights = np.zeros((nbunch * self.nslice, self._dturns), order="F")
 
     @property
     def stds(self):
@@ -730,11 +788,11 @@ def insert(
         """
         frac, elements = zip(*insert_list)
         lg = [0.0 if el is None else el.Length for el in elements]
-        fr = numpy.asarray(frac, dtype=float)
-        lg = 0.5 * numpy.asarray(lg, dtype=float) / self.Length
-        drfrac = numpy.hstack((fr - lg, 1.0)) - numpy.hstack((0.0, fr + lg))
+        fr = np.asarray(frac, dtype=float)
+        lg = 0.5 * np.asarray(lg, dtype=float) / self.Length
+        drfrac = np.hstack((fr - lg, 1.0)) - np.hstack((0.0, fr + lg))
         long_elems = drfrac != 0.0
-        drifts = numpy.ndarray((len(drfrac),), dtype="O")
+        drifts = np.ndarray((len(drfrac),), dtype="O")
         drifts[long_elems] = self.divide(drfrac[long_elems])
         nline = len(drifts) + len(elements)
         line = [None] * nline  # type: list[Optional[Element]]
@@ -783,12 +841,12 @@ def __init__(self, family_name: str, poly_a, poly_b, **kwargs):
         """
 
         def getpol(poly):
-            nonzero = numpy.flatnonzero(poly != 0.0)
+            nonzero = np.flatnonzero(poly != 0.0)
             return poly, len(poly), nonzero[-1] if len(nonzero) > 0 else -1
 
         def lengthen(poly, dl):
             if dl > 0:
-                return numpy.concatenate((poly, numpy.zeros(dl)))
+                return np.concatenate((poly, np.zeros(dl)))
             else:
                 return poly
 
@@ -1198,7 +1256,7 @@ def __init__(self, family_name: str, m66=None, **kwargs):
         Default PassMethod: ``Matrix66Pass``
         """
         if m66 is None:
-            m66 = numpy.identity(6)
+            m66 = np.identity(6)
         kwargs.setdefault("PassMethod", "Matrix66Pass")
         kwargs.setdefault("M66", m66)
         super().__init__(family_name, **kwargs)
@@ -1311,24 +1369,24 @@ def __init__(
         if taux == 0.0:
             dampx = 1
         else:
-            dampx = numpy.exp(-1 / taux)
+            dampx = np.exp(-1 / taux)
 
         assert tauy >= 0.0, "tauy must be greater than or equal to 0"
         if tauy == 0.0:
             dampy = 1
         else:
-            dampy = numpy.exp(-1 / tauy)
+            dampy = np.exp(-1 / tauy)
 
         assert tauz >= 0.0, "tauz must be greater than or equal to 0"
         if tauz == 0.0:
             dampz = 1
         else:
-            dampz = numpy.exp(-1 / tauz)
+            dampz = np.exp(-1 / tauz)
 
         kwargs.setdefault("PassMethod", self.default_pass[True])
         kwargs.setdefault("U0", U0)
         kwargs.setdefault(
-            "damp_mat_diag", numpy.array([dampx, dampx, dampy, dampy, dampz, dampz])
+            "damp_mat_diag", np.array([dampx, dampx, dampy, dampy, dampz, dampz])
         )
 
         super().__init__(family_name, **kwargs)
@@ -1449,7 +1507,7 @@ class QuantumDiffusion(_DictLongtMotion, Element):
     default_pass = {False: "IdentityPass", True: "QuantDiffPass"}
     _conversions = dict(Element._conversions, Lmatp=_array66)
 
-    def __init__(self, family_name: str, lmatp: numpy.ndarray, **kwargs):
+    def __init__(self, family_name: str, lmatp: np.ndarray, **kwargs):
         """Quantum diffusion element
 
         Args:

pyat/test/test_lattice_utils.py

@@ -1,6 +1,8 @@
 import numpy
 import pytest
 
+import math
+from at.lattice import elements as elt
 from at.lattice import checktype, get_cells, refpts_iterator, get_elements
 from at.lattice import elements, uint32_refpts, bool_refpts, checkattr
 from at.lattice import get_s_pos, tilt_elem, shift_elem, set_tilt, set_shift
@@ -176,20 +178,38 @@ def test_get_s_pos_returns_all_pts_for_lat_with_2_elements_using_bool_refpts():
     )
 
 
-def test_tilt_elem(simple_ring):
-    tilt_elem(simple_ring[0], (numpy.pi / 4))
-    v = 1 / 2**0.5
+def test_tilt_elem():
+    elem = elt.Drift("Drift", 1.0)
+    # Test tilt_elem function
+    tilt_elem(elem, math.pi / 4.0)
+    v = math.sqrt(2.0) / 2.0
     a = numpy.diag([v, v, v, v, 1.0, 1.0])
     a[0, 2], a[1, 3], a[2, 0], a[3, 1] = v, v, -v, -v
-    numpy.testing.assert_allclose(simple_ring[0].R1, a)
-    numpy.testing.assert_allclose(simple_ring[0].R2, a.T)
-
-
-def test_shift_elem(simple_ring):
-    shift_elem(simple_ring[2], 1.0, 0.5)
+    numpy.testing.assert_allclose(elem.R1, a, atol=1.0e-15)
+    numpy.testing.assert_allclose(elem.R2, a.T, atol=1.0e-15)
+    numpy.testing.assert_allclose(elem.tilt, numpy.pi / 4.0)
+    # Test tilt property
+    elem.tilt = math.pi / 2.0
+    a = numpy.diag([0.0, 0.0, 0.0, 0.0, 1.0, 1.0])
+    a[0, 2], a[1, 3], a[2, 0], a[3, 1] = 1.0, 1.0, -1.0, -1.0
+    numpy.testing.assert_allclose(elem.R1, a, atol=1.0e-15)
+    numpy.testing.assert_allclose(elem.R2, a.T, atol=1.0e-15)
+
+
+def test_shift_elem():
+    elem = elt.Drift("Drift", 1.0)
+    # Test shift_elem function
+    shift_elem(elem, 1.0, 0.5)
     a = numpy.array([1.0, 0.0, 0.5, 0.0, 0.0, 0.0])
-    numpy.testing.assert_equal(simple_ring[2].T1, -a)
-    numpy.testing.assert_equal(simple_ring[2].T2, a)
+    numpy.testing.assert_equal(elem.T1, -a)
+    numpy.testing.assert_equal(elem.T2, a)
+    numpy.testing.assert_equal(elem.dx, a[0])
+    numpy.testing.assert_equal(elem.dy, a[2])
+    # Test dx, dy properties
+    elem.dx = -2.0
+    elem.dy = -1.0
+    numpy.testing.assert_equal(elem.T1, 2 * a)
+    numpy.testing.assert_equal(elem.T2, -2 * a)
 
 
 def test_set_tilt(simple_ring):