diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py index 442388b94..5868bf73a 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, ) @@ -1242,6 +1243,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"],