diff --git a/src/pals/kinds/mixin/BaseElement.py b/src/pals/kinds/mixin/BaseElement.py index 4f07653..72970d4 100644 --- a/src/pals/kinds/mixin/BaseElement.py +++ b/src/pals/kinds/mixin/BaseElement.py @@ -6,6 +6,7 @@ BodyShiftParameters, FloorParameters, MetaParameters, + PlacementParameters, ReferenceParameters, ReferenceChangeParameters, TrackingParameters, @@ -22,6 +23,7 @@ class BaseElement(BaseModel, validate_assignment=True): name: str # Common parameter groups (optional for all elements) + placement: PlacementParameters | None = None ApertureP: ApertureParameters | None = None BodyShiftP: BodyShiftParameters | None = None FloorP: FloorParameters | None = None diff --git a/src/pals/parameters/PlacementParameters.py b/src/pals/parameters/PlacementParameters.py new file mode 100644 index 0000000..fd4ca0f --- /dev/null +++ b/src/pals/parameters/PlacementParameters.py @@ -0,0 +1,28 @@ +from typing import Literal, Optional, TypeAlias + +from pydantic import BaseModel, Field + + +ReferencePoint: TypeAlias = Literal["ENTRANCE_END", "CENTER", "EXIT_END", "ZERO_POINT"] + + +class PlacementParameters(BaseModel): + offset: float = Field( + default=0.0, + description="Real [m]. Longitudinal offset of the line item. Default is zero.", + allow_inf_nan=False, + ) + + to_point: ReferencePoint = Field( + default="ENTRANCE_END", + description="Line item offset end point. Default is ENTRANCE_END.", + ) + + base_item: Optional[str] = Field( + default=None, description="Line item containing the `from_point`." + ) + + from_point: ReferencePoint = Field( + default="EXIT_END", + description="Base line item offset beginning point. Default is EXIT_END.", + ) diff --git a/src/pals/parameters/__init__.py b/src/pals/parameters/__init__.py index 6334a7c..9b6ae22 100644 --- a/src/pals/parameters/__init__.py +++ b/src/pals/parameters/__init__.py @@ -13,6 +13,7 @@ from .MagneticMultipoleParameters import MagneticMultipoleParameters # noqa: F401 from .MetaParameters import MetaParameters # noqa: F401 from .PatchParameters import PatchParameters # noqa: F401 +from .PlacementParameters import PlacementParameters # noqa: F401 from .ReferenceChangeParameters import ReferenceChangeParameters # noqa: F401 from .ReferenceParameters import ReferenceParameters # noqa: F401 from .RFParameters import RFParameters # noqa: F401 diff --git a/tests/test_elements.py b/tests/test_elements.py index 5a49665..8ba1e8e 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -624,3 +624,71 @@ def test_PlaceholderName_direct(): # Modify the original and verify the reference sees the change drift.length = 3.0 assert ref1.element.length == 3.0 # Change is visible through reference + + +class TestElementsWithPlacementParameters: + @pytest.fixture(autouse=True) + def _create_line(self): + import yaml + + data = yaml.safe_load(""" + test_line: + kind: BeamLine + line: + - marker1: + kind: Marker + - quad1: + kind: Quadrupole + length: 1.0 + MagneticMultipoleP: + Bn1: 1.0 + placement: + offset: 1.0 + base_item: marker1 + to_point: CENTER + from_point: ENTRANCE_END + - marker2: + kind: Marker + placement: + offset: 3.0 + base_item: marker1 + to_point: CENTER + from_point: ENTRANCE_END + """) + + self._beamline = pals.BeamLine(**data) + + def test_valid_beamline(self): + assert isinstance(self._beamline, pals.BeamLine) + assert len(self._beamline.line) == 3 + + def test_marker_constructed_without_placement_parameters(self): + marker = self._beamline.line[0] + + assert isinstance(marker, pals.Marker) + assert marker.name == "marker1" + assert marker.placement is None + + def test_quadrupole_constructed_with_placement_parameters(self): + quad = self._beamline.line[1] + + assert isinstance(quad, pals.Quadrupole) + assert quad.name == "quad1" + assert quad.length == 1.0 + assert quad.MagneticMultipoleP.Bn1 == 1.0 + assert quad.placement is not None + assert quad.placement.offset == 1.0 + assert quad.placement.base_item == "marker1" + assert quad.placement.to_point == "CENTER" + assert quad.placement.from_point == "ENTRANCE_END" + + def test_marker_constructed_with_placement_parameters(self): + marker = self._beamline.line[2] + + assert isinstance(marker, pals.Marker) + assert marker.name == "marker2" + assert marker.placement is not None + assert marker.placement.offset == 3.0 + assert marker.placement.base_item == "marker1" + assert marker.placement.to_point == "CENTER" + assert marker.placement.from_point == "ENTRANCE_END" diff --git a/tests/test_parameters.py b/tests/test_parameters.py index 09416bc..dcf2399 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -1,3 +1,5 @@ +import math + import pytest from pydantic import ValidationError @@ -12,6 +14,7 @@ MagneticMultipoleParameters, MetaParameters, PatchParameters, + PlacementParameters, ReferenceChangeParameters, ReferenceParameters, RFParameters, @@ -134,3 +137,48 @@ def test_ParameterClasses(): # Test BeamBeamParameters beambeam = BeamBeamParameters() assert beambeam is not None + + +class TestPlacementParameters: + def test_default_construction(self): + placement = PlacementParameters() + + assert placement.offset == 0.0 + assert placement.to_point == "ENTRANCE_END" + assert placement.base_item is None + assert placement.from_point == "EXIT_END" + + def test_fully_specified_construction(self): + placement = PlacementParameters( + offset=10.0, to_point="CENTER", base_item="q1", from_point="CENTER" + ) + + assert placement.offset == 10.0 + assert placement.to_point == "CENTER" + assert placement.base_item == "q1" + assert placement.from_point == "CENTER" + + def test_construction_with_all_reference_points(self): + from itertools import combinations + + ref_pts = ("ENTRANCE_END", "CENTER", "EXIT_END", "ZERO_POINT") + + for p1, p2 in combinations(ref_pts, 2): + PlacementParameters(to_point=p1, from_point=p2) + + def test_construction_with_invalid_reference_points(self): + with pytest.raises(ValidationError): + PlacementParameters(to_point="INVALID") + + with pytest.raises(ValidationError): + PlacementParameters(from_point="INVALID") + + def test_construction_raises_error_with_nonnumeric_offset(self): + with pytest.raises(ValidationError): + PlacementParameters(offset="error") + + with pytest.raises(ValidationError): + PlacementParameters(offset=math.nan) + + with pytest.raises(ValidationError): + PlacementParameters(offset=math.inf)