rm_lite.utils.dask_io¶

Dask-backed I/O helpers for chunked 3D RM-synthesis/CLEAN.

Attributes¶

DEFAULT_TARGET_CHUNK_MB

Functions¶

`_channel_mad_lazy`(→ dask.array.Array)	Lazy per-channel MAD std over every spatial pixel of a cube.
`_channel_mad_std_block`(...)
`_chunk_bounds`(→ list[tuple[int, int]])
`_read_fits_block`(→ numpy.typing.NDArray[Any])
`complex_pol_dask`(→ dask.array.Array)	Combine chunked Stokes Q and U dask arrays into a complex Q + iU array.
`estimate_channel_noise_mad`(...)	Robust per-channel noise from Stokes Q/U cubes, for auto-masking/thresholding.
`estimate_stokes_i_channel_noise`(...)	Robust per-channel noise from a single Stokes I cube.
`freq_arr_hz_from_header`(...)	Derive the frequency array in Hz from a FITS header's spectral WCS.
`read_fits_cube_dask`(→ tuple[dask.array.Array, ...)	Lazily read a Stokes FITS cube as a chunked dask array.
`spatial_chunk_size`(→ tuple[int, int])	Pick a square spatial chunk size for a fixed target chunk memory footprint.
`write_zarr_group`(→ None)	Write a set of dask arrays lazily/incrementally to a shared zarr store.

Module Contents¶

rm_lite.utils.dask_io._channel_mad_lazy(cube: dask.array.Array) → dask.array.Array¶

Lazy per-channel MAD std over every spatial pixel of a cube.

Rechunks the spatial axes to one block per channel (a robust median can’t be combined incrementally across separate spatial chunks) and reduces each channel plane to a scalar. Not computed – so several cubes can be reduced in one compute call.

rm_lite.utils.dask_io._channel_mad_std_block(block: numpy.typing.NDArray[numpy.float64]) → numpy.typing.NDArray[numpy.float64]¶

rm_lite.utils.dask_io._chunk_bounds(size: int, chunk: int) → list[tuple[int, int]]¶

rm_lite.utils.dask_io._read_fits_block(path: str | pathlib.Path, y_bounds: tuple[int, int], x_bounds: tuple[int, int]) → numpy.typing.NDArray[Any]¶

rm_lite.utils.dask_io.complex_pol_dask(stokes_q: dask.array.Array, stokes_u: dask.array.Array) → dask.array.Array¶

Combine chunked Stokes Q and U dask arrays into a complex Q + iU array.

Parameters:

stokes_q (da.Array) – Stokes Q dask array.
stokes_u (da.Array) – Stokes U dask array.

Returns:

Complex Q + iU dask array, same chunks as the inputs.

Return type:

da.Array

rm_lite.utils.dask_io.estimate_channel_noise_mad(stokes_q: dask.array.Array, stokes_u: dask.array.Array) → numpy.typing.NDArray[numpy.float64]¶

Robust per-channel noise from Stokes Q/U cubes, for auto-masking/thresholding.

Computes astropy.stats.mad_std over every spatial pixel in each channel plane, then combines the Q and U estimates the same way rm_lite.utils.synthesis.compute_rmsynth_params combines a per-channel complex error into real_qu_error (abs(real + imag) / 2).

Unlike the RM-synth/CLEAN chunking convention, this rechunks the spatial axes to a single block per channel: a robust statistic like the median can’t be combined incrementally across separate spatial chunks the way a sum or mean can, so each channel’s full spatial plane has to be brought together to compute it. This is the same “per channel, one full image plane at a time” access pattern classic per-channel RMS estimation uses.

The per-channel noise this returns can be turned into a weight array (weight_arr = 1 / noise**2) for rm_lite.tools_3d.rmsynth.rmsynth_3d, and from there rm_lite.utils.synthesis.compute_theoretical_noise gives the FDF-domain noise used to set rm_lite.tools_3d.rmclean.rmclean_3d’s mask/threshold (mirroring the 1D run_rmclean_from_synth auto-mask/ auto-threshold convention).

Parameters:

stokes_q (da.Array) – Stokes Q dask array, shape (n_freq, ny, nx).
stokes_u (da.Array) – Stokes U dask array, same shape as stokes_q.

Returns:

Per-channel noise estimate, shape (n_freq,). A plain numpy array, not lazy; computed once, explicitly, here.

Return type:

NDArray[np.float64]

rm_lite.utils.dask_io.estimate_stokes_i_channel_noise(stokes_i: dask.array.Array) → numpy.typing.NDArray[numpy.float64]¶

Robust per-channel noise from a single Stokes I cube.

Same MAD-based per-channel estimator as estimate_channel_noise_mad, but for one cube (no Q/U combination). Useful as the stokes_i_error fed to rm_lite.tools_3d.rmsynth.rmsynth_3d’s per-pixel Stokes I fit when the I cube carries no separate error cube.

Parameters:: stokes_i (da.Array) – Stokes I dask array, shape (n_freq, ny, nx).
Returns:: Per-channel noise estimate, shape (n_freq,). A plain numpy array, not lazy.
Return type:: NDArray[np.float64]

rm_lite.utils.dask_io.freq_arr_hz_from_header(header: astropy.io.fits.Header, n_freq: int) → numpy.typing.NDArray[numpy.float64]¶

Derive the frequency array in Hz from a FITS header’s spectral WCS.

Parameters:

header (Header) – FITS header containing a spectral axis.
n_freq (int) – Number of channels along the spectral axis.

Returns:

Frequency array in Hz.

Return type:

NDArray[np.float64]

rm_lite.utils.dask_io.read_fits_cube_dask(path: str | pathlib.Path, target_chunk_mb: float = DEFAULT_TARGET_CHUNK_MB) → tuple[dask.array.Array, astropy.io.fits.Header]¶

Lazily read a Stokes FITS cube as a chunked dask array.

Each spatial block is read by its own dask.delayed task (reopening the file with memmap=True and slicing out just that block), assembled into one dask array with dask.array.from_delayed + dask.array.block. Actual reads from disk are deferred until a block is computed.

Degenerate length-1 axes (e.g. a dummy Stokes axis, common in ASKAP/EMU cutout cubes) are squeezed out first.

Parameters:

path (str | Path) – Path to the FITS cube. Assumed axis order (freq, y, x) once degenerate axes are squeezed out, i.e. the frequency axis is first in numpy order.
target_chunk_mb (float, optional) – Target chunk memory footprint in MB, see spatial_chunk_size. Defaults to 256.

Returns:

Lazy dask array and the FITS header.

Return type:

tuple[da.Array, Header]

rm_lite.utils.dask_io.spatial_chunk_size(n_freq: int, ny: int, nx: int, itemsize: int, target_chunk_mb: float = DEFAULT_TARGET_CHUNK_MB) → tuple[int, int]¶

Pick a square spatial chunk size for a fixed target chunk memory footprint.

The frequency/Faraday-depth axis is never chunked, so a chunk’s memory footprint is n_freq * cy * cx * itemsize. cy and cx are solved for that footprint to equal target_chunk_mb, then clipped to the image dimensions.

Parameters:

n_freq (int) – Size of the (unchunked) spectral axis.
ny (int) – Full image height in pixels.
nx (int) – Full image width in pixels.
itemsize (int) – Size in bytes of one array element.
target_chunk_mb (float, optional) – Target memory footprint of a single chunk, in MB. Defaults to 256.

Returns:

Spatial chunk size (cy, cx).

Return type:

tuple[int, int]

rm_lite.utils.dask_io.write_zarr_group(store: str | pathlib.Path, arrays: collections.abc.Mapping[str, dask.array.Array], overwrite: bool = True) → None¶

Write a set of dask arrays lazily/incrementally to a shared zarr store.

Each array is written chunk-by-chunk via dask.array.Array.to_zarr; the full array is never materialised in memory before writing. All arrays are written in a single dask.compute() call rather than one to_zarr() call per array: if two arrays share upstream graph nodes (e.g. the four outputs of rm_lite.tools_3d.rmclean.rmclean_3d, which all come from one per-chunk dask.delayed call), computing them separately would silently redo that shared work once per array.

Parameters:

store (str | Path) – Path to the zarr store (a group containing one array per key in arrays).
arrays (Mapping[str, da.Array]) – Name -> dask array to write.
overwrite (bool, optional) – Overwrite existing arrays. Defaults to True.

rm_lite.utils.dask_io.DEFAULT_TARGET_CHUNK_MB = 256¶