rm_lite.utils.dask_io ===================== .. py:module:: rm_lite.utils.dask_io .. autoapi-nested-parse:: Dask-backed I/O helpers for chunked 3D RM-synthesis/CLEAN. Attributes ---------- .. autoapisummary:: rm_lite.utils.dask_io.DEFAULT_TARGET_CHUNK_MB Functions --------- .. autoapisummary:: rm_lite.utils.dask_io._channel_mad_lazy rm_lite.utils.dask_io._channel_mad_std_block rm_lite.utils.dask_io._chunk_bounds rm_lite.utils.dask_io._read_fits_block rm_lite.utils.dask_io.complex_pol_dask rm_lite.utils.dask_io.estimate_channel_noise_mad rm_lite.utils.dask_io.estimate_stokes_i_channel_noise rm_lite.utils.dask_io.freq_arr_hz_from_header rm_lite.utils.dask_io.read_fits_cube_dask rm_lite.utils.dask_io.spatial_chunk_size rm_lite.utils.dask_io.write_zarr_group Module Contents --------------- .. py:function:: _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. .. py:function:: _channel_mad_std_block(block: numpy.typing.NDArray[numpy.float64]) -> numpy.typing.NDArray[numpy.float64] .. py:function:: _chunk_bounds(size: int, chunk: int) -> list[tuple[int, int]] .. py:function:: _read_fits_block(path: str | pathlib.Path, y_bounds: tuple[int, int], x_bounds: tuple[int, int]) -> numpy.typing.NDArray[Any] .. py:function:: 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. :param stokes_q: Stokes Q dask array. :type stokes_q: da.Array :param stokes_u: Stokes U dask array. :type stokes_u: da.Array :returns: Complex Q + iU dask array, same chunks as the inputs. :rtype: da.Array .. py:function:: 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). :param stokes_q: Stokes Q dask array, shape (n_freq, ny, nx). :type stokes_q: da.Array :param stokes_u: Stokes U dask array, same shape as `stokes_q`. :type stokes_u: da.Array :returns: Per-channel noise estimate, shape (n_freq,). A plain numpy array, not lazy; computed once, explicitly, here. :rtype: NDArray[np.float64] .. py:function:: 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. :param stokes_i: Stokes I dask array, shape (n_freq, ny, nx). :type stokes_i: da.Array :returns: Per-channel noise estimate, shape (n_freq,). A plain numpy array, not lazy. :rtype: NDArray[np.float64] .. py:function:: 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. :param header: FITS header containing a spectral axis. :type header: Header :param n_freq: Number of channels along the spectral axis. :type n_freq: int :returns: Frequency array in Hz. :rtype: NDArray[np.float64] .. py:function:: 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. :param 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. :type path: str | Path :param target_chunk_mb: Target chunk memory footprint in MB, see `spatial_chunk_size`. Defaults to 256. :type target_chunk_mb: float, optional :returns: Lazy dask array and the FITS header. :rtype: tuple[da.Array, Header] .. py:function:: 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. :param n_freq: Size of the (unchunked) spectral axis. :type n_freq: int :param ny: Full image height in pixels. :type ny: int :param nx: Full image width in pixels. :type nx: int :param itemsize: Size in bytes of one array element. :type itemsize: int :param target_chunk_mb: Target memory footprint of a single chunk, in MB. Defaults to 256. :type target_chunk_mb: float, optional :returns: Spatial chunk size (cy, cx). :rtype: tuple[int, int] .. py:function:: 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. :param store: Path to the zarr store (a group containing one array per key in `arrays`). :type store: str | Path :param arrays: Name -> dask array to write. :type arrays: Mapping[str, da.Array] :param overwrite: Overwrite existing arrays. Defaults to True. :type overwrite: bool, optional .. py:data:: DEFAULT_TARGET_CHUNK_MB :value: 256