From 15afa74cd67de125869048d73779fe4e726e06eb Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Thu, 6 Aug 2026 17:49:32 -0400 Subject: [PATCH] Add index_start and index_stop to envelope tracker --- py/orbit/envelope/track.py | 15 +++++++++++---- tests/py/orbit/test_env.py | 20 ++++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/py/orbit/envelope/track.py b/py/orbit/envelope/track.py index d17f908a..f3776b6d 100644 --- a/py/orbit/envelope/track.py +++ b/py/orbit/envelope/track.py @@ -3,6 +3,7 @@ from orbit.core.bunch import Bunch from orbit.core.bunch import SyncParticle + from orbit.lattice import AccNode from orbit.lattice import AccLattice from orbit.teapot import BendTEAPOT @@ -47,12 +48,15 @@ def __init__(self, lattice: AccLattice, sc: str | None = None) -> None: node.setParam("ea1", 0.0) node.setParam("ea2", 0.0) - def track(self, envelope: Envelope) -> None: + def track(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> None: """Track envelope through lattice. This is not recursive, so grandchild nodes are not tracked. """ - for node_index, node in enumerate(self.lattice.getNodes()): + nodes = self.lattice.getNodes() + nodes = nodes[index_start : index_stop] + + for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): matrix = get_matrix(child_node, envelope=envelope) if matrix is not None: @@ -91,7 +95,7 @@ def track(self, envelope: Envelope) -> None: if matrix is not None: envelope.transform(matrix) - def track_history(self, envelope: Envelope) -> dict[str, list]: + def track_history(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> dict[str, list]: """Track and return envelope parameters vs. position in lattice.""" history_keys = [ "s", @@ -129,7 +133,10 @@ def update_history(envelope: Envelope, position: float) -> None: path_length = 0.0 update_history(envelope, path_length) - for node_index, node in enumerate(self.lattice.getNodes()): + nodes = self.lattice.getNodes() + nodes = nodes[index_start : index_stop] + + for node_index, node in enumerate(nodes): for child_node in node.getChildNodes(ENTRANCE): matrix = get_matrix(child_node, envelope=envelope) if matrix is not None: diff --git a/tests/py/orbit/test_env.py b/tests/py/orbit/test_env.py index 16f1e5c7..660dbe5a 100644 --- a/tests/py/orbit/test_env.py +++ b/tests/py/orbit/test_env.py @@ -381,3 +381,23 @@ def test_sc_3d_cold_expansion(): # twice the initial size. (See examples from A. Shishlo or # from the ImpactX repo.) pass + + +def test_track_sublattice_no_error(): + bunch = Bunch() + bunch.mass(mass_proton) + bunch.getSyncParticle().kinEnergy(0.001) + + cov_matrix = np.diag(np.square([1e-3, 0, 1e-3, 0.0, 1e-3, 0.0])) + envelope = Envelope(bunch, cov_matrix=cov_matrix) + + lattice = TEAPOT_Lattice() + + n = 5 + for _ in range(n): + lattice.addNode(DriftTEAPOT(length=0.1)) + + tracker = EnvelopeTracker(lattice) + for i in range(n): + tracker.track(envelope, index_start=i) + tracker.track(envelope, index_stop=-i)