Add convolve_weights() to fold a point-spread function into a set of weights

README.md

@@ -52,6 +52,9 @@ pip install regridding
 * [`transpose_weights()`](https://regridding.readthedocs.io/en/latest/_autosummary/regridding.transpose_weights.html)
   and [`transpose_weights_conservative()`](https://regridding.readthedocs.io/en/latest/_autosummary/regridding.transpose_weights_conservative.html),
   which reverse a saved resampling, as needed by iterative inversions.
+* [`convolve_weights()`](https://regridding.readthedocs.io/en/latest/_autosummary/regridding.convolve_weights.html),
+  which folds a convolution of the output, such as a point-spread function,
+  into a saved resampling, so that it is applied and transposed along with it.
 * [`fill()`](https://regridding.readthedocs.io/en/latest/_autosummary/regridding.fill.html),
   which fills the missing values of an array by interpolating from the valid points.
 * [`find_indices()`](https://regridding.readthedocs.io/en/latest/_autosummary/regridding.find_indices.html),

docs/index.rst

@@ -49,10 +49,14 @@ Features
 *   :func:`regridding.transpose_weights` and
     :func:`regridding.transpose_weights_conservative`,
     which reverse a saved resampling, as needed by iterative inversions.
+*   :func:`regridding.convolve_weights`, which folds a convolution of the
+    output, such as a point-spread function, into a saved resampling, so that
+    it is applied and transposed along with it.
 *   Building the weights on a CUDA device and applying them there, with the
     ``device`` argument of :func:`regridding.weights`.
     The result is left in device memory, so a scene which is already on the
-    card is never brought back.
+    card is never brought back, and :func:`regridding.convolve_weights`
+    convolves weights which are there without bringing them back either.
 *   :func:`regridding.fill`, which fills the missing values of an array by
     interpolating from the valid points.
 *   :func:`regridding.find_indices`, which locates the input cell containing

regridding/__init__.py

@@ -14,6 +14,7 @@
     weights,
     transpose_weights,
     transpose_weights_conservative,
+    convolve_weights,
 )
 from ._interp_ndarray import ndarray_linear_interpolation
 from ._regrid import regrid_from_weights, regrid
@@ -26,6 +27,7 @@
     "weights",
     "transpose_weights",
     "transpose_weights_conservative",
+    "convolve_weights",
     "ndarray_linear_interpolation",
     "regrid_from_weights",
     "regrid",

regridding/_cuda.py

@@ -1,7 +1,7 @@
 """
 Allocating and filling arrays which live on a CUDA device.
 
-Both device kernels need the same few operations on device memory, and
+The device kernels need the same few operations on device memory, and
 :mod:`numba` gives its allocation and its kernels annotations which describe
 how the compiler calls them rather than how a caller does, so the waivers
 for that are kept here instead of at each use.
@@ -15,9 +15,11 @@
 __all__ = [
     "threads",
     "available",
+    "on_device",
     "allocate",
     "fill",
     "zeros",
+    "prefix_sum",
 ]
 
 threads = 256
@@ -43,6 +45,21 @@ def available() -> bool:
         return False
 
 
+def on_device(weights: np.ndarray) -> bool:
+    """
+    Test whether a set of weights lives in device memory.
+
+    Parameters
+    ----------
+    weights
+        Weights built by :func:`regridding.weights`.
+    """
+    flat = np.asarray(weights).reshape(-1)
+    if not flat.size:  # pragma: nocover
+        return False
+    return cuda.is_cuda_array(flat[0][2])
+
+
 # this runs on the device, where `coverage` cannot follow it, so it reports
 # the body as missed even when it does the work.  `a` is a device array,
 # which `numba` ships no type for, so it is annotated as `Any`
@@ -125,3 +142,54 @@ def zeros(shape: Any, dtype: np.typing.DTypeLike) -> Any:
         The type of the array's elements.
     """
     return fill(allocate(shape, dtype), 0)
+
+
+def prefix_sum(counts: Any, num: int) -> tuple[Any, int]:
+    """
+    Compute the exclusive prefix sum of a device array of counts, on the
+    device.
+
+    :mod:`numba` has no scan, so this borrows :func:`torch.cumsum`, writing
+    into memory which :mod:`numba` allocated and owns.  The result can then
+    be read by a kernel launched after this returns without anything having
+    to keep a tensor alive: :mod:`torch` returns memory to its own cache when
+    a tensor is dropped and hands it out again without waiting for kernels
+    it does not know about, whereas memory :mod:`numba` frees is only
+    released once the device is done with it.
+
+    The sum runs on the default stream, which is where :mod:`numba` launches
+    the kernels which write `counts` and read the result, so it is ordered
+    after the one and before the other whatever stream :mod:`torch` has been
+    told to use.
+
+    Parameters
+    ----------
+    counts
+        The counts, on the device.
+    num
+        The number of counts.
+
+    Returns
+    -------
+    The ``num + 1`` offsets, as a device array whose last element is the
+    total, and that total.
+    """
+    try:
+        # an optional dependency, so it is absent from the environment the
+        # type checker runs in
+        import torch  # type: ignore[import-not-found]
+    except ImportError as error:  # pragma: nocover
+        raise ImportError(
+            "weights on a device need `torch`, which provides the prefix sum; "
+            "install `regridding[cuda]`"
+        ) from error
+
+    offset = allocate(num + 1, np.int64)
+
+    with torch.cuda.stream(torch.cuda.default_stream()):
+        view = torch.as_tensor(offset, device="cuda")
+        view[0] = 0
+        torch.cumsum(torch.as_tensor(counts, device="cuda"), dim=0, out=view[1:])
+        total = int(view[~0].item())
+
+    return offset, total

regridding/_regrid/_regrid_from_weights.py

@@ -1,9 +1,8 @@
 from typing import Sequence
 import numpy as np
 import numba
-from numba import cuda
 from numba.typed.typedlist import List as TypedList
-from regridding import _util
+from regridding import _util, _cuda
 from ._regrid_from_weights_cuda import regrid_from_weights_cuda
 
 __all__ = [
@@ -103,7 +102,7 @@ def regrid_from_weights(
         axis_output=axis_output,
     )
 
-    if _on_device(weights):
+    if _cuda.on_device(weights):
         return regrid_from_weights_cuda(
             weights=weights,
             shape_orthogonal=shape_orthogonal,
@@ -257,21 +256,6 @@ def _normalize(
     return axis_input, axis_output, shape_orthogonal, shape_input, shape_output
 
 
-def _on_device(weights: np.ndarray) -> bool:
-    """
-    Test whether a set of weights lives in device memory.
-
-    Parameters
-    ----------
-    weights
-        Weights built by :func:`regridding.weights`.
-    """
-    flat = np.asarray(weights).reshape(-1)
-    if not flat.size:  # pragma: nocover
-        return False
-    return cuda.is_cuda_array(flat[0][2])
-
-
 @numba.njit(cache=True, parallel=True)
 def _regrid_from_weights(
     weights: TypedList,

regridding/_regrid/_regrid_from_weights_cuda.py

@@ -58,20 +58,22 @@ def _scatter(  # pragma: nocover
     The indices address a flattened grid of the resampled axes, so each is
     split back into a pair and applied to the strides of those axes.  Slots
     which saw no overlap carry an index of ``-1`` and are skipped, so the
-    weights do not have to be compacted first.
+    weights do not have to be compacted first.  The ``-1`` is on the input
+    side of weights as they were built, and on the output side once they
+    have been transposed, so either side marks a slot as empty.
     """
     w = cuda.grid(1)  # type: ignore[call-arg]
     if w >= values.size:
         return
 
     index_input = indices_input[w]
-    if index_input < 0:
+    index_output = indices_output[w]
+    if index_input < 0 or index_output < 0:
         return
 
     i0 = index_input // num_input
     i1 = index_input - i0 * num_input
 
-    index_output = indices_output[w]
     o0 = index_output // num_output
     o1 = index_output - o0 * num_output
 

regridding/_regrid/_tests/test_regrid_from_weights_cuda.py

@@ -1,3 +1,4 @@
+from typing import Any
 import pytest
 import numpy as np
 import astropy.units as u
@@ -431,3 +432,45 @@ def test_values_output_wrong_shape(self) -> None:
                 values_input=_scene(),
                 values_output=cuda.to_device(np.zeros((3, 3))),
             )
+
+
+@requires_cuda
+def test_transposed_empty_slots() -> None:
+    """
+    Transposing weights built on a device moves the ``-1`` of their empty
+    slots to the output side, where the scatter has to skip it as well.
+
+    A slot which is not skipped adds its input value times zero just before
+    the start of its element of the output, which is the last cell of the
+    element before, so a NaN in the input shows up there.
+    """
+    num = 2
+    (x, y), lattice = _grids()
+    grid_input = tuple(np.broadcast_to(c, (num,) + c.shape) for c in (x, y))
+    grid_output = tuple(np.broadcast_to(c, (num,) + c.shape) for c in lattice)
+
+    weights = regridding.weights(
+        coordinates_input=grid_input,
+        coordinates_output=grid_output,
+        axis_input=(1, 2),
+        axis_output=(1, 2),
+        method="conservative",
+        device="cuda",
+    )
+
+    indices_input = weights[0].reshape(-1)[0][0].copy_to_host()
+    assert np.any(indices_input < 0), "the test needs weights with empty slots"
+
+    transposed = regridding.transpose_weights(weights)
+
+    image = np.random.default_rng(17).random((num, 12, 12))
+    image[1, 0, 0] = np.nan
+
+    result: Any = regridding.regrid_from_weights(
+        *transposed,
+        values_input=image,
+        axis_input=(1, 2),
+        axis_output=(1, 2),
+    )
+
+    assert np.all(np.isfinite(result.copy_to_host()[0]))

regridding/_weights/__init__.py

@@ -3,9 +3,11 @@
     transpose_weights,
     transpose_weights_conservative,
 )
+from ._weights_convolved import convolve_weights
 
 __all__ = [
     "weights",
     "transpose_weights",
     "transpose_weights_conservative",
+    "convolve_weights",
 ]

regridding/_weights/_weights_conservative_2d/_clipping/_cuda.py

@@ -172,35 +172,6 @@ def _allocate_result(
     )
 
 
-def _prefix_sum(counts: Any, num_cell: int) -> tuple[Any, int]:
-    """
-    Compute the exclusive prefix sum of the per-cell counts, on the device.
-
-    :mod:`numba` has no scan, so this borrows :func:`torch.cumsum`, which
-    shares the memory rather than copying it.
-
-    Parameters
-    ----------
-    counts
-        The per-cell counts, on the device.
-    num_cell
-        The number of input cells.
-    """
-    try:
-        # an optional dependency, so it is absent from the environment the
-        # type checker runs in
-        import torch  # type: ignore[import-not-found]
-    except ImportError as error:  # pragma: nocover
-        raise ImportError(
-            "building weights on a device needs `torch`, which provides the "
-            "prefix sum; install `regridding[cuda]`"
-        ) from error
-
-    offset = torch.zeros(num_cell + 1, dtype=torch.int64, device="cuda")
-    torch.cumsum(torch.as_tensor(counts, device="cuda"), dim=0, out=offset[1:])
-    return offset, int(offset[~0].item())
-
-
 def weights_conservative_2d_clipping_cuda(
     grid_input: tuple[np.ndarray, np.ndarray],
     grid_output: tuple[np.ndarray, np.ndarray],
@@ -279,7 +250,7 @@ def weights_conservative_2d_clipping_cuda(
     counts = _cuda.allocate(num_cell, np.int64)
     count_cells[blocks, threads](x, y, num_output_x, num_output_y, counts)  # type: ignore[index]
 
-    offset, num_total = _prefix_sum(counts, num_cell)
+    offset, num_total = _cuda.prefix_sum(counts, num_cell)
 
     indices_input, indices_output, values = _allocate_result(
         num_total,
@@ -296,7 +267,7 @@ def weights_conservative_2d_clipping_cuda(
         factor,
         num_output_x,
         num_output_y,
-        cuda.as_cuda_array(offset),
+        offset,
         indices_input,
         indices_output,
         values,

regridding/_weights/_weights_convolved/__init__.py

@@ -0,0 +1,5 @@
+from ._weights_convolved import convolve_weights
+
+__all__ = [
+    "convolve_weights",
+]