3D RM-CLEAN

We start with the same synthetic cube as the 3D RM-synthesis example (see that page for details) and run RM-synthesis again here to get the dirty FDF and RMSF cubes.

from __future__ import annotations

import tempfile
from pathlib import Path

import astropy.units as u
import dask
import matplotlib.pyplot as plt
import numpy as np
import zarr
from astropy.io import fits
from astropy.visualization import quantity_support

plt.rcParams["figure.dpi"] = 150

_ = quantity_support()
rng = np.random.default_rng(42)

bw_low = 288
freqs = np.linspace(943.5 - bw_low / 2, 943.5 + bw_low / 2, 36) * u.MHz
freq_hz = freqs.to(u.Hz).value

ny, nx = 64, 64
y_grid, x_grid = np.mgrid[0:ny, 0:nx]
blob_y, blob_x = ny * 0.3, nx * 0.7
blob2_y, blob2_x = ny * 0.7, nx * 0.3
rm_map = (
    80.0 * (x_grid / nx - 0.5) * 2
    + np.exp(-((x_grid - blob_x) ** 2 + (y_grid - blob_y) ** 2) / (2 * 4.0**2))
    + np.exp(-((x_grid - blob2_x) ** 2 + (y_grid - blob2_y) ** 2) / (2 * 4.0**2))
)
radius_grid = np.hypot(x_grid - blob_x, y_grid - blob_y)
radius2_grid = np.hypot(x_grid - blob2_x, y_grid - blob2_y)
psi0_deg = 30.0
rms_noise = 0.03

from rm_lite.utils.fitting import gaussian
from rm_lite.utils.synthesis import faraday_simple_spectrum

# Two compact polarised sources (2D Gaussian blobs, reusing the existing
# 1D gaussian() on a radial distance grid), one bottom-right and one
# top-left, on an otherwise unpolarised background. Mostly-empty sky
# is what the noise estimator below needs.
frac_pol_map = gaussian(radius_grid, amplitude=0.6, mean=0.0, fwhm=6.0) + gaussian(
    radius2_grid, amplitude=0.6, mean=0.0, fwhm=6.0
)

stokes_q = np.empty((freq_hz.size, ny, nx), dtype=np.float32)
stokes_u = np.empty((freq_hz.size, ny, nx), dtype=np.float32)
for j in range(ny):
    for i in range(nx):
        complex_spectrum = faraday_simple_spectrum(
            freq_hz,
            frac_pol=frac_pol_map[j, i],
            psi0_deg=psi0_deg,
            rm_radm2=rm_map[j, i],
        )
        stokes_q[:, j, i] = complex_spectrum.real + rng.normal(
            0, rms_noise, freq_hz.size
        )
        stokes_u[:, j, i] = complex_spectrum.imag + rng.normal(
            0, rms_noise, freq_hz.size
        )

tmpdir = Path(tempfile.mkdtemp())


def write_stokes_fits(path: Path, data: np.ndarray, stokes: int) -> None:
    header = fits.Header()
    header["CTYPE1"] = "RA---SIN"
    header["CTYPE2"] = "DEC--SIN"
    header["CTYPE3"] = "FREQ"
    header["CRVAL3"] = freq_hz[0]
    header["CDELT3"] = float(np.diff(freq_hz)[0])
    header["CRPIX3"] = 1.0
    header["CUNIT3"] = "Hz"
    header["CTYPE4"] = "STOKES"
    header["CRVAL4"] = stokes
    header["CDELT4"] = 1.0
    header["CRPIX4"] = 1.0
    fits.PrimaryHDU(data=data[np.newaxis, ...], header=header).writeto(
        path, overwrite=True
    )


stokes_q_fits = tmpdir / "cutout.q.fits"
stokes_u_fits = tmpdir / "cutout.u.fits"
write_stokes_fits(stokes_q_fits, stokes_q, stokes=2)
write_stokes_fits(stokes_u_fits, stokes_u, stokes=3)

from rm_lite.tools_3d.rmsynth import rmsynth_3d
from rm_lite.utils.dask_io import estimate_channel_noise_mad, read_fits_cube_dask

q_dask, _ = read_fits_cube_dask(stokes_q_fits, target_chunk_mb=32 * 1024 / 1024**2)
u_dask, _ = read_fits_cube_dask(stokes_u_fits, target_chunk_mb=32 * 1024 / 1024**2)

channel_noise = estimate_channel_noise_mad(q_dask, u_dask)
weight_arr = 1.0 / channel_noise**2

synth = rmsynth_3d(
    q_dask,
    u_dask,
    freq_hz,
    weight_arr=weight_arr,
    phi_max_radm2=150.0,
    d_phi_radm2=10.0,
)
synth.fdf_dirty_cube

INFO dask_io.estimate_channel_noise_mad: Rechunking Stokes Q/U to one spatial block per channel for noise estimation.
INFO dask_io.estimate_channel_noise_mad: Per-channel noise estimation completed in 0.247 seconds.

Array Chunk Bytes 1.94 MiB 108.98 kiB Shape (31, 64, 64) (31, 15, 15) Dask graph 25 chunks in 165 graph layers Data type complex128 numpy.ndarray

Now run 3D RM-CLEAN with rmclean_3d. The Hogbom clean loop (rm_lite.utils.clean.rmclean) is inherently per-pixel, so it runs once per spatial chunk via dask.delayed. It returns four outputs (clean FDF, model FDF, residual FDF, iteration count) from that one call, split out of the same delayed computation rather than re-running the clean loop per output.

Unlike the 1D convenience wrapper run_rmclean_from_synth, rmclean_3d takes raw flux thresholds directly rather than deriving an auto-mask/auto-threshold itself. A whole cube doesn’t have one noise value, so it needs one explicit uniform threshold rather than a per-pixel one. We derive it here from synth.theoretical_noise, the per-channel-noise-derived FDF noise rmsynth_3d already computed, the same value rmclean_3d_from_synth (below) scales automatically via its own auto_mask/auto_threshold.

from rm_lite.tools_3d.rmclean import rmclean_3d

help(rmclean_3d)

Help on function rmclean_3d in module rm_lite.tools_3d.rmclean:

rmclean_3d(fdf_dirty_cube: 'da.Array', rmsf_cube: 'da.Array', phi_arr_radm2: 'NDArray[np.float64]', phi_double_arr_radm2: 'NDArray[np.float64]', fwhm_rmsf_radm2: 'float', mask: 'float', threshold: 'float', max_iter: 'int' = 1000, gain: 'float' = 0.1, moment_threshold: 'float | None' = None, log_level: 'int' = 40) -> 'RMClean3DResults'
    Run RM-CLEAN on chunked dirty FDF and RMSF cubes.

    Args:
        fdf_dirty_cube (da.Array): Dirty FDF cube, shape (n_phi, ny, nx),
            chunked spatially only (as produced by `rm_lite.tools_3d.rmsynth.rmsynth_3d`).
        rmsf_cube (da.Array): RMSF cube, shape (n_phi_double, ny, nx), with the
            same spatial chunking as `fdf_dirty_cube`.
        phi_arr_radm2 (NDArray[np.float64]): Faraday depth values in rad/m^2.
        phi_double_arr_radm2 (NDArray[np.float64]): Double-length Faraday depth
            values in rad/m^2, for the RMSF.
        fwhm_rmsf_radm2 (float): RMSF FWHM, shared by every pixel (3D RM-CLEAN
            here does not support a per-pixel FWHM map).
        mask (float): Masking threshold. Pixels below this value are not cleaned.
        threshold (float): Cleaning threshold. Stop when all pixels are below this value.
        max_iter (int, optional): Maximum CLEAN iterations. Defaults to 1000.
        gain (float, optional): CLEAN loop gain. Defaults to 0.1.
        moment_threshold (float | None, optional): Amplitude cut (in FDF
            amplitude units) applied to the clean FDF before computing the
            Faraday moment maps, passed to `calc_faraday_moments`. None includes
            all amplitudes (noise-biased). Defaults to None.
        log_level (int, optional): Log level applied to `rm_lite`'s logger while
            each chunk runs. `rmclean`'s Hogbom loop logs at INFO and WARNING
            per pixel (e.g. "Starting minor loop...", "All channels masked...
            performed N iterations"). These are routine per-pixel loop
            termination conditions, not anomalies, and at cube scale they're
            just noise, so this defaults to ERROR (silencing both). Pass
            `logging.WARNING` or `logging.INFO` to restore progressively more
            per-pixel verbosity, e.g. while debugging a specific chunk.
            Defaults to `logging.ERROR`.

    Returns:
        RMClean3DResults: Lazy clean/model/residual FDF cubes and iteration-count map.

auto_mask, auto_threshold = 8.0, 2.0
mask = auto_mask * synth.theoretical_noise.fdf_error_noise
threshold = auto_threshold * synth.theoretical_noise.fdf_error_noise

clean = rmclean_3d(
    synth.fdf_dirty_cube,
    synth.rmsf_cube,
    synth.phi_arr_radm2,
    synth.phi_double_arr_radm2,
    synth.fwhm_rmsf_radm2,
    mask=mask,
    threshold=threshold,
)
clean.iter_count_map

Array Chunk Bytes 32.00 kiB 1.76 kiB Shape (64, 64) (15, 15) Dask graph 25 chunks in 83 graph layers Data type int64 numpy.ndarray

rmclean_3d above takes the dirty FDF/RMSF cubes and phi arrays unpacked from synth one by one. rmclean_3d_from_synth takes the RMSynth3DResults itself and unpacks it internally, mirroring the 1D run_rmclean_from_synth convenience, including deriving mask/threshold itself from auto_mask/auto_threshold scaling synth.theoretical_noise (the same computation as above). It also takes moment_threshold_snr (default 5), the SNR cut applied to the clean FDF before the Faraday moment maps are computed (see below).

from rm_lite.tools_3d.rmclean import rmclean_3d_from_synth

clean_from_synth = rmclean_3d_from_synth(
    synth, auto_mask=auto_mask, auto_threshold=auto_threshold
)
np.array_equal(
    clean_from_synth.iter_count_map.compute(), clean.iter_count_map.compute()
)

INFO rmclean.rmclean_3d_from_synth: Theoretical FDF noise: 0.00527. Auto mask: 0.0421, auto threshold: 0.0105.

True

Everything is still lazy. Compute all four outputs together with dask.compute so the shared per-chunk clean loop only runs once, then compare the dirty and clean peak polarised intensity maps and look at where CLEAN actually did work.

clean_fdf_cube, model_fdf_cube, resid_fdf_cube, iter_count_map = dask.compute(
    clean.clean_fdf_cube,
    clean.model_fdf_cube,
    clean.resid_fdf_cube,
    clean.iter_count_map,
)

dirty_fdf_cube = synth.fdf_dirty_cube.compute()
dirty_pi_map = np.max(np.abs(dirty_fdf_cube), axis=0)
clean_pi_map = np.max(np.abs(clean_fdf_cube), axis=0)

fig, axs = plt.subplots(1, 3, figsize=(14, 4))
im0 = axs[0].imshow(dirty_pi_map, origin="lower")
fig.colorbar(im0, ax=axs[0])
axs[0].set(title="Dirty peak PI")
im1 = axs[1].imshow(clean_pi_map, origin="lower")
fig.colorbar(im1, ax=axs[1])
axs[1].set(title="Clean peak PI")
im2 = axs[2].imshow(iter_count_map, origin="lower")
fig.colorbar(im2, ax=axs[2])
axs[2].set(title="CLEAN iterations")

[Text(0.5, 1.0, 'CLEAN iterations')]

Serialisation: zarr vs FITS

As with the RM-synth outputs, the clean/model/residual FDF cubes are complex, so zarr again is the natural fit, writing lazily chunk by chunk. The iteration-count map is different: it is real-valued (an integer per pixel), so it has no complex-split problem and writes to FITS directly, like any ordinary 2D image.

from rm_lite.utils.dask_io import write_zarr_group

zarr_store = tmpdir / "rmclean3d.zarr"
write_zarr_group(
    zarr_store,
    {
        "fdf_clean": clean.clean_fdf_cube,
        "fdf_model": clean.model_fdf_cube,
        "fdf_resid": clean.resid_fdf_cube,
        "iter_count": clean.iter_count_map,
    },
)

group = zarr.open(zarr_store)
group["fdf_clean"].shape, group["fdf_clean"].dtype, group["iter_count"].dtype

[########################################] | 100% Completed | 5.53 s

INFO dask_io.write_zarr_group: Wrote ['fdf_clean', 'fdf_model', 'fdf_resid', 'iter_count'] to /tmp/tmp092w77zl/rmclean3d.zarr in 5.55 seconds.

((31, 64, 64), dtype('<c16'), dtype('<i8'))

iter_count_fits = tmpdir / "CLEAN_nIter.fits"
fits.PrimaryHDU(iter_count_map.astype(np.float32)).writeto(
    iter_count_fits, overwrite=True
)

roundtrip = fits.getdata(iter_count_fits)
np.array_equal(roundtrip, iter_count_map.astype(np.float32))

True

Faraday moments

rmclean_3d (and rmclean_3d_from_synth) compute the zeroth, first, and second moments of the Faraday depth spectrum (see Dickey et al. 2019) directly on the clean FDF cube, and return them as lazy 2D maps on the result: clean.mom0_map, clean.mom1_map, clean.mom2_map.

• mom0: total polarised intensity. The FDF amplitude is in units per RMSF, so the sum over Faraday depth is divided by the RMSF area (a Gaussian of FWHM fwhm_rmsf_radm2); an unresolved component of peak amplitude \(P\) gives \(\mathrm{mom0} = P\), in the input flux units.

• mom1: intensity-weighted mean Faraday depth, in rad m\(^{-2}\).

• mom2: intensity-weighted Faraday depth dispersion (standard deviation), in rad m\(^{-2}\).

By default no amplitude threshold is applied (clean.mom*_map below), so the noise floor biases the moments, especially mom2. The underlying function is rm_lite.utils.synthesis.calc_faraday_moments; call it directly on clean.clean_fdf_cube to re-threshold without re-cleaning, or pass moment_threshold_snr to rmclean_3d_from_synth (moment_threshold to rmclean_3d) to bake a cut into the result maps.

from rm_lite.utils.synthesis import calc_faraday_moments

# The moment maps come straight off the clean result. `clean` above was run
# without a moment threshold, so these are the raw (noise-biased) moments.
mom0_map, mom1_map, mom2_map = dask.compute(
    clean.mom0_map, clean.mom1_map, clean.mom2_map
)

fig, axs = plt.subplots(1, 3, figsize=(14, 4))
im0 = axs[0].imshow(mom0_map, origin="lower")
fig.colorbar(im0, ax=axs[0])
axs[0].set(title="mom0 (total polarised intensity)")
im1 = axs[1].imshow(mom1_map, origin="lower", cmap="RdBu_r", vmin=-80, vmax=80)
fig.colorbar(im1, ax=axs[1])
axs[1].set(title="mom1 (rad m$^{-2}$)")
im2 = axs[2].imshow(mom2_map, origin="lower", cmap="magma")
fig.colorbar(im2, ax=axs[2])
axs[2].set(title="mom2 (rad m$^{-2}$)")

[Text(0.5, 1.0, 'mom2 (rad m$^{-2}$)')]

To suppress the noise-floor bias, exclude amplitudes below a threshold, here 5x the theoretical FDF noise that rmsynth_3d already computed. The intended way is to ask for it up front, e.g. rmclean_3d_from_synth(synth, moment_threshold_snr=5.0), and read clean.mom*_map. To explore thresholds without re-cleaning, call calc_faraday_moments directly on the clean cube, as below:

moments = calc_faraday_moments(
    clean.clean_fdf_cube,
    synth.phi_arr_radm2,
    synth.fwhm_rmsf_radm2,
    threshold=5.0 * synth.theoretical_noise.fdf_error_noise,
)
mom0_map, mom1_map, mom2_map = dask.compute(*moments)

fig, axs = plt.subplots(1, 3, figsize=(14, 4))
im0 = axs[0].imshow(mom0_map, origin="lower")
fig.colorbar(im0, ax=axs[0])
axs[0].set(title="mom0 (total polarised intensity)")
im1 = axs[1].imshow(mom1_map, origin="lower", cmap="RdBu_r", vmin=-80, vmax=80)
fig.colorbar(im1, ax=axs[1])
axs[1].set(title="mom1 (rad m$^{-2}$)")
im2 = axs[2].imshow(mom2_map, origin="lower", cmap="magma")
fig.colorbar(im2, ax=axs[2])
axs[2].set(title="mom2 (rad m$^{-2}$)")

[Text(0.5, 1.0, 'mom2 (rad m$^{-2}$)')]

Pixels with no amplitude above the threshold have mom0 = 0 and NaN mom1/mom2. At the faint edges of the sources, where only a sample or two clears the cut, mom1 is still noisy, so compare against the input RM map only where mom0 is comfortably detected:

detected = np.isfinite(mom1_map) & (
    mom0_map > 10 * synth.theoretical_noise.fdf_error_noise
)
fig, ax = plt.subplots()
ax.scatter(rm_map[detected], mom1_map[detected], s=10)
ax.axline((0, 0), slope=1, color="k", ls="--")
ax.set(xlabel="Input RM (rad m$^{-2}$)", ylabel="mom1 (rad m$^{-2}$)", aspect="equal")

assert np.allclose(rm_map[detected], mom1_map[detected], atol=10.0)

Unlike the complex FDF cubes, the moment maps are real-valued 2D images, so they save straight to FITS:

for name, data, unit in [
    ("FDF_mom0", mom0_map, ""),
    ("FDF_mom1", mom1_map, "rad m-2"),
    ("FDF_mom2", mom2_map, "rad m-2"),
]:
    hdu = fits.PrimaryHDU(np.asarray(data, dtype=np.float32))
    if unit:
        hdu.header["BUNIT"] = unit
    hdu.writeto(tmpdir / f"{name}.fits", overwrite=True)

sorted(p.name for p in tmpdir.glob("FDF_mom*.fits"))

['FDF_mom0.fits', 'FDF_mom1.fits', 'FDF_mom2.fits']

Debiased moments

np.abs() of a noisy complex FDF is Ricean-distributed, giving polarised intensity a positive noise bias. A >5-sigma cut (as above) sidesteps most of it, but the bias accumulates whenever you integrate without a cut, exactly what mom0 does, so the unthresholded mom0 of an empty sightline is a positive floor, not zero.

rm_lite.utils.synthesis.debias_fdf implements the bias suppression of Mueller, Beck & Krause (2017), adapted for Faraday cubes. Their method median-filters the polarisation angle over the spatial axes (via its cos/sin components, dodging the angle wrap) and projects Q + iU onto the filtered direction, \(P^* = U \sin\theta_m + Q \cos\theta_m\): in signal regions \(P^* \approx |F|\), while in noise regions \(P^*\) is zero-mean, signed Gaussian noise that cancels when summed, instead of a positive Rayleigh floor.

Applied naively per Faraday depth plane, that filtering fails wherever RM varies across the sky: the FDF angle is \(2\psi_0 + 2\lambda_0^2(\mathrm{RM} - \phi)\), so this scene’s RM gradient spins it by \(\sim\) 1 rad per pixel and the projection would destroy real signal. debias_fdf removes that deterministic ramp first, derotating each spectrum by a per-pixel peak Faraday depth estimate (half-max centroid), so only the intrinsic angle \(2\psi_0\), smooth for physical sources, is filtered. This makes it robust to any background RM structure; the remaining caveat is sightlines with several components at very different Faraday depths, where only the dominant peak is fully derotated.

calc_faraday_moments treats real input as-is (no abs()), so the signed debiased amplitudes integrate correctly, or pass debias=True and let it call debias_fdf internally:

from rm_lite.utils.synthesis import debias_fdf

fdf_debias_cube = debias_fdf(
    clean.clean_fdf_cube,
    phi_arr_radm2=synth.phi_arr_radm2,
    lam_sq_0_m2=synth.lam_sq_0_m2,
)

# Unthresholded mom0: |FDF| accumulates the noise floor, the signed
# debiased amplitudes cancel it. `debias=True` is the one-call equivalent.
mom0_unthresh_map, mom0_debias_unthresh_map = dask.compute(
    calc_faraday_moments(
        clean.clean_fdf_cube, synth.phi_arr_radm2, synth.fwhm_rmsf_radm2
    ).mom0,
    calc_faraday_moments(
        clean.clean_fdf_cube,
        synth.phi_arr_radm2,
        synth.fwhm_rmsf_radm2,
        debias=True,
        lam_sq_0_m2=synth.lam_sq_0_m2,
    ).mom0,
)

fig, axs = plt.subplots(1, 3, figsize=(16, 4))
vmax = np.nanmax(mom0_unthresh_map)
im0 = axs[0].imshow(mom0_unthresh_map, origin="lower", vmin=0, vmax=vmax)
fig.colorbar(im0, ax=axs[0])
axs[0].set(title="mom0, no threshold (biased)")
im1 = axs[1].imshow(mom0_debias_unthresh_map, origin="lower", vmin=0, vmax=vmax)
fig.colorbar(im1, ax=axs[1])
axs[1].set(title="mom0, no threshold (debiased)")
im2 = axs[2].imshow(
    mom0_debias_unthresh_map - mom0_unthresh_map,
    origin="lower",
    cmap="coolwarm",
    vmin=-0.1,
    vmax=0.1,
)
fig.colorbar(im2, ax=axs[2])
axs[2].set(title="difference (debiased - biased)")

# Empty sky: the debiased floor collapses...
empty = frac_pol_map < 0.01
floor_raw = np.nanmedian(mom0_unthresh_map[empty])
floor_debias = np.nanmedian(np.abs(mom0_debias_unthresh_map[empty]))
print(f"empty-sky mom0 floor: {floor_raw:.4f} -> {floor_debias:.4f}")
assert floor_debias < 0.2 * floor_raw

# ...while detected sources keep their flux (compare the thresholded mom0)
ratio = mom0_debias_unthresh_map[detected] / mom0_map[detected]
assert 0.85 < np.nanmedian(ratio) < 1.15

fits.PrimaryHDU(mom0_debias_unthresh_map.astype(np.float32)).writeto(
    tmpdir / "FDF_mom0_debias.fits", overwrite=True
)

empty-sky mom0 floor: 0.0309 -> 0.0053