Repository: slaclab/lume-genesis

Issue:
BUG: Archiving is not saving particles
Using the `genesi4_example.ipynb` example:

<img width="582" alt="image" src="https://github.com/user-attachments/assets/8a694ab3-4264-4e2d-bcda-1efb8b2ca7a2">


Discussion:

Patch:
diff --git a/genesis/version4/input/lattice.py b/genesis/version4/input/lattice.py
index d207d87..e5d853a 100644
--- a/genesis/version4/input/lattice.py
+++ b/genesis/version4/input/lattice.py
@@ -708,11 +708,7 @@ class Lattice(BaseModel):
         *,
         count: int = 1,
         start: int = 0,
-        prefix: str = "",
-        suffix: str = "",
-        delimiter: str = "_",
-        line_format: str = "{name}{index}",
-        last_format: str = "",
+        format: str = "L{index}_{name}",
         in_place: bool = True,
         new_beamline: str = "",
     ):
@@ -720,19 +716,11 @@ class Lattice(BaseModel):
         Flatten a beamline such that any nested beamlines are made into distinct elements.
 
         For example, take a beamline ``LN`` which contains 2 nested beamlines
-        "LA" (containing UND) and "LB" (containing QUAD).
-
-        Flattening this with the default line format of ``{name}{index}`` and
-        the default label format of the same ``{name}{index}`` would result
-        with replacing ``LN`` as follows.
-
-        Using start=0, count=1:
-
-        * ``LN0_LA0_UND0`` and ``LN0_LB0_QUAD0``
-
-        Using start=1, count=2:
-
-        * ``LN1_LA1_UND1``, ``LN1_LB1_QUAD1``, ``LN2_LA1_UND1``, ``LN2_LB1_QUAD1``
+        "L1" (containing UND) and "L2" (containing QUAD).
+        Flattening this with the default format of ``L{index}_{name}``
+        would result with replacing ``LN`` with:
+        * ``L0_UND`` and ``L0_QUAD`` (wtih start=0, count=1)
+        * ``L1_UND``, ``L1_QUAD``, ``L2_UND``, ``L2_QUAD`` (with start=1, count=2)
 
         Parameters
         ----------
@@ -742,19 +730,10 @@ class Lattice(BaseModel):
             The number of times to duplicate the flattened beamline.
         start : int, default=0
             The starting index for relabeling elements.
-        prefix : str, default=""
-            A prefix to add to the renamed elements.
-        suffix : str, default=""
-            A suffix to add to the renamed elements.
-        delimiter : str, default="_"
-            The delimiter to use between the name and index in the new element names.
-        line_format : str, default="{name}{index}"
+        format : str, default="L{index}_{name}"
             The string format to use for renaming beamline elements.
-        last_format : str, default=line_format
-            An optional format string to use for the last element in each segment.
-            Defaults to using the 'line_format'.
         in_place : bool, default=True
-            Perform the flattening in-place. Remove unused elements and add
+            Perform the flattening in-place.  Remove unused elements and add
             the new elements directly to the lattice.
         new_beamline : str
             For in-place mode, use this as the name for the new beamline.
@@ -766,66 +745,34 @@ class Lattice(BaseModel):
             from ``start`` to ``start + count - 1``. Each segment is a list of
             created elements based on that index.
         """
-        self.fix_labels()
-        last_format = last_format or line_format
-        line_stack = []
-
-        def _inspect(
-            name: str,
-            seen: Dict[str, int],
-        ) -> Generator[Tuple[str, AnyBeamlineElement], None, None]:
-            if name in line_stack:
-                raise RecursiveLineError(
-                    f"Recursion of beamline elements in line {beamline} detected: {name}"
-                )
-
-            line_stack.append(name)
-
-            element = self.elements[name]
-            seen_count = seen.get(name, start - 1) + 1
-            seen[name] = seen_count
-
-            format = line_format if isinstance(element, Line) else last_format
-
-            line_prefix = format.format(
-                name=name,
-                index=seen_count,
-                beamline=beamline,
-                delimiter=delimiter,
-            )
-
-            if isinstance(element, Line):
-                seen = {}
-                for _, el in element.flattened:
-                    for inner, sub_ele in _inspect(el, seen=seen):
-                        yield f"{line_prefix}{delimiter}{inner}", sub_ele
-            else:
-                yield line_prefix, element
-            line_stack.pop(-1)
-
         to_duplicate = self.inspect_line_by_name(beamline)
         new_elements = {}
         sections = []
-        seen: Dict[str, int] = {}
-        for _ in range(start, start + count):
+        for index in range(start, start + count):
             section = []
-            for new_name, elem in list(_inspect(beamline, seen=seen)):
-                new_name = f"{prefix}{new_name}{suffix}"
+            for name in to_duplicate:
+                new_name = format.format(
+                    name=name,
+                    index=index,
+                    start=start,
+                    count=count,
+                    beamline=beamline,
+                )
+                new_element = copy.deepcopy(self.elements[name])
+                new_element.label = new_name
+
                 if new_name in new_elements:
                     raise ValueError(
-                        f"Corrected name for element conflicts with another new element: {new_name}"
+                        f"Name for element conflicts with another new element: {new_name}"
                     )
-
-                new_element = copy.deepcopy(elem)
-                new_element.label = new_name
                 new_elements[new_name] = new_element
                 section.append(new_element)
             sections.append(section)
 
         if in_place:
-            for label in to_duplicate:
-                if label not in new_elements:
-                    self.elements.pop(label, None)
+            for name in to_duplicate:
+                if name not in new_elements:
+                    self.elements.pop(name, None)
             if new_beamline != beamline and beamline not in new_elements:
                 self.elements.pop(beamline, None)
 
diff --git a/genesis/version4/output.py b/genesis/version4/output.py
index 4ec96db..9d9544d 100644
--- a/genesis/version4/output.py
+++ b/genesis/version4/output.py
@@ -1383,7 +1383,6 @@ class Genesis4Output(Mapping, BaseModel, arbitrary_types_allowed=True):
     )
     particles: Dict[FileKey, ParticleGroup] = pydantic.Field(
         default_factory=dict,
-        exclude=True,
         description="Loaded particle data, keyed by integration step number or filename base.",
     )
     alias: Dict[str, str] = pydantic.Field(


Test patch:
diff --git a/genesis/tests/genesis4/test_archive.py b/genesis/tests/genesis4/test_archive.py
index 06aa0ff..c57a822 100644
--- a/genesis/tests/genesis4/test_archive.py
+++ b/genesis/tests/genesis4/test_archive.py
@@ -162,6 +162,38 @@ def test_hdf_archive(
     assert json_for_comparison(orig_output) == json_for_comparison(genesis4.output)
 
 
+def test_hdf_archive_particles(
+    genesis4: Genesis4,
+    hdf5_filename: pathlib.Path,
+) -> None:
+    genesis4.input.main.namelists.append(Write(beam="end"))
+    orig_output = genesis4.run(raise_on_error=True)
+    assert orig_output.run.success
+
+    orig_output.load_particles()
+
+    orig_particles = orig_output.particles
+    assert len(orig_output.particles)
+
+    t0 = time.monotonic()
+    genesis4.archive(hdf5_filename)
+
+    t1 = time.monotonic()
+    genesis4.load_archive(hdf5_filename)
+    new_output = genesis4.output
+
+    t2 = time.monotonic()
+    print("Took", t1 - t0, "s to archive")
+    print("Took", t2 - t1, "s to restore")
+    assert new_output is not None
+    assert list(new_output.particles) == list(orig_particles)
+
+    for key in new_output.particles:
+        particles = orig_particles[key]
+        print("Checking particles", particles)
+        assert particles == new_output.particles[key]
+
+
 def json_for_comparison(model: BaseModel) -> str:
     # Assuming dictionary keys can't be assumed to be sorted
     data = json.loads(model.model_dump_json())
diff --git a/genesis/tests/genesis4/test_taper.py b/genesis/tests/genesis4/test_taper.py
index 1539143..efece9b 100644
--- a/genesis/tests/genesis4/test_taper.py
+++ b/genesis/tests/genesis4/test_taper.py
@@ -4,7 +4,7 @@ import pytest
 
 from ..conftest import test_artifacts
 from ...errors import NotFlatError, RecursiveLineError
-from ...version4 import Drift, DuplicatedLineItem, Lattice, Undulator, Line, Quadrupole
+from ...version4 import Lattice, Undulator, Line, Quadrupole
 
 
 @pytest.fixture
@@ -247,165 +247,52 @@ def test_plot(
     plt.savefig(test_artifacts / f"{request.node.name}.png")
 
 
-def test_flatten_doc_example() -> None:
-    UND = Undulator(label="UND")
-    UNDA = Undulator(label="UNDA")
-    QUAD = Quadrupole(label="QUAD")
-    LA = Line(elements=["UND"], label="LA")
-    LB = Line(elements=["QUAD"], label="LB")
-    LN = Line(elements=["LA", "LB"], label="LN")
-
-    lattice = Lattice.from_elements([LA, LB, LN, UND, QUAD, UNDA])
-    flattened = lattice.flatten_line(
-        "LN",
-        start=0,
-        count=1,
-        in_place=False,
-    )
-    assert flattened == [
-        [
-            Undulator(label="LN0_LA0_UND0"),
-            Quadrupole(label="LN0_LB0_QUAD0"),
-        ]
-    ]
-
-    flattened = lattice.flatten_line(
-        "LN",
-        start=1,
-        count=2,
-        in_place=False,
-    )
-    assert flattened == [
-        [
-            Undulator(label="LN1_LA1_UND1"),
-            Quadrupole(label="LN1_LB1_QUAD1"),
-        ],
-        [
-            Undulator(label="LN2_LA1_UND1"),
-            Quadrupole(label="LN2_LB1_QUAD1"),
-        ],
-    ]
-
-
 def test_flatten_example() -> None:
     UND = Undulator(label="UND")
-    UNDA = Undulator(label="UNDA")
+    UND1 = Undulator(label="UND1")
     QUAD = Quadrupole(label="QUAD")
-    LA = Line(elements=["UND"], label="LA")
-    LB = Line(elements=["QUAD"], label="LB")
-    LN = Line(elements=["LA", "LB", "UNDA"], label="LN")
-    lattice = Lattice.from_elements([LA, LB, LN, UND, QUAD, UNDA])
-    flattened = lattice.flatten_line(
-        "LN",
-        count=1,
-        start=1,
-        in_place=False,
-        # line_format="{name}_{index}",
-        # last_format="{name}_{index}",
-    )
+    L1 = Line(elements=["UND"], label="L1")
+    L2 = Line(elements=["QUAD"], label="L2")
+    LN = Line(elements=["L1", "L2", "UND1"], label="LN")
+    lattice = Lattice.from_elements([L1, L2, LN, UND, QUAD, UND1])
+    flattened = lattice.flatten_line("LN", count=1, in_place=False)
     assert flattened == [
         [
-            Undulator(label="LN1_LA1_UND1"),
-            Quadrupole(label="LN1_LB1_QUAD1"),
-            Undulator(label="LN1_UNDA1"),
+            Undulator(label="L0_UND"),
+            Quadrupole(label="L0_QUAD"),
+            Undulator(label="L0_UND1"),
         ]
     ]
 
     flattened = lattice.flatten_line("LN", start=1, count=2, in_place=False)
     assert flattened == [
         [
-            Undulator(label="LN1_LA1_UND1"),
-            Quadrupole(label="LN1_LB1_QUAD1"),
-            Undulator(label="LN1_UNDA1"),
+            Undulator(label="L1_UND"),
+            Quadrupole(label="L1_QUAD"),
+            Undulator(label="L1_UND1"),
         ],
         [
-            Undulator(label="LN2_LA1_UND1"),
-            Quadrupole(label="LN2_LB1_QUAD1"),
-            Undulator(label="LN2_UNDA1"),
+            Undulator(label="L2_UND"),
+            Quadrupole(label="L2_QUAD"),
+            Undulator(label="L2_UND1"),
         ],
     ]
 
-    lattice = Lattice.from_elements([LA, LB, LN, UND, QUAD, UNDA])
+    lattice = Lattice.from_elements([L1, L2, LN, UND, QUAD, UND1])
     lattice.flatten_line("LN", count=1, in_place=True)
     assert lattice.elements["LN"].elements == [
-        "LN0_LA0_UND0",
-        "LN0_LB0_QUAD0",
-        "LN0_UNDA0",
+        "L0_UND",
+        "L0_QUAD",
+        "L0_UND1",
     ]
 
-    lattice = Lattice.from_elements([LA, LB, LN, UND, QUAD, UNDA])
+    lattice = Lattice.from_elements([L1, L2, LN, UND, QUAD, UND1])
     lattice.flatten_line("LN", start=1, count=2, in_place=True)
     assert lattice.elements["LN"].elements == [
-        "LN1_LA1_UND1",
-        "LN1_LB1_QUAD1",
-        "LN1_UNDA1",
-        "LN2_LA1_UND1",
-        "LN2_LB1_QUAD1",
-        "LN2_UNDA1",
+        "L1_UND",
+        "L1_QUAD",
+        "L1_UND1",
+        "L2_UND",
+        "L2_QUAD",
+        "L2_UND1",
     ]
-
-
-def test_flatten_example1():
-    fodo_count = 6
-
-    lat0 = Lattice(
-        elements={
-            "D1": Drift(label="D1", L=0.44),
-            "D2": Drift(label="D2", L=0.24),
-            "QF": Quadrupole(label="QF", L=0.08, k1=2.0),
-            "QD": Quadrupole(label="QD", L=0.08, k1=-2.0),
-            "UND": Undulator(
-                label="UND", aw=0.84853, lambdau=0.015, nwig=266, helical=True
-            ),
-            "FODO": Line(
-                label="FODO",
-                elements=["UND", "D1", "QF", "D2", "UND", "D1", "QD", "D2"],
-            ),
-            "FEL": Line(
-                label="FEL",
-                elements=[DuplicatedLineItem(label="FODO", count=fodo_count)],
-            ),
-        },
-    )
-
-    flattened = lat0.flatten_line(
-        "FEL",
-        count=1,
-        start=1,
-        in_place=False,
-        line_format="{name}{index}",
-        last_format="{name}_{index}",
-    )
-    flattened_names = [ele.label for section in flattened for ele in section]
-    for ele in lat0.elements["FODO"].elements:
-        assert f"FEL1_FODO1_{ele}_1" in flattened_names
-
-    assert "FEL1_FODO1_D1_1" in flattened_names
-    assert "FEL2_FODO1_UND_1" not in flattened_names
-
-    assert len(flattened_names) == fodo_count * len(lat0.elements["FODO"].elements)
-
-    print("Flattened:")
-    for name in flattened_names:
-        print(name)
-
-    flattened = lat0.flatten_line(
-        "FEL",
-        count=2,
-        start=1,
-        in_place=False,
-        line_format="{name}{index}",
-        last_format="{name}_{index}",
-    )
-    flattened_names = [ele.label for section in flattened for ele in section]
-    assert "FEL1_FODO1_UND_1" in flattened_names
-    assert "FEL2_FODO1_UND_1" in flattened_names
-    assert "FEL2_FODO1_UND_2" in flattened_names
-    assert "FEL2_FODO2_UND_1" in flattened_names
-
-    for ele in lat0.elements["FODO"].elements:
-        for idx in range(1, 6):
-            assert f"FEL1_FODO{idx}_{ele}_1" in flattened_names
-            assert f"FEL2_FODO{idx}_{ele}_1" in flattened_names
-
-    assert len(flattened_names) == fodo_count * len(lat0.elements["FODO"].elements) * 2