pals-python/examples/fodo.py at d2411edeff36313b166ab81ce4cb9dfd37ca23bc · EZoni/pals-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
from pals import MagneticMultipoleParameters
from pals import Drift
from pals import Quadrupole
from pals import BeamLine
from pals import Lattice


def main():
    drift1 = Drift(
        name="drift1",
        length=0.25,
    )
    quad1 = Quadrupole(
        name="quad1",
        length=1.0,
        MagneticMultipoleP=MagneticMultipoleParameters(
            Bn1=1.0,
        ),
    )
    drift2 = Drift(
        name="drift2",
        length=0.5,
    )
    quad2 = Quadrupole(
        name="quad2",
        length=1.0,
        MagneticMultipoleP=MagneticMultipoleParameters(
            Bn1=-1.0,
        ),
    )
    drift3 = Drift(
        name="drift3",
        length=0.5,
    )
    # Create line with all elements
    line = BeamLine(
        name="fodo_cell",
        line=[
            drift1,
            quad1,
            drift2,
            quad2,
            drift3,
        ],
    )
    # Create lattice with the line as a branch
    lattice = Lattice(
        name="fodo_lattice",
        branches=[line],
    )

    # Serialize to YAML
    yaml_file = "examples_fodo.pals.yaml"
    lattice.to_file(yaml_file)

    # Read YAML data from file
    loaded_lattice = Lattice.from_file(yaml_file)

    # Validate loaded data
    assert lattice == loaded_lattice

    # Serialize to JSON
    json_file = "examples_fodo.pals.json"
    lattice.to_file(json_file)

    # Read JSON data from file
    loaded_lattice = Lattice.from_file(json_file)

    # Validate loaded data
    assert lattice == loaded_lattice


if __name__ == "__main__":
    main()