Repository: desy-ml/cheetah
Pull request #436: Convenience function for computing properties along lattice + track `s`
Author: login: jank324
Labels: name: enhancement
description: New feature or request
Created: 2025-05-02T15:52:39Z | Merged: 2025-05-08T11:51:25Z

<!--- Provide a general summary of your changes in the Title above -->

## Description

<!--- Describe your changes in detail -->

Implements a convenience function `Segment.longitudinal_beam_generator` that generates `Beam` objects along the beamline for easy computation of the longitudinal parameter trajectories. Plot functions have been refactored to make use of this function. To enable this change, `Beam` subclasses now track their `s` position along the beamline.

## Motivation and Context

<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here. -->
<!--- You can use the syntax `closes #100` if this solves the issue #100 -->

- [ ] I have raised an issue to propose this change (required for new features and bug fixes)

I needed this functionality for a project, and we have been talking about wanting it for probably years at this point.

## Types of changes

<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation (update in the documentation)

## Checklist

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->

<!-- - [ ] I've read the [CONTRIBUTION](https://github.com/DLR-RM/stable-baselines3/blob/master/CONTRIBUTING.md) guide (**required**) -->

- [x] I have updated the changelog accordingly (**required**).
- [ ] My change requires a change to the documentation.
- [x] I have updated the tests accordingly (_required for a bug fix or a new feature_).
- [ ] I have updated the documentation accordingly.
  <!-- - [ ] I have opened an associated PR on the [SB3-Contrib repository](https://github.com/Stable-Baselines-Team/stable-baselines3-contrib) (if necessary) -->
  <!-- - [ ] I have opened an associated PR on the [RL-Zoo3 repository](https://github.com/DLR-RM/rl-baselines3-zoo) (if necessary) -->
- [x] I have reformatted the code and checked that formatting passes (**required**).
- [x] I have have fixed all issues found by `flake8` (**required**).
- [x] I have ensured that all `pytest` tests pass (**required**).
- [ ] I have run `pytest` on a machine with a CUDA GPU and made sure all tests pass (**required**).
- [x] I have checked that the documentation builds (**required**).

Note: We are using a maximum length of 88 characters per line.

<!--- This Template is an edited version of the one from https://github.com/DLR-RM/stable-baselines3/ -->


Linked issue #100: Batched execution
It would be in the spirit of the Cheetah's applications, to enable batched execution. This would be, for example, that a quadrupole's strength is set to vector of strengths, which would then result in a 3D tensor of transfer maps and allow parallel evaluations at much higher speed.
Linked issue comments: user: login: cr-xu
created_at: 2023-11-14T15:07:23Z
body: This is a tricky problem, which will be very helpful to have.

There are still two cases to be considered, e.g. batch evaluation of `incoming_beam` and batch evaluation of `segment_setting`

The `incoming_beam` situation is easier to consider and implement: we have a fixed `Segment` and `transfer_map`, we would just need to broadcast the transfer map multiplication properly with the batched/stacked incoming beam (which is already halfway done in the `ParticleBeam` case anyway).

The changing `segment_setting` case is more tricky, but also **more important**. Ultimately, we're modifying the network (transfer map) during the forward pass. One concern is whether we'll be able to keep the execution speed (somewhat) fast enough.

For bookkeeping purposes, the general feature I would like is the ability to perform fast calculation of outgoing beam parameters for one static incoming beam, and vectorized accelerator settings with the general form `[..., n_batch, n_dim]` where `n_dim` being the number of accelerator settings to be changed. 

Currently, I have to do something like: 
```python
def forward(self, X: torch.Tensor) -> torch.Tensor:
    input_shape = X.shape
    X = X.reshape(-1, n_dim)
    Y = torch.zeros(X.shape[:-1])
    for i, x in enumerate(X):
        self.segment.Q1.k1 = x[0] # Set the input parameters 
        ...
        # Track the beam
        out_beam = self.segment(self.incoming_beam)
        # Compute the objective
        obj = ...
        Y[i] = obj
    return Y.reshape(input_shape[:-1])
```

which more or less scale linearly (poorly):
```
i 1.0
CPU times: user 1.55 ms, sys: 1.06 ms, total: 2.61 ms
Wall time: 1.85 ms
i 10.0
CPU times: user 4.6 ms, sys: 14 µs, total: 4.61 ms
Wall time: 4.65 ms
i 100.0
CPU times: user 38.6 ms, sys: 266 µs, total: 38.9 ms
Wall time: 39.4 ms
i 1000.0
CPU times: user 335 ms, sys: 1.22 ms, total: 336 ms
Wall time: 347 ms
i 10000.0
CPU times: user 3.23 s, sys: 15.6 ms, total: 3.25 s
Wall time: 3.26 s
```
author_association: MEMBER
user: login: jank324
created_at: 2023-11-17T12:05:57Z
body: Maybe we can take inspiration from how this was done here: https://github.com/UM-ARM-Lab/pytorch_kinematics
author_association: MEMBER

Diff:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f0ad6e9a3..8e3f14396 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,9 @@
 - Add `KQUAD` and `CSRCSBEND` element names to Elegant converter (see #409) (@amylizzle)
 - Add `Sextupole` to Bmad, Elegant, and Ocelot converters (see #430) (@Hespe)
 - Implement convenience method for quickly setting attributes for all elements of a type in a `Segment` (see #431) (@jank324)
+- Add a method to `ParticleBeam` that lets you subsample a particle beam with fewer particles and the same distribution (see #432) (@jank324)
+- `Segment` now has new functions `beam_along_segment_generator` and `get_beam_attrs_along_segment` for easily retrieving beam objects and their properties. The plot functions have been refactored to use these, and two functions `plot_beam_attrs` and `plot_beam_attrs_over_lattice` were added for straightforward plotting of different beam attributes in a single line of code. (see #436) (@jank324)
+- `Beam` subclasses now track their `s` position along the beamline (see #436) (@jank324)
 
 ### 🐛 Bug fixes
 
@@ -18,6 +21,7 @@
 
 - Bmad is no longer actively run in the test workflows, and comparisons to Bmad are now done on static pre-computed results from Bmad. This also removes the use of Anaconda in the test workflow. (see #429, #431) (@jank324)
 - The PyTorch pin to `<=2.6` was removed, as the issue with `abort trap: 6` was caused by Bmad is no longer actively used in the test workflow (see #429, #431) (@jank324)
+- Temporary pin `snowballstemmer<3.0` for the docs build because of an issue with the latest release. Will be unpinned when yanked or updated. Refer to https://github.com/sphinx-doc/sphinx/issues/13533 and https://github.com/snowballstem/snowball/issues/229. (see #436) (@jank324)
 
 ### 🌟 First Time Contributors
 
@@ -33,7 +37,6 @@
 - Implement `split` method for the `Solenoid` element (see #380) (@cr-xu)
 - Implement a more robust RPN parser, fixing a bug where short strings in an Elegant variable definition would cause parsing to fail. (see #387, #417) (@amylizzle, @Hespe, @jank324)
 - Add a `Sextupole` element (see #406) (@jank324, @Hespe)
-- Add a method to `ParticleBeam` that lets you subsample a particle beam with fewer particles and the same distribution (see #432) (@jank324)
 
 ### 🐛 Bug fixes
 
diff --git a/cheetah/accelerator/aperture.py b/cheetah/accelerator/aperture.py
index 37bead1a7..adda2423f 100644
--- a/cheetah/accelerator/aperture.py
+++ b/cheetah/accelerator/aperture.py
@@ -100,6 +100,7 @@ def track(self, incoming: Beam) -> Beam:
             energy=incoming.energy,
             particle_charges=incoming.particle_charges,
             survival_probabilities=incoming.survival_probabilities * survived_mask,
+            s=incoming.s,
             species=incoming.species.clone(),
         )
 
@@ -107,7 +108,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         # TODO: Implement splitting for aperture properly, for now just return self
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
 
         alpha = 1 if self.is_active else 0.2
diff --git a/cheetah/accelerator/bpm.py b/cheetah/accelerator/bpm.py
index 4b9a3bc9b..5f1a9ca20 100644
--- a/cheetah/accelerator/bpm.py
+++ b/cheetah/accelerator/bpm.py
@@ -57,7 +57,11 @@ def track(self, incoming: Beam) -> Beam:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
 
         alpha = 1 if self.is_active else 0.2
diff --git a/cheetah/accelerator/cavity.py b/cheetah/accelerator/cavity.py
index 6e34893c5..7554f5837 100644
--- a/cheetah/accelerator/cavity.py
+++ b/cheetah/accelerator/cavity.py
@@ -230,6 +230,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
                 cov=outgoing_cov,
                 energy=outgoing_energy,
                 total_charge=incoming.total_charge,
+                s=incoming.s + self.length,
                 device=outgoing_mu.device,
                 dtype=outgoing_mu.dtype,
             )
@@ -240,6 +241,7 @@ def _track_beam(self, incoming: Beam) -> Beam:
                 energy=outgoing_energy,
                 particle_charges=incoming.particle_charges,
                 survival_probabilities=incoming.survival_probabilities,
+                s=incoming.s + self.length,
                 device=outgoing_particles.device,
                 dtype=outgoing_particles.dtype,
             )
@@ -341,7 +343,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         # element itself
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
 
diff --git a/cheetah/accelerator/custom_transfer_map.py b/cheetah/accelerator/custom_transfer_map.py
index a861e273c..e17b89e7e 100644
--- a/cheetah/accelerator/custom_transfer_map.py
+++ b/cheetah/accelerator/custom_transfer_map.py
@@ -102,7 +102,11 @@ def defining_features(self) -> list[str]:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
 
diff --git a/cheetah/accelerator/dipole.py b/cheetah/accelerator/dipole.py
index ffe8bcf49..1e6967ef6 100644
--- a/cheetah/accelerator/dipole.py
+++ b/cheetah/accelerator/dipole.py
@@ -248,6 +248,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             energy=ref_energy,
             particle_charges=incoming.particle_charges,
             survival_probabilities=incoming.survival_probabilities,
+            s=incoming.s + self.length,
             species=incoming.species,
         )
         return outgoing_beam
@@ -498,7 +499,11 @@ def defining_features(self) -> list[str]:
             "tracking_method",
         ]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
         plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
diff --git a/cheetah/accelerator/drift.py b/cheetah/accelerator/drift.py
index 5bf701402..70da858e2 100644
--- a/cheetah/accelerator/drift.py
+++ b/cheetah/accelerator/drift.py
@@ -119,6 +119,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             energy=ref_energy,
             particle_charges=incoming.particle_charges,
             survival_probabilities=incoming.survival_probabilities,
+            s=incoming.s + self.length,
             species=incoming.species,
         )
         return outgoing_beam
@@ -139,8 +140,12 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for i in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
-        pass
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+        # This does nothing on purpose, because drift sections are visualised as gaps.
+        return ax
 
     @property
     def defining_features(self) -> list[str]:
diff --git a/cheetah/accelerator/element.py b/cheetah/accelerator/element.py
index e90a094a8..37577b923 100644
--- a/cheetah/accelerator/element.py
+++ b/cheetah/accelerator/element.py
@@ -68,23 +68,27 @@ def track(self, incoming: Beam) -> Beam:
         """
         if isinstance(incoming, ParameterBeam):
             tm = self.transfer_map(incoming.energy, incoming.species)
-            mu = (tm @ incoming.mu.unsqueeze(-1)).squeeze(-1)
-            cov = tm @ incoming.cov @ tm.transpose(-2, -1)
+            new_mu = (tm @ incoming.mu.unsqueeze(-1)).squeeze(-1)
+            new_cov = tm @ incoming.cov @ tm.transpose(-2, -1)
+            new_s = incoming.s + self.length
             return ParameterBeam(
-                mu,
-                cov,
+                new_mu,
+                new_cov,
                 incoming.energy,
                 total_charge=incoming.total_charge,
+                s=new_s,
                 species=incoming.species.clone(),
             )
         elif isinstance(incoming, ParticleBeam):
             tm = self.transfer_map(incoming.energy, incoming.species)
             new_particles = incoming.particles @ tm.transpose(-2, -1)
+            new_s = incoming.s + self.length
             return ParticleBeam(
                 new_particles,
                 incoming.energy,
                 particle_charges=incoming.particle_charges,
                 survival_probabilities=incoming.survival_probabilities,
+                s=new_s,
                 species=incoming.species.clone(),
             )
         else:
@@ -159,16 +163,18 @@ def split(self, resolution: torch.Tensor) -> list["Element"]:
         raise NotImplementedError
 
     @abstractmethod
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
         """
         Plot a representation of this element into a `matplotlib` Axes at position `s`.
 
-        :param ax: Axes to plot the representation into.
         :param s: Position of the object along s in meters.
         :param vector_idx: Index of the vector dimension to plot. If the model has more
             than one vector dimension, this can be used to select a specific one. In the
             case of present vector dimension but no index provided, the first one is
             used by default.
+        :param ax: Axes to plot the representation into.
         """
         raise NotImplementedError
 
diff --git a/cheetah/accelerator/horizontal_corrector.py b/cheetah/accelerator/horizontal_corrector.py
index 28666736d..7781f290c 100644
--- a/cheetah/accelerator/horizontal_corrector.py
+++ b/cheetah/accelerator/horizontal_corrector.py
@@ -81,7 +81,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for i in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
         plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
diff --git a/cheetah/accelerator/marker.py b/cheetah/accelerator/marker.py
index 66ed1412a..45bd6e472 100644
--- a/cheetah/accelerator/marker.py
+++ b/cheetah/accelerator/marker.py
@@ -40,10 +40,13 @@ def is_skippable(self) -> bool:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
-        # Do nothing on purpose. Maybe later we decide markers should be shown, but for
-        # now they are invisible.
-        pass
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+        # TODO: Implement a better visualisation for markers. At the moment they are
+        # invisible.
+        return ax
 
     @property
     def defining_features(self) -> list[str]:
diff --git a/cheetah/accelerator/quadrupole.py b/cheetah/accelerator/quadrupole.py
index 413a3047e..0e44d5a58 100644
--- a/cheetah/accelerator/quadrupole.py
+++ b/cheetah/accelerator/quadrupole.py
@@ -7,7 +7,12 @@
 from cheetah.accelerator.element import Element
 from cheetah.particles import Beam, ParticleBeam, Species
 from cheetah.track_methods import base_rmatrix, misalignment_matrix
-from cheetah.utils import UniqueNameGenerator, bmadx, verify_device_and_dtype
+from cheetah.utils import (
+    UniqueNameGenerator,
+    bmadx,
+    squash_index_for_unavailable_dims,
+    verify_device_and_dtype,
+)
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
@@ -176,6 +181,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             energy=ref_energy,
             particle_charges=incoming.particle_charges,
             survival_probabilities=incoming.survival_probabilities,
+            s=incoming.s + self.length,
             species=incoming.species,
         )
         return outgoing_beam
@@ -204,10 +210,27 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for i in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
-        plot_k1 = self.k1[vector_idx] if self.k1.dim() > 0 else self.k1
-        plot_s = s[vector_idx] if s.dim() > 0 else s
-        plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
+        plot_k1 = (
+            self.k1[squash_index_for_unavailable_dims(vector_idx, self.k1.shape)]
+            if len(self.k1.shape) > 0
+            else self.k1
+        )
+
+        plot_s = (
+            s[squash_index_for_unavailable_dims(vector_idx, s.shape)]
+            if len(s.shape) > 0
+            else s
+        )
+        plot_length = (
+            self.length[squash_index_for_unavailable_dims(vector_idx, s.shape)]
+            if len(self.length.shape) > 0
+            else self.length
+        )
 
         alpha = 1 if self.is_active else 0.2
         height = 0.8 * (torch.sign(plot_k1) if self.is_active else 1)
diff --git a/cheetah/accelerator/screen.py b/cheetah/accelerator/screen.py
index 4d3fc15dd..74556374a 100644
--- a/cheetah/accelerator/screen.py
+++ b/cheetah/accelerator/screen.py
@@ -200,6 +200,8 @@ def track(self, incoming: Beam) -> Beam:
                     cov=incoming.cov,
                     energy=incoming.energy,
                     total_charge=torch.zeros_like(incoming.total_charge),
+                    s=incoming.s,
+                    species=incoming.species.clone(),
                 )
             elif isinstance(incoming, ParticleBeam):
                 return ParticleBeam(
@@ -209,6 +211,8 @@ def track(self, incoming: Beam) -> Beam:
                     survival_probabilities=torch.zeros_like(
                         incoming.survival_probabilities
                     ),
+                    s=incoming.s,
+                    species=incoming.species.clone(),
                 )
         else:
             return incoming.clone()
@@ -325,7 +329,11 @@ def set_read_beam(self, value: Beam) -> None:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
 
         alpha = 1 if self.is_active else 0.2
diff --git a/cheetah/accelerator/segment.py b/cheetah/accelerator/segment.py
index 28ad7c7e9..86fd368a1 100644
--- a/cheetah/accelerator/segment.py
+++ b/cheetah/accelerator/segment.py
@@ -1,6 +1,6 @@
 from functools import reduce
 from pathlib import Path
-from typing import Any
+from typing import Any, Iterator
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -14,7 +14,7 @@
 from cheetah.converters import bmad, elegant, nxtables
 from cheetah.latticejson import load_cheetah_model, save_cheetah_model
 from cheetah.particles import Beam, Species
-from cheetah.utils import UniqueNameGenerator
+from cheetah.utils import UniqueNameGenerator, squash_index_for_unavailable_dims
 
 generate_unique_name = UniqueNameGenerator(prefix="unnamed_element")
 
@@ -454,6 +454,77 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for split_element in element.split(resolution)
         ]
 
+    def beam_along_segment_generator(
+        self, incoming: Beam, resolution: float | None = None
+    ) -> Iterator[Beam]:
+        """
+        Generator for beam objects along the segment either at the end of each element
+        or at a given resolution.
+
+        :param incoming: Beam that is entering the segment from upstream for which the
+            trajectory is computed.
+        :param resolution: Requested resolution of trajectory. Note that not all
+            elements can be split at arbitrary resolutions, which can lead to deviations
+            from the requested resolution. If `None` is passed, samples are taken at the
+            end of each element.
+        :return: Generator that yields the beam objects along the segment.
+        """
+        # If a resolution is passed, run this method for the split Segment
+        if resolution is not None:
+            yield from self.__class__(
+                elements=self.split(resolution), name=f"{self.name}_split"
+            ).beam_along_segment_generator(incoming)
+        else:
+            yield incoming
+            for element in self.elements:
+                outgoing = element.track(incoming)
+                yield outgoing
+                incoming = outgoing
+
+    def get_beam_attrs_along_segment(
+        self,
+        attr_names: tuple[str, ...] | str,
+        incoming: Beam,
+        resolution: float | None = None,
+    ) -> tuple[torch.Tensor, ...] | torch.Tensor:
+        """
+        Get metrics along the segment either at the end of each element or at a given
+        resolution.
+
+        :param attr_names: Metrics to compute. Can be a single metric or a tuple of
+            metrics. Supported metrics are any property of beam class of `incoming`.
+        :param incoming: Beam that is entering the segment from upstream for which the
+            trajectory is computed.
+        :param resolution: Requested resolution of trajectory. Note that not all
+            elements can be split at arbitrary resolutions, which can lead to deviations
+            from the requested resolution. If `None` is passed, samples are taken at the
+            end of each element.
+        :return: Tuple of tensors containing the requested metrics along the segment.
+        """
+        attr_name_tuple = attr_names if isinstance(attr_names, tuple) else (attr_names,)
+
+        results = zip(
+            *(
+                tuple(getattr(beam, attr_name) for attr_name in attr_name_tuple)
+                for beam in self.beam_along_segment_generator(
+                    incoming, resolution=resolution
+                )
+            )
+        )
+        broadcasted_results = tuple(
+            torch.stack(
+                torch.broadcast_tensors(*attr_tensor),
+                dim=-(incoming.UNVECTORIZED_NUM_ATTR_DIMS[attr_name] + 1),
+            )
+            for attr_tensor, attr_name in zip(results, attr_name_tuple)
+        )
+
+        return (
+            broadcasted_results
+            if isinstance(attr_names, tuple)
+            else broadcasted_results[0]
+        )
+
     def set_attrs_on_every_element_of_type(
         self,
         element_type: type[Element] | tuple[type[Element]],
@@ -478,31 +549,47 @@ def set_attrs_on_every_element_of_type(
                     element_type, is_recursive=True, **kwargs
                 )
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         element_lengths = [element.length for element in self.elements]
-        element_ss = [torch.tensor(0.0)] + [
-            sum(element_lengths[: i + 1]) for i, _ in enumerate(element_lengths)
-        ]
-        element_ss = [s + element_s for element_s in element_ss]
-        broadcast_ss = torch.broadcast_tensors(*element_ss)
-        stacked_ss = torch.stack(broadcast_ss)
-        dimension_reordered_ss = stacked_ss.movedim(0, -1)  # Place vector dims first
+        broadcast_element_lengths = torch.broadcast_tensors(*element_lengths)
+        stacked_element_lengths = torch.stack(broadcast_element_lengths, dim=-1)
+        element_end_s_positions = torch.cumsum(stacked_element_lengths, dim=-1)
+        s_positions = torch.cat(
+            (
+                torch.zeros_like(element_end_s_positions[..., :1]),
+                element_end_s_positions,
+            ),
+            dim=-1,
+        )
 
+        # The element lengths might not capture the entire vector shape with the
+        # incoming beam used in plotting functions that might use this function. The
+        # following tries to extract the correct vector index for just the s positions
+        # of the elements, while preserving the ability to use this element plotting
+        # function without an incoming beam.
         plot_ss = (
-            dimension_reordered_ss[vector_idx]
-            if stacked_ss.dim() > 1
-            else dimension_reordered_ss
+            s_positions[
+                squash_index_for_unavailable_dims(vector_idx, s_positions.shape)
+            ]
+            if s_positions.dim() > 1
+            else s_positions
         ).detach()
 
         ax.plot([0, plot_ss[-1]], [0, 0], "--", color="black")
 
         for element, s in zip(self.elements, plot_ss[:-1]):
-            element.plot(ax, s, vector_idx)
+            element.plot(s, vector_idx, ax)
 
         ax.set_ylim(-1, 1)
         ax.set_xlabel("s (m)")
         ax.set_yticks([])
 
+        return ax
+
     def plot_mean_and_std(
         self,
         axx: plt.Axes,
@@ -525,73 +612,23 @@ def plot_mean_and_std(
             case of present vector dimension but no index provided, the first one is
             used by default.
         """
-        reference_segment = self.clone()
-        splits = reference_segment.split(resolution=torch.tensor(resolution))
+        reference_segment = self.clone()  # Prevent side effects when plotting
 
-        split_lengths = [split.length for split in splits]
-        ss = [torch.tensor(0.0)] + [
-            sum(split_lengths[: i + 1]) for i, _ in enumerate(split_lengths)
-        ]
-        broadcast_ss = torch.broadcast_tensors(*ss)
-        stacked_ss = torch.stack(broadcast_ss)
-        dimension_reordered_ss = stacked_ss.movedim(0, -1)  # Place vector dims first
-
-        references = [incoming]
-        for split in splits:
-            sample = split.track(references[-1])
-            references.append(sample)
-
-        x_means = [reference_beam.mu_x for reference_beam in references]
-        broadcast_x_means = torch.broadcast_tensors(*x_means)
-        stacked_x_means = torch.stack(broadcast_x_means)
-        dimension_reordered_x_means = stacked_x_means.movedim(
-            0, -1
-        )  # Place vector dims first
-        x_stds = [reference_beam.sigma_x for reference_beam in references]
-        broadcast_x_stds = torch.broadcast_tensors(*x_stds)
-        stacked_x_stds = torch.stack(broadcast_x_stds)
-        dimension_reordered_x_stds = stacked_x_stds.movedim(
-            0, -1
-        )  # Place vector dims first
-
-        y_means = [reference_beam.mu_y for reference_beam in references]
-        broadcast_y_means = torch.broadcast_tensors(*y_means)
-        stacked_y_means = torch.stack(broadcast_y_means)
-        dimension_reordered_y_means = stacked_y_means.movedim(
-            0, -1
-        )  # Place vector dims first
-        y_stds = [reference_beam.sigma_y for reference_beam in references]
-        broadcast_y_stds = torch.broadcast_tensors(*y_stds)
-        stacked_y_stds = torch.stack(broadcast_y_stds)
-        dimension_reordered_y_stds = stacked_y_stds.movedim(
-            0, -1
-        )  # Place vector dims first
+        ss, x_means, x_stds, y_means, y_stds = (
+            reference_segment.get_beam_attrs_along_segment(
+                ("s", "mu_x", "sigma_x", "mu_y", "sigma_y"),
+                incoming,
+                resolution=resolution,
+            )
+        )
+        ss, x_means, x_stds, y_means, y_stds = torch.broadcast_tensors(
+            ss, x_means, x_stds, y_means, y_stds
+        )
 
-        plot_ss = (
-            dimension_reordered_ss[vector_idx]
-            if stacked_ss.dim() > 1
-            else dimension_reordered_ss
-        ).detach()
-        plot_x_means = (
-            dimension_reordered_x_means[vector_idx]
-            if stacked_x_means.dim() > 2
-            else dimension_reordered_x_means
-        ).detach()
-        plot_x_stds = (
-            dimension_reordered_x_stds[vector_idx]
-            if stacked_x_stds.dim() > 2
-            else dimension_reordered_x_stds
-        ).detach()
-        plot_y_means = (
-            dimension_reordered_y_means[vector_idx]
-            if stacked_y_means.dim() > 2
-            else dimension_reordered_y_means
-        ).detach()
-        plot_y_stds = (
-            dimension_reordered_y_stds[vector_idx]
-            if stacked_y_stds.dim() > 2
-            else dimension_reordered_y_stds
-        ).detach()
+        plot_ss, plot_x_means, plot_x_stds, plot_y_means, plot_y_stds = (
+            (metric[vector_idx] if metric.dim() > 1 else metric).detach()
+            for metric in (ss, x_means, x_stds, y_means, y_stds)
+        )
 
         axx.plot(plot_ss, plot_x_means)
         axx.fill_between(
@@ -646,69 +683,117 @@ def plot_overview(
 
         plt.tight_layout()
 
-    def plot_twiss(
-        self, incoming: Beam, ax: Any | None = None, vector_idx: tuple | None = None
-    ) -> None:
-        """Plot twiss parameters along the segment."""
-        longitudinal_beams = [incoming]
-        s_positions = [torch.tensor(0.0)]
-        for element in self.flattened().elements:
-            if torch.all(element.length == 0):
-                continue
+    def plot_beam_attrs(
+        self,
+        incoming: Beam,
+        attr_names: tuple[str, ...] | str,
+        resolution: float | None = None,
+        vector_idx: tuple | None = None,
+        ax: plt.Axes | None = None,
+    ) -> plt.Axes:
+        """
+        Plot beam attributes along the segment.
+
+        :param incoming: Beam that is entering the segment from upstream for which the
+            trajectory is computed.
+        :param attr_names: Metrics to compute. Can be a single metric or a tuple of
+            metrics. Supported metrics are any property of beam class of `incoming`.
+        :param resolution: Requested resolution of trajectory. Note that not all
+            elements can be split at arbitrary resolutions, which can lead to deviations
+            from the requested resolution. If `None` is passed, samples are taken at the
+            end of each element.
+        :param vector_idx: Index of the vector dimension to plot. If the model has more
+            than one vector dimension, this can be used to select a specific one. In the
+            case of present vector dimension but no index provided, the first one is
+            used by default.
+        :param ax: Axes to plot into.
+        :return: Axes with the plotted beam attributes.
+        """
+        attr_names_with_s = ("s",) + (
+            attr_names if isinstance(attr_names, tuple) else (attr_names,)
+        )
+        beam_attrs = self.get_beam_attrs_along_segment(
+            attr_names_with_s, incoming, resolution=resolution
+        )
 
-            outgoing = element.track(longitudinal_beams[-1])
+        ax = ax or plt.subplot(111)
 
-            longitudinal_beams.append(outgoing)
-            s_positions.append(s_positions[-1] + element.length)
+        s = beam_attrs[0]
+        for attr, attr_name in zip(beam_attrs[1:], attr_names_with_s[1:]):
+            ax.plot(
+                (s[vector_idx] if s.dim() > 1 else s).detach(),
+                (attr[vector_idx] if attr.dim() > 1 else attr).detach(),
+                label=attr_name,
+            )
 
-        beta_x = [beam.beta_x for beam in longitudinal_beams]
-        beta_y = [beam.beta_y for beam in longitudinal_beams]
+        ax.legend()
 
-        # Extraction of the correct vector element for plotting
-        broadcast_s_positions = torch.broadcast_tensors(*s_positions)
-        stacked_s_positions = torch.stack(broadcast_s_positions)
-        dimension_reordered_s_positions = stacked_s_positions.movedim(0, -1)
-        plot_s_positions = (
-            dimension_reordered_s_positions[vector_idx]
-            if stacked_s_positions.dim() > 1
-            else dimension_reordered_s_positions
-        ).detach()
+        return ax
 
-        broadcast_beta_x = torch.broadcast_tensors(*beta_x)
-        stacked_beta_x = torch.stack(broadcast_beta_x)
-        dimension_reordered_beta_x = stacked_beta_x.movedim(0, -1)
-        plot_beta_x = (
-            dimension_reordered_beta_x[vector_idx]
-            if stacked_beta_x.dim() > 2
-            else dimension_reordered_beta_x
-        ).detach()
+    def plot_beam_attrs_over_lattice(
+        self,
+        incoming: Beam,
+        attr_names: tuple[str, ...] | str,
+        figsize=(8, 4),
+        resolution: float | None = None,
+        vector_idx: tuple | None = None,
+    ) -> None:
+        """
+        Plot beam attributes in a plot over a plot of the lattice.
 
-        broadcast_beta_y = torch.broadcast_tensors(*beta_y)
-        stacked_beta_y = torch.stack(broadcast_beta_y)
-        dimension_reordered_beta_y = stacked_beta_y.movedim(0, -1)
-        plot_beta_y = (
-            dimension_reordered_beta_y[vector_idx]
-            if stacked_beta_y.dim() > 2
-            else dimension_reordered_beta_y
-        ).detach()
+        :param incoming: Beam that is entering the segment from upstream for which the
+            trajectory is computed.
+        :param attr_names: Metrics to compute. Can be a single metric or a tuple of
+            metrics. Supported metrics are any property of beam class of `incoming`.
+        :param figsize: Size of the figure.
+        :param resolution: Minimum resolution of the tracking of the beam position and
+            beam size in the plot.
+        :param vector_idx: Index of the vector dimension to plot. If the model has more
+            than one vector dimension, this can be used to select a specific one. In the
+            case of present vector dimension but no index provided, the first one is
+            used by default.
+        """
+        fig = plt.figure(figsize=figsize)
+        gs = fig.add_gridspec(2, hspace=0, height_ratios=[3, 1])
+        axs = gs.subplots(sharex=True)
 
-        if ax is None:
-            fig = plt.figure()
-            ax = fig.add_subplot(111)
+        self.plot_beam_attrs(
+            incoming=incoming,
+            attr_names=attr_names,
+            resolution=resolution,
+            vector_idx=vector_idx,
+            ax=axs[0],
+        )
+        self.plot(s=0.0, ax=axs[1])
+
+        plt.tight_layout()
+
+    def plot_twiss(
+        self, incoming: Beam, vector_idx: tuple | None = None, ax: Any | None = None
+    ) -> plt.Axes:
+        """Plot twiss parameters along the segment."""
+        ax = self.plot_beam_attrs(
+            incoming,
+            ("beta_x", "beta_y"),
+            resolution=None,
+            vector_idx=vector_idx,
+            ax=ax,
+        )
+
+        beta_x_line = ax.get_lines()[0]
+        beta_y_line = ax.get_lines()[1]
+
+        beta_x_line.set_label(r"$\beta_x$")
+        beta_x_line.set_color("tab:red")
+        beta_y_line.set_label(r"$\beta_y$")
+        beta_y_line.set_color("tab:green")
 
         ax.set_title("Twiss Parameters")
         ax.set_xlabel("s (m)")
         ax.set_ylabel(r"$\beta$ (m)")
-
-        ax.plot(plot_s_positions, plot_beta_x, label=r"$\beta_x$", c="tab:red")
-        ax.plot(plot_s_positions, plot_beta_y, label=r"$\beta_y$", c="tab:green")
-
         ax.legend()
-        plt.tight_layout()
 
-    @property
-    def defining_features(self) -> list[str]:
-        return super().defining_features + ["elements"]
+        return ax
 
     def plot_twiss_over_lattice(self, incoming: Beam, figsize=(8, 4)) -> None:
         """Plot twiss parameters in a plot over a plot of the lattice."""
@@ -717,10 +802,14 @@ def plot_twiss_over_lattice(self, incoming: Beam, figsize=(8, 4)) -> None:
         axs = gs.subplots(sharex=True)
 
         self.plot_twiss(incoming, ax=axs[0])
-        self.plot(axs[1], 0)
+        self.plot(s=0.0, ax=axs[1])
 
         plt.tight_layout()
 
+    @property
+    def defining_features(self) -> list[str]:
+        return super().defining_features + ["elements"]
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}(elements={repr(self.elements)}, "
diff --git a/cheetah/accelerator/sextupole.py b/cheetah/accelerator/sextupole.py
index a3b08ec2d..6e16b7056 100644
--- a/cheetah/accelerator/sextupole.py
+++ b/cheetah/accelerator/sextupole.py
@@ -106,6 +106,7 @@ def track(self, incoming: Beam) -> Beam:
                 energy=incoming.energy,
                 particle_charges=incoming.particle_charges,
                 survival_probabilities=incoming.survival_probabilities,
+                s=incoming.s + self.length,
                 species=incoming.species,
             )
         else:
@@ -125,7 +126,9 @@ def is_active(self) -> bool:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         raise NotImplementedError
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
         raise NotImplementedError
 
     @property
diff --git a/cheetah/accelerator/solenoid.py b/cheetah/accelerator/solenoid.py
index fe26c081c..c9ae29d44 100644
--- a/cheetah/accelerator/solenoid.py
+++ b/cheetah/accelerator/solenoid.py
@@ -126,7 +126,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for _ in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
 
diff --git a/cheetah/accelerator/space_charge_kick.py b/cheetah/accelerator/space_charge_kick.py
index dc3b03852..ed6652324 100644
--- a/cheetah/accelerator/space_charge_kick.py
+++ b/cheetah/accelerator/space_charge_kick.py
@@ -635,6 +635,7 @@ def track(self, incoming: ParticleBeam) -> ParticleBeam:
                 energy=incoming.energy,
                 particle_charges=incoming.particle_charges,
                 survival_probabilities=incoming.survival_probabilities,
+                s=incoming.s,
                 species=incoming.species,
             )
 
@@ -651,7 +652,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
     def is_skippable(self) -> bool:
         return False
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
 
         ax.axvline(plot_s, ymin=0.01, ymax=0.99, color="orange", linestyle="-")
diff --git a/cheetah/accelerator/transverse_deflecting_cavity.py b/cheetah/accelerator/transverse_deflecting_cavity.py
index 671bcd5ba..5b40f17f9 100644
--- a/cheetah/accelerator/transverse_deflecting_cavity.py
+++ b/cheetah/accelerator/transverse_deflecting_cavity.py
@@ -192,6 +192,7 @@ def _track_bmadx(self, incoming: ParticleBeam) -> ParticleBeam:
             energy=ref_energy,
             particle_charges=incoming.particle_charges,
             survival_probabilities=incoming.survival_probabilities,
+            s=incoming.s + self.length,
             species=incoming.species,
         )
         return outgoing_beam
@@ -201,7 +202,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         # element itself
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
 
diff --git a/cheetah/accelerator/undulator.py b/cheetah/accelerator/undulator.py
index 99c5f69b0..1b4621b70 100644
--- a/cheetah/accelerator/undulator.py
+++ b/cheetah/accelerator/undulator.py
@@ -59,7 +59,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
         # TODO: Implement splitting for undulator properly, for now just return self
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
 
diff --git a/cheetah/accelerator/vertical_corrector.py b/cheetah/accelerator/vertical_corrector.py
index 5eb78069a..fe1b31b32 100644
--- a/cheetah/accelerator/vertical_corrector.py
+++ b/cheetah/accelerator/vertical_corrector.py
@@ -16,8 +16,8 @@
 class VerticalCorrector(Element):
     """
     Verticle corrector magnet in a particle accelerator.
-    Note: This is modeled as a drift section with
-        a thin-kick in the vertical plane.
+
+    NOTE: This is modeled as a drift section with a thin-kick in the vertical plane.
 
     :param length: Length in meters.
     :param angle: Particle deflection angle in the vertical plane in rad.
@@ -81,7 +81,11 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for _ in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+
         plot_s = s[vector_idx] if s.dim() > 0 else s
         plot_length = self.length[vector_idx] if self.length.dim() > 0 else self.length
         plot_angle = self.angle[vector_idx] if self.angle.dim() > 0 else self.angle
diff --git a/cheetah/particles/beam.py b/cheetah/particles/beam.py
index fe5a48bd0..cc8386a66 100644
--- a/cheetah/particles/beam.py
+++ b/cheetah/particles/beam.py
@@ -31,6 +31,38 @@ class directly, but use one of the subclasses.
         :math:`\Delta E = E - E_0`
     """
 
+    UNVECTORIZED_NUM_ATTR_DIMS = {
+        "energy": 0,
+        "total_charge": 0,
+        "s": 0,
+        "mu_x": 0,
+        "sigma_x": 0,
+        "mu_px": 0,
+        "sigma_px": 0,
+        "mu_y": 0,
+        "sigma_y": 0,
+        "mu_py": 0,
+        "sigma_py": 0,
+        "mu_tau": 0,
+        "sigma_tau": 0,
+        "mu_p": 0,
+        "sigma_p": 0,
+        "relativistic_gamma": 0,
+        "relativistic_beta": 0,
+        "p0c": 0,
+        "cov_xpx": 0,
+        "cov_ypy": 0,
+        "cov_taup": 0,
+        "emittance_x": 0,
+        "normalized_emittance_x": 0,
+        "beta_x": 0,
+        "alpha_x": 0,
+        "emittance_y": 0,
+        "normalized_emittance_y": 0,
+        "beta_y": 0,
+        "alpha_y": 0,
+    }
+
     @classmethod
     @abstractmethod
     def from_parameters(
@@ -52,6 +84,7 @@ def from_parameters(
         cov_taup: torch.Tensor | None = None,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -80,6 +113,7 @@ def from_parameters(
         :param cov_taup: Covariance between tau and p.
         :param energy: Reference energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
+        :param s: Position along the beamline of the reference particle in meters.
         :param species: Particle species of the beam. Defaults to electron.
         :param device: Device to create the beam on. If set to `"auto"` a CUDA GPU is
             selected if available. The CPU is used otherwise.
@@ -102,6 +136,7 @@ def from_twiss(
         cov_taup: torch.Tensor | None = None,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -123,6 +158,7 @@ def from_twiss(
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
         :param species: Particle species of the beam. Defaults to electron.
+        :param s: Position along the beamline of the reference particle in meters.
         :param device: Device to create the beam on. If set to `"auto"` a CUDA GPU is
             selected if available. The CPU is used otherwise.
         :param dtype: Data type of the beam.
diff --git a/cheetah/particles/parameter_beam.py b/cheetah/particles/parameter_beam.py
index 2cfcbcf36..d8ee49e72 100644
--- a/cheetah/particles/parameter_beam.py
+++ b/cheetah/particles/parameter_beam.py
@@ -15,24 +15,31 @@ class ParameterBeam(Beam):
     :param cov: Covariance matrix of the beam with shape `(..., 7, 7)`.
     :param energy: Reference energy of the beam in eV.
     :param total_charge: Total charge of the beam in C.
+    :param s: Position along the beamline of the reference particle in meters.
     :param species: Particle species of the beam. Defaults to electron.
     :param device: Device to use for the beam. If "auto", use CUDA if available.
         Note: Compuationally it would be faster to use CPU for ParameterBeam.
     :param dtype: Data type of the beam.
     """
 
+    UNVECTORIZED_NUM_ATTR_DIMS = Beam.UNVECTORIZED_NUM_ATTR_DIMS | {
+        "mu": 1,
+        "cov": 2,
+    }
+
     def __init__(
         self,
         mu: torch.Tensor,
         cov: torch.Tensor,
         energy: torch.Tensor,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
         device, dtype = verify_device_and_dtype(
-            [mu, cov, energy, total_charge], device, dtype
+            [mu, cov, energy, total_charge, s], device, dtype
         )
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -50,6 +57,9 @@ def __init__(
                 total_charge if total_charge is not None else 0.0, **factory_kwargs
             ),
         )
+        self.register_buffer_or_parameter(
+            "s", torch.as_tensor(s if s is not None else 0.0, **factory_kwargs)
+        )
 
     @classmethod
     def from_parameters(
@@ -71,6 +81,7 @@ def from_parameters(
         cov_taup: torch.Tensor | None = None,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -95,6 +106,7 @@ def from_parameters(
                 cov_taup,
                 energy,
                 total_charge,
+                s,
             ],
             device,
             dtype,
@@ -190,6 +202,7 @@ def from_parameters(
             cov=cov,
             energy=energy,
             total_charge=total_charge,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -209,6 +222,7 @@ def from_twiss(
         cov_taup: torch.Tensor | None = None,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -227,6 +241,7 @@ def from_twiss(
                 cov_taup,
                 energy,
                 total_charge,
+                s,
             ],
             device,
             dtype,
@@ -293,6 +308,7 @@ def from_twiss(
             cov_xpx=cov_xpx,
             cov_ypy=cov_ypy,
             total_charge=total_charge,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -435,6 +451,7 @@ def transformed_to(
             sigma_p=sigma_p,
             energy=energy,
             total_charge=total_charge,
+            s=self.s,
             species=species,
             device=device,
             dtype=dtype,
@@ -469,6 +486,8 @@ def as_particle_beam(self, num_particles: int) -> "ParticleBeam":  # noqa: F821
             cov_taup=self.cov_taup,
             energy=self.energy,
             total_charge=self.total_charge,
+            s=self.s,
+            species=self.species,
             device=self.mu.device,
             dtype=self.mu.dtype,
         )
@@ -499,6 +518,7 @@ def linspaced(self, num_particles: int) -> "ParticleBeam":  # noqa: F821
             sigma_p=self.sigma_p,
             energy=self.energy,
             total_charge=self.total_charge,
+            s=self.s,
             species=self.species,
             device=self.mu.device,
             dtype=self.mu.dtype,
@@ -570,6 +590,7 @@ def clone(self) -> "ParameterBeam":
             cov=self.cov.clone(),
             energy=self.energy.clone(),
             total_charge=self.total_charge.clone(),
+            s=self.s.clone(),
             species=self.species.clone(),
         )
 
@@ -579,5 +600,6 @@ def __repr__(self) -> str:
             + f"cov={repr(self.cov)}, "
             + f"energy={repr(self.energy)}, "
             + f"total_charge={repr(self.total_charge)}, "
+            + f"s={repr(self.s)}, "
             + f"species={repr(self.species)})"
         )
diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py
index e9a863b83..0a2f27b66 100644
--- a/cheetah/particles/particle_beam.py
+++ b/cheetah/particles/particle_beam.py
@@ -29,6 +29,7 @@ class ParticleBeam(Beam):
     :param survival_probabilities: Vector of probabilities that each particle has
         survived (i.e. not been lost), where 1.0 means the particle has survived and
         0.0 means the particle has been lost. Defaults to ones.
+    :param s: Position along the beamline of the reference particle in meters.
     :param species: Particle species of the beam. Defaults to electron.
     :param device: Device to move the beam's particle array to. If set to `"auto"` a
         CUDA GPU is selected if available. The CPU is used otherwise.
@@ -43,6 +44,17 @@ class ParticleBeam(Beam):
         "tau": r"$\tau$",
         "p": r"$\delta$",
     }
+    UNVECTORIZED_NUM_ATTR_DIMS = Beam.UNVECTORIZED_NUM_ATTR_DIMS | {
+        "particles": 2,
+        "particle_charges": 1,
+        "survival_probabilities": 1,
+        "x": 1,
+        "px": 1,
+        "y": 1,
+        "py": 1,
+        "tau": 1,
+        "p": 1,
+    }
 
     def __init__(
         self,
@@ -50,12 +62,15 @@ def __init__(
         energy: torch.Tensor,
         particle_charges: torch.Tensor | None = None,
         survival_probabilities: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
         device, dtype = verify_device_and_dtype(
-            [particles, energy, particle_charges], device, dtype
+            [particles, energy, particle_charges, survival_probabilities, s],
+            device,
+            dtype,
         )
         factory_kwargs = {"device": device, "dtype": dtype}
         super().__init__()
@@ -64,7 +79,7 @@ def __init__(
             particles.shape[-2] > 0 and particles.shape[-1] == 7
         ), "Particle vectors must be 7-dimensional."
 
-        species = species if species is not None else Species("electron")
+        self.species = species if species is not None else Species("electron")
 
         self.register_buffer_or_parameter(
             "particles", torch.as_tensor(particles, **factory_kwargs)
@@ -90,8 +105,9 @@ def __init__(
                 else torch.ones(particles.shape[-2], **factory_kwargs)
             ),
         )
-
-        self.species = species
+        self.register_buffer_or_parameter(
+            "s", torch.as_tensor(s if s is not None else 0.0, **factory_kwargs)
+        )
 
     @classmethod
     def from_parameters(
@@ -114,6 +130,7 @@ def from_parameters(
         cov_taup: torch.Tensor | None = None,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -144,6 +161,7 @@ def from_parameters(
         :param cov_taup: Correlation between s and p.
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
+        :param s: Position along the beamline of the reference particle in meters.
         :param species: Particle species of the beam. Defaults to electron.
         :param device: Device to move the beam's particle array to. If set to `"auto"` a
             CUDA GPU is selected if available. The CPU is used otherwise.
@@ -170,6 +188,7 @@ def from_parameters(
                 cov_taup,
                 energy,
                 total_charge,
+                s,
             ],
             device,
             dtype,
@@ -257,6 +276,7 @@ def from_parameters(
             num_particles=num_particles,
             energy=energy,
             total_charge=total_charge,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -270,6 +290,7 @@ def from_distribution(
         num_particles: int = 100_000,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -283,6 +304,7 @@ def from_distribution(
         :param cov: Covariance matrix of the multivariate normal distribution.
         :param energy: Energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
+        :param s: Position along the beamline of the reference particle in meters.
         :param species: Particle species of the beam. Defaults to electron.
         :param device: Device to move the beam's particle array to. If set to `"auto"` a
             CUDA GPU is selected if available. The CPU is used otherwise.
@@ -291,7 +313,7 @@ def from_distribution(
         """
         # Extract device and dtype from given arguments
         device, dtype = verify_device_and_dtype(
-            [mu, cov, energy, total_charge], device, dtype
+            [mu, cov, energy, total_charge, s], device, dtype
         )
         factory_kwargs = {"device": device, "dtype": dtype}
 
@@ -327,6 +349,7 @@ def from_distribution(
             particles,
             energy,
             particle_charges=particle_charges,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -347,6 +370,7 @@ def from_twiss(
         sigma_p: torch.Tensor | None = None,
         cov_taup: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -365,6 +389,7 @@ def from_twiss(
                 sigma_p,
                 cov_taup,
                 total_charge,
+                s,
             ],
             device,
             dtype,
@@ -425,6 +450,7 @@ def from_twiss(
             cov_xpx=cov_xpx,
             cov_ypy=cov_ypy,
             total_charge=total_charge,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -442,6 +468,7 @@ def uniform_3d_ellipsoid(
         sigma_p: torch.Tensor | None = None,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -466,6 +493,7 @@ def uniform_3d_ellipsoid(
         :param sigma_p: Sigma of the particle distribution in p, dimensionless.
         :param energy: Reference energy of the beam in eV.
         :param total_charge: Total charge of the beam in C.
+        :param s: Position along the beamline of the reference particle in meters.
         :param species: Particle species of the beam. Defaults to electron.
         :param device: Device to move the beam's particle array to. If set to `"auto"` a
             CUDA GPU is selected if available. The CPU is used otherwise.
@@ -483,6 +511,7 @@ def uniform_3d_ellipsoid(
                 sigma_p,
                 energy,
                 total_charge,
+                s,
             ],
             device,
             dtype,
@@ -517,6 +546,7 @@ def uniform_3d_ellipsoid(
             sigma_p=sigma_p,
             energy=energy,
             total_charge=total_charge,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -564,6 +594,7 @@ def make_linspaced(
         sigma_p: torch.Tensor | None = None,
         energy: torch.Tensor | None = None,
         total_charge: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
@@ -589,6 +620,8 @@ def make_linspaced(
             in meters.
         :param sigma_p: Sigma of the particle distribution in p, dimensionless.
         :param energy: Energy of the beam in eV.
+        :param total_charge: Total charge of the beam in C.
+        :param s: Position along the beamline of the reference particle in meters.
         :param species: Particle species of the beam. Defaults to electron.
         :param device: Device to move the beam's particle array to. If set to `"auto"` a
             CUDA GPU is selected if available. The CPU is used otherwise.
@@ -610,6 +643,7 @@ def make_linspaced(
                 sigma_p,
                 energy,
                 total_charge,
+                s,
             ],
             device,
             dtype,
@@ -694,6 +728,7 @@ def make_linspaced(
             particles=particles,
             energy=energy,
             particle_charges=particle_charges,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -982,6 +1017,8 @@ def transformed_to(
             particles=particles,
             energy=energy,
             particle_charges=particle_charges,
+            survival_probabilities=self.survival_probabilities,
+            s=self.s,
             species=species,
             device=device,
             dtype=dtype,
@@ -1028,6 +1065,9 @@ def linspaced(self, num_particles: int) -> "ParticleBeam":
             sigma_p=self.sigma_p,
             energy=self.energy,
             total_charge=self.total_charge,
+            particle_charges=self.particle_charges,
+            survival_probabilities=self.survival_probabilities,
+            s=self.s,
             species=self.species,
             device=self.particles.device,
             dtype=self.particles.dtype,
@@ -1093,6 +1133,7 @@ def from_xyz_pxpypz(
         energy: torch.Tensor,
         particle_charges: torch.Tensor | None = None,
         survival_probabilities: torch.Tensor | None = None,
+        s: torch.Tensor | None = None,
         species: Species | None = None,
         device: torch.device | None = None,
         dtype=torch.float32,
@@ -1107,6 +1148,7 @@ def from_xyz_pxpypz(
             energy=energy,
             particle_charges=particle_charges,
             survival_probabilities=survival_probabilities,
+            s=s,
             species=species,
             device=device,
             dtype=dtype,
@@ -1695,6 +1737,7 @@ def clone(self) -> "ParticleBeam":
             energy=self.energy.clone(),
             particle_charges=self.particle_charges.clone(),
             survival_probabilities=self.survival_probabilities.clone(),
+            s=self.s.clone(),
             species=self.species.clone(),
         )
 
@@ -1731,5 +1774,6 @@ def __repr__(self) -> str:
             + f"energy={self.energy}, "
             + f"particle_charges={self.particle_charges}, "
             + f"survival_probabilities={self.survival_probabilities}, "
+            + f"s={self.s}, "
             + f"species={repr(self.species)})"
         )
diff --git a/cheetah/utils/__init__.py b/cheetah/utils/__init__.py
index f015bd0c1..c1668f688 100644
--- a/cheetah/utils/__init__.py
+++ b/cheetah/utils/__init__.py
@@ -14,3 +14,4 @@
     unbiased_weighted_variance,
 )
 from .unique_name_generator import UniqueNameGenerator  # noqa: F401
+from .vector import squash_index_for_unavailable_dims  # noqa: F401
diff --git a/cheetah/utils/vector.py b/cheetah/utils/vector.py
new file mode 100644
index 000000000..9be9dfe41
--- /dev/null
+++ b/cheetah/utils/vector.py
@@ -0,0 +1,32 @@
+def squash_index_for_unavailable_dims(index: tuple, shape: tuple) -> tuple:
+    """
+    Helper function to create an index that works on only partially broadcasted tensors.
+
+    You might need this when you vectorized one parameter of your model as `(n,)` and
+    another as `(m, 1)`. Fully broadcasted, this would result in vector shapes of
+    `(m, n)`. Some functions in Cheetah expect an index into the fully broadcasted
+    vector shape. However, sometimes this index is passed onto a function that only
+    handles a result from the model that was only affected by one of these
+    vectorizations, e.g., the second one. This function lets you squash the vector index
+    meant for a vectorised result of shape `(m, n)` into a vector index that works on
+    a result with vector shape `(m, 1)` or `(n,)`.
+
+    Example: The input vector shapes are (a) `(3,)` and (b) `(2, 1)`, resulting in a
+    broadcasted shape of `(2, 3)`. A valid vector index for the fully broadcasted model
+    would be `(1, 2)`. If the result is only affected by vectorisation (b) with shape
+    `(2, 1)`, this function would squash the second dimension of the index to `0`,
+    resulting in `(1, 0)`. Similarly, if the result is only affected by vectorisation
+    (a) with shape `(3,)`, this function would remove the first dimension of the index,
+    resulting in `(2,)`.
+
+    :param index: The index to squash.
+    :param shape: The shape of the result that was only affected by part of the
+        vectorisations.
+    :return: The squashed index that is valid for indexing into a tensor of `shape`.
+    """
+    length_squashed = index[-len(shape) :]
+    unavailable_dims_squashed = tuple(
+        0 if s == 1 else i for i, s in zip(length_squashed, shape)
+    )
+
+    return unavailable_dims_squashed
diff --git a/docs/requirements.txt b/docs/requirements.txt
index 3a00d45c3..4d9872d20 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -2,3 +2,4 @@ furo
 ipykernel
 myst-parser
 nbsphinx
+snowballstemmer<3.0
diff --git a/docs/utils.rst b/docs/utils.rst
index 2375b036f..59e0b4d90 100644
--- a/docs/utils.rst
+++ b/docs/utils.rst
@@ -3,6 +3,10 @@
 Utils
 =====
 
+.. automodule:: utils.argument_verification
+    :members:
+    :undoc-members:
+
 .. automodule:: utils.bmadx
     :members:
     :undoc-members:
@@ -11,6 +15,10 @@ Utils
     :members:
     :undoc-members:
 
+.. automodule:: utils.elementwise_linspace
+    :members:
+    :undoc-members:
+
 .. automodule:: utils.kde
     :members:
     :undoc-members:
@@ -19,6 +27,14 @@ Utils
     :members:
     :undoc-members:
 
+.. automodule:: utils.statistics
+    :members:
+    :undoc-members:
+
 .. automodule:: utils.unique_name_generator
     :members:
     :undoc-members:
+
+.. automodule:: utils.vector
+    :members:
+    :undoc-members:
diff --git a/tests/resources/bmadx/incoming.pt b/tests/resources/bmadx/incoming.pt
index 0db286542..5eeaa346e 100644
Binary files a/tests/resources/bmadx/incoming.pt and b/tests/resources/bmadx/incoming.pt differ
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index cc56f11d5..4f7198967 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -90,12 +90,12 @@ def test_reference_particle_plot_vectorized_2d():
     """
     segment = cheetah.Segment.from_ocelot(ares.cell).subcell("AREASOLA1", "AREABSCR1")
     segment.AREAMQZM1.k1 = torch.tensor(5.0)
-    segment.AREAMQZM2.k1 = torch.tensor([[-5.0, -2.0, -1.0], [1.0, 2.0, 5.0]])
+    segment.AREAMQZM2.k1 = torch.tensor([1.0, 2.0, 5.0])
     segment.AREAMCVM1.k1 = torch.tensor(1e-3)
     segment.AREAMQZM3.k1 = torch.tensor(5.0)
     segment.AREAMCHM1.k1 = torch.tensor(-2e-3)
     segment.Drift_AREAMCHM1.length = (
-        torch.FloatTensor(2, 3).uniform_(0.9, 1.1) * segment.Drift_AREAMCHM1.length
+        torch.FloatTensor(2, 1).uniform_(0.9, 1.1) * segment.Drift_AREAMCHM1.length
     )
 
     incoming = cheetah.ParticleBeam.from_astra(
diff --git a/tests/test_segment.py b/tests/test_segment.py
index b2ca456b5..53be374e2 100644
--- a/tests/test_segment.py
+++ b/tests/test_segment.py
@@ -83,3 +83,91 @@ def test_attr_setting_by_element_type_convenience_method(is_recursive):
             assert element.length == 0.6
         else:
             raise ValueError(f"Unexpected element type: {type(element)}")
+
+
+def test_elementwise_longitudinal_beam_generator():
+    """
+    Test that the longitudinal beam generator works as expected and generates the
+    correct number of beams.
+    """
+    segment = cheetah.Segment(
+        elements=[
+            cheetah.Drift(length=torch.tensor(0.5)),
+            cheetah.Quadrupole(length=torch.tensor(0.3)),
+            cheetah.Drift(length=torch.tensor(0.2)),
+        ]
+    )
+    incoming_beam = cheetah.ParameterBeam.from_astra(
+        "tests/resources/ACHIP_EA1_2021.1351.001"
+    )
+
+    longitudinal_beam_generator = segment.beam_along_segment_generator(
+        incoming=incoming_beam
+    )
+    longitudinal_beam_list = list(longitudinal_beam_generator)
+
+    assert len(longitudinal_beam_list) == 4
+    for longitudinal_beam in longitudinal_beam_list:
+        assert isinstance(longitudinal_beam, incoming_beam.__class__)
+
+
+def test_resolution_longitudinal_beam_generator():
+    """
+    Test that the longitudinal beam generator works as expected and generates the
+    correct number of beams with a specified resolution.
+    """
+    segment = cheetah.Segment(
+        elements=[
+            cheetah.Drift(length=torch.tensor(0.5)),
+            cheetah.Quadrupole(length=torch.tensor(0.3)),
+            cheetah.Drift(length=torch.tensor(0.2)),
+        ]
+    )
+    incoming_beam = cheetah.ParameterBeam.from_astra(
+        "tests/resources/ACHIP_EA1_2021.1351.001"
+    )
+
+    longitudinal_beam_generator = segment.beam_along_segment_generator(
+        incoming=incoming_beam, resolution=0.1
+    )
+    longitudinal_beam_list = list(longitudinal_beam_generator)
+
+    assert len(longitudinal_beam_list) == 11
+    for longitudinal_beam in longitudinal_beam_list:
+        assert isinstance(longitudinal_beam, incoming_beam.__class__)
+
+
+@pytest.mark.parametrize(
+    "attr_names", ["beta_x", ("beta_x",), ("s", "beta_x"), ("x", "mu_x")]
+)
+def test_longitudinal_beam_metric(attr_names):
+    """
+    Test that the convenience method for computing a attributes along the lattice works
+    as expected. Focus is put on the return being a single tensor when an attribute
+    string is passed, and a tuple of tensors when a tuple of metric strings is passed
+    (regardless of the length of the tuple).
+    """
+    segment = cheetah.Segment(
+        elements=[
+            cheetah.Drift(length=torch.tensor(0.5)),
+            cheetah.Quadrupole(length=torch.tensor(0.3)),
+            cheetah.Drift(length=torch.tensor(0.2)),
+        ]
+    )
+    incoming_beam = cheetah.ParticleBeam.from_astra(
+        "tests/resources/ACHIP_EA1_2021.1351.001"
+    )
+
+    result = segment.get_beam_attrs_along_segment(attr_names, incoming_beam)
+
+    if isinstance(attr_names, str):
+        assert isinstance(result, torch.Tensor)
+        assert len(result) == 4
+    else:
+        assert isinstance(result, tuple)
+        assert len(result) == len(attr_names)
+        for attr_result, attr_name in zip(result, attr_names):
+            assert isinstance(attr_result, torch.Tensor)
+            assert attr_result.shape == (
+                (4, incoming_beam.num_particles) if attr_name == "x" else (4,)
+            )


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR introduces a convenience function, Segment.longitudinal_beam_generator, to generate Beam objects along the beamline for convenient computation of longitudinal trajectories. Additionally, the Beam subclasses are updated to track the s position along the beamline, and various accelerator element tracking and plotting methods have been refactored to use the new s parameter.

### Reviewed Changes

Copilot reviewed 20 out of 20 changed files in this pull request and generated no comments.

<details>
<summary>Show a summary per file</summary>

| File                                    | Description                                                                                                  |
| --------------------------------------- | ------------------------------------------------------------------------------------------------------------ |
| tests/test_segment.py                   | Added tests for the new longitudinal_beam_generator functionality.                                           |
| tests/test_plotting.py                  | Updated test cases to reflect changes in plotting and vector handling.                                       |
| cheetah/utils/vector.py                 | Introduced the squash_index_for_unavailable_dims helper with examples of its use in plotting functions.        |
| cheetah/particles/particle_beam.py      | Updated ParticleBeam initialization and cloning to include the s position.                                   |
| cheetah/accelerator/*                   | Modified track methods in Drift, Quadrupole, Dipole, etc., to update the s position; refactored plotting.      |
| cheetah/accelerator/segment.py          | Added longitudinal_beam_generator and refactored plotting functions to use vector indexing via squash_index.    |
| CHANGELOG.md                            | Documented the new feature and changes to s tracking in Beam subclasses.                                     |
</details>



<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**cheetah/accelerator/segment.py:523**
* [nitpick] Consider adding a check to verify that 'vector_idx' is not None before calling 'squash_index_for_unavailable_dims' to prevent potential runtime errors when no vector index is provided.
```
plot_ss = ( element_ss[squash_index_for_unavailable_dims(vector_idx, element_ss.shape)] if element_ss.dim() > 1 else element_ss ).detach()
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-06T15:00:06Z
user: login: Hespe
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-07T12:31:09Z
user: login: Hespe
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-07T13:09:58Z
user: login: jank324
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-07T13:13:33Z
user: login: jank324
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-07T13:14:51Z
user: login: jank324
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-07T13:32:26Z
user: login: Hespe
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-08T07:52:25Z
user: login: jank324
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-08T08:50:14Z
user: login: jank324
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-08T08:52:24Z
user: login: Hespe
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-08T11:46:58Z

Review comments:
path: cheetah/particles/parameter_beam.py
user: login: Hespe
body: The `transformed_to` and `linspaced` methods should also preserve the new `s` property in my opinion. In case of `transformed_to` I would opt to also add `s` as an argument.
created_at: 2025-05-07T08:38:00Z
author_association: MEMBER
path: cheetah/particles/particle_beam.py
user: login: Hespe
body: Same as for `ParameterBeam` a couple of methods are not preserving`s`: `transformed_to`, `linspaced`, `as_parameter_beam`.

Also, should we add `s` to `__repr__` for both beams?
created_at: 2025-05-07T08:42:48Z
author_association: MEMBER
diff_hunk: @@ -70,21 +70,25 @@ def track(self, incoming: Beam) -> Beam:
             tm = self.transfer_map(incoming.energy, incoming.species)
             mu = (tm @ incoming.mu.unsqueeze(-1)).squeeze(-1)
             cov = tm @ incoming.cov @ tm.transpose(-2, -1)
+            s = incoming.s + self.length
             return ParameterBeam(
                 mu,
                 cov,
                 incoming.energy,
                 total_charge=incoming.total_charge,
+                s=s,
                 species=incoming.species.clone(),
             )
         elif isinstance(incoming, ParticleBeam):
             tm = self.transfer_map(incoming.energy, incoming.species)
             new_particles = incoming.particles @ tm.transpose(-2, -1)
+            new_s = incoming.s + self.length
path: cheetah/accelerator/element.py
user: login: Hespe
body: Should we be consistent with the variable naming? Its called just `s` above. Also, we could just do this kind of simple calculation right in the call to the constructor.
created_at: 2025-05-07T08:59:47Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,30 @@
+def squash_index_for_unavailable_dims(index: tuple, shape: tuple) -> tuple:
+    """
+    Helper function to create an index that works on only partially broadcasted tensors.
+
+    You might need this, when you vectorised one parameter of your model as `(n,)` and
+    another as `(m, 1)`. Fully broadcasted, this would result in vector shapes of
+    `(m, n)`. Some functions in Cheetah expect an index into the fully broadcasted
+    vector shape. However, sometimes this index is passed onto a function that only
+    handles a result from the model that was only affected by one of these
+    vectorisations, e.g. the second one. This function let's you squash the vector index
+    meant for a vectorised result of shape `(m, n)` into a vector index that works on
+    a result with vector shape `(m, 1)`.
+
+    Example: The input vector shapes are `(3,)` and `(2, 1)`, resulting in a broadcasted
+    shape of `(3, 2)`. A valid vector index for fully-broadcasted model would be
+    `(0, 2)`, which is not valid for results only affected by the second vectorisation.
+    This function would then squash the second dimension of the index, where the result
+    is not vectorised, to `0`, resulting a valid index of `(0, 0)`.
+
+    :param index: The index to squash.
+    :param shape: The shape of the result that was only affected part of the
path: cheetah/utils/vector.py
user: login: Hespe
body: ```suggestion
    :param shape: The shape of the result that was only affected by parts of the
```
created_at: 2025-05-07T11:45:26Z
author_association: MEMBER
path: cheetah/utils/vector.py
user: login: Hespe
body: It might be helpful to also consider the case where the index is applied to the "shorter" vector shape, i.e. `(3,)` in the current example.
created_at: 2025-05-07T11:46:57Z
author_association: MEMBER
diff_hunk: @@ -454,6 +454,75 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for split_element in element.split(resolution)
         ]
 
+    def longitudinal_beam_generator(
+        self, incoming: Beam, resolution: float | None = None
+    ) -> Iterator[Beam]:
+        """
+        Generator for beam objects along the segment either at the end of each element
+        or at a given resolution.
+
+        :param incoming: Beam that is entering the segment from upstream for which the
+            trajectory is computed.
+        :param resolution: Requested resolution of trajectory. Note that not all
+            elements can be split at arbitrary resolutions, which can lead to deviations
+            from the requested resolution. If `None` is passed, samples are taken at the
+            end of each element.
+        :return: Generator that yields the beam objects along the segment.
+        """
+        # If a resolution is passed, run this method for the split Segment
+        if resolution is not None:
+            yield from self.__class__(
+                elements=self.split(resolution), name=f"{self.name}_split"
+            ).longitudinal_beam_generator(incoming)
+        else:
+            yield incoming
+            for element in self.elements:
+                outgoing = element.track(incoming)
+                yield outgoing
+                incoming = outgoing
+
+    def get_longitudinal_metrics(
+        self,
+        metric_names: tuple[str, ...] | str,
+        incoming: Beam,
+        resolution: float | None = None,
+    ) -> tuple[torch.Tensor, ...]:
path: cheetah/accelerator/segment.py
user: login: Hespe
body: ```suggestion
    ) -> tuple[torch.Tensor, ...] | torch.Tensor:
```
created_at: 2025-05-07T13:09:58Z
author_association: MEMBER
path: cheetah/particles/parameter_beam.py
user: login: jank324
body: Decided to not pass `s` as an argument to `transformed_to` because it can easily be set directly with `beam.s = new_value`.
created_at: 2025-05-07T13:13:33Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,30 @@
+def squash_index_for_unavailable_dims(index: tuple, shape: tuple) -> tuple:
+    """
+    Helper function to create an index that works on only partially broadcasted tensors.
+
+    You might need this, when you vectorised one parameter of your model as `(n,)` and
+    another as `(m, 1)`. Fully broadcasted, this would result in vector shapes of
+    `(m, n)`. Some functions in Cheetah expect an index into the fully broadcasted
+    vector shape. However, sometimes this index is passed onto a function that only
+    handles a result from the model that was only affected by one of these
+    vectorisations, e.g. the second one. This function let's you squash the vector index
+    meant for a vectorised result of shape `(m, n)` into a vector index that works on
+    a result with vector shape `(m, 1)`.
+
+    Example: The input vector shapes are `(3,)` and `(2, 1)`, resulting in a broadcasted
+    shape of `(3, 2)`. A valid vector index for fully-broadcasted model would be
+    `(0, 2)`, which is not valid for results only affected by the second vectorisation.
+    This function would then squash the second dimension of the index, where the result
+    is not vectorised, to `0`, resulting a valid index of `(0, 0)`.
+
+    :param index: The index to squash.
+    :param shape: The shape of the result that was only affected part of the
path: cheetah/utils/vector.py
user: login: jank324
body: Actually, `part` sounds more correct to me.
created_at: 2025-05-07T13:14:51Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,30 @@
+def squash_index_for_unavailable_dims(index: tuple, shape: tuple) -> tuple:
+    """
+    Helper function to create an index that works on only partially broadcasted tensors.
+
+    You might need this, when you vectorised one parameter of your model as `(n,)` and
+    another as `(m, 1)`. Fully broadcasted, this would result in vector shapes of
+    `(m, n)`. Some functions in Cheetah expect an index into the fully broadcasted
+    vector shape. However, sometimes this index is passed onto a function that only
+    handles a result from the model that was only affected by one of these
+    vectorisations, e.g. the second one. This function let's you squash the vector index
+    meant for a vectorised result of shape `(m, n)` into a vector index that works on
+    a result with vector shape `(m, 1)`.
+
+    Example: The input vector shapes are `(3,)` and `(2, 1)`, resulting in a broadcasted
+    shape of `(3, 2)`. A valid vector index for fully-broadcasted model would be
+    `(0, 2)`, which is not valid for results only affected by the second vectorisation.
+    This function would then squash the second dimension of the index, where the result
+    is not vectorised, to `0`, resulting a valid index of `(0, 0)`.
+
+    :param index: The index to squash.
+    :param shape: The shape of the result that was only affected part of the
path: cheetah/utils/vector.py
user: login: jank324
body: ```suggestion
    :param shape: The shape of the result that was only affected by part of the
```
created_at: 2025-05-07T13:32:25Z
author_association: MEMBER
diff_hunk: @@ -139,8 +140,12 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for i in range(num_splits)
         ]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
-        pass
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
+        # Yes, this does nothing on purpose
path: cheetah/accelerator/drift.py
user: login: Hespe
body: Maybe we should rather explain why this does nothing? I.e., because `Drift` is visualized as a gap?
created_at: 2025-05-08T07:29:00Z
author_association: MEMBER
diff_hunk: @@ -40,10 +40,13 @@ def is_skippable(self) -> bool:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         return [self]
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
+        ax = ax or plt.subplot(111)
         # Do nothing on purpose. Maybe later we decide markers should be shown, but for
         # now they are invisible.
path: cheetah/accelerator/marker.py
user: login: Hespe
body: `Drift` and `Marker` should have a similar comment for their plot function doing nothing
created_at: 2025-05-08T07:30:56Z
author_association: MEMBER
diff_hunk: @@ -125,7 +126,9 @@ def is_active(self) -> bool:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         raise NotImplementedError
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
         raise NotImplementedError
path: cheetah/accelerator/sextupole.py
user: login: Hespe
body: That is of course not relevant to the scope of this PR, but we should really provide a plotting implementation of `Sextupole`. Crashing on plotting is not a good sight.
created_at: 2025-05-08T07:33:55Z
author_association: MEMBER
path: cheetah/particles/parameter_beam.py
user: login: Hespe
body: Should `s` be part of `__repr__`?
created_at: 2025-05-08T07:39:01Z
author_association: MEMBER
diff_hunk: @@ -83,3 +83,87 @@ def test_attr_setting_by_element_type_convenience_method(is_recursive):
             assert element.length == 0.6
         else:
             raise ValueError(f"Unexpected element type: {type(element)}")
+
+
+def test_elementwise_longitudinal_beam_generator():
+    """
+    Test that the longitudinal beam generator works as expected and generates the
+    correct number of beams.
+    """
+    segment = cheetah.Segment(
+        elements=[
+            cheetah.Drift(length=torch.tensor(0.5)),
+            cheetah.Quadrupole(length=torch.tensor(0.3)),
+            cheetah.Drift(length=torch.tensor(0.2)),
+        ]
+    )
+    incoming_beam = cheetah.ParameterBeam.from_astra(
+        "tests/resources/ACHIP_EA1_2021.1351.001"
+    )
+
+    longitudinal_beam_generator = segment.beam_along_segment_generator(
+        incoming=incoming_beam
+    )
+    longitudinal_beam_list = list(longitudinal_beam_generator)
+
+    assert len(longitudinal_beam_list) == 4
+    for longitudinal_beam in longitudinal_beam_list:
+        assert isinstance(longitudinal_beam, incoming_beam.__class__)
+
+
+def test_resolution_longitudinal_beam_generator():
+    """
+    Test that the longitudinal beam generator works as expected and generates the
+    correct number of beams with a specified resolution.
+    """
+    segment = cheetah.Segment(
+        elements=[
+            cheetah.Drift(length=torch.tensor(0.5)),
+            cheetah.Quadrupole(length=torch.tensor(0.3)),
+            cheetah.Drift(length=torch.tensor(0.2)),
+        ]
+    )
+    incoming_beam = cheetah.ParameterBeam.from_astra(
+        "tests/resources/ACHIP_EA1_2021.1351.001"
+    )
+
+    longitudinal_beam_generator = segment.beam_along_segment_generator(
+        incoming=incoming_beam, resolution=0.1
+    )
+    longitudinal_beam_list = list(longitudinal_beam_generator)
+
+    assert len(longitudinal_beam_list) == 11
+    for longitudinal_beam in longitudinal_beam_list:
+        assert isinstance(longitudinal_beam, incoming_beam.__class__)
+
+
+@pytest.mark.parametrize("metric_names", ["beta_x", ("beta_x",), ("s", "beta_x")])
+def test_longitudinal_beam_metric(metric_names):
path: tests/test_segment.py
user: login: Hespe
body: Since we decided to rename the function from `metric` to `attrs`, the test and its parameter should probably also be renamed.
created_at: 2025-05-08T07:42:09Z
author_association: MEMBER
diff_hunk: @@ -454,6 +454,76 @@ def split(self, resolution: torch.Tensor) -> list[Element]:
             for split_element in element.split(resolution)
         ]
 
+    def beam_along_segment_generator(
+        self, incoming: Beam, resolution: float | None = None
+    ) -> Iterator[Beam]:
+        """
+        Generator for beam objects along the segment either at the end of each element
+        or at a given resolution.
+
+        :param incoming: Beam that is entering the segment from upstream for which the
+            trajectory is computed.
+        :param resolution: Requested resolution of trajectory. Note that not all
+            elements can be split at arbitrary resolutions, which can lead to deviations
+            from the requested resolution. If `None` is passed, samples are taken at the
+            end of each element.
+        :return: Generator that yields the beam objects along the segment.
+        """
+        # If a resolution is passed, run this method for the split Segment
+        if resolution is not None:
+            yield from self.__class__(
+                elements=self.split(resolution), name=f"{self.name}_split"
+            ).beam_along_segment_generator(incoming)
+        else:
+            yield incoming
+            for element in self.elements:
+                outgoing = element.track(incoming)
+                yield outgoing
+                incoming = outgoing
+
+    def get_beam_attrs_along_segment(
+        self,
+        attr_names: tuple[str, ...] | str,
+        incoming: Beam,
+        resolution: float | None = None,
+    ) -> tuple[torch.Tensor, ...] | torch.Tensor:
+        """
+        Get metrics along the segment either at the end of each element or at a given
+        resolution.
+
+        :param attr_names: Metrics to compute. Can be a single metric or a tuple of
+            metrics. Supported metrics are any property of beam class of `incoming`.
+        :param incoming: Beam that is entering the segment from upstream for which the
+            trajectory is computed.
+        :param resolution: Requested resolution of trajectory. Note that not all
+            elements can be split at arbitrary resolutions, which can lead to deviations
+            from the requested resolution. If `None` is passed, samples are taken at the
+            end of each element.
+        :return: Tuple of tensors containing the requested metrics along the segment.
+        """
+        attr_name_tuple = attr_names if isinstance(attr_names, tuple) else (attr_names,)
+
+        results = zip(
+            *[
+                tuple(getattr(beam, attr_name) for attr_name in attr_name_tuple)
+                for beam in self.beam_along_segment_generator(
+                    incoming, resolution=resolution
+                )
+            ]
path: cheetah/accelerator/segment.py
user: login: Hespe
body: Just to be sure. Using `*[...]` does have the consequence of having to allocate an intermediate list, right? While using `*(...)` would just create another generator? If that is correct, the question would be if that list is actually visible from outside the function or just an implementation detail.
created_at: 2025-05-08T07:51:27Z
author_association: MEMBER
diff_hunk: @@ -125,7 +126,9 @@ def is_active(self) -> bool:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         raise NotImplementedError
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
         raise NotImplementedError
path: cheetah/accelerator/sextupole.py
user: login: jank324
body: 🥸
created_at: 2025-05-08T08:50:14Z
author_association: MEMBER
diff_hunk: @@ -125,7 +126,9 @@ def is_active(self) -> bool:
     def split(self, resolution: torch.Tensor) -> list[Element]:
         raise NotImplementedError
 
-    def plot(self, ax: plt.Axes, s: float, vector_idx: tuple | None = None) -> None:
+    def plot(
+        self, s: float, vector_idx: tuple | None = None, ax: plt.Axes | None = None
+    ) -> plt.Axes:
         raise NotImplementedError
path: cheetah/accelerator/sextupole.py
user: login: jank324
body: I opened an issue #437.
created_at: 2025-05-08T08:52:24Z
author_association: MEMBER

PR comments:
user: login: jank324
created_at: 2025-05-02T16:11:24Z
body: ### TODOs
- [ ] ~~Modify the new function to broadcast~~
- [x] Refactor existing plotting function like `Segment.plot_twiss` to use the new function
author_association: MEMBER
user: login: jank324
created_at: 2025-05-07T12:00:37Z
body: At this stage of this PR, to get a bunch of properties from the lattice, you would do this:
```python
cheetah_s_positions, cheetah_beta_x, cheetah_beta_y = zip(
    *(
        (beam.s, beam.beta_x, beam.beta_y)
        for beam in segment.longitudinal_beam_generator(
            subsampled_design_incoming_particle_beam
        )
    )
)
```
and then probably plot them with 
```python
plt.figure(figsize=(30, 10))
plt.subplot(3, 1, 1)
plt.plot(cheetah_s_positions, cheetah_beta_x, label="Cheetah")
plt.legend()
plt.subplot(3, 1, 2)
plt.plot(cheetah_s_positions, cheetah_beta_y, label="Cheetah")
plt.legend()
plt.subplot(3, 1, 3)
ax = plt.gca()
segment.plot(ax=ax, s=0.0)
plt.show()
```

That's a little more convenient, but not really that convenient. I guess this PR should add another function such that you could do
```python
cheetah_s_positions, cheetah_beta_x, cheetah_beta_y = segment.the_cool_new_function(("s", "beta_x", "beta_y"))
```

Maybe this PR should even extend/refactor the `plot_overview` function on top of that?
```python
segment.plot_overview(metrics=("beta_x", "beta_y")
```
author_association: MEMBER
user: login: jank324
created_at: 2025-05-07T12:47:11Z
body: Some questions:
 - [x] Is `get_longitudinal_metrics` really the right name?
 - [x] Right now the convenient method broadcasts the samples of each metric intra-metric to allow it to build tensors. Is this the interface we want? Or should broadcasting also happen inter-metric (and even build one large tensor?)?
author_association: MEMBER
user: login: jank324
created_at: 2025-05-07T15:02:50Z
body: > Some questions:
> 
> * [x]  Is `get_longitudinal_metrics` really the right name?
> * [ ]  Right now the convenient method broadcasts the samples of each metric intra-metric to allow it to build tensors. Is this the interface we want? Or should broadcasting also happen inter-metric (and even build one large tensor?)?

It might actually make sense that if you request multiple attributes in one call, they will be broadcasted together for your convenience.
author_association: MEMBER
user: login: Hespe
created_at: 2025-05-08T08:22:00Z
body: > > Some questions:
> > 
> > * [x]   Is `get_longitudinal_metrics` really the right name?
> > * [ ]   Right now the convenient method broadcasts the samples of each metric intra-metric to allow it to build tensors. Is this the interface we want? Or should broadcasting also happen inter-metric (and even build one large tensor?)?
> 
> It might actually make sense that if you request multiple attributes in one call, they will be broadcasted together for your convenience.

I played around with the function a bit and noticed that the common broadcasting of attributes can be problematic. The following code crashes right now:

```python
import cheetah
import torch

segment = cheetah.Segment(elements=[cheetah.Drift(length=torch.tensor([1.0, 2.0]))])
beam = cheetah.ParticleBeam.from_parameters()

segment.get_beam_attrs_along_segment(("x", "mu_x"), beam)
``` 

The problem is that `x` by itself is already a vector even for non-vectorized segments. And even if the segment would not be vectorized, `mu_x` would get broadcasted somewhat unexpectedly.
author_association: MEMBER
user: login: jank324
created_at: 2025-05-08T11:38:45Z
body: Temporarily pinning `snowballstemmer<3.0` for the docs build because of an issue with the latest release. Will be unpinned when yanked or updated. Refer to https://github.com/sphinx-doc/sphinx/issues/13533 and https://github.com/snowballstem/snowball/issues/229.
author_association: MEMBER