Start FODO example

Author: EZoni

.github/workflows/unit_tests.yml

@@ -27,3 +27,19 @@ jobs:
       - name: Test
         run: |
           pytest tests -v
+  examples:
+    name: examples
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+      - name: Test
+        run: |
+          python examples/fodo.py

examples/fodo.py

@@ -0,0 +1,84 @@
+import json
+import os
+import sys
+import yaml
+
+# Add the parent directory to sys.path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from schema.MagneticMultipoleParameters import MagneticMultipoleParameters
+
+from schema.DriftElement import DriftElement
+from schema.QuadrupoleElement import QuadrupoleElement
+
+from schema.Item import Item
+from schema.Line import Line
+
+
+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=[
+            Item(item=drift1),
+            Item(item=quad1),
+            Item(item=drift2),
+            Item(item=quad2),
+            Item(item=drift3),
+        ]
+    )
+    # Serialize to YAML
+    yaml_data = yaml.dump(line.model_dump(), default_flow_style=False)
+    # Write YAML data to file
+    yaml_file = "examples_fodo.yaml"
+    with open(yaml_file, "w") as file:
+        file.write(yaml_data)
+    # Read YAML data from file
+    with open(yaml_file, "r") as file:
+        yaml_data = yaml.safe_load(file)
+    # Parse YAML data
+    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)
+    # Write JSON data to file
+    json_file = "examples_fodo.json"
+    with open(json_file, "w") as file:
+        file.write(json_data)
+    # Read JSON data from file
+    with open(json_file, "r") as file:
+        json_data = json.loads(file.read())
+    # Parse JSON data
+    loaded_line = Line(**json_data)
+    # Validate loaded data
+    assert line == loaded_line
+
+
+if __name__ == "__main__":
+    main()

schema/DriftElement.py

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

schema/Item.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel, ConfigDict, Field
-from typing import Annotated, Literal, Union
+from typing import Annotated, Union
 
 from schema.BaseElement import BaseElement
 from schema.ThickElement import ThickElement
@@ -10,9 +10,6 @@
 class Item(BaseModel):
     """An element of a line or a line itself"""
 
-    # Discriminator field
-    kind: Literal["Item"] = "Item"
-
     # Validate every time a new value is assigned to an attribute,
     # not only when an instance of Line is created
     model_config = ConfigDict(validate_assignment=True)

schema/Line.py

@@ -1,5 +1,5 @@
-from pydantic import BaseModel, ConfigDict, Field
-from typing import Annotated, List, Literal
+from pydantic import BaseModel, ConfigDict
+from typing import List
 
 
 from schema.Item import Item
@@ -8,14 +8,11 @@
 class Line(BaseModel):
     """A line of elements and/or other lines"""
 
-    # Discriminator field
-    kind: Literal["Line"] = "Line"
-
     # Validate every time a new value is assigned to an attribute,
     # not only when an instance of Line is created
     model_config = ConfigDict(validate_assignment=True)
 
-    line: List[Annotated[Item, Field(discriminator="kind")]]
+    line: List[Item]
 
 
 # Avoid circular import issues

schema/QuadrupoleElement.py

@@ -8,7 +8,7 @@ class QuadrupoleElement(ThickElement):
     """A quadrupole element"""
 
     # Discriminator field
-    kind: Literal["QuadrupoleElement"] = "QuadrupoleElement"
+    kind: Literal["Quadrupole"] = "Quadrupole"
 
     # Magnetic multipole parameters
     MagneticMultipoleP: MagneticMultipoleParameters