Comprehensive Element Update (#33)

Update to include all currently defined elements and parameter classes.

- [x] rebase out preparatory PRs #34 #36
- [x] vibe-code draft & iterations 
- [x] careful author review of
  - [x] elements
  - [x] parameters
  - [x] tests

---------

Co-authored-by: Edoardo Zoni <[email protected]>

Author: ax3l

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

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:

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

src/pals/kinds/BeamLine.py

@@ -3,8 +3,38 @@
 
 from .BaseElement import BaseElement
 from .ThickElement import ThickElement
+
+from .ACKicker import ACKicker
+from .BeamBeam import BeamBeam
+from .BeginningEle import BeginningEle
+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 .UnionEle import UnionEle
+from .Wiggler import Wiggler
 
 
 class BeamLine(BaseElement):
@@ -19,11 +49,42 @@ class BeamLine(BaseElement):
     line: List[
         Annotated[
             Union[
+                # Base classes (for testing compatibility)
                 BaseElement,
                 ThickElement,
+                # User-Facing element kinds
+                "BeamLine",
+                ACKicker,
+                BeamBeam,
+                BeginningEle,
+                Converter,
+                CrabCavity,
                 Drift,
+                EGun,
+                Feedback,
+                Fiducial,
+                FloorShift,
+                Foil,
+                Fork,
+                Girder,
+                Instrument,
+                Kicker,
+                Marker,
+                Mask,
+                Match,
+                Multipole,
+                NullEle,
+                Octupole,
+                Patch,
                 Quadrupole,
-                "BeamLine",
+                RBend,
+                RFCavity,
+                SBend,
+                Sextupole,
+                Solenoid,
+                Taylor,
+                UnionEle,
+                Wiggler,
             ],
             Field(discriminator="kind"),
         ]

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"

src/pals/kinds/Converter.py

@@ -0,0 +1,17 @@
+from typing import Literal, Optional
+
+from .BaseElement import BaseElement
+from ..parameters import ElectricMultipoleParameters, 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
+    ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None
+    MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None

src/pals/kinds/CrabCavity.py

@@ -0,0 +1,17 @@
+from typing import Literal, Optional
+
+from .ThickElement import ThickElement
+from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters
+from ._warnings import under_construction
+
+
+@under_construction("CrabCavity")
+class CrabCavity(ThickElement):
+    """RF crab cavity"""
+
+    # Discriminator field
+    kind: Literal["CrabCavity"] = "CrabCavity"
+
+    # CrabCavity-specific parameters
+    ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None
+    MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None

src/pals/kinds/Drift.py

@@ -4,7 +4,7 @@
 
 
 class Drift(ThickElement):
-    """A field free region"""
+    """Field free region"""
 
     # Discriminator field
     kind: Literal["Drift"] = "Drift"

src/pals/kinds/EGun.py

@@ -0,0 +1,17 @@
+from typing import Literal, Optional
+
+from .ThickElement import ThickElement
+from ..parameters import ElectricMultipoleParameters, MagneticMultipoleParameters
+from ._warnings import under_construction
+
+
+@under_construction("EGun")
+class EGun(ThickElement):
+    """Electron gun"""
+
+    # Discriminator field
+    kind: Literal["EGun"] = "EGun"
+
+    # EGun-specific parameters
+    ElectricMultipoleP: Optional[ElectricMultipoleParameters] = None
+    MagneticMultipoleP: Optional[MagneticMultipoleParameters] = None

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"