Add Lattice class (#59)

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

Author: EZoni

examples/fodo.py

@@ -2,6 +2,7 @@
 from pals import Drift
 from pals import Quadrupole
 from pals import BeamLine
+from pals import Lattice
 
 
 def main():
@@ -42,26 +43,31 @@ def main():
             drift3,
         ],
     )
+    # Create lattice with the line as a branch
+    lattice = Lattice(
+        name="fodo_lattice",
+        branches=[line],
+    )
 
     # Serialize to YAML
     yaml_file = "examples_fodo.pals.yaml"
-    line.to_file(yaml_file)
+    lattice.to_file(yaml_file)
 
     # Read YAML data from file
-    loaded_line = BeamLine.from_file(yaml_file)
+    loaded_lattice = Lattice.from_file(yaml_file)
 
     # Validate loaded data
-    assert line == loaded_line
+    assert lattice == loaded_lattice
 
     # Serialize to JSON
     json_file = "examples_fodo.pals.json"
-    line.to_file(json_file)
+    lattice.to_file(json_file)
 
     # Read JSON data from file
-    loaded_line = BeamLine.from_file(json_file)
+    loaded_lattice = Lattice.from_file(json_file)
 
     # Validate loaded data
-    assert line == loaded_line
+    assert lattice == loaded_lattice
 
 
 if __name__ == "__main__":

src/pals/kinds/Lattice.py

@@ -0,0 +1,39 @@
+from pydantic import model_validator, Field
+from typing import Annotated, List, Literal, Union
+
+from .BeamLine import BeamLine
+from .mixin import BaseElement
+from ..functions import load_file_to_dict, store_dict_to_file
+
+
+class Lattice(BaseElement):
+    """A lattice combines beamlines"""
+
+    kind: Literal["Lattice"] = "Lattice"
+
+    branches: List[Annotated[Union[BeamLine], Field(discriminator="kind")]]
+
+    @model_validator(mode="before")
+    @classmethod
+    def unpack_json_structure(cls, data):
+        """Deserialize the JSON/YAML/...-like dict for Lattice elements"""
+        from pals.kinds.mixin.all_element_mixin import unpack_element_list_structure
+
+        return unpack_element_list_structure(data, "branches", "branches")
+
+    def model_dump(self, *args, **kwargs):
+        """Custom model dump for Lattice to handle element list formatting"""
+        from pals.kinds.mixin.all_element_mixin import dump_element_list
+
+        return dump_element_list(self, "branches", *args, **kwargs)
+
+    @staticmethod
+    def from_file(filename: str) -> "Lattice":
+        """Load a Lattice from a text file"""
+        pals_dict = load_file_to_dict(filename)
+        return Lattice(**pals_dict)
+
+    def to_file(self, filename: str):
+        """Save a Lattice to a text file"""
+        pals_dict = self.model_dump()
+        store_dict_to_file(filename, pals_dict)

src/pals/kinds/__init__.py

@@ -5,6 +5,7 @@
 from .ACKicker import ACKicker  # noqa: F401
 from .BeamBeam import BeamBeam  # noqa: F401
 from .BeamLine import BeamLine  # noqa: F401
+from .Lattice import Lattice  # noqa: F401
 from .BeginningEle import BeginningEle  # noqa: F401
 from .Converter import Converter  # noqa: F401
 from .CrabCavity import CrabCavity  # noqa: F401
@@ -39,3 +40,4 @@
 # Rebuild pydantic models that depend on other classes
 UnionEle.model_rebuild()
 BeamLine.model_rebuild()
+Lattice.model_rebuild()

src/pals/kinds/all_elements.py

@@ -43,6 +43,7 @@
 def get_all_element_types(extra_types: tuple = None):
     """Return a tuple of all element types that can be used in BeamLine or UnionEle."""
     element_types = (
+        "Lattice",  # Forward reference to handle circular import
         "BeamLine",  # Forward reference to handle circular import
         "UnionEle",  # Forward reference to handle circular import
         ACKicker,

tests/test_elements.py

@@ -511,3 +511,29 @@ def test_UnionEle():
     assert len(element_with_children.elements) == 2
     assert element_with_children.elements[0].name == "m1"
     assert element_with_children.elements[1].name == "d1"
+
+
+def test_Lattice():
+    # Create first line with one base element
+    element1 = pals.Marker(name="element1")
+    line1 = pals.BeamLine(name="line1", line=[element1])
+    assert line1.line == [element1]
+    # Extend first line with one thick element
+    element2 = pals.Drift(name="element2", length=2.0)
+    line1.line.extend([element2])
+    assert line1.line == [element1, element2]
+    # Create second line with one drift element
+    element3 = pals.Drift(name="element3", length=3.0)
+    line2 = pals.BeamLine(name="line2", line=[element3])
+    # Extend first line with second line
+    line1.line.extend(line2.line)
+    assert line1.line == [element1, element2, element3]
+    # Build lattice with two lines
+    lattice = pals.Lattice(name="lattice", branches=[line1, line2])
+    assert lattice.name == "lattice"
+    assert lattice.kind == "Lattice"
+    assert len(lattice.branches) == 2
+    assert lattice.branches[0].name == "line1"
+    assert lattice.branches[1].name == "line2"
+    assert lattice.branches[0].line == [element1, element2, element3]
+    assert lattice.branches[1].line == [element3]

tests/test_serialization.py

@@ -171,40 +171,45 @@ def test_comprehensive_lattice():
     wiggler = pals.Wiggler(name="wiggler1", length=2.0)
 
     # Create comprehensive lattice
-    lattice = pals.BeamLine(
+    lattice = pals.Lattice(
         name="comprehensive_lattice",
-        line=[
-            beginning,  # Start with beginning element
-            fiducial,  # Global coordinate reference
-            marker,  # Mark position
-            drift,  # Field-free region
-            quadrupole,  # Focusing element
-            sextupole,  # Chromatic correction
-            octupole,  # Higher order correction
-            multipole,  # General multipole
-            rbend,  # Rectangular bend
-            sbend,  # Sector bend
-            solenoid,  # Longitudinal focusing
-            rfcavity,  # RF acceleration
-            crabcavity,  # RF crab cavity
-            kicker,  # Transverse kick
-            ackicker,  # AC kicker
-            patch,  # Coordinate transformation
-            floorshift,  # Global coordinate shift
-            instrument,  # Measurement device
-            mask,  # Collimation
-            match,  # Matching element
-            egun,  # Electron source
-            converter,  # Species conversion
-            foil,  # Electron stripping
-            beambeam,  # Colliding beams
-            feedback,  # Feedback system
-            girder,  # Support structure
-            fork,  # Branch connection
-            taylor,  # Taylor map
-            unionele,  # Overlapping elements
-            wiggler,  # Undulator
-            nullele,  # Placeholder
+        branches=[
+            pals.BeamLine(
+                name="comprehensive_beamline",
+                line=[
+                    beginning,  # Start with beginning element
+                    fiducial,  # Global coordinate reference
+                    marker,  # Mark position
+                    drift,  # Field-free region
+                    quadrupole,  # Focusing element
+                    sextupole,  # Chromatic correction
+                    octupole,  # Higher order correction
+                    multipole,  # General multipole
+                    rbend,  # Rectangular bend
+                    sbend,  # Sector bend
+                    solenoid,  # Longitudinal focusing
+                    rfcavity,  # RF acceleration
+                    crabcavity,  # RF crab cavity
+                    kicker,  # Transverse kick
+                    ackicker,  # AC kicker
+                    patch,  # Coordinate transformation
+                    floorshift,  # Global coordinate shift
+                    instrument,  # Measurement device
+                    mask,  # Collimation
+                    match,  # Matching element
+                    egun,  # Electron source
+                    converter,  # Species conversion
+                    foil,  # Electron stripping
+                    beambeam,  # Colliding beams
+                    feedback,  # Feedback system
+                    girder,  # Support structure
+                    fork,  # Branch connection
+                    taylor,  # Taylor map
+                    unionele,  # Overlapping elements
+                    wiggler,  # Undulator
+                    nullele,  # Placeholder
+                ],
+            )
         ],
     )
 
@@ -217,10 +222,13 @@ def test_comprehensive_lattice():
         print(f"\nComprehensive lattice YAML:\n{file.read()}")
 
     # Deserialize back to Python object using Pydantic model logic
-    loaded_lattice = pals.BeamLine.from_file(yaml_file)
+    loaded_lattice = pals.Lattice.from_file(yaml_file)
 
     # Verify the loaded lattice has the correct structure and parameter groups
-    assert len(loaded_lattice.line) == 31  # Should have 31 elements
+    assert len(loaded_lattice.branches) == 1  # Should have 1 branch
+    assert (
+        len(loaded_lattice.branches[0].line) == 31
+    )  # Should have 31 elements in the branch
 
     # Verify specific elements with parameter groups are correctly loaded
     sextupole_loaded = None
@@ -229,7 +237,7 @@ def test_comprehensive_lattice():
     rfcavity_loaded = None
     unionele_loaded = None
 
-    for elem in loaded_lattice.line:
+    for elem in loaded_lattice.branches[0].line:
         if elem.name == "sextupole1":
             sextupole_loaded = elem
         elif elem.name == "octupole1":
@@ -272,10 +280,13 @@ def test_comprehensive_lattice():
         print(f"\nComprehensive lattice JSON:\n{file.read()}")
 
     # Deserialize back to Python object using Pydantic model logic
-    loaded_lattice_json = pals.BeamLine.from_file(json_file)
+    loaded_lattice_json = pals.Lattice.from_file(json_file)
 
     # Verify the loaded lattice has the correct structure and parameter groups
-    assert len(loaded_lattice_json.line) == 31  # Should have 31 elements
+    assert len(loaded_lattice_json.branches) == 1  # Should have 1 branch
+    assert (
+        len(loaded_lattice_json.branches[0].line) == 31
+    )  # Should have 31 elements in the branch
 
     # Verify specific elements with parameter groups are correctly loaded
     sextupole_loaded_json = None
@@ -284,7 +295,7 @@ def test_comprehensive_lattice():
     rfcavity_loaded_json = None
     unionele_loaded_json = None
 
-    for elem in loaded_lattice_json.line:
+    for elem in loaded_lattice_json.branches[0].line:
         if elem.name == "sextupole1":
             sextupole_loaded_json = elem
         elif elem.name == "octupole1":