fix: handle binned data in max res finder

packages/essimaging/src/ess/imaging/tools/resolution.py

@@ -1,118 +1,69 @@
+from itertools import zip_longest
+
 import numpy as np
 import scipp as sc
 from numpy.typing import NDArray
 
 
+def _all_divisors(n):
+    for i in range(1, n + 1):
+        if n % i == 0:
+            yield i
+
+
 def maximum_resolution_achievable(
     events: sc.DataArray,
-    coarse_x_bin_edges: sc.Variable,
-    coarse_y_bin_edges: sc.Variable,
-    time_bin_edges: sc.Variable,
-    max_tries: int = 10,
-    max_pixels_x: int = 2048,
-    max_pixels_y: int = 2048,
-    raise_if_not_maximum: bool = False,
+    image_dims: tuple[str, str],
 ):
     """
     Estimates the maximum resolution achievable
-    given a desired binning in time.
+    given a desired binning.
     The maximum achievable resolution is defined
     as the resolution in ``xy`` such that
-    there is at least one event in every ``xyt`` pixel.
+    there is at least one event in every ``xy...`` pixel,
+    where ``...`` represents the rest of the binning dimenensions
+    of the event data, typically wavelength or time-of-flight.
 
     Parameters
     -------------
     events:
-        1D DataArray containing events with associated x, y, and t coordinates.
-        The names of the coordinates must not be `x`, `y` and `t`,
-        the names of the coordinates are taken from the provided ``bin_edges``
-        for each respective dimension.
-    coarse_x_bin_edges:
-        Minimum acceptable resolution in ``x``.
-    coarse_y_bin_edges:
-        Minimum acceptable resolution in ``y``.
-    time_bin_edges:
-        Desired resolution in ``t``.
-    max_tries:
-        The maximum number of iterations before giving up.
-    max_pixels_x:
-        The maximum number of pixels in ``x``.
-    max_pixels_y:
-        The maximum number of pixels in ``y``.
-    raise_if_not_maximum:
-        Often it is not important to find the exact maximum resolution.
-        Therefore this parameter is ``False`` by default, and the function
-        returns an estimate of the maximum resolution.
-        If you want the returned resolution to be exactly the maximum resolution,
-        set the value of this parameter to ``True``.
+        DataArray binned in at least the ``image_dims`` dimensions and optionally more.
+        The method works best when the number of bins in each ``image_dims`` dimension
+        is a power of ``2``.
 
     Returns
     -------------
-        The bin edges in x respectively y that define the
-        maximum achievable resolution.
+        The event data array folded to the finest resolution where there is at least
+        one event in every pixel.
     """
+    if events.bins is None:
+        raise ValueError(
+            'Input data must be binned.'
+            ' For best result, number of bins in each image dimension should'
+            ' be a power of 2.'
+        )
 
-    lower_nx = coarse_x_bin_edges.size
-    lower_ny = coarse_y_bin_edges.size
-    upper_nx = max_pixels_x
-    upper_ny = max_pixels_y
+    x, y = image_dims
 
-    nx = int(2**0.5 * lower_nx) + 1
-    ny = int(2**0.5 * lower_ny) + 1
-    events = events.bin({time_bin_edges.dim: time_bin_edges})
+    xdivisors = list(_all_divisors(events.sizes[x]))
+    ydivisors = list(_all_divisors(events.sizes[y]))
 
-    for _ in range(max_tries):
-        xbins = sc.linspace(
-            coarse_x_bin_edges.dim, coarse_x_bin_edges[0], coarse_x_bin_edges[-1], nx
-        )
-        ybins = sc.linspace(
-            coarse_y_bin_edges.dim, coarse_y_bin_edges[0], coarse_y_bin_edges[-1], ny
+    for i, j in zip_longest(xdivisors, ydivisors):
+        i = xdivisors[-1] if i is None else i
+        j = ydivisors[-1] if j is None else j
+        tmp_events = (
+            events.fold(x, sizes={x: -1, '_x_aux': i})
+            .fold(y, sizes={y: -1, '_y_aux': j})
+            .bins.concat(['_x_aux', '_y_aux'])
         )
-        min_counts_per_pixel = (
-            events.bin(
-                {
-                    coarse_x_bin_edges.dim: xbins,
-                    coarse_y_bin_edges.dim: ybins,
-                }
-            )
-            .bins.size()
-            .min()
-        )
-
+        min_counts_per_pixel = tmp_events.bins.size().min()
         if min_counts_per_pixel.value > 0:
-            lower_nx = nx
-            lower_ny = ny
-            nx = max(min(round((upper_nx * nx) ** 0.5), nx * 2), lower_nx + 1)
-            ny = max(min(round((upper_ny * ny) ** 0.5), ny * 2), lower_ny + 1)
-        else:
-            upper_nx = nx
-            upper_ny = ny
-            nx = min(round((lower_nx * nx) ** 0.5), upper_nx - 1)
-            ny = min(round((lower_ny * ny) ** 0.5), upper_nx - 1)
-
-        if upper_nx - lower_nx < 2 and upper_ny - lower_ny < 2:
-            break
-
-    if raise_if_not_maximum and upper_nx - lower_nx >= 2 and upper_ny - lower_ny >= 2:
-        raise RuntimeError(
-            'Maximal resolution was not found. Increase `max_tries` to search longer. '
-            'Or set `raise_if_not_maximum=False` if it is not necessary to locate the '
-            'maximum exactly.'
-        )
+            return tmp_events
 
-    return (
-        sc.linspace(
-            coarse_x_bin_edges.dim,
-            coarse_x_bin_edges[0],
-            coarse_x_bin_edges[-1],
-            lower_nx,
-        ),
-        sc.linspace(
-            coarse_y_bin_edges.dim,
-            coarse_y_bin_edges[0],
-            coarse_y_bin_edges[-1],
-            lower_ny,
-        ),
+    raise ValueError(
+        'Even at the coarsest pixel binning there are some pixels that have no events.'
+        ' Probably because the wavelength/time-of-arrival binning is too fine,'
+        ' or because number of bins in each image dimension is not a power of 2.'
     )
 
 

packages/essimaging/tests/imaging/tools/maximum_resolution_monitor_test.py

@@ -21,18 +21,12 @@ def test_finds_maximum_resolution():
         },
     )
 
-    x_be, y_be = maximum_resolution_achievable(
-        events,
-        sc.linspace('x', 0, 1, 2),
-        sc.linspace('x', 0, 1, 2),
-        sc.linspace('t', 0, 2, 2),
+    rebinned = maximum_resolution_achievable(
+        events.bin(x=1024, y=1024),
+        ('x', 'y'),
     )
-    assert len(x_be) == 11
-    assert len(y_be) == 11
-    assert x_be[0] == 0
-    assert x_be[-1] == 1
-    assert y_be[0] == 0
-    assert y_be[-1] == 1
+    assert rebinned.sizes['x'] == 8
+    assert rebinned.sizes['y'] == 8
 
 
 @pytest.mark.parametrize('seed', [0, 1, 2])
@@ -47,26 +41,118 @@ def test_finds_maximum_resolution_random(seed):
             't': sc.ones(dims=['events'], shape=(n,)),
         },
     )
-    x_be, y_be = maximum_resolution_achievable(
+    events = events.bin(x=2**10, y=2**10)
+    del events.bins.coords['x']
+    del events.bins.coords['y']
+    rebinned = maximum_resolution_achievable(
         events,
-        sc.linspace('x', 0, 1, 2),
-        sc.linspace('y', 0, 1, 2),
-        sc.linspace('t', 0, 2, 2),
-        # Need enough tries to be sure we find the optimum
-        max_tries=100,
+        ('x', 'y'),
     )
+
+    assert rebinned.bins.size().min().value > 0
     assert (
-        events.bin(x=x_be, y=y_be, t=sc.linspace('t', 0, 2, 2)).bins.size().min().value
-        > 0
+        events.fold('x', sizes={'x': 2 * rebinned.sizes['x'], '_x_aux': -1})
+        .fold('y', sizes={'y': 2 * rebinned.sizes['y'], '_y_aux': -1})
+        .bins.concat(['_x_aux', '_y_aux'])
+        .bins.size()
+        .min()
+        .value
+    ) == 0
+
+
+def test_finds_maximum_resolution_time_binned_input():
+    np.random.seed(0)
+    n = 100_000
+    events = sc.DataArray(
+        sc.ones(dims=['events'], shape=(n,)),
+        coords={
+            'x': sc.array(dims=['events'], values=np.random.random(n)),
+            'y': sc.array(dims=['events'], values=np.random.random(n)),
+            't': sc.array(dims=['events'], values=np.random.random(n)),
+        },
+    )
+    events = events.bin(x=128, y=128, t=500)
+    rebinned = maximum_resolution_achievable(events, ('x', 'y'))
+
+    assert rebinned.bins.size().min().value > 0
+    assert (
+        events.fold('x', sizes={'x': 2 * rebinned.sizes['x'], '_x_aux': -1})
+        .fold('y', sizes={'y': 2 * rebinned.sizes['y'], '_y_aux': -1})
+        .bins.concat(['_x_aux', '_y_aux'])
+        .bins.size()
+        .min()
+        .value
+    ) == 0
+
+
+def test_finds_maximum_resolution_non_even_number_of_bins():
+    np.random.seed(0)
+    n = 100_000
+    events = sc.DataArray(
+        sc.ones(dims=['events'], shape=(n,)),
+        coords={
+            'x': sc.array(dims=['events'], values=np.random.random(n)),
+            'y': sc.array(dims=['events'], values=np.random.random(n)),
+            't': sc.array(dims=['events'], values=np.random.random(n)),
+        },
     )
+    events = events.bin(x=3**4 * 5, y=3**3 * 5, t=101)
+    rebinned = maximum_resolution_achievable(events, ('x', 'y'))
+
+    assert rebinned.bins.size().min().value > 0
     assert (
-        events.bin(
-            x=sc.linspace('x', 0, 1, len(x_be) + 1),
-            y=sc.linspace('y', 0, 1, len(y_be) + 1),
-            t=sc.linspace('t', 0, 2, 2),
-        )
+        events.fold('x', sizes={'x': 45, '_x_aux': -1})
+        .fold('y', sizes={'y': 15, '_y_aux': -1})
+        .bins.concat(['_x_aux', '_y_aux'])
         .bins.size()
         .min()
         .value
-        == 0
+    ) == 0
+
+
+def test_raises_if_not_binned():
+    np.random.seed(0)
+    n = 100_000
+    events = sc.DataArray(
+        sc.ones(dims=['events'], shape=(n,)),
+        coords={
+            'x': sc.array(dims=['events'], values=np.random.random(n)),
+            'y': sc.array(dims=['events'], values=np.random.random(n)),
+            't': sc.array(dims=['events'], values=np.random.random(n)),
+        },
+    )
+    with pytest.raises(ValueError, match='Input data must be binned'):
+        maximum_resolution_achievable(events, ('x', 'y'))
+
+
+def test_raises_if_reaches_coarsest_grid_without_success():
+    np.random.seed(0)
+    n = 100
+    events = sc.DataArray(
+        sc.ones(dims=['events'], shape=(n,)),
+        coords={
+            'x': sc.array(dims=['events'], values=np.random.random(n)),
+            'y': sc.array(dims=['events'], values=np.random.random(n)),
+            't': sc.array(dims=['events'], values=np.random.random(n)),
+        },
+    )
+    events = events.bin(x=128, y=128, t=101)
+    with pytest.raises(ValueError, match='Even at the coarsest'):
+        maximum_resolution_achievable(events, ('x', 'y'))
+
+
+def test_one_dimension_denser_than_the_other():
+    np.random.seed(0)
+    n = 100
+    events = sc.DataArray(
+        sc.ones(dims=['events'], shape=(n,)),
+        coords={
+            'x': sc.array(dims=['events'], values=np.random.random(n)),
+            'y': sc.array(dims=['events'], values=np.random.random(n)),
+            't': sc.array(dims=['events'], values=np.random.random(n)),
+        },
     )
+    events = events.bin(x=2, y=128, t=2)
+    rebinned = maximum_resolution_achievable(events, ('x', 'y'))
+    assert rebinned.bins.size().min().value > 0
+    assert rebinned.sizes['y'] > 4