{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 3D RM-CLEAN" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We start with the same synthetic cube as the [3D RM-synthesis](rmsynth_3d.ipynb) example (see that page for details) and run RM-synthesis again here to get the dirty FDF and RMSF cubes." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from __future__ import annotations\n", "\n", "import tempfile\n", "from pathlib import Path\n", "\n", "import astropy.units as u\n", "import dask\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import zarr\n", "from astropy.io import fits\n", "from astropy.visualization import quantity_support\n", "\n", "plt.rcParams[\"figure.dpi\"] = 150\n", "\n", "_ = quantity_support()\n", "rng = np.random.default_rng(42)\n", "\n", "bw_low = 288\n", "freqs = np.linspace(943.5 - bw_low / 2, 943.5 + bw_low / 2, 36) * u.MHz\n", "freq_hz = freqs.to(u.Hz).value\n", "\n", "ny, nx = 64, 64\n", "y_grid, x_grid = np.mgrid[0:ny, 0:nx]\n", "blob_y, blob_x = ny * 0.3, nx * 0.7\n", "blob2_y, blob2_x = ny * 0.7, nx * 0.3\n", "rm_map = (\n", " 80.0 * (x_grid / nx - 0.5) * 2\n", " + np.exp(-((x_grid - blob_x) ** 2 + (y_grid - blob_y) ** 2) / (2 * 4.0**2))\n", " + np.exp(-((x_grid - blob2_x) ** 2 + (y_grid - blob2_y) ** 2) / (2 * 4.0**2))\n", ")\n", "radius_grid = np.hypot(x_grid - blob_x, y_grid - blob_y)\n", "radius2_grid = np.hypot(x_grid - blob2_x, y_grid - blob2_y)\n", "psi0_deg = 30.0\n", "rms_noise = 0.03" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from rm_lite.utils.fitting import gaussian\n", "from rm_lite.utils.synthesis import faraday_simple_spectrum\n", "\n", "# Two compact polarised sources (2D Gaussian blobs, reusing the existing\n", "# 1D gaussian() on a radial distance grid), one bottom-right and one\n", "# top-left, on an otherwise unpolarised background. Mostly-empty sky\n", "# is what the noise estimator below needs.\n", "frac_pol_map = gaussian(radius_grid, amplitude=0.6, mean=0.0, fwhm=6.0) + gaussian(\n", " radius2_grid, amplitude=0.6, mean=0.0, fwhm=6.0\n", ")\n", "\n", "stokes_q = np.empty((freq_hz.size, ny, nx), dtype=np.float32)\n", "stokes_u = np.empty((freq_hz.size, ny, nx), dtype=np.float32)\n", "for j in range(ny):\n", " for i in range(nx):\n", " complex_spectrum = faraday_simple_spectrum(\n", " freq_hz,\n", " frac_pol=frac_pol_map[j, i],\n", " psi0_deg=psi0_deg,\n", " rm_radm2=rm_map[j, i],\n", " )\n", " stokes_q[:, j, i] = complex_spectrum.real + rng.normal(\n", " 0, rms_noise, freq_hz.size\n", " )\n", " stokes_u[:, j, i] = complex_spectrum.imag + rng.normal(\n", " 0, rms_noise, freq_hz.size\n", " )\n", "\n", "tmpdir = Path(tempfile.mkdtemp())\n", "\n", "\n", "def write_stokes_fits(path: Path, data: np.ndarray, stokes: int) -> None:\n", " header = fits.Header()\n", " header[\"CTYPE1\"] = \"RA---SIN\"\n", " header[\"CTYPE2\"] = \"DEC--SIN\"\n", " header[\"CTYPE3\"] = \"FREQ\"\n", " header[\"CRVAL3\"] = freq_hz[0]\n", " header[\"CDELT3\"] = float(np.diff(freq_hz)[0])\n", " header[\"CRPIX3\"] = 1.0\n", " header[\"CUNIT3\"] = \"Hz\"\n", " header[\"CTYPE4\"] = \"STOKES\"\n", " header[\"CRVAL4\"] = stokes\n", " header[\"CDELT4\"] = 1.0\n", " header[\"CRPIX4\"] = 1.0\n", " fits.PrimaryHDU(data=data[np.newaxis, ...], header=header).writeto(\n", " path, overwrite=True\n", " )\n", "\n", "\n", "stokes_q_fits = tmpdir / \"cutout.q.fits\"\n", "stokes_u_fits = tmpdir / \"cutout.u.fits\"\n", "write_stokes_fits(stokes_q_fits, stokes_q, stokes=2)\n", "write_stokes_fits(stokes_u_fits, stokes_u, stokes=3)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from rm_lite.tools_3d.rmsynth import rmsynth_3d\n", "from rm_lite.utils.dask_io import estimate_channel_noise_mad, read_fits_cube_dask\n", "\n", "q_dask, _ = read_fits_cube_dask(stokes_q_fits, target_chunk_mb=32 * 1024 / 1024**2)\n", "u_dask, _ = read_fits_cube_dask(stokes_u_fits, target_chunk_mb=32 * 1024 / 1024**2)\n", "\n", "channel_noise = estimate_channel_noise_mad(q_dask, u_dask)\n", "weight_arr = 1.0 / channel_noise**2\n", "\n", "synth = rmsynth_3d(\n", " q_dask,\n", " u_dask,\n", " freq_hz,\n", " weight_arr=weight_arr,\n", " phi_max_radm2=150.0,\n", " d_phi_radm2=10.0,\n", ")\n", "synth.fdf_dirty_cube" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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.\n", "\n", "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`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from rm_lite.tools_3d.rmclean import rmclean_3d\n", "\n", "help(rmclean_3d)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "auto_mask, auto_threshold = 8.0, 2.0\n", "mask = auto_mask * synth.theoretical_noise.fdf_error_noise\n", "threshold = auto_threshold * synth.theoretical_noise.fdf_error_noise\n", "\n", "clean = rmclean_3d(\n", " synth.fdf_dirty_cube,\n", " synth.rmsf_cube,\n", " synth.phi_arr_radm2,\n", " synth.phi_double_arr_radm2,\n", " synth.fwhm_rmsf_radm2,\n", " mask=mask,\n", " threshold=threshold,\n", ")\n", "clean.iter_count_map" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "`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)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from rm_lite.tools_3d.rmclean import rmclean_3d_from_synth\n", "\n", "clean_from_synth = rmclean_3d_from_synth(\n", " synth, auto_mask=auto_mask, auto_threshold=auto_threshold\n", ")\n", "np.array_equal(\n", " clean_from_synth.iter_count_map.compute(), clean.iter_count_map.compute()\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "clean_fdf_cube, model_fdf_cube, resid_fdf_cube, iter_count_map = dask.compute(\n", " clean.clean_fdf_cube,\n", " clean.model_fdf_cube,\n", " clean.resid_fdf_cube,\n", " clean.iter_count_map,\n", ")\n", "\n", "dirty_fdf_cube = synth.fdf_dirty_cube.compute()\n", "dirty_pi_map = np.max(np.abs(dirty_fdf_cube), axis=0)\n", "clean_pi_map = np.max(np.abs(clean_fdf_cube), axis=0)\n", "\n", "fig, axs = plt.subplots(1, 3, figsize=(14, 4))\n", "im0 = axs[0].imshow(dirty_pi_map, origin=\"lower\")\n", "fig.colorbar(im0, ax=axs[0])\n", "axs[0].set(title=\"Dirty peak PI\")\n", "im1 = axs[1].imshow(clean_pi_map, origin=\"lower\")\n", "fig.colorbar(im1, ax=axs[1])\n", "axs[1].set(title=\"Clean peak PI\")\n", "im2 = axs[2].imshow(iter_count_map, origin=\"lower\")\n", "fig.colorbar(im2, ax=axs[2])\n", "axs[2].set(title=\"CLEAN iterations\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Serialisation: zarr vs FITS\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from rm_lite.utils.dask_io import write_zarr_group\n", "\n", "zarr_store = tmpdir / \"rmclean3d.zarr\"\n", "write_zarr_group(\n", " zarr_store,\n", " {\n", " \"fdf_clean\": clean.clean_fdf_cube,\n", " \"fdf_model\": clean.model_fdf_cube,\n", " \"fdf_resid\": clean.resid_fdf_cube,\n", " \"iter_count\": clean.iter_count_map,\n", " },\n", ")\n", "\n", "group = zarr.open(zarr_store)\n", "group[\"fdf_clean\"].shape, group[\"fdf_clean\"].dtype, group[\"iter_count\"].dtype" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "iter_count_fits = tmpdir / \"CLEAN_nIter.fits\"\n", "fits.PrimaryHDU(iter_count_map.astype(np.float32)).writeto(\n", " iter_count_fits, overwrite=True\n", ")\n", "\n", "roundtrip = fits.getdata(iter_count_fits)\n", "np.array_equal(roundtrip, iter_count_map.astype(np.float32))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Faraday moments\n", "\n", "`rmclean_3d` (and `rmclean_3d_from_synth`) compute the zeroth, first, and second moments of the Faraday depth spectrum (see [Dickey et al. 2019](https://arxiv.org/abs/1812.05399)) 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`.\n", "\n", "- **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.\n", "- **mom1**: intensity-weighted mean Faraday depth, in rad m$^{-2}$.\n", "- **mom2**: intensity-weighted Faraday depth dispersion (standard deviation), in rad m$^{-2}$.\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from rm_lite.utils.synthesis import calc_faraday_moments" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# The moment maps come straight off the clean result. `clean` above was run\n", "# without a moment threshold, so these are the raw (noise-biased) moments.\n", "mom0_map, mom1_map, mom2_map = dask.compute(\n", " clean.mom0_map, clean.mom1_map, clean.mom2_map\n", ")\n", "\n", "fig, axs = plt.subplots(1, 3, figsize=(14, 4))\n", "im0 = axs[0].imshow(mom0_map, origin=\"lower\")\n", "fig.colorbar(im0, ax=axs[0])\n", "axs[0].set(title=\"mom0 (total polarised intensity)\")\n", "im1 = axs[1].imshow(mom1_map, origin=\"lower\", cmap=\"RdBu_r\", vmin=-80, vmax=80)\n", "fig.colorbar(im1, ax=axs[1])\n", "axs[1].set(title=\"mom1 (rad m$^{-2}$)\")\n", "im2 = axs[2].imshow(mom2_map, origin=\"lower\", cmap=\"magma\")\n", "fig.colorbar(im2, ax=axs[2])\n", "axs[2].set(title=\"mom2 (rad m$^{-2}$)\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "moments = calc_faraday_moments(\n", " clean.clean_fdf_cube,\n", " synth.phi_arr_radm2,\n", " synth.fwhm_rmsf_radm2,\n", " threshold=5.0 * synth.theoretical_noise.fdf_error_noise,\n", ")\n", "mom0_map, mom1_map, mom2_map = dask.compute(*moments)\n", "\n", "fig, axs = plt.subplots(1, 3, figsize=(14, 4))\n", "im0 = axs[0].imshow(mom0_map, origin=\"lower\")\n", "fig.colorbar(im0, ax=axs[0])\n", "axs[0].set(title=\"mom0 (total polarised intensity)\")\n", "im1 = axs[1].imshow(mom1_map, origin=\"lower\", cmap=\"RdBu_r\", vmin=-80, vmax=80)\n", "fig.colorbar(im1, ax=axs[1])\n", "axs[1].set(title=\"mom1 (rad m$^{-2}$)\")\n", "im2 = axs[2].imshow(mom2_map, origin=\"lower\", cmap=\"magma\")\n", "fig.colorbar(im2, ax=axs[2])\n", "axs[2].set(title=\"mom2 (rad m$^{-2}$)\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "detected = np.isfinite(mom1_map) & (\n", " mom0_map > 10 * synth.theoretical_noise.fdf_error_noise\n", ")\n", "fig, ax = plt.subplots()\n", "ax.scatter(rm_map[detected], mom1_map[detected], s=10)\n", "ax.axline((0, 0), slope=1, color=\"k\", ls=\"--\")\n", "ax.set(xlabel=\"Input RM (rad m$^{-2}$)\", ylabel=\"mom1 (rad m$^{-2}$)\", aspect=\"equal\")\n", "\n", "assert np.allclose(rm_map[detected], mom1_map[detected], atol=10.0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Unlike the complex FDF cubes, the moment maps are real-valued 2D images, so they save straight to FITS:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for name, data, unit in [\n", " (\"FDF_mom0\", mom0_map, \"\"),\n", " (\"FDF_mom1\", mom1_map, \"rad m-2\"),\n", " (\"FDF_mom2\", mom2_map, \"rad m-2\"),\n", "]:\n", " hdu = fits.PrimaryHDU(np.asarray(data, dtype=np.float32))\n", " if unit:\n", " hdu.header[\"BUNIT\"] = unit\n", " hdu.writeto(tmpdir / f\"{name}.fits\", overwrite=True)\n", "\n", "sorted(p.name for p in tmpdir.glob(\"FDF_mom*.fits\"))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Debiased moments\n", "\n", "`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.\n", "\n", "`rm_lite.utils.synthesis.debias_fdf` implements the bias suppression of [Mueller, Beck & Krause (2017)](https://ui.adsabs.harvard.edu/abs/2017A%26A...600A..63M/abstract), 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.\n", "\n", "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.\n", "\n", "`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:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from rm_lite.utils.synthesis import debias_fdf\n", "\n", "fdf_debias_cube = debias_fdf(\n", " clean.clean_fdf_cube,\n", " phi_arr_radm2=synth.phi_arr_radm2,\n", " lam_sq_0_m2=synth.lam_sq_0_m2,\n", ")\n", "\n", "# Unthresholded mom0: |FDF| accumulates the noise floor, the signed\n", "# debiased amplitudes cancel it. `debias=True` is the one-call equivalent.\n", "mom0_unthresh_map, mom0_debias_unthresh_map = dask.compute(\n", " calc_faraday_moments(\n", " clean.clean_fdf_cube, synth.phi_arr_radm2, synth.fwhm_rmsf_radm2\n", " ).mom0,\n", " calc_faraday_moments(\n", " clean.clean_fdf_cube,\n", " synth.phi_arr_radm2,\n", " synth.fwhm_rmsf_radm2,\n", " debias=True,\n", " lam_sq_0_m2=synth.lam_sq_0_m2,\n", " ).mom0,\n", ")\n", "\n", "fig, axs = plt.subplots(1, 3, figsize=(16, 4))\n", "vmax = np.nanmax(mom0_unthresh_map)\n", "im0 = axs[0].imshow(mom0_unthresh_map, origin=\"lower\", vmin=0, vmax=vmax)\n", "fig.colorbar(im0, ax=axs[0])\n", "axs[0].set(title=\"mom0, no threshold (biased)\")\n", "im1 = axs[1].imshow(mom0_debias_unthresh_map, origin=\"lower\", vmin=0, vmax=vmax)\n", "fig.colorbar(im1, ax=axs[1])\n", "axs[1].set(title=\"mom0, no threshold (debiased)\")\n", "im2 = axs[2].imshow(\n", " mom0_debias_unthresh_map - mom0_unthresh_map,\n", " origin=\"lower\",\n", " cmap=\"coolwarm\",\n", " vmin=-0.1,\n", " vmax=0.1,\n", ")\n", "fig.colorbar(im2, ax=axs[2])\n", "axs[2].set(title=\"difference (debiased - biased)\")\n", "\n", "# Empty sky: the debiased floor collapses...\n", "empty = frac_pol_map < 0.01\n", "floor_raw = np.nanmedian(mom0_unthresh_map[empty])\n", "floor_debias = np.nanmedian(np.abs(mom0_debias_unthresh_map[empty]))\n", "print(f\"empty-sky mom0 floor: {floor_raw:.4f} -> {floor_debias:.4f}\")\n", "assert floor_debias < 0.2 * floor_raw\n", "\n", "# ...while detected sources keep their flux (compare the thresholded mom0)\n", "ratio = mom0_debias_unthresh_map[detected] / mom0_map[detected]\n", "assert 0.85 < np.nanmedian(ratio) < 1.15\n", "\n", "fits.PrimaryHDU(mom0_debias_unthresh_map.astype(np.float32)).writeto(\n", " tmpdir / \"FDF_mom0_debias.fits\", overwrite=True\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "rm-lite (3.11.13)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.13" } }, "nbformat": 4, "nbformat_minor": 2 }