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

[roytsmart]

Adds regridding.convolve_weights(weights, kernel, axis_input=None, axis_output=None). It folds a convolution of the output, such as a point-spread function, into a saved set of weights.

The weights W from weights() map input cells onto output cells. This returns the sparse product P W, where P[i', i] = K_i(i' - i) spreads whatever lands in output cell i over its neighbors. Applying the result resamples and blurs in one regrid_from_weights() call. transpose_weights_conservative() transposes the whole operation without any change, which is what ctis needs to model a PSF in MART: its backprojection is the transpose of the forward model.

Kernel

• Layout: the last len(axis_output) axes are the kernel itself, in the order of axis_output, centered at n // 2 as scipy.ndimage.convolve centers it. It must be dimensionless.
• Variation: leading axes broadcast against the output grid, so a kernel can differ per orthogonal element (per wavelength or channel, say) or vary across the output grid. A varying kernel is indexed by the cell the light lands in before it is spread.
• Storage: a kernel that varies along only some resampled axes is stored with one row per cell along those axes, not one per output cell.
• Edges: light spread past the edge of the grid is lost, as off the edge of a sensor.
• Docs: the docstring explains that the kernel acts on cells, so for a PSF it should be the PSF convolved with a cell twice. That is the convention optika.sensors.kernel_diffusion already follows.

Algorithm

Weights are divided into runs: consecutive weights sharing an input cell. Each run scatters its weights through the kernel into a scratch box covering its footprint, which is the bounding box of its output cells grown by the kernel. The box is then read out in order.

• Cost: one multiply-add per weight and kernel element, plus one visit per box cell, however many output cells the input cell covers.
• Ordering: weights ordered by input, as weights() returns them, come out ordered and with each pair once, with no sort.
• Unordered weights: still correct, just less compact.

The algorithm is shared source compiled for the host with numba.njit and for a CUDA device with numba.cuda.jit, as the clipping kernel is.

• Host: each thread keeps one box and grows it as needed. Negative (wraparound) indices from descending grids are handled.
• Device: a first pass measures every run's box, and a prefix sum gives each run its own slice of one scratch allocation. Count and write passes follow.
  • Weights built with device="cuda" are convolved on the device and the result stays there.
  • Their -1 empty slots are dropped, on either side.

Checks against misuse

• Wrong axis_output: if it doesn't describe the grid the weights were built for, the call raises rather than convolving along the wrong axes. The weights can't have more orthogonal axes than it leaves, and every output index must lie in the grid it selects. On the device this is a flag read once at the end.
• New orthogonal axis: a kernel that varies along an orthogonal axis the weights are broadcast along now returns weights and shapes broadcast along it. That case uses axis_input to place the input grid's orthogonal axes.
• Transpose: transpose_weights_conservative() now broadcasts the grids to weights with more orthogonal elements than the grids have.

Related device fixes

• Offsets memory: regridding._cuda.prefix_sum() (moved here from the clipping kernel, with on_device()) now writes into memory numba owns, on the default stream. A kernel reading the offsets after it returns can't race torch reusing the memory. The clipping kernel shares the fix.
• Empty slots in the scatter: the device scatter in regrid_from_weights() now skips a slot whose -1 is on the output side, as it is in transposed device weights. Before, it added zero times an input value to the cell just before its element of the output, which shows up as NaN when the input has one. There is a test for it.

Verification

Tests cover:

• Forward model: equal to scipy.ndimage.convolve(mode="constant") of the resampled array, for odd, even, 1×1 and zero-containing kernels.
• Axis order: axis_output given in reverse order.
• 1D grids: including the descending grids whose indices count from the end.
• Long runs: a coarse grid onto a fine one.
• Varying kernels: along an orthogonal axis, x, y, both, and orthogonal plus x, each equal to an explicit scipy.sparse P @ W.
• Broadcasting: the kernel varying along an axis the weights are broadcast along, applied and transposed.
• Ordering: output ordered and merged.
• Transpose: transpose_weights_conservative of P W equals the geometric transpose after correlating with the kernel.
• Types and units: narrow dtypes and units kept.
• Edge cases: empty elements, the shared kernel body run as plain Python on empty slots, and the axis_output and axis_input errors.

Device tests (-m cuda) check that the device result:

• has the same indices as the host and values to 1e-14, for every kernel shape and for long runs;
• stays on the device and resamples there into the convolved array;
• keeps narrow dtypes;
• skips -1 slots on either side;
• raises on a wrong axis_output;
• handles elements with nothing to spread.

A further test checks the transposed-scatter fix; it fails without the fix.

Locally:

• Full suite: 449 passed and 2 xfailed, including the CUDA tests on an RTX 4090.
• JIT disabled: the changed modules' tests pass with NUMBA_DISABLE_JIT=1.
• Coverage: 100% on every changed module.
• Linters: black and ruff pass, and pyright 1.1.411 is clean on the new code including its tests.

Cost

Measured warm. The first call with a new dtype compiles.

geometry	kernel	weights	host	device
ctis tutorial: 64² scene, 4 channels × 20 wavelengths, onto 128×64	3×3	4.3×	0.08 s	0.18 s
	5×5	9.9×	0.14 s	0.24 s
	7×7	17.9×	0.19 s	0.30 s
	11×11	40.9×	0.30 s	0.49 s
512² scene onto 1024×512, float32/int32 (72M weights)	3×3	4.3×	1.64 s	0.76 s
	5×5	9.9×	2.41 s	1.22 s
8×8 grid onto 400×400 (runs of ~2500)	7×7		0.006 s	0.094 s

• Small elements: at tutorial scale the device pays about 2 ms of launches and allocations for each of 80 small elements, so the host is faster there.
• Long runs: the upsampling case was 1.4 s on the host and 12.9 s on the device before the scratch box.
• Varying kernel: a 7×7 kernel varying along x on 2048×1024 now peaks at 132 MB, mostly the result itself. Before, it copied out an 822 MB table.

Not in this PR

transpose_weights_conservative() still has no device path: it brings the values to the host. That is needed before a whole iterative inversion can stay on the device, and is left for a separate PR.

🤖 Generated with Claude Code

https://claude.ai/code/session_01GqceVSp8s2msxj1Dr9MrGU

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.68%. Comparing base (c76e994) to head (fb07ea0).

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #67      +/-   ##
==========================================
+ Coverage   99.62%   99.68%   +0.06%     
==========================================
  Files          55       62       +7     
  Lines        3192     3835     +643     
==========================================
+ Hits         3180     3823     +643     
  Misses         12       12              

Flag	Coverage Δ
unittests	99.68% <100.00%> (+0.06%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.