diff --git a/CHANGELOG.md b/CHANGELOG.md index ac3824bd..b5c19bb4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ - Add a `skip_superimposed` flag to `segment.flattened()` to allow 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 diff --git a/cheetah/accelerator/screen.py b/cheetah/accelerator/screen.py index 2d5fc8e1..18341225 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,19 @@ 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, + 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 aee51570..e9726db4 100644 --- a/cheetah/utils/__init__.py +++ b/cheetah/utils/__init__.py @@ -3,6 +3,7 @@ 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 .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 diff --git a/cheetah/utils/histogram.py b/cheetah/utils/histogram.py new file mode 100644 index 00000000..5a4f0efb --- /dev/null +++ b/cheetah/utils/histogram.py @@ -0,0 +1,99 @@ +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 diff --git a/tests/test_vectorized.py b/tests/test_vectorized.py index be72138e..b30a7640 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", ["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