[roytsmart]

A prototype of regridding.Regridder: a sparse linear operator that resamples arrays from one grid onto another. It holds the weights of every orthogonal element in one CSR matrix, which it applies in a single pass over its rows. This is M3 of the CUDA MART roadmap.

Builds on #69 (now merged), whose cell volumes, grid check and unit handling it reuses.

Why

Running MART with regrid_from_weights() launches one kernel per orthogonal element, 80 per resample in ctis's tutorial. The GPU spends most of an iteration waiting on those launches, and its scatter adds in no fixed order. A single CSR matrix needs one launch per product and sums every row in a fixed order. The result is then the same on every run, on the host and on a device.

API

regridder = regridding.Regridder.from_weights(
    weights,                    # from regridding.weights()
    axis_input, axis_output,
    shape_input=None,           # the shape of the values it accepts
    shape_output=None,          # the shape of the values it returns
    device=None,                # or "cuda"
    dtype=np.float64,           # or np.float32
)
regridder(values)               # numpy on the host; torch or numba arrays on a device
regridder.T                     # the exact transpose, kept once built
regridder.transpose_conservative(coordinates_input, coordinates_output, weights_input=None)
regridder.to("cuda")

Broadcasting is declared by shapes. Each orthogonal axis's role follows from its length in the weights, the input and the output:

weights	input	output	the operator
N	N	N	applies each element to its own input
N	1	N	applies every element to the same input
N	N	1	sums what the elements give
N	1	1	applies the sum of the elements
1	M	M	applies the same matrix to every input (a batch axis)

• Defaults: when the shapes are left out, each element gets its own input, which is what regrid_from_weights() does today.
• Transposing: swaps the two shapes, so shared inputs become sums and sums become shared inputs.
• For ctis: this makes the forward model one product, from (wavelength, x, y) to (channel, x′, y′).

The conservative transpose scales each entry of the exact transpose by its own element's cells. It isn't two diagonal scalings, because a shared input can have cells of a different size in each element: ctis's scene, distorted, differs by channel. The result shares the indices of .T.

How it is built

• Assembly is a parallel counting sort. The first pass counts each row's entries; the second places them at the rows' cursors. A row keeps its entries in the order of the elements, so the host (numba, per-thread histograms) and the device (torch: integer counts, then a stable sort of each block) build the same matrix bit for bit.
• The product is row-parallel. On the host, each thread takes whole rows. On a device, each row gets a group of threads, as many as an average row has entries, up to a warp: 2 for a backprojection's 3 or so entries, 32 for a forward operator's tens to thousands. The group adds its partial sums in a fixed tree. Its size depends only on the matrix, so a product repeats exactly, and a column gives the same numbers alone or in a batch.
• transpose_weights_conservative()'s cell factors move into a shared _factors(), which both it and the operator use.

Tests

There are 52 host tests and 27 device tests:

• Every role, in two axis layouts: products are compared with regrid_from_weights() followed by a sum. When no axis is shared or summed, they are equal bit for bit.
• Transposes: .T passes a dot-product test, and T.T is self. The conservative transpose equals transpose_weights_conservative() bit for bit.
• Device: assembly, .T and the conservative transpose equal the host's bit for bit. Device products agree with the host's to rounding and repeat exactly. Each group size, from 1 to 32 threads, matches a dense product and gives each column of a batch exactly what it gives alone.
• Also covered: batch and leading axes, units (including a percentage in the conservative transpose), weights_input, float32, weights built on the device, elements with no entries, and invalid shapes and grids. Codecov reports every changed line covered.

The full suite passes locally, with 99 device tests. ruff and black are clean. pyright passes in CI.

Performance

These are MART iterations on an RTX 4090, against ctis's MartInverter, with every path matching ctis to 1e-10 or better:

Problem	ctis	Per-element GPU	Regridder, host	Regridder, GPU
Tutorial, 1.1M weights	44.1 ms	12.4 ms	5.5 ms	1.2 ms
PSF tutorial, 148M weights	183 ms	21.4 ms	105 ms	5.2 ms
256 × 256 scene, 17M weights	524 ms	12.7 ms	73.5 ms	4.9 ms

Assembling 145M weights takes about 0.5 s on the GPU or on 48 host threads, and the transpose another 0.5 s.

Sizing the thread groups to the rows (03f9e67) halves the larger problems. These are the same MART runs, timed back to back on 2026-10-08 with the GPU 10 to 17% busy, so they run slower than the table above:

Problem	A warp per row	Groups sized to rows	torch (cuSPARSE)
Tutorial, 23 iterations	1.3 ms	1.2 ms	0.8 ms
PSF tutorial, 47 iterations	7.2 ms	7.3 ms	8.2 ms
256 × 256 scene, 50 iterations	7.3 ms	3.4 ms	2.7 ms
512 × 512 scene, 20 iterations	26.0 ms	13.9 ms	12.7 ms

• The PSF tutorial is unchanged: its rows are long, so both products still get a full warp.
• The products alone now beat cuSPARSE. At 512 × 512, the backprojection takes 2.0 ms against 3.0, and 14.8 before; the forward product takes 1.8 ms against 2.7.
• torch stays ahead on the whole iteration because of host time. Each device call spends about 0.2 ms on the host: numba's launch, wrapping the two arrays, and torch's layout. That matters only on small problems.

Why not cuSPARSE

torch's sparse product could replace the kernel, but three things rule it out:

• It doesn't repeat exactly. On rows of 4,500 entries, all 100 repeated products differed, by 1 to 4 ulp in about half the values. That held even under torch.use_deterministic_algorithms(True), which neither raised nor fixed it. On rows of any length, a column's result changes with the batch width.
• It can't sum float32 entries in float64. It rejects a float32 matrix times float64 values.
• Its indices must match. int64 row pointers, which ESIS's 2.2e9 entries need, force int64 column indices: 8.8 GB more per direction.

The kernel here repeats exactly in all of those cases.

Not done yet

• A streaming build for weights too large to hold at once (roadmap M6). ESIS's production grid needs it.
• Host time per device call, about 0.2 ms, could shrink by launching without numba's array wrappers. It only matters on small problems.
• A shared-and-summed axis in the conservative transpose: when an axis is both shared and summed and the grids vary along it, transpose_conservative() raises NotImplementedError.
• The device path needs torch (the existing cuda extra) to lay out the values.
• Not yet built: a named-arrays wrapper, a docstring example, and dropping exact zeros during assembly.

🤖 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.77%. Comparing base (91ef979) to head (03f9e67).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #70      +/-   ##
==========================================
+ Coverage   99.71%   99.77%   +0.05%     
==========================================
  Files          65       70       +5     
  Lines        4251     5266    +1015     
==========================================
+ Hits         4239     5254    +1015     
  Misses         12       12              

Flag	Coverage Δ
unittests	99.77% <100.00%> (+0.05%)	⬆️

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.