[WIP] Add method for differentiable 1D histogram

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,54 @@ 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
+        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, **factory_kwargs)
+
+        # Compute histogram
+        if method == "histogram":
+            histogram, _ = torch.histogram(
+                x_array,
+                weight=weights,
+                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]
+                ),
+                **factory_kwargs
+            )
+            histogram = kde_histogram_1d(
+                x=x_array,
+                weights=weights,
+                bins=bin_edges,
+                bandwidth=kde_bandwidth,
+            )
+
+        # Normalize the histogram to get a pdf
+        pdf = histogram / histogram.sum() * (bin_edges[1] - bin_edges[0])
+        return pdf
+
     def plot_1d_distribution(
         self,
         dimension: Literal["x", "px", "y", "py", "tau", "p"],