rm_lite.tools_3d.rmsynth¶

RM-synthesis on chunked 3D Stokes Q/U cubes via dask.

Classes¶

`PixelFit`	One pixel's Stokes I fit within a chunk (see _iter_pixel_fits).
`RMSynth3DResults`	Results of chunked 3D RM-synthesis.

Functions¶

`_alpha_at_ref`(→ float)	Spectral index (d ln I / d ln nu) of spec at the reference frequency.
`_alpha_from_model_block`(...)	Spectral index alpha at the reference frequency from a supplied model cube.
`_compute_global_params`(...)	Compute phi_arr/lam_sq_0_m2/weight_arr and theoretical FDF noise, once for the whole cube.
`_fit_stokes_i_block`(→ numpy.typing.NDArray[numpy.float64])	Fit a Stokes I model per pixel over one spatial chunk, in a single pass.
`_iter_pixel_fits`(→ collections.abc.Iterator[PixelFit])	Yield a PixelFit for every pixel in a chunk.
`_pixel_stokes_i_error`(...)	Per-pixel Stokes I error spectrum from a 3D error cube, 1D per-channel
`_ref_flux_block`(→ numpy.typing.NDArray[numpy.float64])	Interpolate a Stokes I model cube block at the reference frequency ->
`_split_stokes_i_error`(...)	Split a Stokes I error into (per-channel 1D array, per-pixel dask cube).
`_stokes_i_model_cube`(→ tuple[dask.array.Array, ...)	Lazy per-pixel Stokes I model cube, spectral-index map, fitted-order map,
`rmsynth_3d`(→ RMSynth3DResults)	Run RM-synthesis on chunked Stokes Q/U cubes.
`rmsynth_3d_from_fits`(→ RMSynth3DResults)	Run RM-synthesis directly on Stokes Q/U FITS cubes on disk.

Module Contents¶

class rm_lite.tools_3d.rmsynth.PixelFit¶

Bases: NamedTuple

One pixel’s Stokes I fit within a chunk (see _iter_pixel_fits).

fit: rm_lite.utils.fitting.FitResult | None¶: The fit, or None if the pixel was skipped (too few channels / low SNR).

good: numpy.typing.NDArray[numpy.bool_]¶: Finite-channel mask, for the flat-model fallback.

i_spec: numpy.typing.NDArray[numpy.float64]¶: The pixel’s Stokes I spectrum (unmasked), for the flat-model fallback.

x: int¶

y: int¶

class rm_lite.tools_3d.rmsynth.RMSynth3DResults¶

Bases: NamedTuple

Results of chunked 3D RM-synthesis.

fdf_dirty_cube: dask.array.Array¶: Dirty FDF cube, lazy dask array of shape (n_phi, ny, nx).

fwhm_rmsf_radm2: float¶: Analytic RMSF FWHM (per-pixel fitting is not performed in 3D).

lam_sq_0_m2: float¶: Reference wavelength^2 value in m^2.

phi_arr_radm2: numpy.typing.NDArray[numpy.float64]¶: Faraday depth values in rad/m^2.

phi_double_arr_radm2: numpy.typing.NDArray[numpy.float64]¶: Double-length Faraday depth values in rad/m^2, for the RMSF.

rmsf_cube: dask.array.Array¶: RMSF cube, lazy dask array of shape (n_phi_double, ny, nx).

stokes_i_alpha_error_map: dask.array.Array | None = None¶: Per-pixel 1-sigma (16th/84th-percentile) uncertainty on stokes_i_alpha_map, shape (ny, nx), from the same Monte-Carlo over the fit covariance as stokes_i_model_error_cube. None unless compute_model_error=True and a Stokes I cube was fitted (a supplied model has no covariance, so no alpha error). NaN where a pixel was not fitted.

stokes_i_alpha_map: dask.array.Array | None = None¶: Stokes I spectral index (d ln I / d ln nu) at the reference frequency, shape (ny, nx). A 2D map, like the moment maps. NaN where a pixel was not fitted (below the SNR cut). None unless a Stokes I cube or model was supplied.

stokes_i_model_cube: dask.array.Array | None = None¶: Per-pixel Stokes I model cube, lazy, shape (n_freq, ny, nx). None unless a Stokes I cube or model was supplied to rmsynth_3d.

stokes_i_model_error_cube: dask.array.Array | None = None¶: Per-pixel Stokes I model 1-sigma error cube, shape (n_freq, ny, nx). None unless compute_model_error=True (computed in the same per-pixel fit pass).

stokes_i_model_order_map: dask.array.Array | None = None¶: Per-pixel fitted polynomial order of the Stokes I model (len(popt) - 1), shape (ny, nx). With a negative fit_order this is the AIC-chosen order per pixel; with a fixed order it is uniform on fitted pixels. NaN where a pixel was not fitted (below the SNR cut or flat fallback). None unless a Stokes I cube was fitted (a supplied model has no fitted order).

stokes_i_ref_flux_map: dask.array.Array | None = None¶: Stokes I model at the reference frequency (lambda2_to_freq(lam_sq_0_m2)), shape (ny, nx). This is the factor the fractional FDF was multiplied by to reach flux units. A 2D map, like the moment maps. None unless a Stokes I cube or model was supplied.

theoretical_noise: rm_lite.utils.synthesis.TheoreticalNoise¶: Theoretical FDF-domain noise from the per-channel weight array. This is a per-channel, not per-pixel, estimate, so it is uniform across the cube. When a Stokes I model is used the FDF is rescaled to flux per pixel; this noise stays in the Q/U-error domain it was computed in, which the rescaling keeps roughly consistent (exactly so for a flat Stokes I spectrum).

rm_lite.tools_3d.rmsynth._alpha_at_ref(spec: numpy.typing.NDArray[numpy.float64], freq_arr_hz: numpy.typing.NDArray[numpy.float64], ref_freq_hz: float) → float¶

Spectral index (d ln I / d ln nu) of spec at the reference frequency.

Sorts by frequency first: np.interp/np.gradient need an ascending axis, and a descending-frequency cube (negative CDELT3) otherwise gives garbage.

rm_lite.tools_3d.rmsynth._alpha_from_model_block(model_block: numpy.typing.NDArray[numpy.float64], freq_arr_hz: numpy.typing.NDArray[numpy.float64], ref_freq_hz: float) → numpy.typing.NDArray[numpy.float64]¶

Spectral index alpha at the reference frequency from a supplied model cube.

For a model given directly (not fitted), every pixel is modelled, so alpha is finite wherever the model is finite (0 for a flat model). NaN only where the model is not finite.

rm_lite.tools_3d.rmsynth._compute_global_params(freq_arr_hz: numpy.typing.NDArray[numpy.float64], weight_arr: numpy.typing.NDArray[numpy.float64], fdf_options: rm_lite.utils.synthesis.FDFOptions) → tuple[rm_lite.utils.synthesis.RMSynthParams, rm_lite.utils.synthesis.TheoreticalNoise]¶

Compute phi_arr/lam_sq_0_m2/weight_arr and theoretical FDF noise, once for the whole cube.

compute_rmsynth_params is written for a single per-pixel spectrum, but its weight-array derivation round-trips exactly from a per-channel error spectrum (weight = 1/error**2), so a synthetic, fully-finite spectrum with error = 1/sqrt(weight_arr) reuses it unmodified for a per-channel (not per-pixel) weight array shared by every spatial chunk. The same reconstructed error feeds compute_theoretical_noise for a per-channel (not per-pixel) theoretical noise estimate.

rm_lite.tools_3d.rmsynth._fit_stokes_i_block(*arrays: numpy.typing.NDArray[numpy.float64], freq_arr_hz: numpy.typing.NDArray[numpy.float64], ref_freq_hz: float, err_1d: numpy.typing.NDArray[numpy.float64] | None, fit_options: rm_lite.utils.fitting.StokesIFitOptions, log_level: int) → numpy.typing.NDArray[numpy.float64]¶

Fit a Stokes I model per pixel over one spatial chunk, in a single pass.

Returns a stacked block of shape (n_out, cy, cx). Without error: planes 0..n_freq-1 are the model cube, plane n_freq is the fitted spectral index alpha at the reference frequency, and plane n_freq+1 is the fitted polynomial order (len(popt) - 1; n_out = n_freq + 2). With fit_options.compute_model_error: planes n_freq+2..2*n_freq+1 are the 1-sigma model error and plane 2*n_freq+2 is the 1-sigma alpha error (n_out = 2*n_freq + 3). Both errors come from one Monte-Carlo over the same per-pixel fit covariance, so they cost no extra fit; when compute_model_error is False no error work is done at all.

arrays is (i_block,) or (i_block, err_block); the error cube is optional (see _pixel_stokes_i_error). A skipped pixel (too few finite channels or SNR below fit_options.snr_cut; see fit_stokes_i_model) falls back to a flat model at its mean Stokes I, so it gets no spectral correction (its FDF is the plain Q/U FDF), and its alpha, order and errors stay NaN. A pixel with no finite channels stays NaN throughout.

rm_lite.tools_3d.rmsynth._iter_pixel_fits(i_block: numpy.typing.NDArray[numpy.float64], err_block: numpy.typing.NDArray[numpy.float64] | None, err_1d: numpy.typing.NDArray[numpy.float64] | None, freq_arr_hz: numpy.typing.NDArray[numpy.float64], ref_freq_hz: float, fit_options: rm_lite.utils.fitting.StokesIFitOptions) → collections.abc.Iterator[PixelFit]¶

Yield a PixelFit for every pixel in a chunk.

fit is None for a skipped pixel (fit_stokes_i_model returns None when there are too few finite channels or the SNR is below fit_options.snr_cut). Shared by the model/alpha pass and the opt-in error pass so the setup lives once.

rm_lite.tools_3d.rmsynth._pixel_stokes_i_error(err_block: numpy.typing.NDArray[numpy.float64] | None, err_1d: numpy.typing.NDArray[numpy.float64] | None, n_freq: int, y: int, x: int) → numpy.typing.NDArray[numpy.float64]¶: Per-pixel Stokes I error spectrum from a 3D error cube, 1D per-channel array, or neither (zeros -> fit_stokes_i_model fits unweighted).

rm_lite.tools_3d.rmsynth._ref_flux_block(model_block: numpy.typing.NDArray[numpy.float64], freq_arr_hz: numpy.typing.NDArray[numpy.float64], ref_freq_hz: float) → numpy.typing.NDArray[numpy.float64]¶: Interpolate a Stokes I model cube block at the reference frequency -> (cy, cx) reference-flux map block.

rm_lite.tools_3d.rmsynth._split_stokes_i_error(stokes_i_error: numpy.typing.NDArray[numpy.float64] | dask.array.Array | None, chunks: tuple[tuple[int, Ellipsis], Ellipsis]) → tuple[numpy.typing.NDArray[numpy.float64] | None, dask.array.Array | None]¶

Split a Stokes I error into (per-channel 1D array, per-pixel dask cube).

A 3D error (numpy or dask) becomes a dask cube rechunked to chunks; a 1D per-channel error becomes a numpy array; None stays None. da.asarray wraps a numpy cube so .rechunk works either way.

rm_lite.tools_3d.rmsynth._stokes_i_model_cube(stokes_i: dask.array.Array, stokes_i_error: numpy.typing.NDArray[numpy.float64] | dask.array.Array | None, freq_arr_hz: numpy.typing.NDArray[numpy.float64], ref_freq_hz: float, fit_options: rm_lite.utils.fitting.StokesIFitOptions, log_level: int) → tuple[dask.array.Array, dask.array.Array, dask.array.Array, dask.array.Array | None, dask.array.Array | None]¶

Lazy per-pixel Stokes I model cube, spectral-index map, fitted-order map, and (optional) model-error cube and alpha-error map.

All come from one per-pixel fit pass (_fit_stokes_i_block): the model cube is chunked like stokes_i; the alpha and order maps are (ny, nx); and the error cube (n_freq, ny, nx) plus the alpha-error map (ny, nx) – returned only when fit_options.compute_model_error, else None – reuse the same fit, so they cost no extra fit. Alpha and order are NaN for pixels that were not fitted. stokes_i_error is a per-channel 1D array (n_freq,), a per-pixel error cube (n_freq, ny, nx), or None.

rm_lite.tools_3d.rmsynth.rmsynth_3d(stokes_q: dask.array.Array, stokes_u: dask.array.Array, freq_arr_hz: numpy.typing.NDArray[numpy.float64], weight_arr: numpy.typing.NDArray[numpy.float64] | None = None, phi_max_radm2: float | None = None, d_phi_radm2: float | None = None, n_samples: float | None = 10.0, weight_type: rm_lite.utils.synthesis.WeightType = 'variance', robust: float | None = None, stokes_i: dask.array.Array | None = None, stokes_i_error: numpy.typing.NDArray[numpy.float64] | dask.array.Array | None = None, stokes_i_model: dask.array.Array | None = None, estimate_stokes_i_noise: bool = False, fit_order: int = 2, fit_function: Literal['log', 'linear'] = 'log', stokes_i_snr_cut: float | None = 5.0, compute_model_error: bool = False, n_error_samples: int = 1000, nufft_nthreads: int = 1, log_level: int = logging.WARNING) → RMSynth3DResults¶

Run RM-synthesis on chunked Stokes Q/U cubes.

Given a Stokes I cube or model, Q/U are divided by a per-pixel Stokes I model (fitted or supplied) and the FDF is rescaled to flux at the reference frequency; otherwise the FDF stays in Q/U flux.

Parameters:

stokes_q (da.Array) – Stokes Q cube (n_freq, ny, nx), chunked spatially only.
stokes_u (da.Array) – Stokes U cube, same shape/chunks as stokes_q.
freq_arr_hz (NDArray[np.float64]) – Frequency array in Hz.
weight_arr (NDArray[np.float64] | None, optional) – Per-channel (not per-pixel) weight array. Defaults to uniform.
phi_max_radm2 (float | None, optional) – Maximum Faraday depth. Defaults to None.
d_phi_radm2 (float | None, optional) – Faraday depth resolution. Defaults to None.
n_samples (float | None, optional) – Samples across the RMSF. Defaults to 10.0.
weight_type (WeightType, optional) – ‘variance’/’natural’ (1/sigma^2), ‘uniform’ (equal per channel), ‘uniform_lsq’ (equal per lambda^2, narrows the RMSF), ‘briggs’ (robust). Defaults to “variance”.
robust (float | None, optional) – Briggs robust parameter, required for weight_type=’briggs’. Defaults to None.
stokes_i (da.Array | None, optional) – Stokes I cube to fit per pixel for the fractional correction. Ignored if stokes_i_model is given. Defaults to None (FDF stays in Q/U flux).
stokes_i_error (NDArray[np.float64] | da.Array | None, optional) – Stokes I error, per-channel (n_freq,) or per-pixel cube (n_freq, ny, nx), to weight the fit. Defaults to None (unweighted, or estimated if estimate_stokes_i_noise).
stokes_i_model (da.Array | None, optional) – Pre-computed Stokes I model cube, used directly (no fitting). Takes precedence over stokes_i. Defaults to None.
estimate_stokes_i_noise (bool, optional) – Derive a per-channel error from stokes_i when no stokes_i_error is given. Defaults to False.
fit_order (int, optional) – Stokes I fit order; negative iterates orders and picks the best by AIC. Defaults to 2.
fit_function ("log", "linear", optional) – “log” = power law, “linear” = polynomial. Defaults to “log”.
stokes_i_snr_cut (float | None, optional) – Below this frequency-averaged Stokes I SNR a pixel falls back to a flat model (no spectral correction, not blanked). None fits every pixel. Fit path only. Defaults to 5.0.
compute_model_error (bool, optional) – Also compute a per-pixel model error cube via Monte-Carlo over the fit covariance, in the same fit pass. Logs a warning about the compute coupling when enabled. Defaults to False.
n_error_samples (int, optional) – Monte-Carlo samples per pixel for compute_model_error. Defaults to 1000.
nufft_nthreads (int, optional) – finufft OpenMP threads per chunk. Defaults to 1 so dask parallelises across chunks without oversubscribing finufft’s own threads (the fast config on many chunks). Set to 0 (finufft default, all cores) only when computing with few chunks on the synchronous scheduler.
log_level (int, optional) – rm_lite logger level while chunks run; defaults to WARNING to silence per-chunk noise.

Returns:

Lazy FDF cube, RMSF cube, and parameters. With a Stokes I: model, also the model cube and the 2D reference-flux/spectral-index maps.

Return type:

RMSynth3DResults

rm_lite.tools_3d.rmsynth.rmsynth_3d_from_fits(stokes_q_file: str | pathlib.Path, stokes_u_file: str | pathlib.Path, weight_arr: numpy.typing.NDArray[numpy.float64] | None = None, phi_max_radm2: float | None = None, d_phi_radm2: float | None = None, n_samples: float | None = 10.0, weight_type: rm_lite.utils.synthesis.WeightType = 'variance', robust: float | None = None, stokes_i_file: str | pathlib.Path | None = None, stokes_i_error_file: str | pathlib.Path | None = None, stokes_i_model_file: str | pathlib.Path | None = None, estimate_stokes_i_noise: bool = False, fit_order: int = 2, fit_function: Literal['log', 'linear'] = 'log', stokes_i_snr_cut: float | None = 5.0, compute_model_error: bool = False, n_error_samples: int = 1000, nufft_nthreads: int = 1, target_chunk_mb: float = DEFAULT_TARGET_CHUNK_MB, log_level: int = logging.WARNING) → RMSynth3DResults¶

Run RM-synthesis directly on Stokes Q/U FITS cubes on disk.

Convenience wrapper around rm_lite.utils.dask_io.read_fits_cube_dask + rmsynth_3d, for the common case where Q/U are FITS files rather than already-loaded dask arrays. The frequency array is derived from the Stokes Q header’s spectral WCS, and, if weight_arr is not given, so is the per-channel weight array (via estimate_channel_noise_mad).

Parameters:

stokes_q_file (str | Path) – Path to the Stokes Q FITS cube.
stokes_u_file (str | Path) – Path to the Stokes U FITS cube, same shape as the Q cube.
weight_arr (NDArray[np.float64] | None, optional) – Per-channel weight array. Defaults to an estimate from the cube noise (see rm_lite.utils.dask_io.estimate_channel_noise_mad).
phi_max_radm2 (float | None, optional) – Maximum Faraday depth. Defaults to None.
d_phi_radm2 (float | None, optional) – Faraday depth resolution. Defaults to None.
n_samples (float | None, optional) – Number of samples across the RMSF. Defaults to 10.0.
weight_type (WeightType, optional) – See rmsynth_3d. Defaults to “variance”.
robust (float | None, optional) – Briggs robust parameter, required for weight_type=’briggs’. Defaults to None.
stokes_i_file (str | Path | None, optional) – Path to a Stokes I FITS cube (measurements) to fit per pixel for fractional-polarization correction. See rmsynth_3d. Defaults to None.
stokes_i_error_file (str | Path | None, optional) – Path to a Stokes I error FITS cube used to weight the per-pixel fit. Defaults to None.
stokes_i_model_file (str | Path | None, optional) – Path to a pre-computed Stokes I model FITS cube, used directly (no fitting). Takes precedence over stokes_i_file. Defaults to None.
estimate_stokes_i_noise (bool, optional) – See rmsynth_3d. Defaults to False.
fit_order (int, optional) – See rmsynth_3d. Defaults to 2.
fit_function ("log", "linear", optional) – See rmsynth_3d. Defaults to “log”.
stokes_i_snr_cut (float | None, optional) – See rmsynth_3d. Defaults to 5.0.
compute_model_error (bool, optional) – See rmsynth_3d. Defaults to False.
n_error_samples (int, optional) – See rmsynth_3d. Defaults to 1000.
nufft_nthreads (int, optional) – See rmsynth_3d. Defaults to 1.
target_chunk_mb (float, optional) – Target per-chunk memory footprint in MB, see read_fits_cube_dask. Defaults to 256.
log_level (int, optional) – See rmsynth_3d. Defaults to logging.WARNING.

Returns:

Lazy dirty FDF cube, RMSF cube, and associated parameters.

Return type:

RMSynth3DResults