From 6087b327d6917452c8ce4c090e33924debf8dd2a Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Jun 2025 10:06:35 +0200
Subject: [PATCH 1/3] Add method for differentiable 1D histogram

---
 cheetah/particles/particle_beam.py | 36 ++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py
index 5d66c7836..d0979cb90 100644
--- a/cheetah/particles/particle_beam.py
+++ b/cheetah/particles/particle_beam.py
@@ -17,6 +17,7 @@
     unbiased_weighted_covariance_matrix,
     unbiased_weighted_std,
     verify_device_and_dtype,
+    kde_histogram_1d,
 )
 
 
@@ -1220,6 +1221,41 @@ def to_xyz_pxpypz(self) -> torch.Tensor:
 
         return xp_coords
 
+    def get_1d_histogram(
+        self,
+        dimension: Literal["x", "px", "y", "py", "tau", "p"],
+        bins: int = 100,
+        bin_range: tuple[float] | None = None,
+        method: Literal["histogram", "kde"] = "kde",
+        kde_bandwidth: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Get a 1D (differentiable) histogram of the given dimension of the particle distribution.
+
+        TODO: Add documentation
+        """
+        # Extract particle quantities
+        x_array = getattr(self, dimension)
+        weights = getattr(self, 'particle_charges').abs() * getattr(self, 'survival_probabilities')
+        bin_edges = torch.linspace(bin_range[0], bin_range[1], bins)
+
+        # Compute histogram
+        if method == "histogram":
+            histogram, _ = torch.histogramdd(
+                x_array,
+                weight=weights,
+                bins=bin_edges
+            )
+        elif method == "kde":
+            histogram = kde_histogram_1d(
+                x=x_array,
+                weights=weights,
+                bins=bin_edges,
+                bandwidth=kde_bandwidth,
+            )
+
+        return histogram
+
     def plot_1d_distribution(
         self,
         dimension: Literal["x", "px", "y", "py", "tau", "p"],

From ae36355421211d0e625f160ac75f32d57e8663b4 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Jun 2025 10:33:45 +0200
Subject: [PATCH 2/3] Add a few fixes

---
 cheetah/particles/particle_beam.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py
index d0979cb90..4d97bd79a 100644
--- a/cheetah/particles/particle_beam.py
+++ b/cheetah/particles/particle_beam.py
@@ -1233,6 +1233,7 @@ def get_1d_histogram(
         Get a 1D (differentiable) histogram of the given dimension of the particle distribution.
 
         TODO: Add documentation
+        Mention that this returns the PDF (i.e. normalized to 1, when multiplied by bin size)
         """
         # Extract particle quantities
         x_array = getattr(self, dimension)
@@ -1241,12 +1242,19 @@ def get_1d_histogram(
 
         # Compute histogram
         if method == "histogram":
-            histogram, _ = torch.histogramdd(
+            histogram, _ = torch.histogram(
                 x_array,
                 weight=weights,
-                bins=bin_edges
+                bins=bin_edges,
             )
         elif method == "kde":
+            kde_bandwidth = torch.as_tensor(
+                (
+                    kde_bandwidth
+                    if kde_bandwidth is not None
+                    else bin_edges[1] - bin_edges[0]
+                ),
+            )
             histogram = kde_histogram_1d(
                 x=x_array,
                 weights=weights,
@@ -1254,7 +1262,9 @@ def get_1d_histogram(
                 bandwidth=kde_bandwidth,
             )
 
-        return histogram
+        # Normalize the histogram to get a pdf
+        pdf = histogram / histogram.sum() * (bin_edges[1] - bin_edges[0])
+        return pdf
 
     def plot_1d_distribution(
         self,

From 715b9d0a989b62a9cb4f97e12b7650153c889a7d Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 8 Jul 2025 09:00:18 +0200
Subject: [PATCH 3/3] Make histogram work on GPU

---
 cheetah/particles/particle_beam.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py
index 4d97bd79a..d2b0fc529 100644
--- a/cheetah/particles/particle_beam.py
+++ b/cheetah/particles/particle_beam.py
@@ -1235,10 +1235,12 @@ def get_1d_histogram(
         TODO: Add documentation
         Mention that this returns the PDF (i.e. normalized to 1, when multiplied by bin size)
         """
+        factory_kwargs = {"device": self.particles.device, "dtype": self.particles.dtype}
+
         # Extract particle quantities
         x_array = getattr(self, dimension)
         weights = getattr(self, 'particle_charges').abs() * getattr(self, 'survival_probabilities')
-        bin_edges = torch.linspace(bin_range[0], bin_range[1], bins)
+        bin_edges = torch.linspace(bin_range[0], bin_range[1], bins, **factory_kwargs)
 
         # Compute histogram
         if method == "histogram":
@@ -1254,6 +1256,7 @@ def get_1d_histogram(
                     if kde_bandwidth is not None
                     else bin_edges[1] - bin_edges[0]
                 ),
+                **factory_kwargs
             )
             histogram = kde_histogram_1d(
                 x=x_array,