Bin dense radial profiles without sorting

sigima/tests/image/profile_physical_unit_test.py

@@ -105,10 +105,13 @@ def _radial_profile_all_bins(data, center, xcoords, ycoords, step):
         (np.arange(53.0), np.arange(37.0), (26.0, 18.0), 1.0),  # pixel center
         (np.arange(53) * 0.1, np.arange(37) * 0.2, (2.6, 3.5), 0.1),  # anisotropic
         (np.cumsum(np.linspace(0.1, 2.0, 53)), np.arange(37.0), (30.0, 9.5), 0.1),
+        # Sparse bins (more bins than pixels): tiny gap between two pixel centers
+        (np.r_[0.0, 1e-4, np.arange(1.0, 19.0)], np.arange(15.0), (7.3, 6.1), 1e-4),
     ],
 )
 def test_radial_profile_same_as_all_bins(xcoords, ycoords, center, step) -> None:
-    """Accumulating only the bins holding pixels gives bit-identical profiles."""
+    """Accumulating the bins directly (dense bins) or only those holding pixels
+    (sparse bins) gives bit-identical profiles."""
     rng = np.random.default_rng(0)
     data = rng.normal(100.0, 20.0, (ycoords.size, xcoords.size))
     for values in (data, data.astype(np.uint16), np.ma.masked_less(data, 90.0)):
@@ -119,6 +122,39 @@ def test_radial_profile_same_as_all_bins(xcoords, ycoords, center, step) -> None
         assert np.array_equal(x, xref) and np.array_equal(y, yref)
 
 
+@pytest.mark.parametrize(
+    "xcoords, step, sparse",
+    [
+        (np.arange(40.0), 1.0, False),  # Dense bins: no more bins than pixels
+        (np.r_[0.0, 1e-4, np.arange(1.0, 39.0)], 1e-4, True),  # Sparse bins
+    ],
+)
+def test_radial_profile_sorts_only_sparse_bins(
+    monkeypatch: pytest.MonkeyPatch, xcoords: np.ndarray, step: float, sparse: bool
+) -> None:
+    """Dense bins are accumulated directly, without sorting (regression: sorting
+    every bin made radial profiles of large images 2 to 3 times slower); only
+    sparse bins are sorted to keep the bins holding pixels."""
+    calls = []
+    unique = np.unique
+
+    def spy(*args, **kwargs):
+        calls.append(args)
+        return unique(*args, **kwargs)
+
+    monkeypatch.setattr(np, "unique", spy)
+    ycoords = np.arange(30.0)
+    data = np.arange(ycoords.size * xcoords.size, dtype=float)
+    sigima.tools.image.get_radial_profile(
+        data.reshape(ycoords.size, xcoords.size),
+        (12.3, 14.6),
+        xcoords=xcoords,
+        ycoords=ycoords,
+        step=step,
+    )
+    assert bool(calls) == sparse
+
+
 def test_radial_profile_tiny_spacing() -> None:
     """The memory used does not depend on the pixel spacing (regression: a tiny
     spacing allocated one bin per step up to the largest distance, i.e. several

sigima/tools/image/extraction.py

@@ -28,9 +28,10 @@ def get_radial_profile(
     """Return radial profile of image data
 
     The pixels are binned by their distance to the center (distance divided by
-    *step*, rounded down), and the values of each bin are averaged. Only the bins
-    holding pixels are allocated, so that the memory used never exceeds the size
-    of the image, whatever the step.
+    *step*, rounded down), and the values of each bin are averaged. The bins are
+    accumulated directly by index when they are no more numerous than the pixels
+    (usual case), and otherwise only the bins holding pixels are allocated, so
+    that the memory used never exceeds the size of the image, whatever the step.
 
     Args:
         data: Input data (2D array)
@@ -64,16 +65,27 @@ def get_radial_profile(
     if not np.all(np.isfinite(r)):
         raise ValueError("The center and the pixel coordinates must be finite")
     # Bin of each pixel, robust to rounding noise when distances are multiples of
-    # the step: only the bins holding pixels are kept (sorted), each pixel being
-    # given the position of its bin among them
-    bins, index = np.unique(np.floor(r / step + 1e-9).ravel(), return_inverse=True)
+    # the step
+    k = np.floor(r / step + 1e-9).ravel()
 
     # Average over the same distance (distances holding no pixel, possible with
     # anisotropic pixels or non-uniform coordinates, are left out of the profile)
-    tbin = np.bincount(index, data.ravel())  # Sum of pixel values at each distance
-    nr = np.bincount(index)  # Number of pixels at each distance
-
-    yhalf = tbin / nr  # this is the half radial profile
+    if k.max() < k.size:
+        # Dense bins: accumulating every bin by index allocates no more entries
+        # than pixels, without sorting (integer conversion only after this test,
+        # bin indices being unbounded for tiny steps)
+        k = k.astype(np.intp)
+        tbin = np.bincount(k, data.ravel())  # Sum of pixel values at each distance
+        nr = np.bincount(k)  # Number of pixels at each distance
+        bins = np.flatnonzero(nr)
+        yhalf = tbin[bins] / nr[bins]  # this is the half radial profile
+    else:
+        # Sparse bins (tiny step): only the bins holding pixels are kept (sorted),
+        # each pixel being given the position of its bin among them
+        bins, index = np.unique(k, return_inverse=True)
+        tbin = np.bincount(index, data.ravel())
+        nr = np.bincount(index)
+        yhalf = tbin / nr
     xhalf = bins * step
     # Mirror it to get the full radial profile (the first element is the center)
     skip = 1 if bins[0] == 0 else 0