Cut peak memory of the robust convolution by ~3.4x

[AlecThomson]

Why

A single 19533x17000 plane needs ~17 GB to convolve through conv_mode="robust". That is more than twice the per-worker share of a 128 GB / 16-process dask job, so wide-field polarisation cubes get OOM killed part way through a flint run.

Peak RSS was measured by replaying convolve's hot path on a float32 plane with 42.6% NaNs (the real blanking fraction), sampling /proc/self/statm:

plane	float32 size	peak RSS	ratio
6000x5200 (31 Mpix)	0.12 GB	1.83 GB	14.2x
9000x7800 (70 Mpix)	0.26 GB	3.71 GB	13.5x

Linear in pixel count, so 332 Mpix lands at ~17 GB.

What was costing it

Two things, both incidental rather than required by the algorithm:

1. The u/v grids were cast to complex64 before gaussft. Every intermediate in there (ur, vr, g_arg, ur_in, vr_in, dg_arg, g_final) became complex, for a quantity that is real throughout. numba then promoted the whole expression to complex128 against the float64 scalars, at 16 bytes per pixel, and convolve cast the result back down to complex64 afterwards.
2. np.fft.fft2 promotes a float32 image to complex128. numpy's fft API is defined at double precision, so a 1.24 GB plane became a 5.3 GB spectrum, with one of those live per transform axis, then .astype(np.complex64) copied it back down. The NaN mask got a second full complex transform pair at the same cost.

What this does

The image is real, so half its spectrum is redundant. So:

• gaussft now returns a real taper in the dtype of u/v, filled in a single pass. Only the taper itself is ever allocated at image size, instead of up to four full-size intermediates.
• convolve transforms with scipy.fft.rfft2/irfft2, which preserves single precision and stores the last axis once rather than twice. The taper is applied in place and the spectrum is consumed by its own inverse, so neither is duplicated.
• The NaN mask goes through the same real half-spectrum path instead of a second complex pair.

Double-precision input still transforms at double precision, so nothing loses accuracy. real_dtype is matched on dtype kind and width rather than identity, because fits.getdata hands back big-endian >f4 and np.dtype(">f4") == np.float32 is False — an identity check there would send every real image down the double path and undo the saving. There is a test pinning exactly that.

Result

Same 9000x7800 float32 plane, 42.6% NaNs, through convolve:

	peak RSS	ratio	time
before	3.88 GB	13.5x	31.7s
after	1.39 GB	4.0x	9.3s

So ~17 GB becomes ~5 GB per plane, which fits an 8 GB-per-worker job without changing the dask config, and the call is 3.4x faster as a side effect.

This looks to be at the floor for a whole-image FFT: I also tried fusing the taper into the multiply so the full-size taper array never exists, and peak RSS did not move (3.0x either way on the image-only path) because the peak is set by the FFT working set, not the taper. Going lower would need a different algorithm (overlap-add or tiled), since the convolving kernel is tiny next to the image. Not attempted here.

Correctness

Checked against a pristine 4.3.2 copy loaded side by side, over {128x96, 129x97} x {float32, float64} x {NaNs, no NaNs}:

• NaN masks bit-identical in every case
• scaling factors match
• surviving pixels agree to 1e-6 of the image rms (float32 rounding)
• the all-NaN and same-beam short-circuits still return early
• odd axis lengths round-trip (the half-spectrum inverse is given the original shape)

gaussft agrees with a double-precision reference of the same expression to rtol=1e-12.

The miriad-gated tests (test_robust, test_scipy, test_astropy, test_smooth, test_robust_3d) cannot run without miriad installed, so instead pristine 4.3.2 and this branch were both driven through beamcon_2D.smooth_fits_files on the same fixture from tests/test_2d.py. Output is bit-identical for all four conv modes on both fixture beams — beamcon_2D writes float32, and the double-precision difference is exactly one float32 ULP (1.192e-07), so it rounds away on write. Those tests assert atol=1e-3, so their margin is untouched. They will still run for real in CI.

21 passed locally (12 pre-existing + 9 new), with the same 5 miriad errors as on master. ruff 0.15.11 check and format both clean.

New tests

tests/test_convolve_uv.py covers the taper's dtype and realness, agreement with a double-precision reference, cross-precision agreement of convolve including odd shapes and NaN handling, that a >f4 FITS array stays in single precision (asserting on the actual spectrum dtype and half-spectrum shape), and that the caller's array is not consumed.

Open question: should workers be exposed?

scipy.fft takes a workers argument. It is a real speedup when cores are idle, but it is worth nothing in flint's current cores: 16 / processes: 16 config, where every allocated core already has its own task. Measured with N concurrent processes on an N-core box, all starting their transform at the same instant:

concurrency	workers	per-plane	throughput
1 (cores idle)	1	1.469s	0.68 planes/s
1 (cores idle)	4	0.473s	2.11 planes/s
4 (= cores)	1	1.773s	2.23 planes/s
4 (= cores)	4	1.629s	2.44 planes/s

3.1x per plane when cores are free, but only +9% throughput at cores == processes, which is within run-to-run noise. Task-level parallelism has already saturated the cores. Note also that workers is pocketfft's own thread pool and does not read OMP_NUM_THREADS, and workers=-1 resolves to os.cpu_count(), i.e. the whole node rather than the SLURM allocation — so it would need to default to 1 regardless.

Left out of this PR on that basis. It only starts paying if processes drops below cores; happy to plumb it through smooth() and both CLIs as a follow-up if you want it.

🤖 Generated with Claude Code

https://claude.ai/code/session_01AN9y36hf41ZmfkLQPnd4AN

---

Generated by Claude Code