From 4b4b015410bd01f469a1d0802e043b3848d841a9 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 22 Oct 2025 09:33:32 -0700 Subject: [PATCH 1/6] [Vibe] Update Elements Based on the prompt > We defined a meta-data schema for particle accelerator lattice lements in pals. > You will find in pals/source the description of elements and their meta-data, the pals/examples has an example. > Ignore the inheritance feature for now. > I want you to find and read the element definitions. > I want you to find and read the implementation of existing PALS Python elements in pals-python/. Follow their structure and style. > Then, do this workflow: > 1. read and understand one element at a time > 2. go to pals-python/schema and add a new element, one at a time, following the existing implementation > 3. go to pals-python/tests/test_schema.py and define a compact test per new element > Repeat this until all elements defined in PALS are implemented in pals-python. And ~20 follow-up prompts. --- src/pals/kinds/ACKicker.py | 14 + src/pals/kinds/BaseElement.py | 23 +- src/pals/kinds/BeamBeam.py | 16 + src/pals/kinds/BeamLine.py | 63 +++ src/pals/kinds/BeginningEle.py | 12 + src/pals/kinds/Converter.py | 18 + src/pals/kinds/CrabCavity.py | 18 + src/pals/kinds/Drift.py | 4 +- src/pals/kinds/EGun.py | 18 + src/pals/kinds/Feedback.py | 12 + src/pals/kinds/Fiducial.py | 12 + src/pals/kinds/FloorShift.py | 14 + src/pals/kinds/Foil.py | 12 + src/pals/kinds/Fork.py | 16 + src/pals/kinds/Girder.py | 12 + src/pals/kinds/Instrument.py | 18 + src/pals/kinds/Kicker.py | 18 + src/pals/kinds/Marker.py | 12 + src/pals/kinds/Mask.py | 18 + src/pals/kinds/Match.py | 12 + src/pals/kinds/Multipole.py | 18 + src/pals/kinds/NullEle.py | 12 + src/pals/kinds/Octupole.py | 18 + src/pals/kinds/Patch.py | 16 + src/pals/kinds/Quadrupole.py | 10 +- src/pals/kinds/RBend.py | 21 + src/pals/kinds/RFCavity.py | 25 + src/pals/kinds/SBend.py | 21 + src/pals/kinds/Sextupole.py | 18 + src/pals/kinds/Solenoid.py | 23 + src/pals/kinds/Taylor.py | 12 + src/pals/kinds/UnionEle.py | 13 + src/pals/kinds/Wiggler.py | 18 + src/pals/kinds/__init__.py | 29 ++ src/pals/kinds/_warnings.py | 63 +++ src/pals/parameters/ApertureParameters.py | 18 + src/pals/parameters/BeamBeamParameters.py | 10 + src/pals/parameters/BendParameters.py | 23 + src/pals/parameters/BodyShiftParameters.py | 15 + .../parameters/ElectricMultipoleParameters.py | 10 + src/pals/parameters/FloorParameters.py | 15 + src/pals/parameters/FloorShiftParameters.py | 16 + src/pals/parameters/ForkParameters.py | 13 + src/pals/parameters/MetaParameters.py | 13 + src/pals/parameters/PatchParameters.py | 18 + src/pals/parameters/RFParameters.py | 17 + .../parameters/ReferenceChangeParameters.py | 12 + src/pals/parameters/ReferenceParameters.py | 14 + src/pals/parameters/SolenoidParameters.py | 11 + src/pals/parameters/TrackingParameters.py | 10 + src/pals/parameters/__init__.py | 15 + tests/test_elements.py | 432 +++++++++++++++++- tests/test_parameters.py | 91 ++++ tests/test_serialization.py | 317 ++++++++++++- 54 files changed, 1714 insertions(+), 15 deletions(-) create mode 100644 src/pals/kinds/ACKicker.py create mode 100644 src/pals/kinds/BeamBeam.py create mode 100644 src/pals/kinds/BeginningEle.py create mode 100644 src/pals/kinds/Converter.py create mode 100644 src/pals/kinds/CrabCavity.py create mode 100644 src/pals/kinds/EGun.py create mode 100644 src/pals/kinds/Feedback.py create mode 100644 src/pals/kinds/Fiducial.py create mode 100644 src/pals/kinds/FloorShift.py create mode 100644 src/pals/kinds/Foil.py create mode 100644 src/pals/kinds/Fork.py create mode 100644 src/pals/kinds/Girder.py create mode 100644 src/pals/kinds/Instrument.py create mode 100644 src/pals/kinds/Kicker.py create mode 100644 src/pals/kinds/Marker.py create mode 100644 src/pals/kinds/Mask.py create mode 100644 src/pals/kinds/Match.py create mode 100644 src/pals/kinds/Multipole.py create mode 100644 src/pals/kinds/NullEle.py create mode 100644 src/pals/kinds/Octupole.py create mode 100644 src/pals/kinds/Patch.py create mode 100644 src/pals/kinds/RBend.py create mode 100644 src/pals/kinds/RFCavity.py create mode 100644 src/pals/kinds/SBend.py create mode 100644 src/pals/kinds/Sextupole.py create mode 100644 src/pals/kinds/Solenoid.py create mode 100644 src/pals/kinds/Taylor.py create mode 100644 src/pals/kinds/UnionEle.py create mode 100644 src/pals/kinds/Wiggler.py create mode 100644 src/pals/kinds/_warnings.py create mode 100644 src/pals/parameters/ApertureParameters.py create mode 100644 src/pals/parameters/BeamBeamParameters.py create mode 100644 src/pals/parameters/BendParameters.py create mode 100644 src/pals/parameters/BodyShiftParameters.py create mode 100644 src/pals/parameters/ElectricMultipoleParameters.py create mode 100644 src/pals/parameters/FloorParameters.py create mode 100644 src/pals/parameters/FloorShiftParameters.py create mode 100644 src/pals/parameters/ForkParameters.py create mode 100644 src/pals/parameters/MetaParameters.py create mode 100644 src/pals/parameters/PatchParameters.py create mode 100644 src/pals/parameters/RFParameters.py create mode 100644 src/pals/parameters/ReferenceChangeParameters.py create mode 100644 src/pals/parameters/ReferenceParameters.py create mode 100644 src/pals/parameters/SolenoidParameters.py create mode 100644 src/pals/parameters/TrackingParameters.py create mode 100644 tests/test_parameters.py diff --git a/src/pals/kinds/ACKicker.py b/src/pals/kinds/ACKicker.py new file mode 100644 index 0000000..4c6628e --- /dev/null +++ b/src/pals/kinds/ACKicker.py @@ -0,0 +1,14 @@ +from typing import Literal + +from .ThickElement import ThickElement +from ._warnings import under_construction + + +@under_construction("ACKicker") +class ACKicker(ThickElement): + """Time varying kicker element""" + + # Discriminator field + kind: Literal["ACKicker"] = "ACKicker" + + # Note: ACKickerP parameter group not yet implemented diff --git a/src/pals/kinds/BaseElement.py b/src/pals/kinds/BaseElement.py index 577e619..d12799e 100644 --- a/src/pals/kinds/BaseElement.py +++ b/src/pals/kinds/BaseElement.py @@ -1,5 +1,15 @@ from pydantic import BaseModel, ConfigDict -from typing import Literal +from typing import Literal, Optional + +from ..parameters import ( + ApertureParameters, + BodyShiftParameters, + FloorParameters, + MetaParameters, + ReferenceParameters, + ReferenceChangeParameters, + TrackingParameters, +) class BaseElement(BaseModel): @@ -15,8 +25,19 @@ class BaseElement(BaseModel): # element name name: str + # Common parameter groups (optional for all elements) + ApertureP: Optional[ApertureParameters] = None + BodyShiftP: Optional[BodyShiftParameters] = None + FloorP: Optional[FloorParameters] = None + MetaP: Optional[MetaParameters] = None + ReferenceP: Optional[ReferenceParameters] = None + ReferenceChangeP: Optional[ReferenceChangeParameters] = None + TrackingP: Optional[TrackingParameters] = None + def model_dump(self, *args, **kwargs): """This makes sure the element name property is moved out and up to a one-key dictionary""" + # Exclude None values from serialization + kwargs.setdefault("exclude_none", True) elem_dict = super().model_dump(*args, **kwargs) name = elem_dict.pop("name", None) if name is None: diff --git a/src/pals/kinds/BeamBeam.py b/src/pals/kinds/BeamBeam.py new file mode 100644 index 0000000..048ac29 --- /dev/null +++ b/src/pals/kinds/BeamBeam.py @@ -0,0 +1,16 @@ +from typing import Literal, Optional + +from .BaseElement import BaseElement +from ..parameters import BeamBeamParameters +from ._warnings import under_construction + + +@under_construction("BeamBeam") +class BeamBeam(BaseElement): + """Element for simulating colliding beams""" + + # Discriminator field + kind: Literal["BeamBeam"] = "BeamBeam" + + # Beam-beam-specific parameters + BeamBeamP: Optional[BeamBeamParameters] = None diff --git a/src/pals/kinds/BeamLine.py b/src/pals/kinds/BeamLine.py index 433df35..846d277 100644 --- a/src/pals/kinds/BeamLine.py +++ b/src/pals/kinds/BeamLine.py @@ -6,6 +6,37 @@ from .Drift import Drift from .Quadrupole import Quadrupole +# Import schema elements +from .Marker import Marker +from .Sextupole import Sextupole +from .Octupole import Octupole +from .Multipole import Multipole +from .RBend import RBend +from .SBend import SBend +from .Solenoid import Solenoid +from .RFCavity import RFCavity +from .Patch import Patch +from .FloorShift import FloorShift +from .Fork import Fork +from .BeamBeam import BeamBeam +from .BeginningEle import BeginningEle +from .Fiducial import Fiducial +from .NullEle import NullEle +from .Kicker import Kicker +from .ACKicker import ACKicker +from .CrabCavity import CrabCavity +from .EGun import EGun +from .Feedback import Feedback +from .Girder import Girder +from .Instrument import Instrument +from .Mask import Mask +from .Match import Match +from .Taylor import Taylor +from .Wiggler import Wiggler +from .Converter import Converter +from .Foil import Foil +from .UnionEle import UnionEle + class BeamLine(BaseElement): """A line of elements and/or other lines""" @@ -19,11 +50,43 @@ class BeamLine(BaseElement): line: List[ Annotated[ Union[ + # Base classes (for testing compatibility) BaseElement, ThickElement, + # Original elements Drift, Quadrupole, "BeamLine", + # New schema elements + Marker, + Sextupole, + Octupole, + Multipole, + RBend, + SBend, + Solenoid, + RFCavity, + Patch, + FloorShift, + Fork, + BeamBeam, + BeginningEle, + Fiducial, + NullEle, + Kicker, + ACKicker, + CrabCavity, + EGun, + Feedback, + Girder, + Instrument, + Mask, + Match, + Taylor, + Wiggler, + Converter, + Foil, + UnionEle, ], Field(discriminator="kind"), ] diff --git a/src/pals/kinds/BeginningEle.py b/src/pals/kinds/BeginningEle.py new file mode 100644 index 0000000..0af00a7 --- /dev/null +++ b/src/pals/kinds/BeginningEle.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("BeginningEle") +class BeginningEle(BaseElement): + """Initial element at start of a branch""" + + # Discriminator field + kind: Literal["BeginningEle"] = "BeginningEle" diff --git a/src/pals/kinds/Converter.py b/src/pals/kinds/Converter.py new file mode 100644 index 0000000..e7cb724 --- /dev/null +++ b/src/pals/kinds/Converter.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .BaseElement import BaseElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Converter") +class Converter(BaseElement): + """Target to produce new species. EG: Positron converter""" + + # Discriminator field + kind: Literal["Converter"] = "Converter" + + # Converter-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/CrabCavity.py b/src/pals/kinds/CrabCavity.py new file mode 100644 index 0000000..55d0de0 --- /dev/null +++ b/src/pals/kinds/CrabCavity.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("CrabCavity") +class CrabCavity(ThickElement): + """RF crab cavity""" + + # Discriminator field + kind: Literal["CrabCavity"] = "CrabCavity" + + # CrabCavity-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Drift.py b/src/pals/kinds/Drift.py index 18bcfa9..46be16b 100644 --- a/src/pals/kinds/Drift.py +++ b/src/pals/kinds/Drift.py @@ -1,10 +1,12 @@ from typing import Literal from .ThickElement import ThickElement +from ._warnings import under_construction +@under_construction("Drift") class Drift(ThickElement): - """A field free region""" + """Field free region""" # Discriminator field kind: Literal["Drift"] = "Drift" diff --git a/src/pals/kinds/EGun.py b/src/pals/kinds/EGun.py new file mode 100644 index 0000000..f153ccf --- /dev/null +++ b/src/pals/kinds/EGun.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("EGun") +class EGun(ThickElement): + """Electron gun""" + + # Discriminator field + kind: Literal["EGun"] = "EGun" + + # EGun-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Feedback.py b/src/pals/kinds/Feedback.py new file mode 100644 index 0000000..e6db64a --- /dev/null +++ b/src/pals/kinds/Feedback.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("Feedback") +class Feedback(BaseElement): + """Element used to simulate a feedback circuit""" + + # Discriminator field + kind: Literal["Feedback"] = "Feedback" diff --git a/src/pals/kinds/Fiducial.py b/src/pals/kinds/Fiducial.py new file mode 100644 index 0000000..cd107e2 --- /dev/null +++ b/src/pals/kinds/Fiducial.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("Fiducial") +class Fiducial(BaseElement): + """Global coordinate system fiducial point""" + + # Discriminator field + kind: Literal["Fiducial"] = "Fiducial" diff --git a/src/pals/kinds/FloorShift.py b/src/pals/kinds/FloorShift.py new file mode 100644 index 0000000..fafac44 --- /dev/null +++ b/src/pals/kinds/FloorShift.py @@ -0,0 +1,14 @@ +from typing import Literal, Optional + +from .BaseElement import BaseElement +from ..parameters import FloorShiftParameters + + +class FloorShift(BaseElement): + """Global coordinates shift element""" + + # Discriminator field + kind: Literal["FloorShift"] = "FloorShift" + + # Floor shift-specific parameters + FloorShiftP: Optional[FloorShiftParameters] = None diff --git a/src/pals/kinds/Foil.py b/src/pals/kinds/Foil.py new file mode 100644 index 0000000..56fc756 --- /dev/null +++ b/src/pals/kinds/Foil.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("Foil") +class Foil(BaseElement): + """Material that can strip electrons from a particle. Will also cause energy loss and diffusion""" + + # Discriminator field + kind: Literal["Foil"] = "Foil" diff --git a/src/pals/kinds/Fork.py b/src/pals/kinds/Fork.py new file mode 100644 index 0000000..e81f48c --- /dev/null +++ b/src/pals/kinds/Fork.py @@ -0,0 +1,16 @@ +from typing import Literal, Optional + +from .BaseElement import BaseElement +from ..parameters import ForkParameters +from ._warnings import under_construction + + +@under_construction("Fork") +class Fork(BaseElement): + """Element used to connect lattice branches together""" + + # Discriminator field + kind: Literal["Fork"] = "Fork" + + # Fork-specific parameters + ForkP: Optional[ForkParameters] = None diff --git a/src/pals/kinds/Girder.py b/src/pals/kinds/Girder.py new file mode 100644 index 0000000..8b3ea91 --- /dev/null +++ b/src/pals/kinds/Girder.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("Girder") +class Girder(BaseElement): + """Element to support in space a group of other elements""" + + # Discriminator field + kind: Literal["Girder"] = "Girder" diff --git a/src/pals/kinds/Instrument.py b/src/pals/kinds/Instrument.py new file mode 100644 index 0000000..5b2811a --- /dev/null +++ b/src/pals/kinds/Instrument.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Instrument") +class Instrument(ThickElement): + """Measurement element""" + + # Discriminator field + kind: Literal["Instrument"] = "Instrument" + + # Instrument-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Kicker.py b/src/pals/kinds/Kicker.py new file mode 100644 index 0000000..d3f698a --- /dev/null +++ b/src/pals/kinds/Kicker.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Kicker") +class Kicker(ThickElement): + """Particle kicker element""" + + # Discriminator field + kind: Literal["Kicker"] = "Kicker" + + # Kicker-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Marker.py b/src/pals/kinds/Marker.py new file mode 100644 index 0000000..b3c9fcd --- /dev/null +++ b/src/pals/kinds/Marker.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("Marker") +class Marker(BaseElement): + """Zero length element to mark a particular position""" + + # Discriminator field + kind: Literal["Marker"] = "Marker" diff --git a/src/pals/kinds/Mask.py b/src/pals/kinds/Mask.py new file mode 100644 index 0000000..31e11c4 --- /dev/null +++ b/src/pals/kinds/Mask.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Mask") +class Mask(ThickElement): + """Collimation element""" + + # Discriminator field + kind: Literal["Mask"] = "Mask" + + # Mask-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Match.py b/src/pals/kinds/Match.py new file mode 100644 index 0000000..c3a1b8c --- /dev/null +++ b/src/pals/kinds/Match.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("Match") +class Match(BaseElement): + """Orbit, Twiss, and dispersion matching element""" + + # Discriminator field + kind: Literal["Match"] = "Match" diff --git a/src/pals/kinds/Multipole.py b/src/pals/kinds/Multipole.py new file mode 100644 index 0000000..c16a2fb --- /dev/null +++ b/src/pals/kinds/Multipole.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Multipole") +class Multipole(ThickElement): + """Multipole element""" + + # Discriminator field + kind: Literal["Multipole"] = "Multipole" + + # Multipole-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/NullEle.py b/src/pals/kinds/NullEle.py new file mode 100644 index 0000000..7a34480 --- /dev/null +++ b/src/pals/kinds/NullEle.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("NullEle") +class NullEle(BaseElement): + """Placeholder element used for bookkeeping""" + + # Discriminator field + kind: Literal["NullEle"] = "NullEle" diff --git a/src/pals/kinds/Octupole.py b/src/pals/kinds/Octupole.py new file mode 100644 index 0000000..fcc34bc --- /dev/null +++ b/src/pals/kinds/Octupole.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Octupole") +class Octupole(ThickElement): + """Octupole element""" + + # Discriminator field + kind: Literal["Octupole"] = "Octupole" + + # Octupole-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Patch.py b/src/pals/kinds/Patch.py new file mode 100644 index 0000000..900a016 --- /dev/null +++ b/src/pals/kinds/Patch.py @@ -0,0 +1,16 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import PatchParameters +from ._warnings import under_construction + + +@under_construction("Patch") +class Patch(ThickElement): + """Crooked drift used to shift the reference curve""" + + # Discriminator field + kind: Literal["Patch"] = "Patch" + + # Patch-specific parameters + PatchP: Optional[PatchParameters] = None diff --git a/src/pals/kinds/Quadrupole.py b/src/pals/kinds/Quadrupole.py index 5702369..3a38dc5 100644 --- a/src/pals/kinds/Quadrupole.py +++ b/src/pals/kinds/Quadrupole.py @@ -1,14 +1,16 @@ -from typing import Literal +from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import MagneticMultipoleParameters +from ..parameters import MagneticMultipoleParameters, ElectricMultipoleParameters +from ._warnings import under_construction +@under_construction("Quadrupole") class Quadrupole(ThickElement): - """A quadrupole element""" + """Quadrupole element""" # Discriminator field kind: Literal["Quadrupole"] = "Quadrupole" - # Magnetic multipole parameters + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: MagneticMultipoleParameters diff --git a/src/pals/kinds/RBend.py b/src/pals/kinds/RBend.py new file mode 100644 index 0000000..b017119 --- /dev/null +++ b/src/pals/kinds/RBend.py @@ -0,0 +1,21 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ( + BendParameters, + ElectricMultipoleParameters, + MagneticMultipoleParameters, +) + + +class RBend(ThickElement): + """A rectangular bend element""" + + # Discriminator field + kind: Literal["RBend"] = "RBend" + + # Bend-specific parameters + BendP: Optional[BendParameters] = None + + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/RFCavity.py b/src/pals/kinds/RFCavity.py new file mode 100644 index 0000000..76a21f2 --- /dev/null +++ b/src/pals/kinds/RFCavity.py @@ -0,0 +1,25 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ( + RFParameters, + SolenoidParameters, + ElectricMultipoleParameters, + MagneticMultipoleParameters, +) +from ._warnings import under_construction + + +@under_construction("RFCavity") +class RFCavity(ThickElement): + """RF cavity element""" + + # Discriminator field + kind: Literal["RFCavity"] = "RFCavity" + + # RF-specific parameters + RFP: Optional[RFParameters] = None + SolenoidP: Optional[SolenoidParameters] = None + + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/SBend.py b/src/pals/kinds/SBend.py new file mode 100644 index 0000000..98a9eb6 --- /dev/null +++ b/src/pals/kinds/SBend.py @@ -0,0 +1,21 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ( + BendParameters, + ElectricMultipoleParameters, + MagneticMultipoleParameters, +) + + +class SBend(ThickElement): + """A sector bend element""" + + # Discriminator field + kind: Literal["SBend"] = "SBend" + + # Bend-specific parameters + BendP: Optional[BendParameters] = None + + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Sextupole.py b/src/pals/kinds/Sextupole.py new file mode 100644 index 0000000..dd74d2e --- /dev/null +++ b/src/pals/kinds/Sextupole.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Sextupole") +class Sextupole(ThickElement): + """Sextupole element""" + + # Discriminator field + kind: Literal["Sextupole"] = "Sextupole" + + # Sextupole-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Solenoid.py b/src/pals/kinds/Solenoid.py new file mode 100644 index 0000000..939db64 --- /dev/null +++ b/src/pals/kinds/Solenoid.py @@ -0,0 +1,23 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ( + SolenoidParameters, + ElectricMultipoleParameters, + MagneticMultipoleParameters, +) +from ._warnings import under_construction + + +@under_construction("Solenoid") +class Solenoid(ThickElement): + """Solenoid element""" + + # Discriminator field + kind: Literal["Solenoid"] = "Solenoid" + + # Solenoid-specific parameters + SolenoidP: Optional[SolenoidParameters] = None + + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Taylor.py b/src/pals/kinds/Taylor.py new file mode 100644 index 0000000..77f4d12 --- /dev/null +++ b/src/pals/kinds/Taylor.py @@ -0,0 +1,12 @@ +from typing import Literal + +from .BaseElement import BaseElement +from ._warnings import under_construction + + +@under_construction("Taylor") +class Taylor(BaseElement): + """Taylor map element""" + + # Discriminator field + kind: Literal["Taylor"] = "Taylor" diff --git a/src/pals/kinds/UnionEle.py b/src/pals/kinds/UnionEle.py new file mode 100644 index 0000000..cf1fade --- /dev/null +++ b/src/pals/kinds/UnionEle.py @@ -0,0 +1,13 @@ +from typing import List, Literal + +from .BaseElement import BaseElement + + +class UnionEle(BaseElement): + """Union element for overlapping elements""" + + # Discriminator field + kind: Literal["UnionEle"] = "UnionEle" + + # Elements in the union + elements: List[BaseElement] = [] diff --git a/src/pals/kinds/Wiggler.py b/src/pals/kinds/Wiggler.py new file mode 100644 index 0000000..a99c67e --- /dev/null +++ b/src/pals/kinds/Wiggler.py @@ -0,0 +1,18 @@ +from typing import Literal, Optional + +from .ThickElement import ThickElement +from ..parameters import ElectricMultipoleParameters +from ..parameters import MagneticMultipoleParameters +from ._warnings import under_construction + + +@under_construction("Wiggler") +class Wiggler(ThickElement): + """Wiggler element""" + + # Discriminator field + kind: Literal["Wiggler"] = "Wiggler" + + # Wiggler-specific parameters (in addition to inherited ones) + ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None + MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/__init__.py b/src/pals/kinds/__init__.py index 8b8b26b..0ff0f75 100644 --- a/src/pals/kinds/__init__.py +++ b/src/pals/kinds/__init__.py @@ -2,8 +2,37 @@ simpler import statements like `from pals import Drift`. """ +from .ACKicker import ACKicker # noqa: F401 from .BaseElement import BaseElement # noqa: F401 +from .BeamBeam import BeamBeam # noqa: F401 from .BeamLine import BeamLine # noqa: F401 +from .BeginningEle import BeginningEle # noqa: F401 +from .Converter import Converter # noqa: F401 +from .CrabCavity import CrabCavity # noqa: F401 from .Drift import Drift # noqa: F401 +from .EGun import EGun # noqa: F401 +from .Feedback import Feedback # noqa: F401 +from .Fiducial import Fiducial # noqa: F401 +from .FloorShift import FloorShift # noqa: F401 +from .Foil import Foil # noqa: F401 +from .Fork import Fork # noqa: F401 +from .Girder import Girder # noqa: F401 +from .Instrument import Instrument # noqa: F401 +from .Kicker import Kicker # noqa: F401 +from .Marker import Marker # noqa: F401 +from .Mask import Mask # noqa: F401 +from .Match import Match # noqa: F401 +from .Multipole import Multipole # noqa: F401 +from .NullEle import NullEle # noqa: F401 +from .Octupole import Octupole # noqa: F401 +from .Patch import Patch # noqa: F401 from .Quadrupole import Quadrupole # noqa: F401 +from .RBend import RBend # noqa: F401 +from .RFCavity import RFCavity # noqa: F401 +from .SBend import SBend # noqa: F401 +from .Sextupole import Sextupole # noqa: F401 +from .Solenoid import Solenoid # noqa: F401 +from .Taylor import Taylor # noqa: F401 from .ThickElement import ThickElement # noqa: F401 +from .UnionEle import UnionEle # noqa: F401 +from .Wiggler import Wiggler # noqa: F401 diff --git a/src/pals/kinds/_warnings.py b/src/pals/kinds/_warnings.py new file mode 100644 index 0000000..056f20d --- /dev/null +++ b/src/pals/kinds/_warnings.py @@ -0,0 +1,63 @@ +""" +Utility module for handling warnings for under construction elements. +""" + +import warnings +from typing import TypeVar + +T = TypeVar("T", bound=type) + + +def under_construction(element_name: str = None): + """ + Compact decorator to mark an element as under construction. + + Usage: + @under_construction("ElementName") + class MyElement(BaseElement): + pass + + Args: + element_name: Optional custom name for the element. If not provided, + uses the class name. + """ + + def decorator(cls: T) -> T: + # Store original __init__ method + original_init = cls.__init__ + + def new_init(self, *args, **kwargs): + # Call original __init__ first + original_init(self, *args, **kwargs) + # Issue warning after initialization + name = element_name or cls.__name__ + warnings.warn( + f"The {name} element is marked as 'Under Construction' in the PALS standard. " + f"Please refer to the PALS documentation for current status and limitations.", + UserWarning, + stacklevel=3, + ) + + # Replace __init__ method + cls.__init__ = new_init + + # Add warning to class docstring if not already present + if ( + not hasattr(cls, "__doc__") + or not cls.__doc__ + or "UNDER CONSTRUCTION" not in cls.__doc__.upper() + ): + original_doc = cls.__doc__ or "" + warning_doc = f""" +**UNDER CONSTRUCTION**: This element is marked as 'Under Construction' in the PALS standard. + +{original_doc.strip()} + +**Warning**: This element implementation may be incomplete or subject to change. +Please refer to the PALS documentation for the current status and any limitations. +""" + cls.__doc__ = warning_doc.strip() + + return cls + + return decorator diff --git a/src/pals/parameters/ApertureParameters.py b/src/pals/parameters/ApertureParameters.py new file mode 100644 index 0000000..54f39d7 --- /dev/null +++ b/src/pals/parameters/ApertureParameters.py @@ -0,0 +1,18 @@ +from typing import List +from pydantic import BaseModel, ConfigDict + + +class ApertureParameters(BaseModel): + """Aperture parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + x_limits: List[float] = [float("nan"), float("nan")] + y_limits: List[float] = [float("nan"), float("nan")] + shape: str = "RECTANGULAR" + location: str = "ENTRANCE_END" + material: str = "" + thickness: float = 0.0 + aperture_shifts_with_body: bool = False + aperture_active: bool = True diff --git a/src/pals/parameters/BeamBeamParameters.py b/src/pals/parameters/BeamBeamParameters.py new file mode 100644 index 0000000..371befd --- /dev/null +++ b/src/pals/parameters/BeamBeamParameters.py @@ -0,0 +1,10 @@ +from pydantic import BaseModel, ConfigDict + + +class BeamBeamParameters(BaseModel): + """Beam-beam parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + # Parameters will be added when construction is complete diff --git a/src/pals/parameters/BendParameters.py b/src/pals/parameters/BendParameters.py new file mode 100644 index 0000000..e01287c --- /dev/null +++ b/src/pals/parameters/BendParameters.py @@ -0,0 +1,23 @@ +from pydantic import BaseModel, ConfigDict + + +class BendParameters(BaseModel): + """Bend parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + rho_ref: float = 0.0 # [radian] Reference bend angle + bend_field_ref: float = 0.0 # [T] Reference bend field + e1: float = 0.0 # [radian] Entrance end pole face rotation with respect to a sector geometry + e2: float = 0.0 # [radian] Exit end pole face rotation with respect to a rectangular geometry + e1_rect: float = 0.0 # [radian] Entrance end pole face rotation with respect to a rectangular geometry + e2_rect: float = 0.0 # [radian] Exit end pole face rotation with respect to a rectangular geometry + edge_int1: float = 0.0 # [T*m] Entrance end fringe field integral + edge_int2: float = 0.0 # [T*m] Exit end fringe field integral + g_ref: float = 0.0 # [1/m] Reference bend strength = 1/radius_ref + h1: float = 0.0 # [TODO] Entrance end pole face curvature + h2: float = 0.0 # [TODO] Exit end pole face curvature + L_chord: float = 0.0 # [m] Chord length + L_sagitta: float = 0.0 # [m] Sagitta length (output parameter) + tilt_ref: float = 0.0 # [radian] Reference tilt diff --git a/src/pals/parameters/BodyShiftParameters.py b/src/pals/parameters/BodyShiftParameters.py new file mode 100644 index 0000000..3845ad3 --- /dev/null +++ b/src/pals/parameters/BodyShiftParameters.py @@ -0,0 +1,15 @@ +from pydantic import BaseModel, ConfigDict + + +class BodyShiftParameters(BaseModel): + """Body shift parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + x_offset: float = 0.0 + y_offset: float = 0.0 + z_offset: float = 0.0 + x_rot: float = 0.0 + y_rot: float = 0.0 + z_rot: float = 0.0 diff --git a/src/pals/parameters/ElectricMultipoleParameters.py b/src/pals/parameters/ElectricMultipoleParameters.py new file mode 100644 index 0000000..3c95065 --- /dev/null +++ b/src/pals/parameters/ElectricMultipoleParameters.py @@ -0,0 +1,10 @@ +from pydantic import BaseModel, ConfigDict + + +class ElectricMultipoleParameters(BaseModel): + """Electric multipole parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + # Parameters will be added when construction is complete diff --git a/src/pals/parameters/FloorParameters.py b/src/pals/parameters/FloorParameters.py new file mode 100644 index 0000000..eb995c2 --- /dev/null +++ b/src/pals/parameters/FloorParameters.py @@ -0,0 +1,15 @@ +from pydantic import BaseModel, ConfigDict + + +class FloorParameters(BaseModel): + """Floor position and orientation parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + x_position: float = 0.0 # [m] Floor position in x + y_position: float = 0.0 # [m] Floor position in y + z_position: float = 0.0 # [m] Floor position in z + x_rotation: float = 0.0 # [radian] Floor rotation around x + y_rotation: float = 0.0 # [radian] Floor rotation around y + z_rotation: float = 0.0 # [radian] Floor rotation around z diff --git a/src/pals/parameters/FloorShiftParameters.py b/src/pals/parameters/FloorShiftParameters.py new file mode 100644 index 0000000..2cb79e5 --- /dev/null +++ b/src/pals/parameters/FloorShiftParameters.py @@ -0,0 +1,16 @@ +from pydantic import BaseModel, ConfigDict + + +class FloorShiftParameters(BaseModel): + """Floor shift parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + x_offset: float = 0.0 + y_offset: float = 0.0 + z_offset: float = 0.0 + t_offset: float = 0.0 + x_rot: float = 0.0 + y_rot: float = 0.0 + z_rot: float = 0.0 diff --git a/src/pals/parameters/ForkParameters.py b/src/pals/parameters/ForkParameters.py new file mode 100644 index 0000000..a5765b5 --- /dev/null +++ b/src/pals/parameters/ForkParameters.py @@ -0,0 +1,13 @@ +from pydantic import BaseModel, ConfigDict + + +class ForkParameters(BaseModel): + """Fork parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + to_line: str = "" + to_ele: str = "" + direction: str = "FORWARDS" # "FORWARDS" or "BACKWARDS" + propagate_reference: bool = True diff --git a/src/pals/parameters/MetaParameters.py b/src/pals/parameters/MetaParameters.py new file mode 100644 index 0000000..68f37e4 --- /dev/null +++ b/src/pals/parameters/MetaParameters.py @@ -0,0 +1,13 @@ +from pydantic import BaseModel, ConfigDict + + +class MetaParameters(BaseModel): + """Meta parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + alias: str = "" + ID: str = "" + label: str = "" + description: str = "" diff --git a/src/pals/parameters/PatchParameters.py b/src/pals/parameters/PatchParameters.py new file mode 100644 index 0000000..1dc9fea --- /dev/null +++ b/src/pals/parameters/PatchParameters.py @@ -0,0 +1,18 @@ +from pydantic import BaseModel, ConfigDict + + +class PatchParameters(BaseModel): + """Patch parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + x_offset: float = 0.0 + y_offset: float = 0.0 + z_offset: float = 0.0 + x_rot: float = 0.0 + y_rot: float = 0.0 + z_rot: float = 0.0 + flexible: bool = False + ref_coords: str = "exit_end" # "entrance_end" or "exit_end" + user_sets_length: bool = False diff --git a/src/pals/parameters/RFParameters.py b/src/pals/parameters/RFParameters.py new file mode 100644 index 0000000..7a27a89 --- /dev/null +++ b/src/pals/parameters/RFParameters.py @@ -0,0 +1,17 @@ +from pydantic import BaseModel, ConfigDict + + +class RFParameters(BaseModel): + """RF parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + frequency: float = 0.0 # [Hz] RF frequency + harmon: int = 0 # [unitless] RF frequency harmonic number + voltage: float = 0.0 # [V] RF voltage + gradient: float = 0.0 # [V/m] RF gradient + phase: float = 0.0 # [unitless] RF phase in 0 to 2*pi + multipass_phase: float = 0.0 # [unitless] RF Phase added to multipass elements + cavity_type: str = "STANDING_WAVE" # [string] Cavity type + n_cell: int = 1 # [unitless] Number of cavity cells diff --git a/src/pals/parameters/ReferenceChangeParameters.py b/src/pals/parameters/ReferenceChangeParameters.py new file mode 100644 index 0000000..951c142 --- /dev/null +++ b/src/pals/parameters/ReferenceChangeParameters.py @@ -0,0 +1,12 @@ +from pydantic import BaseModel, ConfigDict + + +class ReferenceChangeParameters(BaseModel): + """Reference energy change and/or reference time correction parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + delta_e: float = 0.0 # [eV] Energy change + delta_pc: float = 0.0 # [momentum*c] Momentum change + delta_time: float = 0.0 # [s] Time correction diff --git a/src/pals/parameters/ReferenceParameters.py b/src/pals/parameters/ReferenceParameters.py new file mode 100644 index 0000000..723a70c --- /dev/null +++ b/src/pals/parameters/ReferenceParameters.py @@ -0,0 +1,14 @@ +from pydantic import BaseModel, ConfigDict + + +class ReferenceParameters(BaseModel): + """Reference parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + species_ref: str = "" + pc_ref: float = 0.0 # [momentum*c] Reference momentum times speed of light + E_tot_ref: float = 0.0 # [eV] Reference total energy + time_ref: float = 0.0 # [s] Reference time + location: str = "" # Where reference parameters are evaluated diff --git a/src/pals/parameters/SolenoidParameters.py b/src/pals/parameters/SolenoidParameters.py new file mode 100644 index 0000000..dac5927 --- /dev/null +++ b/src/pals/parameters/SolenoidParameters.py @@ -0,0 +1,11 @@ +from pydantic import BaseModel, ConfigDict + + +class SolenoidParameters(BaseModel): + """Solenoid parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + Ksol: float = 0.0 # Normalized solenoid strength + Bsol: float = 0.0 # Solenoid field diff --git a/src/pals/parameters/TrackingParameters.py b/src/pals/parameters/TrackingParameters.py new file mode 100644 index 0000000..bb06da5 --- /dev/null +++ b/src/pals/parameters/TrackingParameters.py @@ -0,0 +1,10 @@ +from pydantic import BaseModel, ConfigDict + + +class TrackingParameters(BaseModel): + """Tracking parameters""" + + # Allow arbitrary fields + model_config = ConfigDict(extra="allow") + + # Parameters will be added when construction is complete diff --git a/src/pals/parameters/__init__.py b/src/pals/parameters/__init__.py index 9a4d433..6334a7c 100644 --- a/src/pals/parameters/__init__.py +++ b/src/pals/parameters/__init__.py @@ -2,4 +2,19 @@ simpler import statements like `from pals import MagneticMultipoleParameters`. """ +from .ApertureParameters import ApertureParameters # noqa: F401 +from .BeamBeamParameters import BeamBeamParameters # noqa: F401 +from .BendParameters import BendParameters # noqa: F401 +from .BodyShiftParameters import BodyShiftParameters # noqa: F401 +from .ElectricMultipoleParameters import ElectricMultipoleParameters # noqa: F401 +from .FloorParameters import FloorParameters # noqa: F401 +from .FloorShiftParameters import FloorShiftParameters # noqa: F401 +from .ForkParameters import ForkParameters # noqa: F401 from .MagneticMultipoleParameters import MagneticMultipoleParameters # noqa: F401 +from .MetaParameters import MetaParameters # noqa: F401 +from .PatchParameters import PatchParameters # noqa: F401 +from .ReferenceChangeParameters import ReferenceChangeParameters # noqa: F401 +from .ReferenceParameters import ReferenceParameters # noqa: F401 +from .RFParameters import RFParameters # noqa: F401 +from .SolenoidParameters import SolenoidParameters # noqa: F401 +from .TrackingParameters import TrackingParameters # noqa: F401 diff --git a/tests/test_elements.py b/tests/test_elements.py index 940bf1c..d37f169 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -1,11 +1,60 @@ from pydantic import ValidationError -from pals import MagneticMultipoleParameters -from pals import BaseElement -from pals import ThickElement -from pals import Drift -from pals import Quadrupole -from pals import BeamLine +from pals.kinds import ( + BaseElement, + ThickElement, + Drift, + Quadrupole, + BeamLine, + CrabCavity, + EGun, + Fork, + Mask, + RBend, + Marker, + ACKicker, + Solenoid, + BeamBeam, + Girder, + Taylor, + FloorShift, + Multipole, + Wiggler, + Octupole, + RFCavity, + Feedback, + BeginningEle, + UnionEle, + Patch, + NullEle, + Foil, + Sextupole, + Instrument, + Match, + SBend, + Fiducial, + Converter, + Kicker, +) + +from pals.parameters import ( + SolenoidParameters, + FloorParameters, + TrackingParameters, + FloorShiftParameters, + BeamBeamParameters, + ApertureParameters, + ElectricMultipoleParameters, + MagneticMultipoleParameters, + ReferenceParameters, + BodyShiftParameters, + ReferenceChangeParameters, + PatchParameters, + MetaParameters, + ForkParameters, + RFParameters, + BendParameters, +) def test_BaseElement(): @@ -112,3 +161,374 @@ def test_BeamLine(): # Extend first line with second line line1.line.extend(line2.line) assert line1.line == [element1, element2, element3] + + +def test_Marker(): + """Test Marker element""" + element = Marker(name="marker1") + assert element.name == "marker1" + assert element.kind == "Marker" + + +def test_Sextupole(): + """Test Sextupole element""" + element = Sextupole( + name="sext1", + length=0.5, + MagneticMultipoleP=MagneticMultipoleParameters(Bn2=1.0), + ApertureP=ApertureParameters(x_limits=[-0.1, 0.1]), + ) + assert element.name == "sext1" + assert element.length == 0.5 + assert element.kind == "Sextupole" + assert element.MagneticMultipoleP.Bn2 == 1.0 + assert element.ApertureP.x_limits == [-0.1, 0.1] + + +def test_Octupole(): + """Test Octupole element""" + element = Octupole( + name="oct1", + length=0.3, + ElectricMultipoleP=ElectricMultipoleParameters(En3=0.5), + MetaP=MetaParameters(alias="octupole_test"), + ) + assert element.name == "oct1" + assert element.length == 0.3 + assert element.kind == "Octupole" + assert element.ElectricMultipoleP.En3 == 0.5 + assert element.MetaP.alias == "octupole_test" + + +def test_Multipole(): + """Test Multipole element""" + element = Multipole( + name="mult1", + length=0.4, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=2.0, Bn2=1.5), + BodyShiftP=BodyShiftParameters(x_offset=0.01), + ) + assert element.name == "mult1" + assert element.length == 0.4 + assert element.kind == "Multipole" + assert element.MagneticMultipoleP.Bn1 == 2.0 + assert element.BodyShiftP.x_offset == 0.01 + + +def test_RBend(): + """Test RBend element""" + bend_params = BendParameters(rho_ref=1.0, bend_field_ref=2.0) + element = RBend( + name="rbend1", + length=1.0, + BendP=bend_params, + ApertureP=ApertureParameters(x_limits=[-0.2, 0.2]), + MetaP=MetaParameters(description="Test bend"), + ) + assert element.name == "rbend1" + assert element.length == 1.0 + assert element.kind == "RBend" + assert element.BendP.rho_ref == 1.0 + assert element.ApertureP.x_limits == [-0.2, 0.2] + assert element.MetaP.description == "Test bend" + + +def test_SBend(): + """Test SBend element""" + bend_params = BendParameters(rho_ref=1.5, bend_field_ref=3.0) + element = SBend( + name="sbend1", + length=1.2, + BendP=bend_params, + ReferenceP=ReferenceParameters(species_ref="proton"), + ) + assert element.name == "sbend1" + assert element.length == 1.2 + assert element.kind == "SBend" + assert element.BendP.rho_ref == 1.5 + assert element.ReferenceP.species_ref == "proton" + + +def test_Solenoid(): + """Test Solenoid element""" + sol_params = SolenoidParameters(Ksol=0.1, Bsol=0.2) + element = Solenoid( + name="sol1", + length=0.8, + SolenoidP=sol_params, + TrackingP=TrackingParameters(is_on=True), + ) + assert element.name == "sol1" + assert element.length == 0.8 + assert element.kind == "Solenoid" + assert element.SolenoidP.Ksol == 0.1 + assert element.TrackingP.is_on + + +def test_RFCavity(): + """Test RFCavity element""" + rf_params = RFParameters(frequency=1e9, voltage=1e6) + element = RFCavity( + name="rf1", + length=0.5, + RFP=rf_params, + SolenoidP=SolenoidParameters(Ksol=0.05), + FloorP=FloorParameters(x_offset=0.1), + ) + assert element.name == "rf1" + assert element.length == 0.5 + assert element.kind == "RFCavity" + assert element.RFP.frequency == 1e9 + assert element.SolenoidP.Ksol == 0.05 + assert element.FloorP.x_offset == 0.1 + + +def test_Patch(): + """Test Patch element""" + patch_params = PatchParameters(x_offset=0.1, y_offset=0.2) + element = Patch( + name="patch1", + length=0.3, + PatchP=patch_params, + ReferenceChangeP=ReferenceChangeParameters(pc_change=1e6), + ) + assert element.name == "patch1" + assert element.length == 0.3 + assert element.kind == "Patch" + assert element.PatchP.x_offset == 0.1 + assert element.ReferenceChangeP.pc_change == 1e6 + + +def test_FloorShift(): + """Test FloorShift element""" + floor_params = FloorShiftParameters(x_offset=0.5, z_offset=1.0) + element = FloorShift( + name="floor1", + FloorShiftP=floor_params, + MetaP=MetaParameters(alias="floor_test"), + ) + assert element.name == "floor1" + assert element.kind == "FloorShift" + assert element.FloorShiftP.x_offset == 0.5 + assert element.MetaP.alias == "floor_test" + + +def test_Fork(): + """Test Fork element""" + fork_params = ForkParameters(to_line="line1", direction="FORWARDS") + element = Fork( + name="fork1", + ForkP=fork_params, + ReferenceP=ReferenceParameters(species_ref="electron"), + ) + assert element.name == "fork1" + assert element.kind == "Fork" + assert element.ForkP.to_line == "line1" + assert element.ReferenceP.species_ref == "electron" + + +def test_BeamBeam(): + """Test BeamBeam element""" + bb_params = BeamBeamParameters() + element = BeamBeam( + name="bb1", + BeamBeamP=bb_params, + ApertureP=ApertureParameters(x_limits=[-0.05, 0.05]), + ) + assert element.name == "bb1" + assert element.kind == "BeamBeam" + assert element.ApertureP.x_limits == [-0.05, 0.05] + + +def test_BeginningEle(): + """Test BeginningEle element""" + element = BeginningEle( + name="begin1", MetaP=MetaParameters(description="Start of lattice") + ) + assert element.name == "begin1" + assert element.kind == "BeginningEle" + assert element.MetaP.description == "Start of lattice" + + +def test_Fiducial(): + """Test Fiducial element""" + element = Fiducial( + name="fid1", ReferenceP=ReferenceParameters(species_ref="proton") + ) + assert element.name == "fid1" + assert element.kind == "Fiducial" + assert element.ReferenceP.species_ref == "proton" + + +def test_NullEle(): + """Test NullEle element""" + element = NullEle(name="null1", TrackingP=TrackingParameters(is_on=False)) + assert element.name == "null1" + assert element.kind == "NullEle" + assert not element.TrackingP.is_on + + +def test_Kicker(): + """Test Kicker element""" + element = Kicker( + name="kick1", + length=0.2, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.5), + ElectricMultipoleP=ElectricMultipoleParameters(En1=0.3), + ) + assert element.name == "kick1" + assert element.length == 0.2 + assert element.kind == "Kicker" + assert element.MagneticMultipoleP.Bn1 == 0.5 + assert element.ElectricMultipoleP.En1 == 0.3 + + +def test_ACKicker(): + """Test ACKicker element""" + element = ACKicker(name="ackick1", length=0.15) + assert element.name == "ackick1" + assert element.length == 0.15 + assert element.kind == "ACKicker" + + +def test_CrabCavity(): + """Test CrabCavity element""" + element = CrabCavity( + name="crab1", + length=0.25, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.8), + ElectricMultipoleP=ElectricMultipoleParameters(En1=0.4), + ) + assert element.name == "crab1" + assert element.length == 0.25 + assert element.kind == "CrabCavity" + assert element.MagneticMultipoleP.Bn1 == 0.8 + assert element.ElectricMultipoleP.En1 == 0.4 + + +def test_EGun(): + """Test EGun element""" + element = EGun( + name="egun1", + length=0.1, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=1.2), + ElectricMultipoleP=ElectricMultipoleParameters(En1=0.6), + ) + assert element.name == "egun1" + assert element.length == 0.1 + assert element.kind == "EGun" + assert element.MagneticMultipoleP.Bn1 == 1.2 + assert element.ElectricMultipoleP.En1 == 0.6 + + +def test_Feedback(): + """Test Feedback element""" + element = Feedback(name="fb1", MetaP=MetaParameters(alias="feedback_test")) + assert element.name == "fb1" + assert element.kind == "Feedback" + assert element.MetaP.alias == "feedback_test" + + +def test_Girder(): + """Test Girder element""" + element = Girder(name="girder1", FloorP=FloorParameters(x_offset=0.1, y_offset=0.2)) + assert element.name == "girder1" + assert element.kind == "Girder" + assert element.FloorP.x_offset == 0.1 + assert element.FloorP.y_offset == 0.2 + + +def test_Instrument(): + """Test Instrument element""" + element = Instrument( + name="inst1", + length=0.05, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.2), + ElectricMultipoleP=ElectricMultipoleParameters(En1=0.1), + ) + assert element.name == "inst1" + assert element.length == 0.05 + assert element.kind == "Instrument" + assert element.MagneticMultipoleP.Bn1 == 0.2 + assert element.ElectricMultipoleP.En1 == 0.1 + + +def test_Mask(): + """Test Mask element""" + element = Mask( + name="mask1", + length=0.02, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.15), + ElectricMultipoleP=ElectricMultipoleParameters(En1=0.08), + ) + assert element.name == "mask1" + assert element.length == 0.02 + assert element.kind == "Mask" + assert element.MagneticMultipoleP.Bn1 == 0.15 + assert element.ElectricMultipoleP.En1 == 0.08 + + +def test_Match(): + """Test Match element""" + element = Match( + name="match1", BodyShiftP=BodyShiftParameters(x_offset=0.01, y_rot=0.02) + ) + assert element.name == "match1" + assert element.kind == "Match" + assert element.BodyShiftP.x_offset == 0.01 + assert element.BodyShiftP.y_rot == 0.02 + + +def test_Taylor(): + """Test Taylor element""" + element = Taylor( + name="taylor1", + ReferenceChangeP=ReferenceChangeParameters(pc_change=1e6, t_change=1e-9), + ) + assert element.name == "taylor1" + assert element.kind == "Taylor" + assert element.ReferenceChangeP.pc_change == 1e6 + assert element.ReferenceChangeP.t_change == 1e-9 + + +def test_Wiggler(): + """Test Wiggler element""" + element = Wiggler( + name="wig1", + length=2.0, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.5), + ElectricMultipoleP=ElectricMultipoleParameters(En1=0.3), + ) + assert element.name == "wig1" + assert element.length == 2.0 + assert element.kind == "Wiggler" + assert element.MagneticMultipoleP.Bn1 == 0.5 + assert element.ElectricMultipoleP.En1 == 0.3 + + +def test_Converter(): + """Test Converter element""" + element = Converter( + name="conv1", + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.4), + ElectricMultipoleP=ElectricMultipoleParameters(En1=0.2), + ) + assert element.name == "conv1" + assert element.kind == "Converter" + assert element.MagneticMultipoleP.Bn1 == 0.4 + assert element.ElectricMultipoleP.En1 == 0.2 + + +def test_Foil(): + """Test Foil element""" + element = Foil(name="foil1") + assert element.name == "foil1" + assert element.kind == "Foil" + + +def test_UnionEle(): + """Test UnionEle element""" + element = UnionEle(name="union1", elements=[]) + assert element.name == "union1" + assert element.kind == "UnionEle" + assert element.elements == [] diff --git a/tests/test_parameters.py b/tests/test_parameters.py new file mode 100644 index 0000000..d0d65bc --- /dev/null +++ b/tests/test_parameters.py @@ -0,0 +1,91 @@ +from pals import MagneticMultipoleParameters + +# Import parameter classes +from pals.parameters import ( + SolenoidParameters, + FloorParameters, + TrackingParameters, + FloorShiftParameters, + BeamBeamParameters, + ApertureParameters, + ElectricMultipoleParameters, + ReferenceParameters, + BodyShiftParameters, + ReferenceChangeParameters, + PatchParameters, + MetaParameters, + ForkParameters, + RFParameters, + BendParameters, +) + + +def test_ParameterClasses(): + """Test parameter classes""" + # Test ApertureParameters + aperture = ApertureParameters(x_limits=[-0.1, 0.1], y_limits=[-0.05, 0.05]) + assert aperture.x_limits == [-0.1, 0.1] + + # Test BodyShiftParameters + body_shift = BodyShiftParameters(x_offset=0.01, y_rot=0.02) + assert body_shift.x_offset == 0.01 + + # Test MetaParameters + meta = MetaParameters(alias="test", description="test element") + assert meta.alias == "test" + + # Test ElectricMultipoleParameters + emp = ElectricMultipoleParameters(En1=1.0, Es1=0.5) + assert emp.En1 == 1.0 + + # Test MagneticMultipoleParameters + mmp = MagneticMultipoleParameters(Bn1=1.0, Bs1=0.5) + assert mmp.Bn1 == 1.0 + assert mmp.Bs1 == 0.5 + + # Test SolenoidParameters + sol = SolenoidParameters(Ksol=0.1, Bsol=0.2) + assert sol.Ksol == 0.1 + + # Test RFParameters + rf = RFParameters(frequency=1e9, voltage=1e6) + assert rf.frequency == 1e9 + + # Test BendParameters + bend = BendParameters(rho_ref=1.0, bend_field_ref=2.0) + assert bend.rho_ref == 1.0 + + # Test PatchParameters + patch = PatchParameters(x_offset=0.1, flexible=True) + assert patch.x_offset == 0.1 + + # Test FloorShiftParameters + floor = FloorShiftParameters(x_offset=0.5, z_offset=1.0) + assert floor.x_offset == 0.5 + + # Test ForkParameters + fork = ForkParameters(to_line="line1", direction="FORWARDS") + assert fork.to_line == "line1" + + # Test ReferenceParameters + ref = ReferenceParameters(species_ref="electron", pc_ref=1e6) + assert ref.species_ref == "electron" + + # Test TrackingParameters + tracking = TrackingParameters(is_on=True) + assert tracking.is_on + + # Test FloorParameters + floor_pos = FloorParameters(x_offset=0.1, y_offset=0.2, z_offset=0.3) + assert floor_pos.x_offset == 0.1 + assert floor_pos.y_offset == 0.2 + assert floor_pos.z_offset == 0.3 + + # Test ReferenceChangeParameters + ref_change = ReferenceChangeParameters(pc_change=1e6, t_change=1e-9) + assert ref_change.pc_change == 1e6 + assert ref_change.t_change == 1e-9 + + # Test BeamBeamParameters + beambeam = BeamBeamParameters() + assert beambeam is not None diff --git a/tests/test_serialization.py b/tests/test_serialization.py index eb718cd..5ae3a9a 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -2,9 +2,56 @@ import os import yaml -from pals import BaseElement -from pals import ThickElement -from pals import BeamLine +from pals.kinds import ( + BaseElement, + ThickElement, + BeamLine, + Drift, + Quadrupole, + CrabCavity, + EGun, + Fork, + Mask, + RBend, + Marker, + ACKicker, + Solenoid, + BeamBeam, + Girder, + Taylor, + FloorShift, + Multipole, + Wiggler, + Octupole, + RFCavity, + Feedback, + BeginningEle, + UnionEle, + Patch, + NullEle, + Foil, + Sextupole, + Instrument, + Match, + SBend, + Fiducial, + Converter, + Kicker, +) + +from pals.parameters import ( + SolenoidParameters, + FloorShiftParameters, + BeamBeamParameters, + ApertureParameters, + ElectricMultipoleParameters, + MagneticMultipoleParameters, + PatchParameters, + MetaParameters, + ForkParameters, + RFParameters, + BendParameters, +) def test_yaml(): @@ -55,3 +102,267 @@ def test_json(): os.remove(test_file) # Validate loaded BeamLine object assert line == loaded_line + + +def test_comprehensive_lattice(): + """Test a comprehensive lattice using every PALS element at least once""" + + # Create elements in alphabetical order for easy maintenance + # ACKicker + ackicker = ACKicker(name="ackicker1", length=0.1) + + # BeamBeam + beambeam = BeamBeam(name="beambeam1", BeamBeamP=BeamBeamParameters()) + + # BeginningEle + beginning = BeginningEle(name="beginning1") + + # Converter + converter = Converter(name="converter1") + + # CrabCavity + crabcavity = CrabCavity(name="crabcavity1", length=0.2) + + # Drift + drift = Drift(name="drift1", length=0.5) + + # EGun + egun = EGun(name="egun1", length=0.15) + + # Feedback + feedback = Feedback(name="feedback1") + + # Fiducial + fiducial = Fiducial(name="fiducial1") + + # FloorShift + floorshift = FloorShift( + name="floorshift1", FloorShiftP=FloorShiftParameters(x_offset=0.1) + ) + + # Foil + foil = Foil(name="foil1") + + # Fork + fork = Fork(name="fork1", ForkP=ForkParameters(to_line="line1")) + + # Girder + girder = Girder(name="girder1") + + # Instrument + instrument = Instrument(name="instrument1", length=0.05) + + # Kicker + kicker = Kicker(name="kicker1", length=0.1) + + # Marker + marker = Marker(name="marker1") + + # Mask + mask = Mask(name="mask1", length=0.02) + + # Match + match = Match(name="match1") + + # Multipole + multipole = Multipole(name="multipole1", length=0.3) + + # NullEle + nullele = NullEle(name="nullele1") + + # Octupole + octupole = Octupole( + name="octupole1", + length=0.25, + ElectricMultipoleP=ElectricMultipoleParameters(En3=0.5), + MetaP=MetaParameters(alias="octupole_test"), + ) + + # Patch + patch = Patch(name="patch1", length=0.4, PatchP=PatchParameters(x_offset=0.05)) + + # Quadrupole + quadrupole = Quadrupole( + name="quadrupole1", + length=0.8, + MagneticMultipoleP=MagneticMultipoleParameters(Bn1=1.0), + ) + + # RBend + rbend = RBend( + name="rbend1", + length=1.0, + BendP=BendParameters(rho_ref=2.0), + ApertureP=ApertureParameters(x_limits=[-0.2, 0.2]), + ) + + # RFCavity + rfcavity = RFCavity( + name="rfcavity1", + length=0.3, + RFP=RFParameters(frequency=1e9), + SolenoidP=SolenoidParameters(Ksol=0.05), + ) + + # SBend + sbend = SBend(name="sbend1", length=1.2, BendP=BendParameters(rho_ref=1.5)) + + # Sextupole + sextupole = Sextupole( + name="sextupole1", + length=0.2, + MagneticMultipoleP=MagneticMultipoleParameters(Bn2=1.0), + ApertureP=ApertureParameters(x_limits=[-0.1, 0.1]), + ) + + # Solenoid + solenoid = Solenoid( + name="solenoid1", length=0.6, SolenoidP=SolenoidParameters(Ksol=0.1) + ) + + # Taylor + taylor = Taylor(name="taylor1") + + # UnionEle + unionele = UnionEle(name="unionele1", elements=[]) + + # Wiggler + wiggler = Wiggler(name="wiggler1", length=2.0) + + # Create comprehensive lattice + lattice = BeamLine( + 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 + ], + ) + + # Test serialization to YAML + yaml_data = yaml.dump(lattice.model_dump(), default_flow_style=False) + print(f"\nComprehensive lattice YAML:\n{yaml_data}") + + # Write to temporary file + yaml_file = "comprehensive_lattice.yaml" + with open(yaml_file, "w") as file: + file.write(yaml_data) + + # Read back from file + with open(yaml_file, "r") as file: + loaded_yaml_data = yaml.safe_load(file) + + # Deserialize back to Python object using Pydantic model logic + loaded_lattice = BeamLine(**loaded_yaml_data) + + # Verify the loaded lattice has the correct structure and parameter groups + assert len(loaded_lattice.line) == 31 # Should have 31 elements + + # Verify specific elements with parameter groups are correctly loaded + sextupole_loaded = None + octupole_loaded = None + rbend_loaded = None + rfcavity_loaded = None + + for elem in loaded_lattice.line: + if elem.name == "sextupole1": + sextupole_loaded = elem + elif elem.name == "octupole1": + octupole_loaded = elem + elif elem.name == "rbend1": + rbend_loaded = elem + elif elem.name == "rfcavity1": + rfcavity_loaded = elem + + # Test that parameter groups are correctly deserialized + assert sextupole_loaded.MagneticMultipoleP.Bn2 == 1.0 + assert sextupole_loaded.ApertureP.x_limits == [-0.1, 0.1] + + assert octupole_loaded.ElectricMultipoleP.En3 == 0.5 + assert octupole_loaded.MetaP.alias == "octupole_test" + + assert rbend_loaded.BendP.rho_ref == 2.0 + assert rbend_loaded.ApertureP.x_limits == [-0.2, 0.2] + + assert rfcavity_loaded.RFP.frequency == 1e9 + assert rfcavity_loaded.SolenoidP.Ksol == 0.05 + + # Test serialization to JSON + json_data = json.dumps(lattice.model_dump(), sort_keys=True, indent=2) + print(f"\nComprehensive lattice JSON:\n{json_data}") + + # Write to temporary file + json_file = "comprehensive_lattice.json" + with open(json_file, "w") as file: + file.write(json_data) + + # Read back from file + with open(json_file, "r") as file: + loaded_json_data = json.loads(file.read()) + + # Deserialize back to Python object using Pydantic model logic + loaded_lattice_json = BeamLine(**loaded_json_data) + + # Verify the loaded lattice has the correct structure and parameter groups + assert len(loaded_lattice_json.line) == 31 # Should have 31 elements + + # Verify specific elements with parameter groups are correctly loaded + sextupole_loaded_json = None + octupole_loaded_json = None + rbend_loaded_json = None + rfcavity_loaded_json = None + + for elem in loaded_lattice_json.line: + if elem.name == "sextupole1": + sextupole_loaded_json = elem + elif elem.name == "octupole1": + octupole_loaded_json = elem + elif elem.name == "rbend1": + rbend_loaded_json = elem + elif elem.name == "rfcavity1": + rfcavity_loaded_json = elem + + # Test that parameter groups are correctly deserialized + assert sextupole_loaded_json.MagneticMultipoleP.Bn2 == 1.0 + assert sextupole_loaded_json.ApertureP.x_limits == [-0.1, 0.1] + + assert octupole_loaded_json.ElectricMultipoleP.En3 == 0.5 + assert octupole_loaded_json.MetaP.alias == "octupole_test" + + assert rbend_loaded_json.BendP.rho_ref == 2.0 + assert rbend_loaded_json.ApertureP.x_limits == [-0.2, 0.2] + + assert rfcavity_loaded_json.RFP.frequency == 1e9 + assert rfcavity_loaded_json.SolenoidP.Ksol == 0.05 + + # Clean up temporary files + os.remove(yaml_file) + os.remove(json_file) From fed09901f0281687d383a04e1aaa2aadb7bb8068 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 23 Oct 2025 11:53:05 -0700 Subject: [PATCH 2/6] Final Review & Cleanup --- src/pals/kinds/BeamLine.py | 43 ++-- src/pals/kinds/Converter.py | 5 +- src/pals/kinds/CrabCavity.py | 5 +- src/pals/kinds/Drift.py | 2 - src/pals/kinds/EGun.py | 5 +- src/pals/kinds/Instrument.py | 5 +- src/pals/kinds/Kicker.py | 5 +- src/pals/kinds/Mask.py | 5 +- src/pals/kinds/Multipole.py | 5 +- src/pals/kinds/Octupole.py | 5 +- src/pals/kinds/Quadrupole.py | 3 +- src/pals/kinds/Sextupole.py | 5 +- src/pals/kinds/UnionEle.py | 1 + src/pals/kinds/Wiggler.py | 5 +- src/pals/parameters/ApertureParameters.py | 5 +- src/pals/parameters/BeamBeamParameters.py | 5 +- src/pals/parameters/BendParameters.py | 5 +- src/pals/parameters/BodyShiftParameters.py | 5 +- .../parameters/ElectricMultipoleParameters.py | 3 +- src/pals/parameters/FloorParameters.py | 12 +- src/pals/parameters/FloorShiftParameters.py | 5 +- src/pals/parameters/ForkParameters.py | 5 +- src/pals/parameters/MetaParameters.py | 5 +- src/pals/parameters/PatchParameters.py | 5 +- src/pals/parameters/RFParameters.py | 5 +- .../parameters/ReferenceChangeParameters.py | 10 +- src/pals/parameters/ReferenceParameters.py | 5 +- src/pals/parameters/SolenoidParameters.py | 5 +- src/pals/parameters/TrackingParameters.py | 5 +- tests/test_elements.py | 231 +++++++----------- tests/test_parameters.py | 43 ++-- tests/test_serialization.py | 161 +++++------- 32 files changed, 226 insertions(+), 393 deletions(-) diff --git a/src/pals/kinds/BeamLine.py b/src/pals/kinds/BeamLine.py index 846d277..de5aed9 100644 --- a/src/pals/kinds/BeamLine.py +++ b/src/pals/kinds/BeamLine.py @@ -53,40 +53,39 @@ class BeamLine(BaseElement): # Base classes (for testing compatibility) BaseElement, ThickElement, - # Original elements - Drift, - Quadrupole, + # User-Facing element kinds "BeamLine", - # New schema elements - Marker, - Sextupole, - Octupole, - Multipole, - RBend, - SBend, - Solenoid, - RFCavity, - Patch, - FloorShift, - Fork, + ACKicker, BeamBeam, BeginningEle, - Fiducial, - NullEle, - Kicker, - ACKicker, + Converter, CrabCavity, + Drift, EGun, Feedback, + Fiducial, + FloorShift, + Foil, + Fork, Girder, Instrument, + Kicker, + Marker, Mask, Match, + Multipole, + NullEle, + Octupole, + Patch, + Quadrupole, + RBend, + RFCavity, + SBend, + Sextupole, + Solenoid, Taylor, - Wiggler, - Converter, - Foil, UnionEle, + Wiggler, ], Field(discriminator="kind"), ] diff --git a/src/pals/kinds/Converter.py b/src/pals/kinds/Converter.py index e7cb724..cdeee8d 100644 --- a/src/pals/kinds/Converter.py +++ b/src/pals/kinds/Converter.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .BaseElement import BaseElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Converter(BaseElement): # Discriminator field kind: Literal["Converter"] = "Converter" - # Converter-specific parameters (in addition to inherited ones) + # Converter-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/CrabCavity.py b/src/pals/kinds/CrabCavity.py index 55d0de0..70cc334 100644 --- a/src/pals/kinds/CrabCavity.py +++ b/src/pals/kinds/CrabCavity.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class CrabCavity(ThickElement): # Discriminator field kind: Literal["CrabCavity"] = "CrabCavity" - # CrabCavity-specific parameters (in addition to inherited ones) + # CrabCavity-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Drift.py b/src/pals/kinds/Drift.py index 46be16b..e38c6f5 100644 --- a/src/pals/kinds/Drift.py +++ b/src/pals/kinds/Drift.py @@ -1,10 +1,8 @@ from typing import Literal from .ThickElement import ThickElement -from ._warnings import under_construction -@under_construction("Drift") class Drift(ThickElement): """Field free region""" diff --git a/src/pals/kinds/EGun.py b/src/pals/kinds/EGun.py index f153ccf..5dd719e 100644 --- a/src/pals/kinds/EGun.py +++ b/src/pals/kinds/EGun.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class EGun(ThickElement): # Discriminator field kind: Literal["EGun"] = "EGun" - # EGun-specific parameters (in addition to inherited ones) + # EGun-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Instrument.py b/src/pals/kinds/Instrument.py index 5b2811a..d7ac343 100644 --- a/src/pals/kinds/Instrument.py +++ b/src/pals/kinds/Instrument.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Instrument(ThickElement): # Discriminator field kind: Literal["Instrument"] = "Instrument" - # Instrument-specific parameters (in addition to inherited ones) + # Instrument-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Kicker.py b/src/pals/kinds/Kicker.py index d3f698a..5ee5f3a 100644 --- a/src/pals/kinds/Kicker.py +++ b/src/pals/kinds/Kicker.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Kicker(ThickElement): # Discriminator field kind: Literal["Kicker"] = "Kicker" - # Kicker-specific parameters (in addition to inherited ones) + # Kicker-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Mask.py b/src/pals/kinds/Mask.py index 31e11c4..4222682 100644 --- a/src/pals/kinds/Mask.py +++ b/src/pals/kinds/Mask.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Mask(ThickElement): # Discriminator field kind: Literal["Mask"] = "Mask" - # Mask-specific parameters (in addition to inherited ones) + # Mask-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Multipole.py b/src/pals/kinds/Multipole.py index c16a2fb..1fac2f6 100644 --- a/src/pals/kinds/Multipole.py +++ b/src/pals/kinds/Multipole.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Multipole(ThickElement): # Discriminator field kind: Literal["Multipole"] = "Multipole" - # Multipole-specific parameters (in addition to inherited ones) + # Multipole-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Octupole.py b/src/pals/kinds/Octupole.py index fcc34bc..3126d6d 100644 --- a/src/pals/kinds/Octupole.py +++ b/src/pals/kinds/Octupole.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Octupole(ThickElement): # Discriminator field kind: Literal["Octupole"] = "Octupole" - # Octupole-specific parameters (in addition to inherited ones) + # Octupole-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/Quadrupole.py b/src/pals/kinds/Quadrupole.py index 3a38dc5..a83b671 100644 --- a/src/pals/kinds/Quadrupole.py +++ b/src/pals/kinds/Quadrupole.py @@ -2,15 +2,14 @@ from .ThickElement import ThickElement from ..parameters import MagneticMultipoleParameters, ElectricMultipoleParameters -from ._warnings import under_construction -@under_construction("Quadrupole") class Quadrupole(ThickElement): """Quadrupole element""" # Discriminator field kind: Literal["Quadrupole"] = "Quadrupole" + # Octupole-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: MagneticMultipoleParameters diff --git a/src/pals/kinds/Sextupole.py b/src/pals/kinds/Sextupole.py index dd74d2e..a09c43b 100644 --- a/src/pals/kinds/Sextupole.py +++ b/src/pals/kinds/Sextupole.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Sextupole(ThickElement): # Discriminator field kind: Literal["Sextupole"] = "Sextupole" - # Sextupole-specific parameters (in addition to inherited ones) + # Sextupole-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/kinds/UnionEle.py b/src/pals/kinds/UnionEle.py index cf1fade..1bfa0d0 100644 --- a/src/pals/kinds/UnionEle.py +++ b/src/pals/kinds/UnionEle.py @@ -10,4 +10,5 @@ class UnionEle(BaseElement): kind: Literal["UnionEle"] = "UnionEle" # Elements in the union + # Note: https://github.com/campa-consortium/pals/issues/89 elements: List[BaseElement] = [] diff --git a/src/pals/kinds/Wiggler.py b/src/pals/kinds/Wiggler.py index a99c67e..a9e1510 100644 --- a/src/pals/kinds/Wiggler.py +++ b/src/pals/kinds/Wiggler.py @@ -1,8 +1,7 @@ from typing import Literal, Optional from .ThickElement import ThickElement -from ..parameters import ElectricMultipoleParameters -from ..parameters import MagneticMultipoleParameters +from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters from ._warnings import under_construction @@ -13,6 +12,6 @@ class Wiggler(ThickElement): # Discriminator field kind: Literal["Wiggler"] = "Wiggler" - # Wiggler-specific parameters (in addition to inherited ones) + # Wiggler-specific parameters ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None diff --git a/src/pals/parameters/ApertureParameters.py b/src/pals/parameters/ApertureParameters.py index 54f39d7..0c593d7 100644 --- a/src/pals/parameters/ApertureParameters.py +++ b/src/pals/parameters/ApertureParameters.py @@ -1,13 +1,10 @@ from typing import List -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class ApertureParameters(BaseModel): """Aperture parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - x_limits: List[float] = [float("nan"), float("nan")] y_limits: List[float] = [float("nan"), float("nan")] shape: str = "RECTANGULAR" diff --git a/src/pals/parameters/BeamBeamParameters.py b/src/pals/parameters/BeamBeamParameters.py index 371befd..ba530c7 100644 --- a/src/pals/parameters/BeamBeamParameters.py +++ b/src/pals/parameters/BeamBeamParameters.py @@ -1,10 +1,7 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class BeamBeamParameters(BaseModel): """Beam-beam parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - # Parameters will be added when construction is complete diff --git a/src/pals/parameters/BendParameters.py b/src/pals/parameters/BendParameters.py index e01287c..4c1359e 100644 --- a/src/pals/parameters/BendParameters.py +++ b/src/pals/parameters/BendParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class BendParameters(BaseModel): """Bend parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - rho_ref: float = 0.0 # [radian] Reference bend angle bend_field_ref: float = 0.0 # [T] Reference bend field e1: float = 0.0 # [radian] Entrance end pole face rotation with respect to a sector geometry diff --git a/src/pals/parameters/BodyShiftParameters.py b/src/pals/parameters/BodyShiftParameters.py index 3845ad3..9624a50 100644 --- a/src/pals/parameters/BodyShiftParameters.py +++ b/src/pals/parameters/BodyShiftParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class BodyShiftParameters(BaseModel): """Body shift parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - x_offset: float = 0.0 y_offset: float = 0.0 z_offset: float = 0.0 diff --git a/src/pals/parameters/ElectricMultipoleParameters.py b/src/pals/parameters/ElectricMultipoleParameters.py index 3c95065..2df8fb3 100644 --- a/src/pals/parameters/ElectricMultipoleParameters.py +++ b/src/pals/parameters/ElectricMultipoleParameters.py @@ -7,4 +7,5 @@ class ElectricMultipoleParameters(BaseModel): # Allow arbitrary fields model_config = ConfigDict(extra="allow") - # Parameters will be added when construction is complete + # TODO: add ElectricMultipoleParameters in a follow-up RP + # https://pals-project.readthedocs.io/en/latest/element-parameters.html#electricmultipolep-electric-multipole-parameters diff --git a/src/pals/parameters/FloorParameters.py b/src/pals/parameters/FloorParameters.py index eb995c2..74f80b8 100644 --- a/src/pals/parameters/FloorParameters.py +++ b/src/pals/parameters/FloorParameters.py @@ -1,15 +1,7 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class FloorParameters(BaseModel): """Floor position and orientation parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - - x_position: float = 0.0 # [m] Floor position in x - y_position: float = 0.0 # [m] Floor position in y - z_position: float = 0.0 # [m] Floor position in z - x_rotation: float = 0.0 # [radian] Floor rotation around x - y_rotation: float = 0.0 # [radian] Floor rotation around y - z_rotation: float = 0.0 # [radian] Floor rotation around z + # Under construction diff --git a/src/pals/parameters/FloorShiftParameters.py b/src/pals/parameters/FloorShiftParameters.py index 2cb79e5..e9bdfae 100644 --- a/src/pals/parameters/FloorShiftParameters.py +++ b/src/pals/parameters/FloorShiftParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class FloorShiftParameters(BaseModel): """Floor shift parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - x_offset: float = 0.0 y_offset: float = 0.0 z_offset: float = 0.0 diff --git a/src/pals/parameters/ForkParameters.py b/src/pals/parameters/ForkParameters.py index a5765b5..836e01c 100644 --- a/src/pals/parameters/ForkParameters.py +++ b/src/pals/parameters/ForkParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class ForkParameters(BaseModel): """Fork parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - to_line: str = "" to_ele: str = "" direction: str = "FORWARDS" # "FORWARDS" or "BACKWARDS" diff --git a/src/pals/parameters/MetaParameters.py b/src/pals/parameters/MetaParameters.py index 68f37e4..ff79506 100644 --- a/src/pals/parameters/MetaParameters.py +++ b/src/pals/parameters/MetaParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class MetaParameters(BaseModel): """Meta parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - alias: str = "" ID: str = "" label: str = "" diff --git a/src/pals/parameters/PatchParameters.py b/src/pals/parameters/PatchParameters.py index 1dc9fea..54c9a9d 100644 --- a/src/pals/parameters/PatchParameters.py +++ b/src/pals/parameters/PatchParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class PatchParameters(BaseModel): """Patch parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - x_offset: float = 0.0 y_offset: float = 0.0 z_offset: float = 0.0 diff --git a/src/pals/parameters/RFParameters.py b/src/pals/parameters/RFParameters.py index 7a27a89..272f2b5 100644 --- a/src/pals/parameters/RFParameters.py +++ b/src/pals/parameters/RFParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class RFParameters(BaseModel): """RF parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - frequency: float = 0.0 # [Hz] RF frequency harmon: int = 0 # [unitless] RF frequency harmonic number voltage: float = 0.0 # [V] RF voltage diff --git a/src/pals/parameters/ReferenceChangeParameters.py b/src/pals/parameters/ReferenceChangeParameters.py index 951c142..b7bcde8 100644 --- a/src/pals/parameters/ReferenceChangeParameters.py +++ b/src/pals/parameters/ReferenceChangeParameters.py @@ -1,12 +1,8 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class ReferenceChangeParameters(BaseModel): """Reference energy change and/or reference time correction parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - - delta_e: float = 0.0 # [eV] Energy change - delta_pc: float = 0.0 # [momentum*c] Momentum change - delta_time: float = 0.0 # [s] Time correction + dE_ref: float = 0.0 # Change in reference energy + extra_dtime_ref: float = 0.0 # Reference time deviation from nominal diff --git a/src/pals/parameters/ReferenceParameters.py b/src/pals/parameters/ReferenceParameters.py index 723a70c..53daeb1 100644 --- a/src/pals/parameters/ReferenceParameters.py +++ b/src/pals/parameters/ReferenceParameters.py @@ -1,12 +1,9 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class ReferenceParameters(BaseModel): """Reference parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - species_ref: str = "" pc_ref: float = 0.0 # [momentum*c] Reference momentum times speed of light E_tot_ref: float = 0.0 # [eV] Reference total energy diff --git a/src/pals/parameters/SolenoidParameters.py b/src/pals/parameters/SolenoidParameters.py index dac5927..2768b32 100644 --- a/src/pals/parameters/SolenoidParameters.py +++ b/src/pals/parameters/SolenoidParameters.py @@ -1,11 +1,8 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class SolenoidParameters(BaseModel): """Solenoid parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - Ksol: float = 0.0 # Normalized solenoid strength Bsol: float = 0.0 # Solenoid field diff --git a/src/pals/parameters/TrackingParameters.py b/src/pals/parameters/TrackingParameters.py index bb06da5..e26139f 100644 --- a/src/pals/parameters/TrackingParameters.py +++ b/src/pals/parameters/TrackingParameters.py @@ -1,10 +1,7 @@ -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel class TrackingParameters(BaseModel): """Tracking parameters""" - # Allow arbitrary fields - model_config = ConfigDict(extra="allow") - # Parameters will be added when construction is complete diff --git a/tests/test_elements.py b/tests/test_elements.py index d37f169..dcc0671 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -1,66 +1,12 @@ from pydantic import ValidationError -from pals.kinds import ( - BaseElement, - ThickElement, - Drift, - Quadrupole, - BeamLine, - CrabCavity, - EGun, - Fork, - Mask, - RBend, - Marker, - ACKicker, - Solenoid, - BeamBeam, - Girder, - Taylor, - FloorShift, - Multipole, - Wiggler, - Octupole, - RFCavity, - Feedback, - BeginningEle, - UnionEle, - Patch, - NullEle, - Foil, - Sextupole, - Instrument, - Match, - SBend, - Fiducial, - Converter, - Kicker, -) - -from pals.parameters import ( - SolenoidParameters, - FloorParameters, - TrackingParameters, - FloorShiftParameters, - BeamBeamParameters, - ApertureParameters, - ElectricMultipoleParameters, - MagneticMultipoleParameters, - ReferenceParameters, - BodyShiftParameters, - ReferenceChangeParameters, - PatchParameters, - MetaParameters, - ForkParameters, - RFParameters, - BendParameters, -) +import pals def test_BaseElement(): # Create one base element with custom name element_name = "base_element" - element = BaseElement(name=element_name) + element = pals.BaseElement(name=element_name) assert element.name == element_name @@ -68,7 +14,7 @@ def test_ThickElement(): # Create one thick element with custom name and length element_name = "thick_element" element_length = 1.0 - element = ThickElement( + element = pals.ThickElement( name=element_name, length=element_length, ) @@ -90,7 +36,7 @@ def test_Drift(): # Create one drift element with custom name and length element_name = "drift_element" element_length = 1.0 - element = Drift( + element = pals.Drift( name=element_name, length=element_length, ) @@ -120,7 +66,7 @@ def test_Quadrupole(): element_magnetic_multipole_Bs2 = 2.2 element_magnetic_multipole_tilt1 = 3.1 element_magnetic_multipole_tilt2 = 3.2 - element_magnetic_multipole = MagneticMultipoleParameters( + element_magnetic_multipole = pals.MagneticMultipoleParameters( Bn1=element_magnetic_multipole_Bn1, Bs1=element_magnetic_multipole_Bs1, tilt1=element_magnetic_multipole_tilt1, @@ -128,7 +74,7 @@ def test_Quadrupole(): Bs2=element_magnetic_multipole_Bs2, tilt2=element_magnetic_multipole_tilt2, ) - element = Quadrupole( + element = pals.Quadrupole( name=element_name, length=element_length, MagneticMultipoleP=element_magnetic_multipole, @@ -148,16 +94,16 @@ def test_Quadrupole(): def test_BeamLine(): # Create first line with one base element - element1 = BaseElement(name="element1") - line1 = BeamLine(name="line1", line=[element1]) + element1 = pals.BaseElement(name="element1") + line1 = pals.BeamLine(name="line1", line=[element1]) assert line1.line == [element1] # Extend first line with one thick element - element2 = ThickElement(name="element2", length=2.0) + element2 = pals.ThickElement(name="element2", length=2.0) line1.line.extend([element2]) assert line1.line == [element1, element2] # Create second line with one drift element - element3 = Drift(name="element3", length=3.0) - line2 = BeamLine(name="line2", line=[element3]) + 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] @@ -165,18 +111,18 @@ def test_BeamLine(): def test_Marker(): """Test Marker element""" - element = Marker(name="marker1") + element = pals.Marker(name="marker1") assert element.name == "marker1" assert element.kind == "Marker" def test_Sextupole(): """Test Sextupole element""" - element = Sextupole( + element = pals.Sextupole( name="sext1", length=0.5, - MagneticMultipoleP=MagneticMultipoleParameters(Bn2=1.0), - ApertureP=ApertureParameters(x_limits=[-0.1, 0.1]), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn2=1.0), + ApertureP=pals.ApertureParameters(x_limits=[-0.1, 0.1]), ) assert element.name == "sext1" assert element.length == 0.5 @@ -187,11 +133,11 @@ def test_Sextupole(): def test_Octupole(): """Test Octupole element""" - element = Octupole( + element = pals.Octupole( name="oct1", length=0.3, - ElectricMultipoleP=ElectricMultipoleParameters(En3=0.5), - MetaP=MetaParameters(alias="octupole_test"), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En3=0.5), + MetaP=pals.MetaParameters(alias="octupole_test"), ) assert element.name == "oct1" assert element.length == 0.3 @@ -202,11 +148,11 @@ def test_Octupole(): def test_Multipole(): """Test Multipole element""" - element = Multipole( + element = pals.Multipole( name="mult1", length=0.4, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=2.0, Bn2=1.5), - BodyShiftP=BodyShiftParameters(x_offset=0.01), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=2.0, Bn2=1.5), + BodyShiftP=pals.BodyShiftParameters(x_offset=0.01), ) assert element.name == "mult1" assert element.length == 0.4 @@ -217,13 +163,13 @@ def test_Multipole(): def test_RBend(): """Test RBend element""" - bend_params = BendParameters(rho_ref=1.0, bend_field_ref=2.0) - element = RBend( + bend_params = pals.BendParameters(rho_ref=1.0, bend_field_ref=2.0) + element = pals.RBend( name="rbend1", length=1.0, BendP=bend_params, - ApertureP=ApertureParameters(x_limits=[-0.2, 0.2]), - MetaP=MetaParameters(description="Test bend"), + ApertureP=pals.ApertureParameters(x_limits=[-0.2, 0.2]), + MetaP=pals.MetaParameters(description="Test bend"), ) assert element.name == "rbend1" assert element.length == 1.0 @@ -235,12 +181,12 @@ def test_RBend(): def test_SBend(): """Test SBend element""" - bend_params = BendParameters(rho_ref=1.5, bend_field_ref=3.0) - element = SBend( + bend_params = pals.BendParameters(rho_ref=1.5, bend_field_ref=3.0) + element = pals.SBend( name="sbend1", length=1.2, BendP=bend_params, - ReferenceP=ReferenceParameters(species_ref="proton"), + ReferenceP=pals.ReferenceParameters(species_ref="proton"), ) assert element.name == "sbend1" assert element.length == 1.2 @@ -251,61 +197,57 @@ def test_SBend(): def test_Solenoid(): """Test Solenoid element""" - sol_params = SolenoidParameters(Ksol=0.1, Bsol=0.2) - element = Solenoid( + sol_params = pals.SolenoidParameters(Ksol=0.1, Bsol=0.2) + element = pals.Solenoid( name="sol1", length=0.8, SolenoidP=sol_params, - TrackingP=TrackingParameters(is_on=True), ) assert element.name == "sol1" assert element.length == 0.8 assert element.kind == "Solenoid" assert element.SolenoidP.Ksol == 0.1 - assert element.TrackingP.is_on def test_RFCavity(): """Test RFCavity element""" - rf_params = RFParameters(frequency=1e9, voltage=1e6) - element = RFCavity( + rf_params = pals.RFParameters(frequency=1e9, voltage=1e6) + element = pals.RFCavity( name="rf1", length=0.5, RFP=rf_params, - SolenoidP=SolenoidParameters(Ksol=0.05), - FloorP=FloorParameters(x_offset=0.1), + SolenoidP=pals.SolenoidParameters(Ksol=0.05), ) assert element.name == "rf1" assert element.length == 0.5 assert element.kind == "RFCavity" assert element.RFP.frequency == 1e9 assert element.SolenoidP.Ksol == 0.05 - assert element.FloorP.x_offset == 0.1 def test_Patch(): """Test Patch element""" - patch_params = PatchParameters(x_offset=0.1, y_offset=0.2) - element = Patch( + patch_params = pals.PatchParameters(x_offset=0.1, y_offset=0.2) + element = pals.Patch( name="patch1", length=0.3, PatchP=patch_params, - ReferenceChangeP=ReferenceChangeParameters(pc_change=1e6), + ReferenceChangeP=pals.ReferenceChangeParameters(dE_ref=1e6), ) assert element.name == "patch1" assert element.length == 0.3 assert element.kind == "Patch" assert element.PatchP.x_offset == 0.1 - assert element.ReferenceChangeP.pc_change == 1e6 + assert element.ReferenceChangeP.dE_ref == 1e6 def test_FloorShift(): """Test FloorShift element""" - floor_params = FloorShiftParameters(x_offset=0.5, z_offset=1.0) - element = FloorShift( + floor_params = pals.FloorShiftParameters(x_offset=0.5, z_offset=1.0) + element = pals.FloorShift( name="floor1", FloorShiftP=floor_params, - MetaP=MetaParameters(alias="floor_test"), + MetaP=pals.MetaParameters(alias="floor_test"), ) assert element.name == "floor1" assert element.kind == "FloorShift" @@ -315,11 +257,11 @@ def test_FloorShift(): def test_Fork(): """Test Fork element""" - fork_params = ForkParameters(to_line="line1", direction="FORWARDS") - element = Fork( + fork_params = pals.ForkParameters(to_line="line1", direction="FORWARDS") + element = pals.Fork( name="fork1", ForkP=fork_params, - ReferenceP=ReferenceParameters(species_ref="electron"), + ReferenceP=pals.ReferenceParameters(species_ref="electron"), ) assert element.name == "fork1" assert element.kind == "Fork" @@ -329,11 +271,11 @@ def test_Fork(): def test_BeamBeam(): """Test BeamBeam element""" - bb_params = BeamBeamParameters() - element = BeamBeam( + bb_params = pals.BeamBeamParameters() + element = pals.BeamBeam( name="bb1", BeamBeamP=bb_params, - ApertureP=ApertureParameters(x_limits=[-0.05, 0.05]), + ApertureP=pals.ApertureParameters(x_limits=[-0.05, 0.05]), ) assert element.name == "bb1" assert element.kind == "BeamBeam" @@ -342,8 +284,8 @@ def test_BeamBeam(): def test_BeginningEle(): """Test BeginningEle element""" - element = BeginningEle( - name="begin1", MetaP=MetaParameters(description="Start of lattice") + element = pals.BeginningEle( + name="begin1", MetaP=pals.MetaParameters(description="Start of lattice") ) assert element.name == "begin1" assert element.kind == "BeginningEle" @@ -352,8 +294,8 @@ def test_BeginningEle(): def test_Fiducial(): """Test Fiducial element""" - element = Fiducial( - name="fid1", ReferenceP=ReferenceParameters(species_ref="proton") + element = pals.Fiducial( + name="fid1", ReferenceP=pals.ReferenceParameters(species_ref="proton") ) assert element.name == "fid1" assert element.kind == "Fiducial" @@ -362,19 +304,18 @@ def test_Fiducial(): def test_NullEle(): """Test NullEle element""" - element = NullEle(name="null1", TrackingP=TrackingParameters(is_on=False)) + element = pals.NullEle(name="null1") assert element.name == "null1" assert element.kind == "NullEle" - assert not element.TrackingP.is_on def test_Kicker(): """Test Kicker element""" - element = Kicker( + element = pals.Kicker( name="kick1", length=0.2, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.5), - ElectricMultipoleP=ElectricMultipoleParameters(En1=0.3), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=0.5), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En1=0.3), ) assert element.name == "kick1" assert element.length == 0.2 @@ -385,7 +326,7 @@ def test_Kicker(): def test_ACKicker(): """Test ACKicker element""" - element = ACKicker(name="ackick1", length=0.15) + element = pals.ACKicker(name="ackick1", length=0.15) assert element.name == "ackick1" assert element.length == 0.15 assert element.kind == "ACKicker" @@ -393,11 +334,11 @@ def test_ACKicker(): def test_CrabCavity(): """Test CrabCavity element""" - element = CrabCavity( + element = pals.CrabCavity( name="crab1", length=0.25, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.8), - ElectricMultipoleP=ElectricMultipoleParameters(En1=0.4), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=0.8), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En1=0.4), ) assert element.name == "crab1" assert element.length == 0.25 @@ -408,11 +349,11 @@ def test_CrabCavity(): def test_EGun(): """Test EGun element""" - element = EGun( + element = pals.EGun( name="egun1", length=0.1, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=1.2), - ElectricMultipoleP=ElectricMultipoleParameters(En1=0.6), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=1.2), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En1=0.6), ) assert element.name == "egun1" assert element.length == 0.1 @@ -423,7 +364,9 @@ def test_EGun(): def test_Feedback(): """Test Feedback element""" - element = Feedback(name="fb1", MetaP=MetaParameters(alias="feedback_test")) + element = pals.Feedback( + name="fb1", MetaP=pals.MetaParameters(alias="feedback_test") + ) assert element.name == "fb1" assert element.kind == "Feedback" assert element.MetaP.alias == "feedback_test" @@ -431,20 +374,18 @@ def test_Feedback(): def test_Girder(): """Test Girder element""" - element = Girder(name="girder1", FloorP=FloorParameters(x_offset=0.1, y_offset=0.2)) + element = pals.Girder(name="girder1") assert element.name == "girder1" assert element.kind == "Girder" - assert element.FloorP.x_offset == 0.1 - assert element.FloorP.y_offset == 0.2 def test_Instrument(): """Test Instrument element""" - element = Instrument( + element = pals.Instrument( name="inst1", length=0.05, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.2), - ElectricMultipoleP=ElectricMultipoleParameters(En1=0.1), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=0.2), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En1=0.1), ) assert element.name == "inst1" assert element.length == 0.05 @@ -455,11 +396,11 @@ def test_Instrument(): def test_Mask(): """Test Mask element""" - element = Mask( + element = pals.Mask( name="mask1", length=0.02, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.15), - ElectricMultipoleP=ElectricMultipoleParameters(En1=0.08), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=0.15), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En1=0.08), ) assert element.name == "mask1" assert element.length == 0.02 @@ -470,8 +411,8 @@ def test_Mask(): def test_Match(): """Test Match element""" - element = Match( - name="match1", BodyShiftP=BodyShiftParameters(x_offset=0.01, y_rot=0.02) + element = pals.Match( + name="match1", BodyShiftP=pals.BodyShiftParameters(x_offset=0.01, y_rot=0.02) ) assert element.name == "match1" assert element.kind == "Match" @@ -481,23 +422,25 @@ def test_Match(): def test_Taylor(): """Test Taylor element""" - element = Taylor( + element = pals.Taylor( name="taylor1", - ReferenceChangeP=ReferenceChangeParameters(pc_change=1e6, t_change=1e-9), + ReferenceChangeP=pals.ReferenceChangeParameters( + dE_ref=1e6, extra_dtime_ref=1e-9 + ), ) assert element.name == "taylor1" assert element.kind == "Taylor" - assert element.ReferenceChangeP.pc_change == 1e6 - assert element.ReferenceChangeP.t_change == 1e-9 + assert element.ReferenceChangeP.dE_ref == 1e6 + assert element.ReferenceChangeP.extra_dtime_ref == 1e-9 def test_Wiggler(): """Test Wiggler element""" - element = Wiggler( + element = pals.Wiggler( name="wig1", length=2.0, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.5), - ElectricMultipoleP=ElectricMultipoleParameters(En1=0.3), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=0.5), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En1=0.3), ) assert element.name == "wig1" assert element.length == 2.0 @@ -508,10 +451,10 @@ def test_Wiggler(): def test_Converter(): """Test Converter element""" - element = Converter( + element = pals.Converter( name="conv1", - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=0.4), - ElectricMultipoleP=ElectricMultipoleParameters(En1=0.2), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=0.4), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En1=0.2), ) assert element.name == "conv1" assert element.kind == "Converter" @@ -521,14 +464,14 @@ def test_Converter(): def test_Foil(): """Test Foil element""" - element = Foil(name="foil1") + element = pals.Foil(name="foil1") assert element.name == "foil1" assert element.kind == "Foil" def test_UnionEle(): """Test UnionEle element""" - element = UnionEle(name="union1", elements=[]) + element = pals.UnionEle(name="union1", elements=[]) assert element.name == "union1" assert element.kind == "UnionEle" assert element.elements == [] diff --git a/tests/test_parameters.py b/tests/test_parameters.py index d0d65bc..63c26bd 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -1,22 +1,21 @@ from pals import MagneticMultipoleParameters -# Import parameter classes -from pals.parameters import ( - SolenoidParameters, - FloorParameters, - TrackingParameters, - FloorShiftParameters, - BeamBeamParameters, +from pals import ( ApertureParameters, - ElectricMultipoleParameters, - ReferenceParameters, + BeamBeamParameters, + BendParameters, BodyShiftParameters, - ReferenceChangeParameters, - PatchParameters, - MetaParameters, + ElectricMultipoleParameters, + # FloorParameters, # not yet tested + FloorShiftParameters, ForkParameters, + MetaParameters, + PatchParameters, + ReferenceChangeParameters, + ReferenceParameters, RFParameters, - BendParameters, + SolenoidParameters, + # TrackingParameters, # not yet tested ) @@ -71,20 +70,16 @@ def test_ParameterClasses(): ref = ReferenceParameters(species_ref="electron", pc_ref=1e6) assert ref.species_ref == "electron" - # Test TrackingParameters - tracking = TrackingParameters(is_on=True) - assert tracking.is_on + # TODO: Test TrackingParameters + # tracking = TrackingParameters(...) + # assert tracking.i.. - # Test FloorParameters - floor_pos = FloorParameters(x_offset=0.1, y_offset=0.2, z_offset=0.3) - assert floor_pos.x_offset == 0.1 - assert floor_pos.y_offset == 0.2 - assert floor_pos.z_offset == 0.3 + # TODO: Test FloorParameters # Test ReferenceChangeParameters - ref_change = ReferenceChangeParameters(pc_change=1e6, t_change=1e-9) - assert ref_change.pc_change == 1e6 - assert ref_change.t_change == 1e-9 + ref_change = ReferenceChangeParameters(extra_dtime_ref=1e6, dE_ref=1e-9) + assert ref_change.extra_dtime_ref == 1e6 + assert ref_change.dE_ref == 1e-9 # Test BeamBeamParameters beambeam = BeamBeamParameters() diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 5ae3a9a..635840c 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -2,65 +2,16 @@ import os import yaml -from pals.kinds import ( - BaseElement, - ThickElement, - BeamLine, - Drift, - Quadrupole, - CrabCavity, - EGun, - Fork, - Mask, - RBend, - Marker, - ACKicker, - Solenoid, - BeamBeam, - Girder, - Taylor, - FloorShift, - Multipole, - Wiggler, - Octupole, - RFCavity, - Feedback, - BeginningEle, - UnionEle, - Patch, - NullEle, - Foil, - Sextupole, - Instrument, - Match, - SBend, - Fiducial, - Converter, - Kicker, -) - -from pals.parameters import ( - SolenoidParameters, - FloorShiftParameters, - BeamBeamParameters, - ApertureParameters, - ElectricMultipoleParameters, - MagneticMultipoleParameters, - PatchParameters, - MetaParameters, - ForkParameters, - RFParameters, - BendParameters, -) +import pals def test_yaml(): # Create one base element - element1 = BaseElement(name="element1") + element1 = pals.BaseElement(name="element1") # Create one thick element - element2 = ThickElement(name="element2", length=2.0) + element2 = pals.ThickElement(name="element2", length=2.0) # Create line with both elements - line = BeamLine(name="line", line=[element1, element2]) + line = pals.BeamLine(name="line", line=[element1, element2]) # Serialize the BeamLine object to YAML yaml_data = yaml.dump(line.model_dump(), default_flow_style=False) print(f"\n{yaml_data}") @@ -72,7 +23,7 @@ def test_yaml(): with open(test_file, "r") as file: yaml_data = yaml.safe_load(file) # Parse the YAML data back into a BeamLine object - loaded_line = BeamLine(**yaml_data) + loaded_line = pals.BeamLine(**yaml_data) # Remove the test file os.remove(test_file) # Validate loaded BeamLine object @@ -81,11 +32,11 @@ def test_yaml(): def test_json(): # Create one base element - element1 = BaseElement(name="element1") + element1 = pals.BaseElement(name="element1") # Create one thick element - element2 = ThickElement(name="element2", length=2.0) + element2 = pals.ThickElement(name="element2", length=2.0) # Create line with both elements - line = BeamLine(name="line", line=[element1, element2]) + line = pals.BeamLine(name="line", line=[element1, element2]) # Serialize the BeamLine object to JSON json_data = json.dumps(line.model_dump(), sort_keys=True, indent=2) print(f"\n{json_data}") @@ -97,7 +48,7 @@ def test_json(): with open(test_file, "r") as file: json_data = json.loads(file.read()) # Parse the JSON data back into a BeamLine object - loaded_line = BeamLine(**json_data) + loaded_line = pals.BeamLine(**json_data) # Remove the test file os.remove(test_file) # Validate loaded BeamLine object @@ -109,128 +60,132 @@ def test_comprehensive_lattice(): # Create elements in alphabetical order for easy maintenance # ACKicker - ackicker = ACKicker(name="ackicker1", length=0.1) + ackicker = pals.ACKicker(name="ackicker1", length=0.1) # BeamBeam - beambeam = BeamBeam(name="beambeam1", BeamBeamP=BeamBeamParameters()) + beambeam = pals.BeamBeam(name="beambeam1", BeamBeamP=pals.BeamBeamParameters()) # BeginningEle - beginning = BeginningEle(name="beginning1") + beginning = pals.BeginningEle(name="beginning1") # Converter - converter = Converter(name="converter1") + converter = pals.Converter(name="converter1") # CrabCavity - crabcavity = CrabCavity(name="crabcavity1", length=0.2) + crabcavity = pals.CrabCavity(name="crabcavity1", length=0.2) # Drift - drift = Drift(name="drift1", length=0.5) + drift = pals.Drift(name="drift1", length=0.5) # EGun - egun = EGun(name="egun1", length=0.15) + egun = pals.EGun(name="egun1", length=0.15) # Feedback - feedback = Feedback(name="feedback1") + feedback = pals.Feedback(name="feedback1") # Fiducial - fiducial = Fiducial(name="fiducial1") + fiducial = pals.Fiducial(name="fiducial1") # FloorShift - floorshift = FloorShift( - name="floorshift1", FloorShiftP=FloorShiftParameters(x_offset=0.1) + floorshift = pals.FloorShift( + name="floorshift1", FloorShiftP=pals.FloorShiftParameters(x_offset=0.1) ) # Foil - foil = Foil(name="foil1") + foil = pals.Foil(name="foil1") # Fork - fork = Fork(name="fork1", ForkP=ForkParameters(to_line="line1")) + fork = pals.Fork(name="fork1", ForkP=pals.ForkParameters(to_line="line1")) # Girder - girder = Girder(name="girder1") + girder = pals.Girder(name="girder1") # Instrument - instrument = Instrument(name="instrument1", length=0.05) + instrument = pals.Instrument(name="instrument1", length=0.05) # Kicker - kicker = Kicker(name="kicker1", length=0.1) + kicker = pals.Kicker(name="kicker1", length=0.1) # Marker - marker = Marker(name="marker1") + marker = pals.Marker(name="marker1") # Mask - mask = Mask(name="mask1", length=0.02) + mask = pals.Mask(name="mask1", length=0.02) # Match - match = Match(name="match1") + match = pals.Match(name="match1") # Multipole - multipole = Multipole(name="multipole1", length=0.3) + multipole = pals.Multipole(name="multipole1", length=0.3) # NullEle - nullele = NullEle(name="nullele1") + nullele = pals.NullEle(name="nullele1") # Octupole - octupole = Octupole( + octupole = pals.Octupole( name="octupole1", length=0.25, - ElectricMultipoleP=ElectricMultipoleParameters(En3=0.5), - MetaP=MetaParameters(alias="octupole_test"), + ElectricMultipoleP=pals.ElectricMultipoleParameters(En3=0.5), + MetaP=pals.MetaParameters(alias="octupole_test"), ) # Patch - patch = Patch(name="patch1", length=0.4, PatchP=PatchParameters(x_offset=0.05)) + patch = pals.Patch( + name="patch1", length=0.4, PatchP=pals.PatchParameters(x_offset=0.05) + ) # Quadrupole - quadrupole = Quadrupole( + quadrupole = pals.Quadrupole( name="quadrupole1", length=0.8, - MagneticMultipoleP=MagneticMultipoleParameters(Bn1=1.0), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn1=1.0), ) # RBend - rbend = RBend( + rbend = pals.RBend( name="rbend1", length=1.0, - BendP=BendParameters(rho_ref=2.0), - ApertureP=ApertureParameters(x_limits=[-0.2, 0.2]), + BendP=pals.BendParameters(rho_ref=2.0), + ApertureP=pals.ApertureParameters(x_limits=[-0.2, 0.2]), ) # RFCavity - rfcavity = RFCavity( + rfcavity = pals.RFCavity( name="rfcavity1", length=0.3, - RFP=RFParameters(frequency=1e9), - SolenoidP=SolenoidParameters(Ksol=0.05), + RFP=pals.RFParameters(frequency=1e9), + SolenoidP=pals.SolenoidParameters(Ksol=0.05), ) # SBend - sbend = SBend(name="sbend1", length=1.2, BendP=BendParameters(rho_ref=1.5)) + sbend = pals.SBend( + name="sbend1", length=1.2, BendP=pals.BendParameters(rho_ref=1.5) + ) # Sextupole - sextupole = Sextupole( + sextupole = pals.Sextupole( name="sextupole1", length=0.2, - MagneticMultipoleP=MagneticMultipoleParameters(Bn2=1.0), - ApertureP=ApertureParameters(x_limits=[-0.1, 0.1]), + MagneticMultipoleP=pals.MagneticMultipoleParameters(Bn2=1.0), + ApertureP=pals.ApertureParameters(x_limits=[-0.1, 0.1]), ) # Solenoid - solenoid = Solenoid( - name="solenoid1", length=0.6, SolenoidP=SolenoidParameters(Ksol=0.1) + solenoid = pals.Solenoid( + name="solenoid1", length=0.6, SolenoidP=pals.SolenoidParameters(Ksol=0.1) ) # Taylor - taylor = Taylor(name="taylor1") + taylor = pals.Taylor(name="taylor1") # UnionEle - unionele = UnionEle(name="unionele1", elements=[]) + unionele = pals.UnionEle(name="unionele1", elements=[]) # Wiggler - wiggler = Wiggler(name="wiggler1", length=2.0) + wiggler = pals.Wiggler(name="wiggler1", length=2.0) # Create comprehensive lattice - lattice = BeamLine( + lattice = pals.BeamLine( name="comprehensive_lattice", line=[ beginning, # Start with beginning element @@ -281,7 +236,7 @@ def test_comprehensive_lattice(): loaded_yaml_data = yaml.safe_load(file) # Deserialize back to Python object using Pydantic model logic - loaded_lattice = BeamLine(**loaded_yaml_data) + loaded_lattice = pals.BeamLine(**loaded_yaml_data) # Verify the loaded lattice has the correct structure and parameter groups assert len(loaded_lattice.line) == 31 # Should have 31 elements @@ -329,7 +284,7 @@ def test_comprehensive_lattice(): loaded_json_data = json.loads(file.read()) # Deserialize back to Python object using Pydantic model logic - loaded_lattice_json = BeamLine(**loaded_json_data) + loaded_lattice_json = pals.BeamLine(**loaded_json_data) # Verify the loaded lattice has the correct structure and parameter groups assert len(loaded_lattice_json.line) == 31 # Should have 31 elements From e5c7ac3843060fcf947c62fbe609af9dca3f8258 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 27 Oct 2025 13:04:46 -0700 Subject: [PATCH 3/6] Clean/Simplify Import Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com> --- tests/test_parameters.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_parameters.py b/tests/test_parameters.py index 63c26bd..486269f 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -1,5 +1,3 @@ -from pals import MagneticMultipoleParameters - from pals import ( ApertureParameters, BeamBeamParameters, @@ -9,6 +7,7 @@ # FloorParameters, # not yet tested FloorShiftParameters, ForkParameters, + MagneticMultipoleParameters, MetaParameters, PatchParameters, ReferenceChangeParameters, From 162d4df87f4949d73af259ea8e08e7e7c9e2a20d Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 27 Oct 2025 13:34:25 -0700 Subject: [PATCH 4/6] More Parameter Validations And more limits on ints and string literal values. --- src/pals/parameters/ApertureParameters.py | 24 ++++++++++----- .../parameters/ElectricMultipoleParameters.py | 2 +- src/pals/parameters/ForkParameters.py | 3 +- src/pals/parameters/PatchParameters.py | 3 +- src/pals/parameters/RFParameters.py | 4 +-- src/pals/parameters/ReferenceParameters.py | 4 +++ tests/test_elements.py | 15 ++-------- tests/test_parameters.py | 29 +++++++++++++++---- 8 files changed, 55 insertions(+), 29 deletions(-) diff --git a/src/pals/parameters/ApertureParameters.py b/src/pals/parameters/ApertureParameters.py index 0c593d7..ff5d956 100644 --- a/src/pals/parameters/ApertureParameters.py +++ b/src/pals/parameters/ApertureParameters.py @@ -1,15 +1,25 @@ -from typing import List -from pydantic import BaseModel +from typing import Literal +from pydantic import BaseModel, Field class ApertureParameters(BaseModel): """Aperture parameters""" - x_limits: List[float] = [float("nan"), float("nan")] - y_limits: List[float] = [float("nan"), float("nan")] - shape: str = "RECTANGULAR" - location: str = "ENTRANCE_END" + x_limits: list[float | None, float | None] = Field( + default=[None, None], + validate_args=lambda x: (x[0] is None or x[1] is None or x[0] < x[1]), + ) + y_limits: list[float | None, float | None] = Field( + default=[None, None], + validate_args=lambda x: (x[0] is None or x[1] is None or x[0] < x[1]), + ) + shape: Literal["RECTANGULAR", "ELLIPTICAL", "VERTICES", "CUSTOM_SHAPE"] = ( + "RECTANGULAR" + ) + location: Literal[ + "ENTRANCE_END", "CENTER", "EXIT_END", "BOTH_ENDS", "NOWHERE", "EVERYWHERE" + ] = "ENTRANCE_END" material: str = "" - thickness: float = 0.0 + thickness: float = Field(default=0.0, ge=0.0) aperture_shifts_with_body: bool = False aperture_active: bool = True diff --git a/src/pals/parameters/ElectricMultipoleParameters.py b/src/pals/parameters/ElectricMultipoleParameters.py index 2df8fb3..52fe74f 100644 --- a/src/pals/parameters/ElectricMultipoleParameters.py +++ b/src/pals/parameters/ElectricMultipoleParameters.py @@ -4,7 +4,7 @@ class ElectricMultipoleParameters(BaseModel): """Electric multipole parameters""" - # Allow arbitrary fields + # Allow arbitrary fields (TODO: remove this) model_config = ConfigDict(extra="allow") # TODO: add ElectricMultipoleParameters in a follow-up RP diff --git a/src/pals/parameters/ForkParameters.py b/src/pals/parameters/ForkParameters.py index 836e01c..0738943 100644 --- a/src/pals/parameters/ForkParameters.py +++ b/src/pals/parameters/ForkParameters.py @@ -1,3 +1,4 @@ +from typing import Literal from pydantic import BaseModel @@ -6,5 +7,5 @@ class ForkParameters(BaseModel): to_line: str = "" to_ele: str = "" - direction: str = "FORWARDS" # "FORWARDS" or "BACKWARDS" + direction: Literal["FORWARDS", "BACKWARDS"] = "FORWARDS" propagate_reference: bool = True diff --git a/src/pals/parameters/PatchParameters.py b/src/pals/parameters/PatchParameters.py index 54c9a9d..c9e946e 100644 --- a/src/pals/parameters/PatchParameters.py +++ b/src/pals/parameters/PatchParameters.py @@ -1,3 +1,4 @@ +from typing import Literal from pydantic import BaseModel @@ -11,5 +12,5 @@ class PatchParameters(BaseModel): y_rot: float = 0.0 z_rot: float = 0.0 flexible: bool = False - ref_coords: str = "exit_end" # "entrance_end" or "exit_end" + ref_coords: Literal["entrance_end", "exit_end"] = "exit_end" user_sets_length: bool = False diff --git a/src/pals/parameters/RFParameters.py b/src/pals/parameters/RFParameters.py index 272f2b5..55c03c6 100644 --- a/src/pals/parameters/RFParameters.py +++ b/src/pals/parameters/RFParameters.py @@ -1,4 +1,4 @@ -from pydantic import BaseModel +from pydantic import BaseModel, Field class RFParameters(BaseModel): @@ -11,4 +11,4 @@ class RFParameters(BaseModel): phase: float = 0.0 # [unitless] RF phase in 0 to 2*pi multipass_phase: float = 0.0 # [unitless] RF Phase added to multipass elements cavity_type: str = "STANDING_WAVE" # [string] Cavity type - n_cell: int = 1 # [unitless] Number of cavity cells + n_cell: int = Field(default=1, gt=0) # [unitless] Number of cavity cells diff --git a/src/pals/parameters/ReferenceParameters.py b/src/pals/parameters/ReferenceParameters.py index 53daeb1..c1e84a3 100644 --- a/src/pals/parameters/ReferenceParameters.py +++ b/src/pals/parameters/ReferenceParameters.py @@ -1,3 +1,4 @@ +from typing import Literal from pydantic import BaseModel @@ -9,3 +10,6 @@ class ReferenceParameters(BaseModel): E_tot_ref: float = 0.0 # [eV] Reference total energy time_ref: float = 0.0 # [s] Reference time location: str = "" # Where reference parameters are evaluated + location: Literal[ + "UPSTREAM_END", "DOWNSTREAM_END" + ] # TODO: undefined default in PALS? diff --git a/tests/test_elements.py b/tests/test_elements.py index dcc0671..0a3c3b4 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -1,3 +1,4 @@ +import pytest from pydantic import ValidationError import pals @@ -23,13 +24,8 @@ def test_ThickElement(): # Try to assign negative length and # detect validation error without breaking pytest element_length = -1.0 - passed = True - try: + with pytest.raises(ValidationError): element.length = element_length - except ValidationError as e: - print(e) - passed = False - assert not passed def test_Drift(): @@ -45,13 +41,8 @@ def test_Drift(): # Try to assign negative length and # detect validation error without breaking pytest element_length = -1.0 - passed = True - try: + with pytest.raises(ValidationError): element.length = element_length - except ValidationError as e: - print(e) - passed = False - assert not passed def test_Quadrupole(): diff --git a/tests/test_parameters.py b/tests/test_parameters.py index 486269f..21eb8f9 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -1,10 +1,11 @@ +import pytest +from pydantic import ValidationError + from pals import ( ApertureParameters, BeamBeamParameters, BendParameters, BodyShiftParameters, - ElectricMultipoleParameters, - # FloorParameters, # not yet tested FloorShiftParameters, ForkParameters, MagneticMultipoleParameters, @@ -24,6 +25,11 @@ def test_ParameterClasses(): aperture = ApertureParameters(x_limits=[-0.1, 0.1], y_limits=[-0.05, 0.05]) assert aperture.x_limits == [-0.1, 0.1] + with pytest.raises(ValidationError): + _ = ApertureParameters( + x_limits=[-0.1, 0.1], y_limits=[-0.05, 0.05, 0.1], shape="wrong" + ) + # Test BodyShiftParameters body_shift = BodyShiftParameters(x_offset=0.01, y_rot=0.02) assert body_shift.x_offset == 0.01 @@ -32,15 +38,23 @@ def test_ParameterClasses(): meta = MetaParameters(alias="test", description="test element") assert meta.alias == "test" - # Test ElectricMultipoleParameters - emp = ElectricMultipoleParameters(En1=1.0, Es1=0.5) - assert emp.En1 == 1.0 + # Test ElectricMultipoleParameters (TODO) + # emp = ElectricMultipoleParameters(En1=1.0, Es1=0.5) + # assert emp.En1 == 1.0 # Test MagneticMultipoleParameters mmp = MagneticMultipoleParameters(Bn1=1.0, Bs1=0.5) assert mmp.Bn1 == 1.0 assert mmp.Bs1 == 0.5 + # catch typos + with pytest.raises(ValidationError): + _ = MagneticMultipoleParameters(Bm1=1.0, Bs1=0.5) + with pytest.raises(ValidationError): + _ = MagneticMultipoleParameters(Bn1=1.0, Bv1=0.5) + with pytest.raises(ValidationError): + _ = MagneticMultipoleParameters(Bn01=1.0, Bs01=0.5) + # Test SolenoidParameters sol = SolenoidParameters(Ksol=0.1, Bsol=0.2) assert sol.Ksol == 0.1 @@ -49,6 +63,11 @@ def test_ParameterClasses(): rf = RFParameters(frequency=1e9, voltage=1e6) assert rf.frequency == 1e9 + with pytest.raises(ValidationError): + _ = RFParameters(frequency=1e9, voltage=1e6, n_cell=0) + with pytest.raises(ValidationError): + _ = RFParameters(frequency=1e9, voltage=1e6, n_cell=-1) + # Test BendParameters bend = BendParameters(rho_ref=1.0, bend_field_ref=2.0) assert bend.rho_ref == 1.0 From 55557e1dca3b6ce8d234806d53f52a55a5eb2a7e Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 27 Oct 2025 13:43:34 -0700 Subject: [PATCH 5/6] Cleanup --- src/pals/kinds/BeamLine.py | 41 +++++++++++------------ src/pals/parameters/ApertureParameters.py | 20 ++++++----- 2 files changed, 31 insertions(+), 30 deletions(-) diff --git a/src/pals/kinds/BeamLine.py b/src/pals/kinds/BeamLine.py index de5aed9..8d87e48 100644 --- a/src/pals/kinds/BeamLine.py +++ b/src/pals/kinds/BeamLine.py @@ -3,39 +3,38 @@ from .BaseElement import BaseElement from .ThickElement import ThickElement -from .Drift import Drift -from .Quadrupole import Quadrupole -# Import schema elements -from .Marker import Marker -from .Sextupole import Sextupole -from .Octupole import Octupole -from .Multipole import Multipole -from .RBend import RBend -from .SBend import SBend -from .Solenoid import Solenoid -from .RFCavity import RFCavity -from .Patch import Patch -from .FloorShift import FloorShift -from .Fork import Fork +from .ACKicker import ACKicker from .BeamBeam import BeamBeam from .BeginningEle import BeginningEle -from .Fiducial import Fiducial -from .NullEle import NullEle -from .Kicker import Kicker -from .ACKicker import ACKicker +from .Converter import Converter from .CrabCavity import CrabCavity +from .Drift import Drift from .EGun import EGun from .Feedback import Feedback +from .Fiducial import Fiducial +from .FloorShift import FloorShift +from .Foil import Foil +from .Fork import Fork from .Girder import Girder from .Instrument import Instrument +from .Kicker import Kicker +from .Marker import Marker from .Mask import Mask from .Match import Match +from .Multipole import Multipole +from .NullEle import NullEle +from .Octupole import Octupole +from .Patch import Patch +from .Quadrupole import Quadrupole +from .RBend import RBend +from .RFCavity import RFCavity +from .SBend import SBend +from .Sextupole import Sextupole +from .Solenoid import Solenoid from .Taylor import Taylor -from .Wiggler import Wiggler -from .Converter import Converter -from .Foil import Foil from .UnionEle import UnionEle +from .Wiggler import Wiggler class BeamLine(BaseElement): diff --git a/src/pals/parameters/ApertureParameters.py b/src/pals/parameters/ApertureParameters.py index ff5d956..f7477ef 100644 --- a/src/pals/parameters/ApertureParameters.py +++ b/src/pals/parameters/ApertureParameters.py @@ -1,18 +1,20 @@ from typing import Literal -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator class ApertureParameters(BaseModel): """Aperture parameters""" - x_limits: list[float | None, float | None] = Field( - default=[None, None], - validate_args=lambda x: (x[0] is None or x[1] is None or x[0] < x[1]), - ) - y_limits: list[float | None, float | None] = Field( - default=[None, None], - validate_args=lambda x: (x[0] is None or x[1] is None or x[0] < x[1]), - ) + @field_validator("x_limits", "y_limits") + @classmethod + def validate_limits(cls, v): + """Validate that limits are None or that min < max""" + if v[0] is not None and v[1] is not None and v[0] >= v[1]: + raise ValueError("Lower limit must be less than upper limit") + return v + + x_limits: list[float | None, float | None] = Field(default=[None, None]) + y_limits: list[float | None, float | None] = Field(default=[None, None]) shape: Literal["RECTANGULAR", "ELLIPTICAL", "VERTICES", "CUSTOM_SHAPE"] = ( "RECTANGULAR" ) From 34afd31530e53200d47b16e2849733fc4448ec26 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 27 Oct 2025 15:47:38 -0700 Subject: [PATCH 6/6] Update Test: w/o Mixin Classes --- tests/test_elements.py | 4 ++-- tests/test_serialization.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_elements.py b/tests/test_elements.py index 0a3c3b4..a115484 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -85,11 +85,11 @@ def test_Quadrupole(): def test_BeamLine(): # Create first line with one base element - element1 = pals.BaseElement(name="element1") + 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.ThickElement(name="element2", length=2.0) + element2 = pals.Drift(name="element2", length=2.0) line1.line.extend([element2]) assert line1.line == [element1, element2] # Create second line with one drift element diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 635840c..b4b7181 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -7,9 +7,9 @@ def test_yaml(): # Create one base element - element1 = pals.BaseElement(name="element1") + element1 = pals.Marker(name="element1") # Create one thick element - element2 = pals.ThickElement(name="element2", length=2.0) + element2 = pals.Drift(name="element2", length=2.0) # Create line with both elements line = pals.BeamLine(name="line", line=[element1, element2]) # Serialize the BeamLine object to YAML @@ -32,9 +32,9 @@ def test_yaml(): def test_json(): # Create one base element - element1 = pals.BaseElement(name="element1") + element1 = pals.Marker(name="element1") # Create one thick element - element2 = pals.ThickElement(name="element2", length=2.0) + element2 = pals.Drift(name="element2", length=2.0) # Create line with both elements line = pals.BeamLine(name="line", line=[element1, element2]) # Serialize the BeamLine object to JSON