Stokes I 3D (#53)

* Stokes I

* Cleanup

* Cleanup

* Self review

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Unify options classes (#54)

* Da roolz

* Use common classes

* Notebook

* Cleanup

* Drop uncertainties

---------

Co-authored-by: Alec Thomson <[email protected]>

* Fix for sigfig

---------

Co-authored-by: Alec Thomson <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

CLAUDE.md

@@ -0,0 +1,55 @@
+# rm-lite
+
+Mini reimplementation of RM-Tools. Fully typed Python, mypy strict on
+`rm_lite/*`.
+
+## Data (this repo only — overrides global caution)
+
+Root FITS cutouts and `scratch.ipynb` are dev/test fixtures, not SKA proprietary
+data. OK to read/run directly. Same for anything under `tests/` or
+`rm_lite/data/`.
+
+## Environment
+
+Use `uv run` for everything — never call `pytest`/`mypy`/`ruff` bare.
+
+## Verify gate (run before calling any task done)
+
+```
+uv run pytest
+uv run prek run --all-files
+```
+
+Full test suite every time, not just touched modules — dask/zarr tests are slow
+but regressions there are easy to miss with partial runs.
+
+## Docs
+
+`docs/` uses sphinx-autoapi + nbsphinx. When a public function/param in
+`rm_lite/` changes shape, update the relevant docstring/.rst/example in the same
+change — don't leave it for a follow-up.
+
+## Code style
+
+- Docstrings and comments: terse, one-liners (`"""Does X"""`), same for
+  NamedTuple field docstrings. No numpydoc/Google sections unless surrounding
+  code already has them.
+- Write like a human, not a manual. No over-explaining, no em/en dashes anywhere
+  (docs, comments, commit messages, prose).
+- `from __future__ import annotations` required in every module (ruff isort
+  enforces).
+- Python must be fully typed. mypy strict + `disallow_untyped_defs` on
+  `rm_lite.*`.
+- Ruff ruleset is wide (see `pyproject.toml`); don't add per-file ignores
+  without asking.
+- Prefer plain Options/Results containers (NamedTuple/dataclass) over full OOP
+  with methods and inheritance. Functions take an Options object, return a
+  Results object. Reach for a class-with-methods only when state genuinely needs
+  to persist across calls.
+
+## Documentation must be self-checking
+
+Example notebooks and docs (`docs/`, `*.ipynb`) must actually run and assert
+their own claims, not just show code with commentary. If a doc says "returns X",
+it should execute and check X, not describe it and move on. Stale docs are worse
+than no docs.

docs/examples/rmclean_3d.ipynb

@@ -11,7 +11,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We'll start with the same synthetic cube as the 3D RM-synth example -- see that page for details -- and run RM-synthesis again here to get the dirty FDF and RMSF cubes."
+    "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."
    ]
   },
   {
@@ -67,8 +67,8 @@
     "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, one\n",
-    "# top-left -- on an otherwise unpolarised background -- mostly-empty sky\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",
@@ -148,9 +148,9 @@
    "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 -- they're split out of the same delayed computation rather than re-running the clean loop per output.\n",
+    "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`."
+    "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`."
    ]
   },
   {
@@ -190,7 +190,7 @@
    "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)."
+    "`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)."
    ]
   },
   {
@@ -251,7 +251,7 @@
    "source": [
     "## Serialisation: zarr vs FITS\n",
     "\n",
-    "As with the RM-synth outputs, the clean/model/residual FDF cubes are complex -- zarr again is the natural fit, writing lazily chunk by chunk. The iteration-count map is different: it's real-valued (an integer per pixel), so it has no complex-split problem and writes to FITS directly, same as any ordinary 2D image."
+    "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."
    ]
   },
   {
@@ -298,13 +298,13 @@
    "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",
+    "`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",
+    "- **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."
+    "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."
    ]
   },
   {
@@ -344,7 +344,7 @@
    "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:"
+    "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:"
    ]
   },
   {
@@ -377,7 +377,7 @@
    "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:"
+    "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:"
    ]
   },
   {
@@ -429,13 +429,13 @@
    "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",
+    "`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",
+    "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:"
+    "`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:"
    ]
   },
   {

docs/examples/rmsynth_1d.ipynb

@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 1D RM-synthesis"
+    "# 1D RM-Synthesis"
    ]
   },
   {