[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.

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

@@ -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"),
         ]

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,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

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

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"

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

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"