From 073edb4e1019ed6cb02136f4aeef1bcc00f4a4c0 Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:07:06 -0700 Subject: [PATCH 1/9] add vectorized histogramming --- cheetah/accelerator/screen.py | 50 +++++++++++------- cheetah/utils/__init__.py | 1 + cheetah/utils/histogram.py | 96 +++++++++++++++++++++++++++++++++++ 3 files changed, 129 insertions(+), 18 deletions(-) create mode 100644 cheetah/utils/histogram.py diff --git a/cheetah/accelerator/screen.py b/cheetah/accelerator/screen.py index 2d5fc8e11..64c4ec80c 100644 --- a/cheetah/accelerator/screen.py +++ b/cheetah/accelerator/screen.py @@ -12,6 +12,7 @@ cache_transfer_map, cloud_in_cell_charge_deposition, kde_histogram_2d, + vectorized_histogram_2d, ) generate_unique_name = UniqueNameGenerator(prefix="unnamed_element") @@ -290,25 +291,38 @@ def reading(self) -> torch.Tensor: image = dist.log_prob(pos).exp().mT elif isinstance(read_beam, ParticleBeam): if self.method == "histogram": - # Catch vectorisation, which is currently not supported by "histogram" - if ( - len(read_beam.particles.shape) > 2 - or len(read_beam.particle_charges.shape) > 1 - or len(read_beam.energy.shape) > 0 - ): - raise NotImplementedError( - "The `'histogram'` method of `Screen` does not support " - "vectorization. Use `'kde'` instead. If this is a feature you " - "would like to see, please open an issue on GitHub." - ) - - image_transposed, _ = torch.histogramdd( - torch.stack((read_beam.x, read_beam.y)).mT, - bins=self.pixel_bin_edges, - weight=read_beam.particle_charges.abs() - * read_beam.survival_probabilities, + # # Catch vectorisation, which is currently not supported by "histogram" + # if ( + # len(read_beam.particles.shape) > 2 + # or len(read_beam.particle_charges.shape) > 1 + # or len(read_beam.energy.shape) > 0 + # ): + # raise NotImplementedError( + # "The `'histogram'` method of `Screen` does not support " + # "vectorization. Use `'kde'` instead. If this is a feature you " + # "would like to see, please open an issue on GitHub." + # ) + + # image_transposed, _ = torch.histogramdd( + # torch.stack((read_beam.x, read_beam.y)).mT, + # bins=self.pixel_bin_edges, + # weight=read_beam.particle_charges.abs() + # * read_beam.survival_probabilities, + # ) + # image = image_transposed.mT + weights = ( + read_beam.particle_charges.abs() * read_beam.survival_probabilities + ) + broadcasted_x, broadcasted_y, broadcasted_weights = ( + torch.broadcast_tensors(read_beam.x, read_beam.y, weights) ) - image = image_transposed.mT + image = vectorized_histogram_2d( + x1=broadcasted_x, + x2=broadcasted_y, + bins1=self.pixel_bin_centers[0], + bins2=self.pixel_bin_centers[1], + weights=broadcasted_weights, + ).mT elif self.method == "kde": weights = ( read_beam.particle_charges.abs() * read_beam.survival_probabilities diff --git a/cheetah/utils/__init__.py b/cheetah/utils/__init__.py index 61c67d49b..8af8c1f02 100644 --- a/cheetah/utils/__init__.py +++ b/cheetah/utils/__init__.py @@ -4,6 +4,7 @@ from .device import is_mps_available_and_functional # noqa: F401 from .elementwise_linspace import elementwise_linspace # noqa: F401 from .kde import kde_histogram_1d, kde_histogram_2d # noqa: F401 +from .histogram import vectorized_histogram_2d # noqa: F401 from .physics import compute_relativistic_factors # noqa: F401 from .plot import ( # noqa: F401 format_axis_as_percentage, diff --git a/cheetah/utils/histogram.py b/cheetah/utils/histogram.py new file mode 100644 index 000000000..880b0d5c9 --- /dev/null +++ b/cheetah/utils/histogram.py @@ -0,0 +1,96 @@ +import torch + +def _bin_centers_to_edges(centers: torch.Tensor) -> torch.Tensor: + """Convert a 1D tensor of evenly-spaced bin centers to bin edges.""" + if centers.numel() < 2: + raise ValueError("Need at least 2 bin centers to infer edges.") + + step = centers[1] - centers[0] + left_edge = centers[0] - step / 2 + right_edge = centers[-1] + step / 2 + midpoints = (centers[1:] + centers[:-1]) / 2 + return torch.cat([left_edge.unsqueeze(0), midpoints, right_edge.unsqueeze(0)]) + + +def vectorized_histogram_2d( + x1: torch.Tensor, + x2: torch.Tensor, + bins1: torch.Tensor, + bins2: torch.Tensor, + weights: torch.Tensor | None = None, +) -> torch.Tensor: + """ + Compute batched 2D histograms for coordinate pairs ``(x1, x2)``. + + Returns a single tensor of counts (rather than a KDE-smoothed density) so + it is a drop-in alternative to :func:`kde_histogram_2d` in + :meth:`Screen.reading`. Like :func:`kde_histogram_2d`, any number of + leading batch dimensions is supported (e.g. an ensemble beam shaped + ``(n_draws, n_samples, n_particles)``), not just a single batch dim. + + :param x1: Input tensor of shape ``(*batch_shape, N)`` (``batch_shape`` may + be empty for a single unbatched set of ``N`` points). + :param x2: Tensor broadcastable with ``x1``. + :param bins1: Bin centers for the first axis, shape ``(N_bins1,)``. + :param bins2: Bin centers for the second axis, shape ``(N_bins2,)``. + :param weights: Optional weights, broadcastable with ``x1`` and ``x2``. + :returns: Histogram of shape ``(*broadcast_batch_shape, N_bins1, N_bins2)``. + + .. note:: + ``x1`` and ``x2`` are paired elementwise (the value at a given index in + ``x1`` is binned jointly with the value at the same index in ``x2``), + unlike :func:`kde_histogram_2d`, which pairs them implicitly via a + batched matmul contraction. Because pairing happens via + :func:`torch.broadcast_tensors`, the trailing (particle-count) dimension + of ``x1`` and ``x2`` is also broadcast, not just the leading batch + dims. So ``N=1`` on one side is technically allowed and will broadcast + that single value against every point on the other side - a degenerate + edge case (pairing one fixed point against many) rather than a + meaningful joint histogram of two particle ensembles. + """ + if weights is None: + weights = torch.ones_like(x1) + # `broadcast_tensors` returns expanded, non-contiguous (stride-0) views, so + # contiguity has to be enforced *after* this call (see the `.reshape` calls + # below), not before it - otherwise `torch.bucketize` silently makes its + # own contiguous copy and warns. + x1, x2, weights = torch.broadcast_tensors(x1, x2, weights) + + batch_shape = x1.shape[:-1] + N = x1.shape[-1] + + # Flatten all leading batch dims into one so a single `bincount` call can + # produce every batch element's histogram at once; `.contiguous()` is + # required here (not just `.reshape`, which is a no-op view when the shape + # doesn't actually change) because `torch.bucketize` silently makes its own + # contiguous copy - and warns - if given a non-contiguous/broadcast input. + x1_flat = x1.reshape(-1, N).contiguous() + x2_flat = x2.reshape(-1, N).contiguous() + weights_flat = weights.reshape(-1, N) + + B = x1_flat.shape[0] + device = x1_flat.device + dtype = x1_flat.dtype + + x1_edges = _bin_centers_to_edges(bins1) + x2_edges = _bin_centers_to_edges(bins2) + bins_x1 = x1_edges.numel() - 1 + bins_x2 = x2_edges.numel() - 1 + + ix1 = (torch.bucketize(x1_flat, x1_edges) - 1).clamp(0, bins_x1 - 1).long() # (B, N) + ix2 = (torch.bucketize(x2_flat, x2_edges) - 1).clamp(0, bins_x2 - 1).long() # (B, N) + + idx_flat = ix1 * bins_x2 + ix2 # (B, N) + + offset = torch.arange(B, device=device, dtype=idx_flat.dtype) * ( + bins_x1 * bins_x2 + ) + idx_flat_offset = (idx_flat + offset.unsqueeze(1)).reshape(-1) + + weights_flat = weights_flat.reshape(-1).to(dtype) + hist_flat = torch.bincount( + idx_flat_offset, weights=weights_flat, minlength=B * bins_x1 * bins_x2 + ).to(dtype) + hist = hist_flat.view(*batch_shape, bins_x1, bins_x2) + + return hist \ No newline at end of file From c73222538ed178924fb0fd38d446176419bab34b Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:59:46 -0700 Subject: [PATCH 2/9] fix formatting --- cheetah/utils/histogram.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/cheetah/utils/histogram.py b/cheetah/utils/histogram.py index 880b0d5c9..5a4f0efba 100644 --- a/cheetah/utils/histogram.py +++ b/cheetah/utils/histogram.py @@ -1,5 +1,6 @@ import torch + def _bin_centers_to_edges(centers: torch.Tensor) -> torch.Tensor: """Convert a 1D tensor of evenly-spaced bin centers to bin edges.""" if centers.numel() < 2: @@ -77,14 +78,16 @@ def vectorized_histogram_2d( bins_x1 = x1_edges.numel() - 1 bins_x2 = x2_edges.numel() - 1 - ix1 = (torch.bucketize(x1_flat, x1_edges) - 1).clamp(0, bins_x1 - 1).long() # (B, N) - ix2 = (torch.bucketize(x2_flat, x2_edges) - 1).clamp(0, bins_x2 - 1).long() # (B, N) + ix1 = ( + (torch.bucketize(x1_flat, x1_edges) - 1).clamp(0, bins_x1 - 1).long() + ) # (B, N) + ix2 = ( + (torch.bucketize(x2_flat, x2_edges) - 1).clamp(0, bins_x2 - 1).long() + ) # (B, N) idx_flat = ix1 * bins_x2 + ix2 # (B, N) - offset = torch.arange(B, device=device, dtype=idx_flat.dtype) * ( - bins_x1 * bins_x2 - ) + offset = torch.arange(B, device=device, dtype=idx_flat.dtype) * (bins_x1 * bins_x2) idx_flat_offset = (idx_flat + offset.unsqueeze(1)).reshape(-1) weights_flat = weights_flat.reshape(-1).to(dtype) @@ -93,4 +96,4 @@ def vectorized_histogram_2d( ).to(dtype) hist = hist_flat.view(*batch_shape, bins_x1, bins_x2) - return hist \ No newline at end of file + return hist From 8c42926ba9a9bd3c0c8eca52cf5c269bafb38a2f Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:03:49 -0700 Subject: [PATCH 3/9] update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f6cc846e..b413d67fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ - Add the flag `skip_superimposed` to `segment.flattened()` which allows the user to flatten segments without flattening superimposed elements (see #664) (@roussel-ryan) - `TransverseDeflectingCavity` now supports `linear` as a `tracking_method`. (see #678) (@cr-xu) +- Add vectorized histogramming (see #680) (@jp-ga) ### 🐛 Bug fixes From a5d2b9eb5b0341157c6182185fa6e91d63bdf857 Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:07:09 -0700 Subject: [PATCH 4/9] isort --- cheetah/utils/__init__.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/cheetah/utils/__init__.py b/cheetah/utils/__init__.py index 0963c108d..e378d7020 100644 --- a/cheetah/utils/__init__.py +++ b/cheetah/utils/__init__.py @@ -3,25 +3,23 @@ from .cloud_in_cell import cloud_in_cell_charge_deposition # noqa: F401 from .device import is_mps_available_and_functional # noqa: F401 from .elementwise_linspace import elementwise_linspace # noqa: F401 -from .kde import kde_histogram_1d, kde_histogram_2d # noqa: F401 from .histogram import vectorized_histogram_2d # noqa: F401 +from .kde import kde_histogram_1d, kde_histogram_2d # noqa: F401 from .names import UniqueNameGenerator, merge_element_names # noqa: F401 from .physics import compute_relativistic_factors # noqa: F401 -from .plot import ( # noqa: F401 - format_axis_as_percentage, - format_axis_with_prefixed_unit, -) -from .statistics import ( # noqa: F401 - match_distribution_moments, +from .plot import format_axis_as_percentage # noqa: F401 +from .plot import format_axis_with_prefixed_unit +from .statistics import match_distribution_moments # noqa: F401 +from .statistics import ( unbiased_weighted_covariance, unbiased_weighted_covariance_matrix, unbiased_weighted_std, unbiased_weighted_variance, ) from .vector import squash_index_for_unavailable_dims # noqa: F401 -from .warnings import ( # noqa: F401 +from .warnings import DirtyNameWarning # noqa: F401 +from .warnings import ( DefaultParameterWarning, - DirtyNameWarning, NoBeamPropertiesInLatticeWarning, NotUnderstoodPropertyWarning, PhysicsWarning, From ee36088248ac8658ef9c6d62fc94ea5444dfb5bd Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:20:10 -0700 Subject: [PATCH 5/9] flake8 --- cheetah/accelerator/screen.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/cheetah/accelerator/screen.py b/cheetah/accelerator/screen.py index 64c4ec80c..f26425a81 100644 --- a/cheetah/accelerator/screen.py +++ b/cheetah/accelerator/screen.py @@ -291,16 +291,18 @@ def reading(self) -> torch.Tensor: image = dist.log_prob(pos).exp().mT elif isinstance(read_beam, ParticleBeam): if self.method == "histogram": - # # Catch vectorisation, which is currently not supported by "histogram" + # # Catch vectorisation, which is currently not supported by + # # "histogram" # if ( # len(read_beam.particles.shape) > 2 # or len(read_beam.particle_charges.shape) > 1 # or len(read_beam.energy.shape) > 0 # ): # raise NotImplementedError( - # "The `'histogram'` method of `Screen` does not support " - # "vectorization. Use `'kde'` instead. If this is a feature you " - # "would like to see, please open an issue on GitHub." + # "The `'histogram'` method of `Screen` does not " + # "support vectorization. Use `'kde'` instead. If " + # "this is a feature you would like to see, please " + # "open an issue on GitHub." # ) # image_transposed, _ = torch.histogramdd( From c5ad880fd1ca93c78503eb4984ea43eadc9562b6 Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:30:45 -0700 Subject: [PATCH 6/9] delete commented code --- cheetah/accelerator/screen.py | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/cheetah/accelerator/screen.py b/cheetah/accelerator/screen.py index f26425a81..18341225a 100644 --- a/cheetah/accelerator/screen.py +++ b/cheetah/accelerator/screen.py @@ -291,27 +291,6 @@ def reading(self) -> torch.Tensor: image = dist.log_prob(pos).exp().mT elif isinstance(read_beam, ParticleBeam): if self.method == "histogram": - # # Catch vectorisation, which is currently not supported by - # # "histogram" - # if ( - # len(read_beam.particles.shape) > 2 - # or len(read_beam.particle_charges.shape) > 1 - # or len(read_beam.energy.shape) > 0 - # ): - # raise NotImplementedError( - # "The `'histogram'` method of `Screen` does not " - # "support vectorization. Use `'kde'` instead. If " - # "this is a feature you would like to see, please " - # "open an issue on GitHub." - # ) - - # image_transposed, _ = torch.histogramdd( - # torch.stack((read_beam.x, read_beam.y)).mT, - # bins=self.pixel_bin_edges, - # weight=read_beam.particle_charges.abs() - # * read_beam.survival_probabilities, - # ) - # image = image_transposed.mT weights = ( read_beam.particle_charges.abs() * read_beam.survival_probabilities ) From 6c90db226b3ada9be3a93c2b9bee2cc6f590cb06 Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:47:24 -0700 Subject: [PATCH 7/9] fix linting --- cheetah/utils/__init__.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/cheetah/utils/__init__.py b/cheetah/utils/__init__.py index e378d7020..e9726db42 100644 --- a/cheetah/utils/__init__.py +++ b/cheetah/utils/__init__.py @@ -7,19 +7,21 @@ from .kde import kde_histogram_1d, kde_histogram_2d # noqa: F401 from .names import UniqueNameGenerator, merge_element_names # noqa: F401 from .physics import compute_relativistic_factors # noqa: F401 -from .plot import format_axis_as_percentage # noqa: F401 -from .plot import format_axis_with_prefixed_unit -from .statistics import match_distribution_moments # noqa: F401 -from .statistics import ( +from .plot import ( # noqa: F401 + format_axis_as_percentage, + format_axis_with_prefixed_unit, +) +from .statistics import ( # noqa: F401 + match_distribution_moments, unbiased_weighted_covariance, unbiased_weighted_covariance_matrix, unbiased_weighted_std, unbiased_weighted_variance, ) from .vector import squash_index_for_unavailable_dims # noqa: F401 -from .warnings import DirtyNameWarning # noqa: F401 -from .warnings import ( +from .warnings import ( # noqa: F401 DefaultParameterWarning, + DirtyNameWarning, NoBeamPropertiesInLatticeWarning, NotUnderstoodPropertyWarning, PhysicsWarning, From 2812cc51b6a3d94db225f6eddd396af211f70718 Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:08:10 -0700 Subject: [PATCH 8/9] add test --- tests/test_vectorized.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_vectorized.py b/tests/test_vectorized.py index be72138e0..b76795836 100644 --- a/tests/test_vectorized.py +++ b/tests/test_vectorized.py @@ -303,7 +303,7 @@ def test_vectorized_solenoid(BeamClass): @pytest.mark.parametrize("BeamClass", [cheetah.ParticleBeam]) -@pytest.mark.parametrize("method", ["kde"]) # Currently only KDE supports vectorisation +@pytest.mark.parametrize("method", ["kde", "histogram"]) def test_vectorized_screen_2d(BeamClass, method): """ Test that a vectorized `Screen` is able to track a particle beam and produce a From bb2831b99e451e048e2bd1f7c5ddca76627311e7 Mon Sep 17 00:00:00 2001 From: jp-ga <97850904+jp-ga@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:05:46 -0700 Subject: [PATCH 9/9] add all histogram types to vectorized test --- tests/test_vectorized.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_vectorized.py b/tests/test_vectorized.py index b76795836..b30a76407 100644 --- a/tests/test_vectorized.py +++ b/tests/test_vectorized.py @@ -303,7 +303,7 @@ def test_vectorized_solenoid(BeamClass): @pytest.mark.parametrize("BeamClass", [cheetah.ParticleBeam]) -@pytest.mark.parametrize("method", ["kde", "histogram"]) +@pytest.mark.parametrize("method", ["histogram", "kde", "cloud-in-cell"]) def test_vectorized_screen_2d(BeamClass, method): """ Test that a vectorized `Screen` is able to track a particle beam and produce a