Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 11 additions & 14 deletions examples/fodo.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,43 +16,39 @@

def main():
drift1 = DriftElement(
name="drift1",
length=0.25,
)
quad1 = QuadrupoleElement(
name="quad1",
length=1.0,
MagneticMultipoleP=MagneticMultipoleParameters(
Bn1=1.0,
),
)
drift2 = DriftElement(
name="drift2",
length=0.5,
)
quad2 = QuadrupoleElement(
name="quad2",
length=1.0,
MagneticMultipoleP=MagneticMultipoleParameters(
Bn1=-1.0,
),
)
drift3 = DriftElement(
name="drift3",
length=0.5,
)
# Create line with all elements
line = Line(
line=[
drift1,
quad1,
drift2,
quad2,
drift3,
]
line={
"drift1": drift1,
"quad1": quad1,
"drift2": drift2,
"quad2": quad2,
"drift3": drift3,
}
)

# Serialize to YAML
yaml_data = yaml.dump(line.model_dump(), default_flow_style=False)
yaml_data = yaml.dump(line.model_dump(), default_flow_style=False, sort_keys=False)
print("Dumping YAML data...")
print(f"{yaml_data}")
# Write YAML data to file
Expand All @@ -66,8 +62,9 @@ def main():
loaded_line = Line(**yaml_data)
# Validate loaded data
assert line == loaded_line

# Serialize to JSON
json_data = json.dumps(line.model_dump(), sort_keys=True, indent=2)
json_data = json.dumps(line.model_dump(), indent=2)
print("Dumping JSON data...")
print(f"{json_data}")
# Write JSON data to file
Expand Down
5 changes: 1 addition & 4 deletions schema/BaseElement.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from pydantic import BaseModel, ConfigDict
from typing import Literal, Optional
from typing import Literal


class BaseElement(BaseModel):
Expand All @@ -11,6 +11,3 @@ class BaseElement(BaseModel):
# Validate every time a new value is assigned to an attribute,
# not only when an instance of BaseElement is created
model_config = ConfigDict(validate_assignment=True)

# Unique element name
name: Optional[str] = None
7 changes: 4 additions & 3 deletions schema/Line.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from pydantic import BaseModel, ConfigDict, Field
from typing import Annotated, List, Literal, Union
from typing import Annotated, Literal, Union, OrderedDict

from schema.BaseElement import BaseElement
from schema.ThickElement import ThickElement
Expand All @@ -16,7 +16,8 @@ class Line(BaseModel):

kind: Literal["Line"] = "Line"

line: List[
line: OrderedDict[
Copy link
Member

@ax3l ax3l Jul 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I love the idea of the PR, but I just noticed an issue with the data structure:

I think it does not work to make Line an ordered dict, or a dictionary at all.

The problem with dictionaries is that their elements are either unsorted (generally random, i.e. independent of name and insertion order because they are usually implemented as hash tables, where a key is hashed to determine which "bucket" an element belongs in) or sorted, by (key) name in our case, which is not what we want for a beamline/segment: we need a datastructure that strictly keeps the insertions order (across programming languages and formats). That format is a list.

As a specialty of Python, since 3.7 Python dicts preserve the insertation order, but that is not very portable to assume for other programming languages (or even YAML libraries).

Copy link
Member

@ax3l ax3l Jul 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But what works is this:

  • Line is a list
  • Element is a dict (with only one key) of dict (the actual element properties)
- thingA:
    kind: Sextupole
    length: 3
    
- thingB:
    kind: Quad
    length: 1

- thingC:
    length: 2

that then looks in Python like this:

[{'thingA': {'kind': 'Sextupole', 'length': 3}}, {'thingB': {'kind': 'Quad', 'length': 1}}, {'thingC': {'length': 2}}]

and in a Python API it stays as now:

Sextupole(name="thingA", length=3)

Copy link
Member

@ax3l ax3l Jul 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just for completeness, one cannot really do:

- thingA:
  kind: Sextupole
  length: 3

- thingB:
  kind: Quad
  length: 1

- thingC:
  length: 2

which represents

[{'thingA': None, 'kind': 'Sextupole', 'length': 3}, {'thingB': None, 'kind': 'Quad', 'length': 1}, {'thingC': None, 'length': 2}]

because then the random key that is the name is unclear how to pick out (again: order is not guaranteed, it might not be first)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ax3l

In your comment above, #12 (comment), did you have in mind that line would be a list of dict, so something like

    line: List[
        Dict[
            str,
            Annotated[
                Union[
                    BaseElement,
                    ThickElement,
                    DriftElement,
                    QuadrupoleElement,
                    "Line",
                ],
                Field(discriminator="kind"),
            ],
        ]
    ]

or did you have in mind that we should implement some sort of new WrapperElement class (also derived from BaseModel, like all other element classes) that has a single attribute, a dictionary, which is the initialized to store one single "key: value" pair, where the value is the original dictionary of properties obtained by a model_dump from the a given element?

I have experimented a bit with the latter but it does not seem very easy to me. It also seems to require custom serialization and deserialization, which makes me think that we could do just that (custom serialization and deserialization, as done in #8) without the layer of the dictionary wrapper.

Copy link
Member

@ax3l ax3l Jul 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, line as a list of dict :)

Copy link
Member

@ax3l ax3l Jul 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(with references to other pre-defined elements, it will then be a list of dict|str, but we can add references and inheritance of elements in a follow-up. The dict part or its attributes we could still derive from a base element.)

str,
Annotated[
Union[
BaseElement,
Expand All @@ -26,7 +27,7 @@ class Line(BaseModel):
"Line",
],
Field(discriminator="kind"),
]
],
]


Expand Down
51 changes: 22 additions & 29 deletions tests/test_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,21 +15,16 @@


def test_BaseElement():
# Create one base element with custom name
element_name = "base_element"
element = BaseElement(name=element_name)
assert element.name == element_name
# nothing to test here
pass


def test_ThickElement():
# Create one thick element with custom name and length
element_name = "thick_element"
element_length = 1.0
element = ThickElement(
name=element_name,
length=element_length,
)
assert element.name == element_name
assert element.length == element_length
# Try to assign negative length and
# detect validation error without breaking pytest
Expand All @@ -45,13 +40,10 @@ def test_ThickElement():

def test_DriftElement():
# Create one drift element with custom name and length
element_name = "drift_element"
element_length = 1.0
element = DriftElement(
name=element_name,
length=element_length,
)
assert element.name == element_name
assert element.length == element_length
# Try to assign negative length and
# detect validation error without breaking pytest
Expand All @@ -67,7 +59,6 @@ def test_DriftElement():

def test_QuadrupoleElement():
# Create one drift element with custom name and length
element_name = "quadrupole_element"
element_length = 1.0
element_magnetic_multipole_Bn1 = 1.1
element_magnetic_multipole_Bn2 = 1.2
Expand All @@ -84,11 +75,9 @@ def test_QuadrupoleElement():
tilt2=element_magnetic_multipole_tilt2,
)
element = QuadrupoleElement(
name=element_name,
length=element_length,
MagneticMultipoleP=element_magnetic_multipole,
)
assert element.name == element_name
assert element.length == element_length
assert element.MagneticMultipoleP.Bn1 == element_magnetic_multipole_Bn1
assert element.MagneticMultipoleP.Bs1 == element_magnetic_multipole_Bs1
Expand All @@ -103,28 +92,32 @@ def test_QuadrupoleElement():

def test_Line():
# Create first line with one base element
element1 = BaseElement(name="element1")
line1 = Line(line=[element1])
assert line1.line == [element1]
element1 = BaseElement()
line1 = Line(line={"element1": element1})
assert line1.line == {"element1": element1}
# Extend first line with one thick element
element2 = ThickElement(name="element2", length=2.0)
line1.line.extend([element2])
assert line1.line == [element1, element2]
element2 = ThickElement(length=2.0)
line1.line.update({"element2": element2})
assert line1.line == {"element1": element1, "element2": element2}
# Create second line with one drift element
element3 = DriftElement(name="element3", length=3.0)
line2 = Line(line=[element3])
element3 = DriftElement(length=3.0)
line2 = Line(line={"element3": element3})
# Extend first line with second line
line1.line.extend(line2.line)
assert line1.line == [element1, element2, element3]
line1.line.update(line2.line)
assert line1.line == {
"element1": element1,
"element2": element2,
"element3": element3,
}


def test_yaml():
# Create one base element
element1 = BaseElement(name="element1")
element1 = BaseElement()
# Create one thick element
element2 = ThickElement(name="element2", length=2.0)
element2 = ThickElement(length=2.0)
# Create line with both elements
line = Line(line=[element1, element2])
line = Line(line={"element1": element1, "element2": element2})
# Serialize the Line object to YAML
yaml_data = yaml.dump(line.model_dump(), default_flow_style=False)
print(f"\n{yaml_data}")
Expand All @@ -145,11 +138,11 @@ def test_yaml():

def test_json():
# Create one base element
element1 = BaseElement(name="element1")
element1 = BaseElement()
# Create one thick element
element2 = ThickElement(name="element2", length=2.0)
element2 = ThickElement(length=2.0)
# Create line with both elements
line = Line(line=[element1, element2])
line = Line(line={"element1": element1, "element2": element2})
# Serialize the Line object to JSON
json_data = json.dumps(line.model_dump(), sort_keys=True, indent=2)
print(f"\n{json_data}")
Expand Down