Skip to content

Commit b39c90e

Browse files
committed
Merge branch 'main' into run_examples_from_pals
2 parents 10d54f3 + ef707cc commit b39c90e

11 files changed

Lines changed: 285 additions & 19 deletions

File tree

examples/fodo.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from pals import Quadrupole
44
from pals import BeamLine
55
from pals import Lattice
6+
from pals import load, store
67

78

89
def main():
@@ -51,23 +52,23 @@ def main():
5152

5253
# Serialize to YAML
5354
yaml_file = "examples_fodo.pals.yaml"
54-
lattice.to_file(yaml_file)
55+
store(yaml_file, lattice)
5556

5657
# Read YAML data from file
57-
loaded_lattice = Lattice.from_file(yaml_file)
58+
loaded_lattice = load(yaml_file)
5859

5960
# Validate loaded data
60-
assert lattice == loaded_lattice
61+
assert lattice == loaded_lattice.facility[0]
6162

6263
# Serialize to JSON
6364
json_file = "examples_fodo.pals.json"
64-
lattice.to_file(json_file)
65+
store(json_file, lattice)
6566

6667
# Read JSON data from file
67-
loaded_lattice = Lattice.from_file(json_file)
68+
loaded_lattice = load(json_file)
6869

6970
# Validate loaded data
70-
assert lattice == loaded_lattice
71+
assert lattice == loaded_lattice.facility[0]
7172

7273

7374
if __name__ == "__main__":

src/pals/PALS.py

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
from pydantic import BaseModel
2+
3+
from pydantic import model_validator
4+
from typing import List, Optional
5+
6+
from .kinds import Lattice
7+
from .kinds.all_elements import get_all_elements_as_annotation
8+
from .functions import load_file_to_dict, store_dict_to_file
9+
10+
11+
Facility = List[get_all_elements_as_annotation()]
12+
13+
14+
class PALSroot(BaseModel):
15+
"""Represent the roo PALS structure"""
16+
17+
version: Optional[str] = None
18+
19+
facility: Facility
20+
21+
@model_validator(mode="before")
22+
@classmethod
23+
def unpack_json_structure(cls, data):
24+
"""Deserialize the JSON/YAML/...-like dict for facility elements"""
25+
from pals.kinds.mixin.all_element_mixin import unpack_element_list_structure
26+
27+
return unpack_element_list_structure(data, "facility", "facility")
28+
29+
def model_dump(self, *args, **kwargs):
30+
"""Custom model dump for facility to handle element list formatting"""
31+
from pals.kinds.mixin.all_element_mixin import dump_element_list
32+
33+
data = {}
34+
data["PALS"] = {}
35+
data["PALS"]["version"] = self.version
36+
data["PALS"] = dump_element_list(self, "facility", *args, **kwargs)
37+
return data
38+
39+
@staticmethod
40+
def from_file(filename: str) -> "PALSroot":
41+
"""Load a facility from a text file"""
42+
pals_dict = load_file_to_dict(filename)
43+
return PALSroot(**pals_dict)
44+
45+
def to_file(self, filename: str):
46+
"""Save a facility to a text file"""
47+
pals_dict = self.model_dump()
48+
store_dict_to_file(filename, pals_dict)
49+
50+
51+
def load(filename: str) -> PALSroot:
52+
"""Load a facility from a text file"""
53+
pals_dict = load_file_to_dict(filename)
54+
return PALSroot(**pals_dict)
55+
56+
57+
def store(filename: str, pals_root: PALSroot | Facility | Lattice):
58+
# wrap single elements in a list, facility in a PALSroot
59+
if isinstance(pals_root, Lattice):
60+
pals_root = PALSroot(version=None, facility=[pals_root])
61+
elif isinstance(pals_root, list):
62+
pals_root = PALSroot(version=None, facility=pals_root)
63+
64+
pals_dict = pals_root.model_dump()
65+
store_dict_to_file(filename, pals_dict)

src/pals/__init__.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,8 @@
66

77
from .kinds import * # noqa
88
from .parameters import * # noqa
9+
from .PALS import PALSroot, load, store # noqa
10+
11+
12+
# Rebuild pydantic models that depend on other classes
13+
PALSroot.model_rebuild()

src/pals/functions.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,14 @@ def store_dict_to_file(filename: str, pals_dict: dict):
6363
if extension == ".json":
6464
import json
6565

66-
json_data = json.dumps(pals_dict, sort_keys=True, indent=2)
66+
json_data = json.dumps(pals_dict, sort_keys=False, indent=2)
6767
with open(filename, "w") as file:
6868
file.write(json_data)
6969

7070
elif extension == ".yaml":
7171
import yaml
7272

73-
yaml_data = yaml.dump(pals_dict, default_flow_style=False)
73+
yaml_data = yaml.dump(pals_dict, default_flow_style=False, sort_keys=False)
7474
with open(filename, "w") as file:
7575
file.write(yaml_data)
7676

src/pals/kinds/Lattice.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
from pydantic import model_validator, Field
2-
from typing import Annotated, List, Literal, Union
1+
from pydantic import model_validator
2+
from typing import List, Literal, Union
33

44
from .BeamLine import BeamLine
5+
from .PlaceholderName import PlaceholderName
56
from .mixin import BaseElement
67
from ..functions import load_file_to_dict, store_dict_to_file
78

@@ -11,7 +12,7 @@ class Lattice(BaseElement):
1112

1213
kind: Literal["Lattice"] = "Lattice"
1314

14-
branches: List[Annotated[Union[BeamLine], Field(discriminator="kind")]]
15+
branches: List[Union[BeamLine, PlaceholderName]]
1516

1617
@model_validator(mode="before")
1718
@classmethod

src/pals/kinds/PlaceholderName.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
from pydantic import BaseModel, Field, model_serializer
2+
from typing import Annotated
3+
4+
from .mixin import BaseElement
5+
6+
7+
class PlaceholderName(BaseModel):
8+
"""Represents a reference to a named element.
9+
10+
This placeholder is replaced with a physically distinct
11+
element during beamline expansion.
12+
13+
This class behaves like a string (via __str__ and __eq__) but stores
14+
a true reference to the actual element object once it's resolved.
15+
16+
The element field holds a reference (not a copy) to the actual element.
17+
18+
Attributes:
19+
name: The name of the referenced element
20+
element: A reference to the resolved element object (None until resolved)
21+
22+
Example:
23+
>>> ref = PlaceholderName(name="drift1")
24+
>>> ref.name
25+
'drift1'
26+
>>> str(ref)
27+
'drift1'
28+
>>> ref == "drift1"
29+
True
30+
>>> ref.element # None until resolved
31+
>>> drift = pals.Drift(name="drift1", length=1.0)
32+
>>> ref.element = drift
33+
>>> ref.is_resolved()
34+
True
35+
>>> ref.element is drift # True - it's a reference, not a copy
36+
True
37+
"""
38+
39+
name: str = Field(..., description="The name of the referenced element")
40+
element: Annotated[
41+
"BaseElement | None",
42+
Field(default=None, description="Reference to the resolved element object"),
43+
] = None
44+
45+
@model_serializer(mode="plain")
46+
def _serialize_as_name(self) -> str:
47+
"""Serialize this reference as just its name.
48+
49+
This makes `model_dump()` return a string (the element name), so nested
50+
serialization (e.g. inside BeamLine.line) produces plain strings too.
51+
"""
52+
return self.name
53+
54+
def __init__(self, name: str | None = None, /, **data):
55+
"""Initialize with either positional name or keyword arguments."""
56+
if name is not None:
57+
super().__init__(name=name, **data)
58+
else:
59+
super().__init__(**data)
60+
61+
def __str__(self) -> str:
62+
"""Return the element name as string."""
63+
return self.name
64+
65+
def __eq__(self, other: object) -> bool:
66+
"""Enable string comparison."""
67+
if isinstance(other, str):
68+
return self.name == other
69+
if isinstance(other, PlaceholderName):
70+
return self.name == other.name and self.element is other.element
71+
return False
72+
73+
def __hash__(self) -> int:
74+
"""Make hashable like a string."""
75+
return hash(self.name)
76+
77+
def is_resolved(self) -> bool:
78+
"""Check if this reference has been resolved to an actual element."""
79+
return self.element is not None
80+
81+
def __repr__(self) -> str:
82+
"""Return a representation of the PlaceholderName."""
83+
resolved = "resolved" if self.is_resolved() else "unresolved"
84+
return f"PlaceholderName('{self.name}', {resolved})"

src/pals/kinds/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from .NullEle import NullEle # noqa: F401
2727
from .Octupole import Octupole # noqa: F401
2828
from .Patch import Patch # noqa: F401
29+
from .PlaceholderName import PlaceholderName # noqa: F401
2930
from .Quadrupole import Quadrupole # noqa: F401
3031
from .RBend import RBend # noqa: F401
3132
from .RFCavity import RFCavity # noqa: F401

src/pals/kinds/all_elements.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,8 @@
44
avoiding duplication between BeamLine.line and UnionEle.elements.
55
"""
66

7-
from typing import Annotated, Union
7+
from typing import Union
88

9-
from pydantic import Field
109

1110
from .ACKicker import ACKicker
1211
from .BeamBeam import BeamBeam
@@ -15,6 +14,7 @@
1514
from .CrabCavity import CrabCavity
1615
from .Drift import Drift
1716
from .EGun import EGun
17+
from .PlaceholderName import PlaceholderName
1818
from .Feedback import Feedback
1919
from .Fiducial import Fiducial
2020
from .FloorShift import FloorShift
@@ -83,6 +83,13 @@ def get_all_element_types(extra_types: tuple = None):
8383

8484

8585
def get_all_elements_as_annotation(extra_types: tuple = None):
86-
"""Return the Union type of all allowed elements with their name as the discriminator field."""
87-
types = get_all_element_types(extra_types)
88-
return Annotated[Union[types], Field(discriminator="kind")]
86+
"""Return the Union type of all allowed elements with their kind as the discriminator field.
87+
88+
Note: PlaceholderName is included to support string references to named elements.
89+
Since PlaceholderName doesn't have a 'kind' field, we cannot use discriminator.
90+
Pydantic will still properly validate the union by trying each type in order in
91+
our unpack_element_list_structure method.
92+
"""
93+
types = get_all_element_types(extra_types) + (PlaceholderName,)
94+
# We can't use discriminator with PlaceholderName in the union since it has no 'kind' field
95+
return Union[types]

src/pals/kinds/mixin/all_element_mixin.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
"""
66

77
from . import BaseElement
8+
from ..PlaceholderName import PlaceholderName
89

910

1011
def unpack_element_list_structure(
@@ -44,7 +45,14 @@ def unpack_element_list_structure(
4445
for item in data[field_name]:
4546
# An element can be a string that refers to another element
4647
if isinstance(item, str):
47-
raise RuntimeError("Reference/alias elements not yet implemented")
48+
# Wrap the string in a Placeholder name object
49+
new_list.append(PlaceholderName(item))
50+
continue
51+
# An element can be a PlaceholderName instance directly
52+
elif isinstance(item, PlaceholderName):
53+
# Keep the PlaceholderName as-is
54+
new_list.append(item)
55+
continue
4856
# An element can be a dict
4957
elif isinstance(item, dict):
5058
if not (len(item) == 1):
@@ -69,7 +77,7 @@ def unpack_element_list_structure(
6977
continue
7078

7179
raise TypeError(
72-
f"Value must be a reference string or a dict, but we got {item!r}"
80+
f"Value must be a reference string, PlaceholderName, or a dict, but we got {item!r}"
7381
)
7482

7583
data[field_name] = new_list
@@ -102,5 +110,8 @@ def dump_element_list(self, field_name: str, *args, **kwargs) -> dict:
102110
elem_dict = elem.model_dump(**kwargs)
103111
new_list.append(elem_dict)
104112

105-
data[self.name][field_name] = new_list
113+
if hasattr(self, "name"): # all but PALSroot have a name
114+
data[self.name][field_name] = new_list
115+
else:
116+
data[field_name] = new_list
106117
return data

src/pals/schema_version.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
from typing import Optional
2+
3+
# PALS schema version - null for now, will be set when version scheme is finalized
4+
PALS_SCHEMA_VERSION: Optional[str] = None

0 commit comments

Comments
 (0)