diff --git a/examples/pals/fodo.pals.yaml b/examples/pals/fodo.pals.yaml index 07ef63f7b..64ab99e28 100644 --- a/examples/pals/fodo.pals.yaml +++ b/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 diff --git a/requirements.txt b/requirements.txt index 09c4adc29..de9830ff1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ numpy>=1.15 -pals-schema~=0.2.0 +pals-schema~=0.3.0 quantiphy~=2.19 diff --git a/src/python/impactx/extensions/KnownElementsList.py b/src/python/impactx/extensions/KnownElementsList.py index a2f0e2025..1a350f2b7 100644 --- a/src/python/impactx/extensions/KnownElementsList.py +++ b/src/python/impactx/extensions/KnownElementsList.py @@ -181,27 +181,9 @@ 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( @@ -209,18 +191,73 @@ def load_file(self, filename, nslice=1): ) +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( @@ -228,12 +265,28 @@ def from_pals(self, pals_beamline, nslice=1): ) ) 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, ) )