Merge pull request #289 from Hespe/clone-methods

Add `clone` method for beams and elements

Author: jank324

CHANGELOG.md

@@ -12,6 +12,8 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 - `Screen` no longer blocks the beam (by default). To return to old behaviour, set `Screen.is_blocking = True`. (see #208) (@jank324, @roussel-ryan)
 - The way `dtype`s are determined is now more in line with PyTorch's conventions. This may cause different-than-expected `dtype`s in old code. (see #254) (@hespe, @jank324)
 - `Beam.parameters()` no longer shadows `torch.nn.Module.parameters()`. The previously returned properties now need to be queried individually. (see #300) (@hespe)
+- `e1` and `e2` in `Dipole` and `RBend` have been renamed and made more consistent between the different magnet types. They now have prefixes `dipole_` and `rbend_` respectively. (see #289) (@hespe, @jank324)
+- The `_transfer_map` property of `CustomTransferMap` has been renamed to `predefined_transfer_map`. (see #289) (@hespe, @jank324)
 
 ### 🚀 Features
 
@@ -27,6 +29,7 @@ This is a major release with significant upgrades under the hood of Cheetah. Des
 - Add `TransverseDeflectingCavity` element (following the Bmad-X implementation) (see #240, #278 #296) (@jp-ga, @cr-xu, @jank324)
 - `Dipole` and `RBend` now take a focusing moment `k1` (see #235, #247) (@hespe)
 - Implement a converter for lattice files imported from Elegant (see #222, #251, #273, #281) (@hespe, @jank324)
+- `Beam` and `Element` objects now have a `.clone()` method to create a deep copy (see #289) (@hespe, @jank324)
 
 ### 🐛 Bug fixes
 

cheetah/accelerator/aperture.py

@@ -126,12 +126,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No
 
     @property
     def defining_features(self) -> list[str]:
-        return super().defining_features + [
-            "x_max",
-            "y_max",
-            "shape",
-            "is_active",
-        ]
+        return super().defining_features + ["x_max", "y_max", "shape", "is_active"]
 
     def __repr__(self) -> str:
         return (

cheetah/accelerator/bpm.py

@@ -1,4 +1,3 @@
-from copy import deepcopy
 from typing import Optional
 
 import matplotlib.pyplot as plt
@@ -44,7 +43,7 @@ def track(self, incoming: Beam) -> Beam:
         else:
             raise TypeError(f"Parameter incoming is of invalid type {type(incoming)}")
 
-        return deepcopy(incoming)
+        return incoming.clone()
 
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
@@ -60,7 +59,7 @@ def plot(self, ax: plt.Axes, s: float, vector_idx: Optional[tuple] = None) -> No
 
     @property
     def defining_features(self) -> list[str]:
-        return super().defining_features
+        return super().defining_features + ["is_active"]
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(name={repr(self.name)})"

cheetah/accelerator/custom_transfer_map.py

@@ -18,28 +18,31 @@ class CustomTransferMap(Element):
 
     def __init__(
         self,
-        transfer_map: torch.Tensor,
+        predefined_transfer_map: torch.Tensor,
         length: Optional[torch.Tensor] = None,
         name: Optional[str] = None,
         device=None,
         dtype=None,
     ) -> None:
-        device, dtype = verify_device_and_dtype([transfer_map, length], device, dtype)
+        device, dtype = verify_device_and_dtype(
+            [predefined_transfer_map, length], device, dtype
+        )
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__(name=name)
 
-        assert isinstance(transfer_map, torch.Tensor)
-        assert transfer_map.shape[-2:] == (7, 7)
+        assert isinstance(predefined_transfer_map, torch.Tensor)
+        assert predefined_transfer_map.shape[-2:] == (7, 7)
 
         self.register_buffer(
-            "_transfer_map", torch.as_tensor(transfer_map, **factory_kwargs)
+            "predefined_transfer_map",
+            torch.as_tensor(predefined_transfer_map, **factory_kwargs),
         )
         self.register_buffer(
             "length",
             (
                 torch.as_tensor(length, **factory_kwargs)
                 if length is not None
-                else torch.zeros(transfer_map.shape[:-2], **factory_kwargs)
+                else torch.zeros(predefined_transfer_map.shape[:-2], **factory_kwargs)
             ),
         )
 
@@ -83,21 +86,23 @@ def from_merging_elements(
         )
 
     def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
-        return self._transfer_map
+        return self.predefined_transfer_map
 
     @property
     def is_skippable(self) -> bool:
         return True
 
     def __repr__(self):
         return (
-            f"{self.__class__.__name__}(transfer_map={repr(self._transfer_map)}, "
+            f"{self.__class__.__name__}("
+            + f"predefined_transfer_map={repr(self.predefined_transfer_map)}, "
             + f"length={repr(self.length)}, "
             + f"name={repr(self.name)})"
         )
 
+    @property
     def defining_features(self) -> list[str]:
-        return super().defining_features + ["transfer_map"]
+        return super().defining_features + ["length", "predefined_transfer_map"]
 
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]

cheetah/accelerator/dipole.py

@@ -23,8 +23,8 @@ class Dipole(Element):
     :param length: Length in meters.
     :param angle: Deflection angle in rad.
     :param k1: Focussing strength in 1/m^-2. Only used with `"cheetah"` tracking method.
-    :param e1: The angle of inclination of the entrance face [rad].
-    :param e2: The angle of inclination of the exit face [rad].
+    :param dipole_e1: The angle of inclination of the entrance face in rad.
+    :param dipole_e2: The angle of inclination of the exit face in rad.
     :param tilt: Tilt of the magnet in x-y plane [rad].
     :param gap: The magnet gap in meters. Note that in MAD and ELEGANT: HGAP = gap/2.
     :param gap_exit: The magnet gap at the exit in meters. Note that in MAD and
@@ -49,8 +49,8 @@ def __init__(
         length: torch.Tensor,
         angle: Optional[torch.Tensor] = None,
         k1: Optional[torch.Tensor] = None,
-        e1: Optional[torch.Tensor] = None,
-        e2: Optional[torch.Tensor] = None,
+        dipole_e1: Optional[torch.Tensor] = None,
+        dipole_e2: Optional[torch.Tensor] = None,
         tilt: Optional[torch.Tensor] = None,
         gap: Optional[torch.Tensor] = None,
         gap_exit: Optional[torch.Tensor] = None,
@@ -68,8 +68,8 @@ def __init__(
                 length,
                 angle,
                 k1,
-                e1,
-                e2,
+                dipole_e1,
+                dipole_e2,
                 tilt,
                 gap,
                 gap_exit,
@@ -88,55 +88,55 @@ def __init__(
             (
                 torch.as_tensor(angle, **factory_kwargs)
                 if angle is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.register_buffer(
             "k1",
             (
                 torch.as_tensor(k1, **factory_kwargs)
                 if k1 is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.register_buffer(
-            "e1",
+            "_e1",
             (
-                torch.as_tensor(e1, **factory_kwargs)
-                if e1 is not None
-                else torch.zeros_like(self.length)
+                torch.as_tensor(dipole_e1, **factory_kwargs)
+                if dipole_e1 is not None
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.register_buffer(
-            "e2",
+            "_e2",
             (
-                torch.as_tensor(e2, **factory_kwargs)
-                if e2 is not None
-                else torch.zeros_like(self.length)
+                torch.as_tensor(dipole_e2, **factory_kwargs)
+                if dipole_e2 is not None
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.register_buffer(
             "fringe_integral",
             (
                 torch.as_tensor(fringe_integral, **factory_kwargs)
                 if fringe_integral is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.register_buffer(
             "fringe_integral_exit",
             (
                 self.fringe_integral
                 if fringe_integral_exit is None
-                else torch.as_tensor(fringe_integral_exit, **factory_kwargs)
+                else torch.tensor(fringe_integral_exit, **factory_kwargs)
             ),
         )
         self.register_buffer(
             "gap",
             (
                 torch.as_tensor(gap, **factory_kwargs)
                 if gap is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.register_buffer(
@@ -152,7 +152,7 @@ def __init__(
             (
                 torch.as_tensor(tilt, **factory_kwargs)
                 if tilt is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.fringe_at = fringe_at
@@ -163,6 +163,22 @@ def __init__(
     def hx(self) -> torch.Tensor:
         return torch.where(self.length == 0.0, 0.0, self.angle / self.length)
 
+    @property
+    def dipole_e1(self) -> torch.Tensor:
+        return self._e1
+
+    @dipole_e1.setter
+    def dipole_e1(self, value: torch.Tensor):
+        self._e1 = value
+
+    @property
+    def dipole_e2(self) -> torch.Tensor:
+        return self._e2
+
+    @dipole_e2.setter
+    def dipole_e2(self, value: torch.Tensor):
+        self._e2 = value
+
     @property
     def is_skippable(self) -> bool:
         return self.tracking_method == "cheetah"
@@ -379,7 +395,7 @@ def _bmadx_fringe_linear(
         :return: px, py final Bmad cannonical coordinates.
         """
         g = self.angle / self.length
-        e = self.e1 * (location == "entrance") + self.e2 * (location == "exit")
+        e = self._e1 * (location == "entrance") + self._e2 * (location == "exit")
         f_int = self.fringe_integral * (
             location == "entrance"
         ) + self.fringe_integral_exit * (location == "exit")
@@ -432,18 +448,18 @@ def _transfer_map_enter(self) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
-        sec_e = 1.0 / torch.cos(self.e1)
+        sec_e = 1.0 / torch.cos(self._e1)
         phi = (
             self.fringe_integral
             * self.hx
             * self.gap
             * sec_e
-            * (1 + torch.sin(self.e1) ** 2)
+            * (1 + torch.sin(self._e1) ** 2)
         )
 
         tm = torch.eye(7, device=device, dtype=dtype).repeat(*phi.shape, 1, 1)
-        tm[..., 1, 0] = self.hx * torch.tan(self.e1)
-        tm[..., 3, 2] = -self.hx * torch.tan(self.e1 - phi)
+        tm[..., 1, 0] = self.hx * torch.tan(self._e1)
+        tm[..., 3, 2] = -self.hx * torch.tan(self._e1 - phi)
 
         return tm
 
@@ -452,18 +468,18 @@ def _transfer_map_exit(self) -> torch.Tensor:
         device = self.length.device
         dtype = self.length.dtype
 
-        sec_e = 1.0 / torch.cos(self.e2)
+        sec_e = 1.0 / torch.cos(self._e2)
         phi = (
             self.fringe_integral_exit
             * self.hx
             * self.gap
             * sec_e
-            * (1 + torch.sin(self.e2) ** 2)
+            * (1 + torch.sin(self._e2) ** 2)
         )
 
         tm = torch.eye(7, device=device, dtype=dtype).repeat(*phi.shape, 1, 1)
-        tm[..., 1, 0] = self.hx * torch.tan(self.e2)
-        tm[..., 3, 2] = -self.hx * torch.tan(self.e2 - phi)
+        tm[..., 1, 0] = self.hx * torch.tan(self._e2)
+        tm[..., 3, 2] = -self.hx * torch.tan(self._e2 - phi)
 
         return tm
 
@@ -477,8 +493,8 @@ def __repr__(self):
             f"{self.__class__.__name__}(length={repr(self.length)}, "
             + f"angle={repr(self.angle)}, "
             + f"k1={repr(self.k1)}, "
-            + f"e1={repr(self.e1)},"
-            + f"e2={repr(self.e2)},"
+            + f"dipole_e1={repr(self.dipole_e1)},"
+            + f"dipole_e2={repr(self.dipole_e2)},"
             + f"tilt={repr(self.tilt)},"
             + f"gap={repr(self.gap)},"
             + f"gap_exit={repr(self.gap_exit)},"
@@ -496,8 +512,8 @@ def defining_features(self) -> list[str]:
             "length",
             "angle",
             "k1",
-            "e1",
-            "e2",
+            "dipole_e1",
+            "dipole_e2",
             "tilt",
             "gap",
             "gap_exit",

cheetah/accelerator/element.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from copy import deepcopy
 from typing import Optional
 
 import matplotlib.pyplot as plt
@@ -110,7 +111,20 @@ def defining_features(self) -> list[str]:
         NOTE: When overriding this property, make sure to call the super method and
         extend the list it returns.
         """
-        return []
+        return ["name"]
+
+    def clone(self) -> "Element":
+        """Create a copy of the element which does not share the underlying memory."""
+        return self.__class__(
+            **{
+                feature: (
+                    getattr(self, feature).clone()
+                    if isinstance(getattr(self, feature), torch.Tensor)
+                    else deepcopy(getattr(self, feature))
+                )
+                for feature in self.defining_features
+            }
+        )
 
     @abstractmethod
     def split(self, resolution: torch.Tensor) -> list["Element"]:

cheetah/accelerator/horizontal_corrector.py

@@ -44,7 +44,7 @@ def __init__(
             (
                 torch.as_tensor(angle, **factory_kwargs)
                 if angle is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
 

cheetah/accelerator/quadrupole.py

@@ -55,7 +55,7 @@ def __init__(
             (
                 torch.as_tensor(k1, **factory_kwargs)
                 if k1 is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.register_buffer(
@@ -71,7 +71,7 @@ def __init__(
             (
                 torch.as_tensor(tilt, **factory_kwargs)
                 if tilt is not None
-                else torch.zeros_like(self.length)
+                else torch.tensor(0.0, **factory_kwargs)
             ),
         )
         self.num_steps = num_steps