Vectorized histogramming

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
 

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

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

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

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