Vectorised beam plotting by jp-ga · Pull Request #584 · desy-ml/cheetah

jp-ga · 2025-10-17T02:28:04Z

Description

Add the following plotting functionalities for beams with vector dimensions:

Plot average 1d projection histogram with lower and upper bounds
Plot average 2d projection histogram with lower and upper bounds
Triangle plot with the above functionalities

Motivation and Context

I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

Bug fix (non-breaking change which fixes an issue)
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

I have updated the changelog accordingly (required).
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.
I have reformatted the code and checked that formatting passes (required).
I have have fixed all issues found by flake8 (required).
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).
I have checked that the documentation builds (required).

Note: We are using a maximum length of 88 characters per line.

…ies`

jp-ga · 2025-10-17T04:16:25Z

Example contour plot:

Example histogram plot:

jp-ga · 2025-10-17T04:18:16Z

@roussel-ryan @cr-xu this PR might be of interest for you. Let me know your thoughts.

cr-xu · 2025-10-20T22:25:14Z

This looks great! Thanks for the feature!

…ality

… support since NumPy 2.0

… is only support since NumPy 2.0" This reverts commit 9312eaf.

Copilot

Pull request overview

This pull request adds ensemble plotting capabilities for vectorized ParticleBeam instances, enabling visualization of mean distributions with uncertainty bounds. The changes introduce new statistical utilities for computing histograms over vectorized distributions and extend the existing plotting methods to display confidence intervals.

Changes:

Added vectorized histogram computation functions (vectorized_histogram_1d, vectorized_histogram_2d) that efficiently compute histograms across multiple beam instances
Added functions to compute mean histograms with confidence bounds (distribution_histogram_and_confidence_1d/2d, histograms_mean_and_confidence) supporting three uncertainty methods: standard deviation (sd), standard error (se), and percentile intervals (pi)
Extended plot_1d_distribution, plot_2d_distribution, and plot_distribution methods to support uncertainty visualization for vectorized beams
Updated minimum dependency versions: NumPy 1.23.3→2.0.0, Matplotlib 3.5.0→3.9.0, SciPy 1.10.1→1.13.0

Reviewed changes

Copilot reviewed 8 out of 9 changed files in this pull request and generated 12 comments.

Show a summary per file

File	Description
cheetah/utils/statistics.py	Adds five new functions for vectorized histogram computation and confidence interval calculation over beam ensembles
cheetah/utils/init.py	Exports the new statistical functions for public API access
cheetah/particles/particle_beam.py	Refactors plotting methods to use new histogram functions and adds support for uncertainty bands via errorbar parameter
tests/test_statistics.py	Adds comprehensive tests for new histogram and confidence functions with various vector shapes and errorbar methods
tests/test_plotting.py	Adds test for vectorized beam distribution plotting with both histogram and contour styles
setup.py	Updates minimum required versions for matplotlib, numpy, and scipy dependencies
test_minimum_requirements_constraints.txt	Updates test constraints to match new minimum dependency versions
CHANGELOG.md	Documents the new plotting functionality in v0.8.0 release notes

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-02-24T15:56:16Z

+            smoothed_lower_bound = gaussian_filter(lower_bound, smoothing)
+            smoothed_upper_bound = gaussian_filter(upper_bound, smoothing)


The gaussian_filter function from scipy.ndimage expects NumPy arrays, but lower_bound and upper_bound are PyTorch tensors returned from distribution_histogram_and_confidence_1d. These tensors need to be converted to NumPy arrays before being passed to gaussian_filter, and the results should be converted back to tensors if needed for consistency with the rest of the code. The same issue applies to histogram on line 1299.

Copilot · 2026-02-24T15:56:17Z

+            ax.fill_between(
+                bin_centers,
+                smoothed_lower_bound,
+                smoothed_upper_bound,
+                **({"color": "C1", "alpha": 0.5} | (fill_between_kws or {})),
+            )

-        ax.plot(
-            centers,
-            histogram / histogram.max(),
-            **{"color": "black"} | (plot_kws or {}),
-        )
-        ax.set_xlabel(f"{self.PRETTY_DIMENSION_LABELS[dimension]}")
+        smoothed_histogram = gaussian_filter(histogram, smoothing)
+        ax.plot(bin_centers, smoothed_histogram, **({"color": "C0"} | (plot_kws or {})))


Matplotlib's plotting functions typically expect NumPy arrays. While bin_centers is a PyTorch tensor, it should be converted to NumPy (using .numpy() or .cpu().numpy()) before being passed to ax.fill_between and ax.plot to ensure compatibility.

Copilot · 2026-02-24T15:56:17Z

            ax.pcolormesh(
-                x_edges,
-                y_edges,
-                clipped_histogram.T / smoothed_histogram.max(),
-                **{"cmap": "rainbow"} | (pcolormesh_kws or {}),
+                bin_centers_x,
+                bin_centers_y,
+                smoothed_histogram.mT,
+                **({"cmap": "rainbow"} | (pcolormesh_kws or {})),
            )
        elif style == "contour":
-            contour_histogram = gaussian_filter(histogram, contour_smoothing)
-
-            ax.contour(
-                x_centers,
-                y_centers,
-                contour_histogram.T / contour_histogram.max(),
-                **{"levels": 3} | (contour_kws or {}),
+            contour_set_of_mean = ax.contour(
+                bin_centers_x,
+                bin_centers_y,
+                smoothed_histogram.mT,
+                **({"levels": 3} | (distribution_contour_kws or {})),
            )

+            if lower_bound is not None and upper_bound is not None:
+                smoothed_lower_bound = gaussian_filter(lower_bound, smoothing)
+                smoothed_upper_bound = gaussian_filter(upper_bound, smoothing)
+
+                ax.contour(
+                    bin_centers_x,
+                    bin_centers_y,
+                    smoothed_lower_bound.mT,
+                    **(
+                        {"levels": contour_set_of_mean.levels, "linestyles": "--"}
+                        | (confidence_contour_kws or {})
+                    ),
+                )
+                ax.contour(
+                    bin_centers_x,
+                    bin_centers_y,
+                    smoothed_upper_bound.mT,
+                    **(
+                        {"levels": contour_set_of_mean.levels, "linestyles": "--"}
+                        | (confidence_contour_kws or {})
+                    ),
+                )


Matplotlib's pcolormesh and contour functions typically expect NumPy arrays. The bin_centers_x, bin_centers_y, and histogram tensors should be converted to NumPy arrays before being passed to these plotting functions to ensure compatibility.

Copilot · 2026-02-24T15:56:19Z

+    bin_indicies_x = torch.bucketize(x_flat.contiguous(), boundaries_x)
+    bin_indicies_y = torch.bucketize(y_flat.contiguous(), boundaries_y)
+
+    # Flatten 2-dimensional bin indices to 1 dimension
+    bin_indicies_flat = bin_indicies_x * bins[1] + bin_indicies_y
+
+    # Flatten batch with offsets
+    vector_offsets = torch.arange(num_vector_elements, device=x.device) * (
+        bins[0] * bins[1]
+    )
+    bin_indicies_flat = (bin_indicies_flat + vector_offsets.unsqueeze(1)).flatten()
+
+    # Count occurrences
+    histogram_flat = torch.bincount(
+        bin_indicies_flat, minlength=num_vector_elements * bins[0] * bins[1]


Typo: 'bin_indicies' should be 'bin_indices' (correct spelling). This typo appears throughout the function.

Suggested change

bin_indicies_x = torch.bucketize(x_flat.contiguous(), boundaries_x)

bin_indicies_y = torch.bucketize(y_flat.contiguous(), boundaries_y)

# Flatten 2-dimensional bin indices to 1 dimension

bin_indicies_flat = bin_indicies_x * bins[1] + bin_indicies_y

# Flatten batch with offsets

vector_offsets = torch.arange(num_vector_elements, device=x.device) * (

bins[0] * bins[1]

)

bin_indicies_flat = (bin_indicies_flat + vector_offsets.unsqueeze(1)).flatten()

# Count occurrences

histogram_flat = torch.bincount(

bin_indicies_flat, minlength=num_vector_elements * bins[0] * bins[1]

bin_indices_x = torch.bucketize(x_flat.contiguous(), boundaries_x)

bin_indices_y = torch.bucketize(y_flat.contiguous(), boundaries_y)

# Flatten 2-dimensional bin indices to 1 dimension

bin_indices_flat = bin_indices_x * bins[1] + bin_indices_y

# Flatten batch with offsets

vector_offsets = torch.arange(num_vector_elements, device=x.device) * (

bins[0] * bins[1]

)

bin_indices_flat = (bin_indices_flat + vector_offsets.unsqueeze(1)).flatten()

# Count occurrences

histogram_flat = torch.bincount(

bin_indices_flat, minlength=num_vector_elements * bins[0] * bins[1]

Copilot · 2026-02-24T15:56:19Z

-            **{"color": "black"} | (plot_kws or {}),
-        )
-        ax.set_xlabel(f"{self.PRETTY_DIMENSION_LABELS[dimension]}")
+        smoothed_histogram = gaussian_filter(histogram, smoothing)


The gaussian_filter function from scipy.ndimage expects NumPy arrays, but histogram is a PyTorch tensor returned from distribution_histogram_and_confidence_1d. This tensor needs to be converted to a NumPy array before being passed to gaussian_filter.

Copilot · 2026-02-24T15:56:19Z

+        smoothed_histogram = gaussian_filter(mean_histogram, smoothing)
+
        if style == "histogram":
            ax.pcolormesh(
-                x_edges,
-                y_edges,
-                clipped_histogram.T / smoothed_histogram.max(),
-                **{"cmap": "rainbow"} | (pcolormesh_kws or {}),
+                bin_centers_x,
+                bin_centers_y,
+                smoothed_histogram.mT,
+                **({"cmap": "rainbow"} | (pcolormesh_kws or {})),
            )
        elif style == "contour":
-            contour_histogram = gaussian_filter(histogram, contour_smoothing)
-
-            ax.contour(
-                x_centers,
-                y_centers,
-                contour_histogram.T / contour_histogram.max(),
-                **{"levels": 3} | (contour_kws or {}),
+            contour_set_of_mean = ax.contour(
+                bin_centers_x,
+                bin_centers_y,
+                smoothed_histogram.mT,
+                **({"levels": 3} | (distribution_contour_kws or {})),
            )

+            if lower_bound is not None and upper_bound is not None:
+                smoothed_lower_bound = gaussian_filter(lower_bound, smoothing)
+                smoothed_upper_bound = gaussian_filter(upper_bound, smoothing)


The gaussian_filter function from scipy.ndimage expects NumPy arrays, but mean_histogram, lower_bound, and upper_bound are PyTorch tensors returned from distribution_histogram_and_confidence_2d. These tensors need to be converted to NumPy arrays before being passed to gaussian_filter.

Copilot · 2026-02-24T15:56:19Z

+        bin_ranges = (
+            (float(x.min()), float(x.max())),
+            (float(y.min()), float(y.max())),


When all x or y values are identical, x.min() equals x.max() (or y.min() equals y.max()), resulting in zero-width bin ranges. This could cause issues with torch.linspace and histogram binning. Consider adding a check to handle this edge case by expanding the range slightly or raising an informative error.

Suggested change

bin_ranges = (

(float(x.min()), float(x.max())),

(float(y.min()), float(y.max())),

x_min = float(x.min())

x_max = float(x.max())

y_min = float(y.min())

y_max = float(y.max())

# Handle edge case where all x or all y values are identical, which would

# otherwise produce zero-width bin ranges.

if x_min == x_max:

if torch.is_floating_point(x):

eps_x = float(torch.finfo(x.dtype).eps) or 1e-6

else:

# For integer types, expand by half a unit on each side.

eps_x = 0.5

x_min -= eps_x

x_max += eps_x

if y_min == y_max:

if torch.is_floating_point(y):

eps_y = float(torch.finfo(y.dtype).eps) or 1e-6

else:

# For integer types, expand by half a unit on each side.

eps_y = 0.5

y_min -= eps_y

y_max += eps_y

bin_ranges = (

(x_min, x_max),

(y_min, y_max),

Copilot · 2026-02-24T15:56:20Z

+    )  # (num_vector_elements,)
+    bin_indicies_flat = (
+        bin_indicies + vector_offsets.unsqueeze(-1)
+    ).flatten()  # (num_vector_elements * num_bins,)


The comment states the shape is (num_vector_elements * num_bins,), but it should be (num_vector_elements * num_samples,) because we're flattening all the bin indices for all samples across all vector elements.

Suggested change

).flatten() # (num_vector_elements * num_bins,)

).flatten() # (num_vector_elements * num_samples,)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

cr-xu · 2026-03-18T22:51:38Z

Hmmm.... just tried out this PR, it doesn't look like the example plot JP showed above.
I couldn't really track what was changed

But for me the confidence_contour_kws doesn't seem to produce contour lines, and the histogram somehow still shows color at zero locations (tried set_under and manually passing the cmap but it didn't help)

jp-ga added 2 commits October 14, 2025 12:52

correct default shapes of particle_charges and `survival_probabilit…

9efcb77

…ies`

add ensemble plot functionality

c4ba854

jp-ga linked an issue Oct 17, 2025 that may be closed by this pull request

add beam ensemble plotting #583

Open

update changelog

48b67e1

jp-ga requested a review from jank324 October 17, 2025 02:34

jp-ga self-assigned this Oct 17, 2025

jp-ga added the enhancement New feature or request label Oct 17, 2025

jp-ga added 3 commits October 16, 2025 22:40

revert survival probabilities and charges default shaping

d8b5f30

move style to plot_2d_kws

7d57136

add ensemble plotting test

1a1bd25

jp-ga marked this pull request as ready for review October 17, 2025 04:13

jp-ga added 5 commits October 16, 2025 23:23

update changelog

04f97a8

update changelog

76bf864

fix docstring

e6ecefc

fix docstring

718bc21

fix docstring

3ed190b

jank324 requested changes Oct 21, 2025

View reviewed changes

Various cleanups in function signatures and docstrings

af21b42

jank324 reviewed Oct 21, 2025

View reviewed changes

Comment thread cheetah/particles/particle_beam.py Outdated

jank324 added 4 commits October 21, 2025 12:13

Cleanup test

42df81b

Cleanup test ammend

8f62df6

Slightly clearer line in test

27aac7a

Better test name

3cd56e6

jank324 reviewed Oct 21, 2025

View reviewed changes

Comment thread CHANGELOG.md Outdated

jank324 added 2 commits October 21, 2025 12:17

Apply suggestions from code review

a473af8

Various small cleanups

c2fb7fc

jank324 reviewed Oct 21, 2025

View reviewed changes

Comment thread cheetah/particles/ensemble_utils.py Outdated

jank324 added 12 commits December 2, 2025 16:27

Remove no longer avaialble method from docstrings and error messages

e7ece47

Fix slightling confusing argument names

1e26b67

Clearer changelog entry

8c35ef5

Fix PR ID in changelog entry

f1c39cf

Test some things

996df93

Minor code cleanup in notebook for testing

a026443

Fix smoothing in 1D histogram

9580388

Fix squiggly confidence contours

0d98b8d

Make bounds clearer as such

75f4604

Remove histogram from contour smoothing to maintain original function…

3c5fb3c

…ality

Make colour behaviour more predictable accordig to theme

31e6c15

Improve overall structure of 2D plot function

ab0bea1

jank324 reviewed Dec 9, 2025

View reviewed changes

Comment thread cheetah/particles/particle_beam.py Outdated

jank324 added 9 commits December 9, 2025 16:43

Automatic levels (slightly broken)

d4190e7

Match new code style on transpose

032fde7

Merge branch 'master' into 583-add-beam-ensemble-plotting

3a89304

Change .mT for NumPy arrays back to .T because the former is only…

9312eaf

… support since NumPy 2.0

Rerun ensemble plot test notebook with current state

17e5582

Revert "Change .mT for NumPy arrays back to .T because the former…

c0bcf07

… is only support since NumPy 2.0" This reverts commit 9312eaf.

Increase minimum NumPy requirement to 2.0.0

f719607

Up minimum scipy version to lowest compatible with NumPy 2

04c7eb7

Increase minimum Matplotlib version for NumPy 2 compatibility

849138d

jank324 requested a review from Copilot February 24, 2026 15:46

Copilot started reviewing on behalf of jank324 February 24, 2026 15:47 View session

Copilot AI reviewed Feb 24, 2026

View reviewed changes

Apply suggestion from Copilot code review

bedf00a

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

jank324 changed the title ~~583 add beam ensemble plotting~~ Vectorised beam plotting Feb 25, 2026

Merge branch 'master' into 583-add-beam-ensemble-plotting

c147f33

cr-xu mentioned this pull request Mar 19, 2026

Distribution Plot is not sharing all axes #615

Open

		smoothed_lower_bound = gaussian_filter(lower_bound, smoothing)
		smoothed_upper_bound = gaussian_filter(upper_bound, smoothing)

-        bin_ranges = (
-            (float(x.min()), float(x.max())),
-            (float(y.min()), float(y.max())),
+        x_min = float(x.min())
+        x_max = float(x.max())
+        y_min = float(y.min())
+        y_max = float(y.max())
+        # Handle edge case where all x or all y values are identical, which would
+        # otherwise produce zero-width bin ranges.
+        if x_min == x_max:
+            if torch.is_floating_point(x):
+                eps_x = float(torch.finfo(x.dtype).eps) or 1e-6
+            else:
+                # For integer types, expand by half a unit on each side.
+                eps_x = 0.5
+            x_min -= eps_x
+            x_max += eps_x
+        if y_min == y_max:
+            if torch.is_floating_point(y):
+                eps_y = float(torch.finfo(y.dtype).eps) or 1e-6
+            else:
+                # For integer types, expand by half a unit on each side.
+                eps_y = 0.5
+            y_min -= eps_y
+            y_max += eps_y
+        bin_ranges = (
+            (x_min, x_max),
+            (y_min, y_max),

	).flatten() # (num_vector_elements * num_bins,)
	).flatten() # (num_vector_elements * num_samples,)

Uh oh!

Conversation

jp-ga commented Oct 17, 2025 • edited by jank324 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

jp-ga commented Oct 17, 2025

Uh oh!

jp-ga commented Oct 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cr-xu commented Oct 20, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

cr-xu commented Mar 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jp-ga commented Oct 17, 2025 •

edited by jank324

Loading

jp-ga commented Oct 17, 2025 •

edited

Loading