PALS Python 0.3.0, Unit Convention for Quad

examples/pals/fodo.pals.yaml

@@ -1,22 +1,33 @@
-fodo_cell:
-  kind: BeamLine
-  line:
-  - drift1:
-      kind: Drift
-      length: 0.25
-  - quad1:
-      MagneticMultipoleP:
-        Bn1: 1.0
-      kind: Quadrupole
-      length: 1.0
-  - drift2:
-      kind: Drift
-      length: 0.5
-  - quad2:
-      MagneticMultipoleP:
-        Bn1: -1.0
-      kind: Quadrupole
-      length: 1.0
-  - drift3:
-      kind: Drift
-      length: 0.25
+PALS:
+  version: null  # the PALS schema is not yet versioned
+
+  facility:
+    - fodo_cell:
+        kind: BeamLine
+        line:
+          - drift1:
+              kind: Drift
+              length: 0.25
+          - quad1:
+              MagneticMultipoleP:
+                Kn1: 1.0
+              kind: Quadrupole
+              length: 1.0
+          - drift2:
+              kind: Drift
+              length: 0.5
+          - quad2:
+              MagneticMultipoleP:
+                Kn1: -1.0
+              kind: Quadrupole
+              length: 1.0
+          - drift3:
+              kind: Drift
+              length: 0.25
+
+    - fodo_lattice:
+        kind: Lattice
+        branches:
+          - fodo_cell
+
+    - use: fodo_lattice

requirements.txt

@@ -1,3 +1,3 @@
 numpy>=1.15
-pals-schema~=0.2.0
+pals-schema~=0.3.0
 quantiphy~=2.19

src/python/impactx/extensions/KnownElementsList.py

@@ -181,59 +181,112 @@ def load_file(self, filename, nslice=1):
         return
 
     elif extension_inner == ".pals":
-        from pals.BeamLine import BeamLine
+        from pals import load as load_pals_file
 
-        # examples: fodo.pals.yaml, fodo.pals.json
-        with open(filename, "r") as file:
-            if extension == ".json":
-                import json
-
-                pals_data = json.loads(file.read())
-            elif extension == ".yaml":
-                import yaml
-
-                pals_data = yaml.safe_load(file)
-            # TODO: toml, xml
-            else:
-                raise RuntimeError(
-                    f"load_file: No support for PALS file {filename} with extension {extension} yet."
-                )
-
-        # Parse the data dictionary back into a PALS `BeamLine` object.
-        # The automatically PALS data validation happens here.
-        self.from_pals(BeamLine(**pals_data), nslice)
+        self.from_pals(load_pals_file(filename), nslice)
         return
 
     raise RuntimeError(
         f"load_file: No support for file {filename} with extension {extension} yet."
     )
 
 
+def flatten_pals(pals_data, registry=None):
+    """Flatten a PALS root, lattice, or beamline to a list of PALS elements.
+
+    Placeholder references are resolved from the root facility definitions.
+    """
+    from pals import BeamLine, Lattice, PALSroot, PlaceholderName
+
+    if registry is None:
+        registry = {}
+
+    if isinstance(pals_data, PALSroot):
+        registry = {
+            item.name: item
+            for item in pals_data.facility
+            if not isinstance(item, PlaceholderName) and hasattr(item, "name")
+        }
+
+        if len(pals_data.facility) == 1:
+            return flatten_pals(pals_data.facility[0], registry)
+
+        active_entry = pals_data.facility[-1]
+        if not isinstance(active_entry, PlaceholderName):
+            raise RuntimeError(
+                "from_pals: PALS roots with multiple facility entries must "
+                "select the active lattice or beamline with a final 'use' entry."
+            )
+        return flatten_pals(active_entry, registry)
+
+    if isinstance(pals_data, PlaceholderName):
+        if pals_data.element is not None:
+            return flatten_pals(pals_data.element, registry)
+        if pals_data.name not in registry:
+            raise RuntimeError(
+                f"from_pals: Cannot resolve PALS element reference {pals_data.name!r}."
+            )
+        return flatten_pals(registry[pals_data.name], registry)
+
+    if isinstance(pals_data, Lattice):
+        if len(pals_data.branches) != 1:
+            raise RuntimeError(
+                "from_pals: ImpactX currently supports PALS lattices with exactly "
+                f"one branch, but got {len(pals_data.branches)}."
+            )
+        return flatten_pals(pals_data.branches[0], registry)
+
+    if isinstance(pals_data, BeamLine):
+        pals_elements = []
+        for element in pals_data.line:
+            pals_elements.extend(flatten_pals(element, registry))
+        return pals_elements
+
+    return [pals_data]
+
+
 def from_pals(self, pals_beamline, nslice=1):
-    """Load and append a lattice from a Particle Accelerator Lattice Standard (PALS) Python BeamLine.
+    """Load and append a lattice from a Particle Accelerator Lattice Standard (PALS) object.
 
     https://github.com/campa-consortium/pals-python
     """
-    from pals.Drift import Drift
-    from pals.Quadrupole import Quadrupole
+    from pals import Drift, Quadrupole
+
+    pals_elements = flatten_pals(pals_beamline)
 
     # Loop over the pals_beamline and create a new ImpactX KnownElementsList from it.
     #       Use self.extend(...) on the latter.
     ix_beamline = []
-    for pals_element in pals_beamline.line:
+    for pals_element in pals_elements:
         if isinstance(pals_element, Drift):
             ix_beamline.append(
                 elements.Drift(
                     name=pals_element.name, ds=pals_element.length, nslice=nslice
                 )
             )
         elif isinstance(pals_element, Quadrupole):
+            magnetic_multipole = pals_element.MagneticMultipoleP
+            if magnetic_multipole is None:
+                raise RuntimeError(
+                    f"from_pals: No magnetic multipole input provided for element of kind {type(pals_element)}."
+                )
+
+            if getattr(magnetic_multipole, "Bn1", None) is not None:
+                k_quad = magnetic_multipole.Bn1
+                unit_quad = 1
+            elif getattr(magnetic_multipole, "Kn1", None) is not None:
+                k_quad = magnetic_multipole.Kn1
+                unit_quad = 0
+            else:
+                raise RuntimeError(
+                    f"from_pals: No gradient input provided for element of kind {type(pals_element)}."
+                )
             ix_beamline.append(
                 elements.ChrQuad(
                     name=pals_element.name,
                     ds=pals_element.length,
-                    k=pals_element.MagneticMultipoleP.Bn1,
-                    unit=0,
+                    k=k_quad,
+                    unit=unit_quad,
                     nslice=nslice,
                 )
             )