From e29fe3bda1cfce5b3c380b0cbbf0e3da010d6aa7 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 08:10:32 +0000 Subject: [PATCH 1/8] Weight Q/U by their variance after the Stokes I division Dividing Q/U by a Stokes I model divides their noise too, but the transform kept 1/sigma^2 weights. Channels where the model is faint got magnified noise at full weight, which blows up PI wherever a model falls towards zero. `stokes_i_weighting` makes the noise-based weights follow the division: - "global" (default): one field-wide power law nu**alpha for every pixel, so the map keeps one RMSF and one reference frequency. `stokes_i_weight_alpha` pins alpha, or "auto" fits it to the field's mean Stokes I. - "per_pixel": each pixel's own model. Most SNR, but the RMSF follows the spectral index. - None: the old weights, unchanged. With the weighting on, pixels without a kept fit divide by the same power law instead of a flat model, so PI is continuous across the Stokes I SNR cut. lam_sq_0 "auto" follows the template weights in both modes. The 1D tool takes the same options. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- docs/examples/rmsynth_3d_stokes_i.ipynb | 243 +++++++++++++++- rm_lite/tools_1d/rmsynth.py | 104 +++++-- rm_lite/tools_3d/rmsynth.py | 135 ++++++++- rm_lite/utils/fitting.py | 92 ++++++- rm_lite/utils/synthesis.py | 89 +++++- tests/test_rmsynth.py | 14 +- tests/test_stokes_i_weighting.py | 351 ++++++++++++++++++++++++ tests/test_tools_3d_stokes_i.py | 2 +- 8 files changed, 968 insertions(+), 62 deletions(-) create mode 100644 tests/test_stokes_i_weighting.py diff --git a/docs/examples/rmsynth_3d_stokes_i.ipynb b/docs/examples/rmsynth_3d_stokes_i.ipynb index c69ce33..f27d9d3 100644 --- a/docs/examples/rmsynth_3d_stokes_i.ipynb +++ b/docs/examples/rmsynth_3d_stokes_i.ipynb @@ -318,7 +318,7 @@ "source": [ "### Masking low-SNR pixels\n", "\n", - "Fitting a Stokes I model to a noise pixel is pointless and can inject fake spectral structure. `rmsynth_3d` fits only pixels whose frequency-averaged Stokes I SNR, `median(I) * sqrt(N) / median(noise)`, is at least `stokes_i_snr_cut` (default 5). Below the cut, or if a fit fails to converge, the pixel gets a flat model, so it has no spectral correction and its FDF is the plain Q/U FDF. The FDF cube is never blanked, but `stokes_i_ref_flux_map` is NaN on those pixels, along with alpha, order and the term cubes: the flat model is a mean of noise, so reporting it as a flux would give a meaningless (often negative) fractional polarisation.\n", + "Fitting a Stokes I model to a noise pixel is pointless and can inject fake spectral structure. `rmsynth_3d` fits only pixels whose frequency-averaged Stokes I SNR, `median(I) * sqrt(N) / median(noise)`, is at least `stokes_i_snr_cut` (default 5). Below the cut, or if a fit fails to converge, the pixel gets a flat model, so it has no spectral correction and its FDF is the plain Q/U FDF. This is with `weight_type=\"uniform\"`, as on this page. With a noise-based weight type the pixel gets a power law instead (see [Weights after the Stokes I division](#Weights-after-the-Stokes-I-division)). The FDF cube is never blanked, but `stokes_i_ref_flux_map` is NaN on those pixels, along with alpha, order and the term cubes: the flat model is a mean of noise, so reporting it as a flux would give a meaningless (often negative) fractional polarisation.\n", "\n", "Medians, not the mean and rms, so one bad channel cannot decide whether a pixel is fitted: an rms over the errors turns one inflated error channel into a near-zero SNR, and a mean over the fluxes lets one negative spike make it negative. Either quietly flattens a good pixel.\n", "\n", @@ -369,7 +369,7 @@ "source": [ "### Blanked Stokes I\n", "\n", - "The fallback has two steps, and neither blanks a pixel. A pixel the fitter cannot use, below the cut, not converged, or with a model that runs away, gets a flat model at its mean Stokes I, which cancels out of the FDF. A pixel with no mean to divide by either, negative or with Stokes I blank there, gets no correction at all: RM-synthesis on the raw Q/U.\n", + "The fallback has two steps, and neither blanks a pixel. As above, this is the fallback with `weight_type=\"uniform\"`. A pixel the fitter cannot use, below the cut, not converged, or with a model that runs away, gets a flat model at its mean Stokes I, which cancels out of the FDF. A pixel with no mean to divide by either, negative or with Stokes I blank there, gets no correction at all: RM-synthesis on the raw Q/U.\n", "\n", "The first step covers more than the cut, which is why it does not need an error of its own. At a mosaic edge the linmos weight goes to zero, so the Stokes I error rm-lite derives from it is infinite and the cut has nothing to measure against, while Q and U still hold real data. A peak PI map has to have a value at those pixels, so the Stokes I pass hands back a usable model whatever it was given." ] @@ -678,6 +678,245 @@ "cell_type": "markdown", "id": "24", "metadata": {}, + "source": [ + "### Weights after the Stokes I division\n", + "\n", + "The correction divides Q and U by the Stokes I model, and that divides their noise too. In channel $j$ the corrected data $q_j = (Q_j + iU_j)\\, I(\\nu_0)/I_j$ have noise $\\sigma_{q,j} = \\sigma_j\\, I(\\nu_0)/I_j$. Inverse-variance weighting of the corrected data therefore means $w_j = 1/\\sigma_{q,j}^2 \\propto I_j^2/\\sigma_j^2$. Weighting them by $1/\\sigma_j^2$ gives full weight to the channels where $I_j$ is small and $\\sigma_{q,j}$ is large, and the FDF noise grows as $\\sum_j \\sigma_j^2 / I_j^2$.\n", + "\n", + "`stokes_i_weighting` sets the weights:\n", + "\n", + "- `\"global\"` (the default): $w_j = T_j^2/\\sigma_j^2$, with $T_j = \\nu_j^{\\alpha_T}$ the same for every pixel. `stokes_i_weight_alpha=\"auto\"` fits $\\alpha_T$ to the Stokes I averaged over the field; a number sets it.\n", + "- `\"per_pixel\"`: $w_j = I_j^2/\\sigma_j^2$, with each pixel's own model.\n", + "- `None`: $w_j = 1/\\sigma_j^2$, as before this option existed.\n", + "\n", + "It changes only the noise-based weight types (`\"variance\"`, `\"natural\"`, `\"uniform_lsq\"`, `\"briggs\"`). In both modes $\\lambda^2_0$ from `lam_sq_0_m2=\"auto\"`, and so the Stokes I reference frequency, is the weighted mean of $\\lambda^2$ under the template weights $T_j^2/\\sigma_j^2$.\n", + "\n", + "Below is a Faraday-thin source with 10% fractional polarisation, in every pixel of a 40 x 40 cube, over an 800 to 1800 MHz band. Its Stokes I falls to 2% of its peak at the top of the band. The model is supplied, so the only difference between the runs is the weights." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "from rm_lite.utils.synthesis import lambda2_to_freq\n", + "\n", + "wide_hz = np.linspace(800e6, 1800e6, 64)\n", + "wide_lsq = freq_to_lambda2(wide_hz)\n", + "wchunks = (-1, 20, 20)\n", + "sigma_w, frac_w = 0.02, 0.1\n", + "\n", + "\n", + "def thin_cube(model_cube, sigma, rm_radm2=30.0):\n", + " # Faraday-thin Q/U that follow the Stokes I model, plus noise.\n", + " angle = (2 * rm_radm2 * wide_lsq)[:, None, None]\n", + " noise = rng.normal(0, sigma, (2, *model_cube.shape))\n", + " return (\n", + " frac_w * model_cube * np.cos(angle) + noise[0],\n", + " frac_w * model_cube * np.sin(angle) + noise[1],\n", + " )\n", + "\n", + "\n", + "def synth_wide(model_cube, q, u, d_phi_radm2=0.5, **kwargs):\n", + " return rmsynth_3d(\n", + " da.from_array(q, chunks=wchunks),\n", + " da.from_array(u, chunks=wchunks),\n", + " wide_hz,\n", + " weight_arr=np.full(wide_hz.size, 1 / sigma_w**2),\n", + " stokes_i_model=da.from_array(model_cube, chunks=wchunks),\n", + " phi_max_radm2=300.0,\n", + " d_phi_radm2=d_phi_radm2,\n", + " **kwargs,\n", + " )\n", + "\n", + "\n", + "t = (wide_hz - wide_hz[0]) / (wide_hz[-1] - wide_hz[0])\n", + "fading_i = np.broadcast_to((1 - 0.98 * t**6)[:, None, None], (wide_hz.size, 40, 40))\n", + "q_w, u_w = thin_cube(fading_i, sigma_w)\n", + "\n", + "snr = {}\n", + "for mode in (None, \"global\", \"per_pixel\"):\n", + " res = synth_wide(fading_i, q_w, u_w, stokes_i_weighting=mode)\n", + " truth = frac_w * res.stokes_i_ref_flux_map.compute()\n", + " pi = np.abs(res.fdf_dirty_cube.compute()).max(axis=0)\n", + " snr[mode] = float(np.median(truth / np.asarray(res.theoretical_noise.fdf_q_noise)))\n", + " print(\n", + " f\"{mode!s:>9}: median PI / truth {np.median(pi / truth):.3f}, SNR {snr[mode]:.1f}\"\n", + " )\n", + " if mode is not None:\n", + " assert abs(np.median(pi / truth) - 1) < 0.02\n", + "\n", + "assert snr[None] < snr[\"global\"] < snr[\"per_pixel\"]" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "Both modes give an unbiased PI and a higher SNR than $1/\\sigma_j^2$. The per-pixel weights are the inverse variance of each pixel's own corrected data, so they give the highest SNR that pixel can have. Global mode gets less here, because a power law cannot follow a spectrum that turns down this sharply.\n", + "\n", + "#### The RMSF\n", + "\n", + "The weights also set the RMSF. In global mode every pixel has the same weights, so every pixel has the same RMSF and the same effective $\\lambda^2$, whatever its spectral index. In per-pixel mode the weights follow each pixel's Stokes I, so the RMSF and the effective $\\lambda^2$ change with the spectral index, and any Faraday-complexity measure made against the RMSF, such as the second moment of the FDF, is partly set by the Stokes I spectrum.\n", + "\n", + "Two noiseless Faraday-thin sources, with $\\alpha = -0.5$ and $\\alpha = -2.5$. The global template is fitted to their mean Stokes I:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "two_alpha = np.stack([(wide_hz / 1e9) ** -0.5, (wide_hz / 1e9) ** -2.5], axis=-1)[\n", + " :, None, :\n", + "]\n", + "q_two, u_two = thin_cube(two_alpha, 0.0)\n", + "\n", + "\n", + "def half_power_width(rmsf, phi):\n", + " # Width of the main lobe at half its peak.\n", + " return float(np.ptp(phi[np.abs(rmsf) >= 0.5 * np.abs(rmsf).max()]))\n", + "\n", + "\n", + "rmsf_width = {}\n", + "for mode in (None, \"global\", \"per_pixel\"):\n", + " res = synth_wide(\n", + " two_alpha,\n", + " q_two,\n", + " u_two,\n", + " d_phi_radm2=0.2,\n", + " stokes_i_weighting=mode,\n", + " per_pixel_rmsf=True,\n", + " )\n", + " rmsf = res.rmsf_cube.compute()\n", + " rmsf_width[mode] = [\n", + " half_power_width(rmsf[:, 0, i], res.phi_double_arr_radm2) for i in range(2)\n", + " ]\n", + " print(\n", + " f\"{mode!s:>9}: RMSF FWHM {rmsf_width[mode][0]:.1f} (alpha -0.5) and \"\n", + " f\"{rmsf_width[mode][1]:.1f} (alpha -2.5) rad/m^2\"\n", + " )\n", + "\n", + "assert np.isclose(rmsf_width[None][0], rmsf_width[None][1])\n", + "assert np.isclose(rmsf_width[\"global\"][0], rmsf_width[\"global\"][1])\n", + "assert rmsf_width[\"per_pixel\"][1] > 1.04 * rmsf_width[\"per_pixel\"][0]" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "#### Pixels below the Stokes I SNR cut\n", + "\n", + "A pixel below `stokes_i_snr_cut` has no fit. With the weighting on, it is divided by the template $T(\\nu)/T(\\nu_0)$ instead of a flat model. Its PI is then the PI at the reference frequency, like its fitted neighbours, to the extent that its spectrum follows the template. With `stokes_i_weighting=None` it keeps the flat model and reports a band-averaged PI instead. For a steep spectrum the two differ, so the PI map steps where the cut falls.\n", + "\n", + "A row of pixels with the same $\\alpha = -3$ spectrum and a brightness that falls through the cut:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "amp = np.geomspace(0.05, 1.0, 12)[None, None, :]\n", + "row_i = 0.02 * (wide_hz / 1.2e9)[:, None, None] ** -3.0 * amp\n", + "q_row, u_row = thin_cube(row_i, 0.0)\n", + "frac_row = {}\n", + "for mode in (None, \"global\"):\n", + " res = rmsynth_3d(\n", + " da.from_array(q_row, chunks=(-1, 1, 12)),\n", + " da.from_array(u_row, chunks=(-1, 1, 12)),\n", + " wide_hz,\n", + " weight_arr=np.full(wide_hz.size, 1 / sigma_w**2),\n", + " stokes_i=da.from_array(row_i, chunks=(-1, 1, 12)),\n", + " stokes_i_error=np.full(wide_hz.size, 1e-3),\n", + " fit_order=1,\n", + " stokes_i_snr_cut=50.0,\n", + " phi_max_radm2=100.0,\n", + " d_phi_radm2=0.5,\n", + " stokes_i_weighting=mode,\n", + " )\n", + " fitted_row = np.isfinite(res.stokes_i_model_order_map.compute())[0]\n", + " ref_hz = float(lambda2_to_freq(res.lam_sq_0_m2))\n", + " i_ref = np.array([np.interp(ref_hz, wide_hz, row_i[:, 0, i]) for i in range(12)])\n", + " frac_row[mode] = np.abs(res.fdf_dirty_cube.compute()).max(axis=0)[0] / i_ref\n", + " print(\n", + " f\"{mode!s:>6}: fractional pol, fitted pixels {np.round(frac_row[mode][fitted_row], 3)}\"\n", + " )\n", + " print(\n", + " f\"{'':>6} fractional pol, below the cut {np.round(frac_row[mode][~fitted_row], 3)}\"\n", + " )\n", + "\n", + "assert fitted_row.any()\n", + "assert (~fitted_row).any()\n", + "np.testing.assert_allclose(frac_row[\"global\"], frac_w, rtol=1e-3)\n", + "assert np.all(np.abs(frac_row[None][~fitted_row] / frac_w - 1) > 0.05)" + ] + }, + { + "cell_type": "markdown", + "id": "30", + "metadata": {}, + "source": [ + "#### Choosing a mode\n", + "\n", + "The numbers from the cells above:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [ + "rows = [\n", + " (\n", + " \"SNR, Stokes I falling at the band top\",\n", + " snr[None],\n", + " snr[\"global\"],\n", + " snr[\"per_pixel\"],\n", + " ),\n", + " (\n", + " \"RMSF FWHM, alpha -0.5 (rad/m^2)\",\n", + " rmsf_width[None][0],\n", + " rmsf_width[\"global\"][0],\n", + " rmsf_width[\"per_pixel\"][0],\n", + " ),\n", + " (\n", + " \"RMSF FWHM, alpha -2.5 (rad/m^2)\",\n", + " rmsf_width[None][1],\n", + " rmsf_width[\"global\"][1],\n", + " rmsf_width[\"per_pixel\"][1],\n", + " ),\n", + "]\n", + "print(f\"{'':40s}{'None':>10}{'global':>10}{'per_pixel':>10}\")\n", + "for label, *values in rows:\n", + " print(f\"{label:40s}\" + \"\".join(f\"{v:10.1f}\" for v in values))" + ] + }, + { + "cell_type": "markdown", + "id": "32", + "metadata": {}, + "source": [ + "- `\"global\"`: one RMSF and one effective frequency for the whole map, so PI, RM and Faraday-complexity maps compare pixel to pixel. A pixel whose spectrum is far from the template has less SNR than it could, and a pixel whose model falls far below the template in some channels is only partly protected from the noise there. The RMSF is set by the template, so it differs from the $1/\\sigma_j^2$ RMSF. Whether it is wider or narrower depends on the band and the noise: here it is narrower, and on a band with gaps it can be wider. With `\"auto\"` the template comes from each field, so pass a number to give several fields one RMSF.\n", + "- `\"per_pixel\"`: the highest SNR in every pixel. The RMSF and the effective frequency change with the spectral index, and it cannot be used with `lam_sq_0_m2=\"per_pixel\"`.\n", + "- `None`: the weights before this option existed. The FDF noise is unbounded where a model falls towards zero." + ] + }, + { + "cell_type": "markdown", + "id": "33", + "metadata": {}, "source": [ "`rmsynth_3d_from_fits` takes the same Stokes I options as file paths (`stokes_i_file`, `stokes_i_error_file`, or `stokes_i_model_file`). See the [3D RM-CLEAN](rmclean_3d.ipynb) page to deconvolve the FDF cube." ] diff --git a/rm_lite/tools_1d/rmsynth.py b/rm_lite/tools_1d/rmsynth.py index 073ed26..893e25c 100644 --- a/rm_lite/tools_1d/rmsynth.py +++ b/rm_lite/tools_1d/rmsynth.py @@ -3,8 +3,10 @@ from __future__ import annotations import time +from dataclasses import replace from typing import Literal, NamedTuple +import dask.array as da import numpy as np import polars as pl from numpy.typing import NDArray @@ -16,13 +18,17 @@ StokesIFitOptions, coefficient_errors, coefficient_names, + field_spectral_index, ) from rm_lite.utils.logging import logger from rm_lite.utils.synthesis import ( + NOISE_WEIGHT_TYPES, FDFOptions, LamSq0Mode, StokesData, + StokesIWeighting, WeightType, + apply_weight_type, compute_rmsynth_params, compute_theoretical_noise, create_fractional_spectra, @@ -32,6 +38,7 @@ get_rmsf_nufft, lambda2_to_freq, rmsynth_nufft, + stokes_i_template, ) @@ -114,6 +121,25 @@ def _stokes_i_terms( ) +def _spectrum_alpha( + stokes_data: StokesData, + fdf_options: FDFOptions, + fit_options: StokesIFitOptions, +) -> float: + """The power-law index for the weights, pinned or fitted to this spectrum.""" + if fdf_options.stokes_i_weight_alpha != "auto": + return float(fdf_options.stokes_i_weight_alpha) + model = stokes_data.stokes_i_model_arr + spectrum = model if model is not None else stokes_data.stokes_i_arr + assert spectrum is not None + return field_spectral_index( + da.from_array(np.asarray(spectrum)[:, np.newaxis, np.newaxis]), + None if model is not None else stokes_data.stokes_i_error_arr, + stokes_data.freq_arr_hz, + fit_options, + ) + + def run_rmsynth( freq_arr_hz: NDArray[np.float64], complex_pol_arr: NDArray[np.complexfloating], @@ -134,6 +160,8 @@ def run_rmsynth( fit_order: int = 2, stokes_i_robust_loss: RobustLoss = "cauchy", stokes_i_f_scale: float = 3.0, + stokes_i_weighting: StokesIWeighting | None = "global", + stokes_i_weight_alpha: float | Literal["auto"] = "auto", ignore_stokes_i: bool = False, moment_threshold_snr: float = 5.0, ) -> RMSynth1DResults: @@ -170,6 +198,14 @@ def run_rmsynth( Defaults to "cauchy". stokes_i_f_scale (float, optional): How far out, in sigma, before a channel is downweighted. Flat from 1 to 10. Defaults to 3.0. + stokes_i_weighting ("global", "per_pixel", None, optional): How the + inverse-variance weights follow the Stokes I division, see + `rm_lite.tools_3d.rmsynth.rmsynth_3d`. "global" weights by a power law + of `stokes_i_weight_alpha`, "per_pixel" by this spectrum's own model, + None by 1/sigma^2. Defaults to "global". + stokes_i_weight_alpha (float | "auto", optional): Spectral index of the + power law. "auto" fits this spectrum, so pin it to give a catalogue + of spectra one RMSF. Defaults to "auto". moment_threshold_snr (float, optional): SNR cut (times the theoretical FDF noise) applied to FDF amplitudes before computing the Faraday moments. Defaults to 5.0. Returns: @@ -200,6 +236,8 @@ def run_rmsynth( lam_sq_0_m2=lam_sq_0_m2, do_fit_rmsf=do_fit_rmsf, do_fit_rmsf_real=do_fit_rmsf_real, + stokes_i_weighting=stokes_i_weighting, + stokes_i_weight_alpha=stokes_i_weight_alpha, ) # snr_cut=None: the 1D fractional fit has only one spectrum, so an SNR cut # would just silently drop fractional polarisation rather than fall back to @@ -254,10 +292,24 @@ def _run_rmsynth( stokes_i_terms (pl.DataFrame): Fitted Stokes I model terms """ + weighting = ( + fdf_options.stokes_i_weighting + if not ignore_stokes_i and fdf_options.weight_type in NOISE_WEIGHT_TYPES + else None + ) + noise_error = stokes_data.complex_pol_error + weight_error = noise_error + if weighting is not None: + weight_alpha = _spectrum_alpha(stokes_data, fdf_options, fit_options) + weight_error = noise_error / stokes_i_template( + stokes_data.freq_arr_hz, weight_alpha + ) + fit_options = replace(fit_options, fallback_alpha=weight_alpha) + rmsynth_params = compute_rmsynth_params( freq_arr_hz=stokes_data.freq_arr_hz, complex_pol_arr=stokes_data.complex_pol_arr, - complex_pol_error=stokes_data.complex_pol_error, + complex_pol_error=weight_error, fdf_options=fdf_options, ) @@ -276,15 +328,43 @@ def _run_rmsynth( no_nan_idx = fractional_stokes_data.no_nan_idx fit_result = fractional_stokes_data.fit_result - # Compute after any fractional spectra have been created - tick = time.time() - # Perform RM-synthesis on the spectrum all_flagged = (~no_nan_idx).all() if all_flagged: msg = "All channels have been masked!" logger.warning(msg) + stokes_i_reference_flux = np.nan + if not ignore_stokes_i: + assert stokes_data.stokes_i_model_arr is not None + assert stokes_data.freq_arr_hz.shape == stokes_data.stokes_i_model_arr.shape + if not all_flagged: + stokes_i_model = interpolate.interp1d( + stokes_data.freq_arr_hz[no_nan_idx], + stokes_data.stokes_i_model_arr[no_nan_idx], + ) + + stokes_i_reference_flux = float(stokes_i_model(ref_freq_hz)) + else: + logger.warning("Using mean as reference flux") + stokes_i_reference_flux = float(np.nanmean(stokes_data.stokes_i_model_arr)) + + if weighting == "per_pixel" and stokes_data.stokes_i_model_arr is not None: + # This spectrum's own model, in place of the power law used up to here. + shape = stokes_data.stokes_i_model_arr / stokes_i_reference_flux + rmsynth_params = rmsynth_params._replace( + weight_arr=apply_weight_type( + lambda_sq_arr_m2=rmsynth_params.lambda_sq_arr_m2, + real_qu_error=np.abs(noise_error.real + noise_error.imag) / 2.0 / shape, + channel_mask=~np.isfinite(stokes_data.complex_pol_arr), + fdf_options=fdf_options, + cell_m2=rmsynth_params.cell_m2, + ) + ) + + # Compute after any fractional spectra have been created + tick = time.time() + fdf_dirty_arr = rmsynth_nufft( complex_pol_arr=stokes_data.complex_pol_arr[no_nan_idx], lambda_sq_arr_m2=rmsynth_params.lambda_sq_arr_m2[no_nan_idx], @@ -314,19 +394,6 @@ def _run_rmsynth( ) if not ignore_stokes_i: - assert stokes_data.stokes_i_model_arr is not None - assert stokes_data.freq_arr_hz.shape == stokes_data.stokes_i_model_arr.shape - if not all_flagged: - stokes_i_model = interpolate.interp1d( - stokes_data.freq_arr_hz[no_nan_idx], - stokes_data.stokes_i_model_arr[no_nan_idx], - ) - - stokes_i_reference_flux = float(stokes_i_model(ref_freq_hz)) - else: - logger.warning("Using mean as reference flux") - stokes_i_reference_flux = float(np.nanmean(stokes_data.stokes_i_model_arr)) - fdf_dirty_arr *= stokes_i_reference_flux theoretical_noise = theoretical_noise._replace( @@ -335,9 +402,6 @@ def _run_rmsynth( fdf_u_noise=theoretical_noise.fdf_u_noise * stokes_i_reference_flux, ) - else: - stokes_i_reference_flux = np.nan - # Measure the parameters of the dirty FDF # Use the theoretical noise to calculate uncertainties logger.warning( diff --git a/rm_lite/tools_3d/rmsynth.py b/rm_lite/tools_3d/rmsynth.py index 69d71ae..a23eec5 100644 --- a/rm_lite/tools_3d/rmsynth.py +++ b/rm_lite/tools_3d/rmsynth.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from dataclasses import replace from functools import partial from pathlib import Path from typing import Any, Literal, NamedTuple, cast @@ -36,14 +37,17 @@ StokesIFitOptions, alpha_from_model_block, coefficient_names, + field_spectral_index, fit_stokes_cube, ref_flux_from_block, ) from rm_lite.utils.logging import logger, quiet_logs from rm_lite.utils.synthesis import ( + NOISE_WEIGHT_TYPES, FDFOptions, LamSq0Mode, RMSynthParams, + StokesIWeighting, TheoreticalNoise, WeightType, apply_weight_type, @@ -58,6 +62,7 @@ lambda2_to_freq, make_double_phi_arr, rmsynth_nufft, + stokes_i_template, ) @@ -160,6 +165,11 @@ class RMSynth3DResults(NamedTuple): `per_pixel_rmsf`, so asking for both computes the transforms twice.""" per_pixel_rmsf: PerPixelRMSF | None = None """How RM-CLEAN rebuilds the RMSF per block. Set with `rmsf_cube`.""" + stokes_i_weighting: StokesIWeighting | None = None + """How the weights followed the Stokes I division; None if they did not.""" + stokes_i_weight_alpha: float | None = None + """Spectral index of the field-wide power law the weights and the fallback + used; None if the weights did not follow the Stokes I division.""" def _compute_global_params( @@ -204,13 +214,20 @@ def noise_weights( pol_cube: da.Array, rmsynth_params: RMSynthParams, fdf_options: FDFOptions, + noise_weight_arr: NDArray[np.float64] | da.Array | None = None, ) -> NoiseWeights: """The error and weights the FDF noise follows from, aligned to `pol_cube`. - Lazy, so nothing is read until the noise is asked for. + `noise_weight_arr` is the Q/U inverse variance, when the transform weights + `weight_arr` are something else. Lazy, so nothing is read until the noise is + asked for. """ aligned = _weight_arr_map_blocks_args(weight_arr, pol_cube) - complex_pol_error = error_from_weight(aligned[0] if aligned else weight_arr) + noise_source = weight_arr if noise_weight_arr is None else noise_weight_arr + noise_aligned = _weight_arr_map_blocks_args(noise_source, pol_cube) + complex_pol_error = error_from_weight( + noise_aligned[0] if noise_aligned else noise_source + ) if np.ndim(complex_pol_error) == 1: # Spatial axes too, or it will not broadcast against the weight cube. complex_pol_error = broadcast_over_channels(complex_pol_error, pol_cube) @@ -638,6 +655,45 @@ def _weight_arr_map_blocks_args( return (weight_da.rechunk({0: -1, 1: target.chunks[1], 2: target.chunks[2]}),) +def _scale_channels( + weight_arr: NDArray[Any] | da.Array, factor: NDArray[np.float64] +) -> NDArray[Any] | da.Array: + """`weight_arr` times a per-channel factor, in its own dtype.""" + if weight_arr.ndim == 3: + factor = factor[:, np.newaxis, np.newaxis] + return (weight_arr * factor).astype(weight_arr.dtype) + + +def _per_pixel_weights( + noise_weight_arr: NDArray[Any] | da.Array, + model_cube: da.Array, + ref_flux_map: da.Array, +) -> da.Array: + """Inverse variance of Q/U once divided by each pixel's own model. + + Dividing by `model / ref_flux` scales the error by its inverse, so the + weight gains its square. + """ + shape = model_cube / ref_flux_map[np.newaxis].astype(model_cube.dtype) + if np.ndim(noise_weight_arr) == 1: + noise = np.asarray(noise_weight_arr)[:, np.newaxis, np.newaxis] + return cast("da.Array", noise * shape**2) + return da.asarray(noise_weight_arr).rechunk(model_cube.chunks) * shape**2 + + +def _stokes_i_weight_alpha( + fdf_options: FDFOptions, + stokes_i: da.Array, + stokes_i_error: NDArray[np.float64] | da.Array | None, + freq_arr_hz: NDArray[np.float64], + fit_options: StokesIFitOptions, +) -> float: + """The field-wide spectral index, pinned or fitted to the summed Stokes I.""" + if fdf_options.stokes_i_weight_alpha != "auto": + return float(fdf_options.stokes_i_weight_alpha) + return field_spectral_index(stokes_i, stokes_i_error, freq_arr_hz, fit_options) + + def rmsynth_3d( stokes_q: da.Array, stokes_u: da.Array, @@ -659,6 +715,8 @@ def rmsynth_3d( stokes_i_model_floor_sigma: float = 0.01, stokes_i_robust_loss: RobustLoss = "cauchy", stokes_i_f_scale: float = 3.0, + stokes_i_weighting: StokesIWeighting | None = "global", + stokes_i_weight_alpha: float | Literal["auto"] = "auto", compute_model_error: bool = False, n_error_samples: int = 1000, per_pixel_rmsf: bool = False, @@ -720,6 +778,20 @@ def rmsynth_3d( path only. Defaults to "cauchy". stokes_i_f_scale (float, optional): How far out, in sigma, before a channel is downweighted. Flat from 1 to 10. Defaults to 3.0. + stokes_i_weighting ("global", "per_pixel", None, optional): Dividing Q/U + by a Stokes I model divides their error too, so the inverse-variance + weights gain a factor of the model shape squared. "global" uses one + field-wide power law for every pixel: one RMSF and one reference + frequency for the whole map. "per_pixel" uses each pixel's own + model: the most sensitive, but the RMSF then varies with the + spectral index. Either way pixels without a kept fit fall back to the + field-wide power law rather than a flat model. None keeps 1/sigma^2 + and the flat fallback. No effect without Stokes I or with + weight_type="uniform". Defaults to "global". + stokes_i_weight_alpha (float | "auto", optional): Spectral index of the + field-wide power law, or "auto" to fit it to the Stokes I (or model) + averaged over the field. Pin it to share one RMSF between fields. + Defaults to "auto". 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. @@ -763,6 +835,8 @@ def rmsynth_3d( weight_type=weight_type, robust=robust, lam_sq_0_m2=lam_sq_0_m2, + stokes_i_weighting=stokes_i_weighting, + stokes_i_weight_alpha=stokes_i_weight_alpha, ) fit_options = StokesIFitOptions( fit_order=fit_order, @@ -777,6 +851,35 @@ def rmsynth_3d( if weight_arr is None: weight_arr = np.ones_like(freq_arr_hz) + stokes_i_source = stokes_i_model if stokes_i_model is not None else stokes_i + weighting = ( + fdf_options.stokes_i_weighting + if stokes_i_source is not None and weight_type in NOISE_WEIGHT_TYPES + else None + ) + noise_weight_arr = weight_arr + weight_alpha: float | None = None + if stokes_i_source is not None and weighting is not None: + weight_alpha = _stokes_i_weight_alpha( + fdf_options, + stokes_i_source, + stokes_i_error if stokes_i_model is None else None, + freq_arr_hz, + fit_options, + ) + # The field-wide shape sets the weights up to the fit, so lam_sq_0_m2 and + # the Stokes I reference frequency follow them in either mode. + weight_arr = _scale_channels( + weight_arr, stokes_i_template(freq_arr_hz, weight_alpha) ** 2 + ) + fit_options = replace(fit_options, fallback_alpha=weight_alpha) + if weighting == "per_pixel" and not per_pixel_rmsf: + logger.info( + "stokes_i_weighting='per_pixel' weights each pixel by its own Stokes I " + "model, so the RMSF differs pixel to pixel; computing the per-pixel " + "RMSF cube, which RM-CLEAN needs to match the FDF it is cleaning." + ) + per_pixel_rmsf = True weight_summary = _summarise_weight(weight_arr) if fdf_options.lam_sq_0_m2 == "per_pixel" and not per_pixel_rmsf: logger.info( @@ -900,9 +1003,20 @@ def rmsynth_3d( ref_freq_hz=ref_freq_hz, ) + if ( + weighting == "per_pixel" + and stokes_i_model_cube is not None + and ref_flux_map is not None + ): + weight_arr = _per_pixel_weights( + noise_weight_arr, stokes_i_model_cube, ref_flux_map + ) + # After any Stokes I division, so the blanking folded in is the blanking the # FDF is actually built from. - weights = noise_weights(weight_arr, pol_cube, rmsynth_params, fdf_options) + weights = noise_weights( + weight_arr, pol_cube, rmsynth_params, fdf_options, noise_weight_arr + ) if stokes_i_model_cube is not None and ref_flux_map is not None: theoretical_noise = fractional_theoretical_noise( weights=weights, @@ -1004,6 +1118,8 @@ def rmsynth_3d( stokes_i_ref_freq_hz=(ref_freq_hz if stokes_i_model_cube is not None else None), rmsf_cube=rmsf_cube, per_pixel_rmsf=rmsf_recipe, + stokes_i_weighting=weighting, + stokes_i_weight_alpha=weight_alpha, ) @@ -1151,6 +1267,8 @@ def rmsynth_3d_from_fits( stokes_i_model_floor_sigma: float = 0.01, stokes_i_robust_loss: RobustLoss = "cauchy", stokes_i_f_scale: float = 3.0, + stokes_i_weighting: StokesIWeighting | None = "global", + stokes_i_weight_alpha: float | Literal["auto"] = "auto", compute_model_error: bool = False, n_error_samples: int = 1000, per_pixel_rmsf: bool = False, @@ -1201,6 +1319,8 @@ def rmsynth_3d_from_fits( stokes_i_model_floor_sigma (float, optional): See `rmsynth_3d`. Defaults to 0.01. stokes_i_robust_loss (RobustLoss, optional): See `rmsynth_3d`. Defaults to "cauchy". stokes_i_f_scale (float, optional): See `rmsynth_3d`. Defaults to 3.0. + stokes_i_weighting ("global", "per_pixel", None, optional): See `rmsynth_3d`. Defaults to "global". + stokes_i_weight_alpha (float | "auto", optional): See `rmsynth_3d`. Defaults to "auto". compute_model_error (bool, optional): See `rmsynth_3d`. Defaults to False. n_error_samples (int, optional): See `rmsynth_3d`. Defaults to 1000. per_pixel_rmsf (bool, optional): See `rmsynth_3d`. Defaults to False. @@ -1292,12 +1412,7 @@ def rmsynth_3d_from_fits( # Noise-based types use 1/sigma^2 as their base (uniform_lsq/briggs then apply # the geometric lambda^2 factor); per-channel `uniform` deliberately ignores noise. - if weight_arr is None and weight_type in ( - "variance", - "natural", - "uniform_lsq", - "briggs", - ): + if weight_arr is None and weight_type in NOISE_WEIGHT_TYPES: weight_arr = get_weight_arr_from_fits( stokes_q_file, stokes_u_file, @@ -1350,6 +1465,8 @@ def rmsynth_3d_from_fits( stokes_i_model_floor_sigma=stokes_i_model_floor_sigma, stokes_i_robust_loss=stokes_i_robust_loss, stokes_i_f_scale=stokes_i_f_scale, + stokes_i_weighting=stokes_i_weighting, + stokes_i_weight_alpha=stokes_i_weight_alpha, compute_model_error=compute_model_error, n_error_samples=n_error_samples, per_pixel_rmsf=per_pixel_rmsf, diff --git a/rm_lite/utils/fitting.py b/rm_lite/utils/fitting.py index 9cdb090..62dc779 100644 --- a/rm_lite/utils/fitting.py +++ b/rm_lite/utils/fitting.py @@ -9,6 +9,7 @@ import dask.array as da import numpy as np from astropy.stats import akaike_info_criterion_lsq +from dask.base import compute from numpy.typing import ArrayLike, NDArray from scipy import optimize, stats @@ -60,6 +61,8 @@ class StokesIFitOptions: model_floor_sigma: float = 0.01 """Reject a model dipping this many sigma below the band-averaged Stokes I noise, falling back to a flat one (see `model_noise_floor`); 0 disables""" + fallback_alpha: float | None = None + """Spectral index of the model used where no fit is kept; None keeps a flat one""" def __post_init__(self) -> None: if self.fit_function not in ("log", "linear"): @@ -847,6 +850,57 @@ def fit_stokes_cube( ) +def field_spectral_index( + stokes_i: da.Array, + stokes_i_error: NDArray[np.float64] | da.Array | None, + freq_arr_hz: NDArray[np.float64], + fit_options: StokesIFitOptions, +) -> float: + """Power-law index of the field's mean Stokes I spectrum. + + One pass over the cube: each channel is averaged over the pixels that have + it, and the errors are carried in quadrature. 0 if the fit gives nothing. + """ + finite = da.isfinite(stokes_i) + total = da.sum(da.where(finite, stokes_i, 0.0), axis=(1, 2)) + count = da.sum(finite, axis=(1, 2)) + variance: da.Array | NDArray[np.float64] | None = None + if stokes_i_error is not None and np.ndim(stokes_i_error) == 3: + variance = da.sum( + da.where(finite, da.asarray(stokes_i_error) ** 2, 0.0), axis=(1, 2) + ) + total, count, variance = compute(total, count, variance) + count = np.asarray(count, dtype=np.float64) + with np.errstate(divide="ignore", invalid="ignore"): + mean = np.asarray(total, dtype=np.float64) / count + if variance is not None: + error = np.sqrt(np.asarray(variance, dtype=np.float64)) / count + elif stokes_i_error is not None: + error = np.asarray(stokes_i_error, dtype=np.float64) / np.sqrt(count) + else: + error = np.zeros_like(mean) + good = np.isfinite(mean) & np.isfinite(error) + if int(good.sum()) < 3: + logger.warning("Too few channels to fit the field's Stokes I; using alpha=0.") + return 0.0 + fit = static_fit( + freq_arr_hz[good], + float(np.exp(np.mean(np.log(freq_arr_hz[good])))), + mean[good], + error[good], + fit_order=1, + fit_function="log", + robust_loss=fit_options.robust_loss, + f_scale=fit_options.f_scale, + ) + alpha = float(np.asarray(fit.popt)[1]) + if not np.isfinite(alpha): + logger.warning("The field's Stokes I gave no spectral index; using alpha=0.") + return 0.0 + logger.info(f"Field-wide Stokes I spectral index for the weights: {alpha:.3f}") + return alpha + + def _pixel_stokes_i_error( err_block: NDArray[np.float64] | None, err_1d: NDArray[np.float64] | None, @@ -1013,15 +1067,21 @@ def _write_error_planes( out[planes.alpha_error, y, x] = abs(a_high - a_low) -def _write_flat_model( - out: NDArray[np.float64], - y: int, - x: int, - planes: BlockPlanes, +def fallback_model( + freq_arr_hz: NDArray[np.float64], + ref_freq_hz: float, mean_flux: float, -) -> None: - """Flat model, so Q/U get no correction; alpha, order and terms stay NaN.""" - out[planes.model, y, x] = flat_model_value(mean_flux) + alpha: float | None, +) -> NDArray[np.float64]: + """Model to divide by where no fit is kept. + + Flat at the mean without `alpha`, so Q/U get no correction. With it, a power + law of that index, 1 at the reference frequency: the amplitude cancels once + the FDF is rescaled, so it needs no flux. + """ + if alpha is None: + return np.full_like(freq_arr_hz, flat_model_value(mean_flux)) + return np.asarray((freq_arr_hz / ref_freq_hz) ** alpha, dtype=np.float64) RefFreqHz: TypeAlias = float | NDArray[np.float64] @@ -1062,8 +1122,9 @@ def _fit_stokes_i_block( finite channels or SNR below `fit_options.snr_cut`) or whose model is unusable (non-finite anywhere, or below `fit_options.model_floor_sigma` times the pixel's band-averaged noise, see `model_is_usable`) falls back to - a flat model at its mean Stokes I, and one whose mean cannot divide either - (negative, or no finite Stokes I at all) to a flat 1.0, leaving Q/U + `fallback_model`: a power law of `fit_options.fallback_alpha` when that is + set, else a flat model at its mean Stokes I, and one whose mean cannot divide + either (negative, or no finite Stokes I at all) to a flat 1.0, leaving Q/U uncorrected. Either way the model is finite, so no pixel of the FDF is blanked by the Stokes I pass alone; alpha, order, terms and errors stay NaN. """ @@ -1086,11 +1147,14 @@ def _fit_stokes_i_block( for y, x, i_spec, e_spec, good, fit in _iter_pixel_fits( i_block, err_block, err_1d, freq_arr_hz, ref_freq_hz, fit_options ): + pixel_ref_hz = ref_freq_for_pixel(ref_freq_hz, y, x) mean_flux = float(np.mean(i_spec[good])) if good.any() else np.nan + fallback = fallback_model( + freq_arr_hz, pixel_ref_hz, mean_flux, fit_options.fallback_alpha + ) if fit is None: - _write_flat_model(out, y, x, planes, mean_flux) + out[planes.model, y, x] = fallback continue - pixel_ref_hz = ref_freq_for_pixel(ref_freq_hz, y, x) model = fit.stokes_i_model_func( freq_arr_hz / pixel_ref_hz, *np.asarray(fit.popt) ) @@ -1099,7 +1163,7 @@ def _fit_stokes_i_block( # flagged and Q/U did not, so all of it has to be usable. if not model_is_usable(model, noise_floor): n_rejected += 1 - _write_flat_model(out, y, x, planes, mean_flux) + out[planes.model, y, x] = fallback continue _write_model_planes( out, y, x, planes, fit, model, freq_arr_hz, pixel_ref_hz @@ -1118,7 +1182,7 @@ def _fit_stokes_i_block( if n_rejected: logger.warning( f"{n_rejected} of {cy * cx} pixels in this chunk fitted an unusable " - "Stokes I model and fell back to a flat one (see " + "Stokes I model and fell back to `fallback_model` (see " "`rm_lite.utils.fitting.model_is_usable`). Expect this on pixels with " "no real Stokes I signal, i.e. when `stokes_i_snr_cut` is None, and " "on artefacts, whose Stokes I passes through zero." diff --git a/rm_lite/utils/synthesis.py b/rm_lite/utils/synthesis.py index 63039be..82e044e 100644 --- a/rm_lite/utils/synthesis.py +++ b/rm_lite/utils/synthesis.py @@ -40,6 +40,7 @@ FitResult, StokesIFitOptions, check_snr_cut_has_error, + fallback_model, fit_fdf, fit_rmsf, fit_sampled_peak, @@ -147,6 +148,20 @@ class TheoreticalNoise(NamedTuple): narrows the RMSF), `briggs` (robust interpolation between natural and uniform_lsq, needs `robust`). """ +StokesIWeighting: TypeAlias = Literal["global", "per_pixel"] +""" How the noise-based weights follow the Stokes I division: `global` scales +every pixel's weights by one field-wide spectrum, `per_pixel` by each pixel's own +model. """ + +NOISE_WEIGHT_TYPES: tuple[WeightType, ...] = ( + "variance", + "natural", + "uniform_lsq", + "briggs", +) +""" Weight types built on the Q/U noise, so the only ones the Stokes I +weighting can change. """ + @dataclass(frozen=True, kw_only=True, slots=True) class FDFOptions: @@ -170,6 +185,12 @@ class FDFOptions: """ Reference lambda^2 in m^2, or "auto"/"per_pixel" to derive one. The Stokes I reference frequency is derived from it, so the phase and flux references always match. """ + stokes_i_weighting: StokesIWeighting | None = "global" + """ Inverse-variance weights for the Q/U that were divided by Stokes I: + "global" or "per_pixel" (see `StokesIWeighting`); None keeps 1/sigma^2 """ + stokes_i_weight_alpha: float | Literal["auto"] = "auto" + """ Spectral index of the field-wide spectrum, or "auto" to fit it to the + summed Stokes I """ def __post_init__(self) -> None: if isinstance(self.lam_sq_0_m2, str): @@ -191,6 +212,30 @@ def __post_init__(self) -> None: if self.weight_type == "briggs" and self.robust is None: msg = "weight_type='briggs' requires a `robust` parameter." raise ValueError(msg) + if self.stokes_i_weighting not in (*get_args(StokesIWeighting), None): + msg = ( + f"stokes_i_weighting must be one of {get_args(StokesIWeighting)} " + f"or None, got {self.stokes_i_weighting!r}." + ) + raise ValueError(msg) + if self.stokes_i_weighting == "per_pixel" and self.lam_sq_0_m2 == "per_pixel": + # A per-pixel reference is taken from the weights before the Stokes I + # fit, so it could not follow weights that come from the fit. + msg = ( + "stokes_i_weighting='per_pixel' cannot be combined with " + "lam_sq_0_m2='per_pixel'." + ) + raise ValueError(msg) + alpha = self.stokes_i_weight_alpha + bad_alpha = ( + alpha != "auto" if isinstance(alpha, str) else not np.isfinite(alpha) + ) + if bad_alpha: + msg = ( + "stokes_i_weight_alpha must be 'auto' or a finite value, " + f"got {self.stokes_i_weight_alpha!r}." + ) + raise ValueError(msg) if self.d_phi_radm2 is None and self.n_samples is None: msg = "Either d_phi_radm2 or n_samples must be provided." raise ValueError(msg) @@ -1153,21 +1198,35 @@ def create_fractional_spectra( noise_floor = model_noise_floor( stokes_data.stokes_i_error_arr[no_nan_idx], fit_options.model_floor_sigma ) + kept_fit: FitResult | None = fit_result if not model_is_usable(model_good, noise_floor): logger.warning( "The fitted Stokes I model cannot safely divide Q/U (see " - "`rm_lite.utils.fitting.model_is_usable`); falling back to a flat " - "model at the mean Stokes I, so Q/U get no spectral correction." - ) - fit_result = flat_fit_result( - flat_model_value(float(np.mean(i_good))), - len(np.asarray(fit_result.popt)) - 1, - fit_options.fit_function, + "`rm_lite.utils.fitting.model_is_usable`); falling back to " + "`rm_lite.utils.fitting.fallback_model`." ) + kept_fit = None + if fit_options.fallback_alpha is None: + kept_fit = flat_fit_result( + flat_model_value(float(np.mean(i_good))), + len(np.asarray(fit_result.popt)) - 1, + fit_options.fit_function, + ) - stokes_i_model_arr, stokes_i_model_error = sample_model_error( - fit_result, stokes_data.freq_arr_hz / ref_freq_hz, fit_options.n_error_samples - ) + if kept_fit is None: + stokes_i_model_arr = fallback_model( + stokes_data.freq_arr_hz, + ref_freq_hz, + float(np.mean(i_good)), + fit_options.fallback_alpha, + ) + stokes_i_model_error = np.zeros_like(stokes_i_model_arr) + else: + stokes_i_model_arr, stokes_i_model_error = sample_model_error( + kept_fit, + stokes_data.freq_arr_hz / ref_freq_hz, + fit_options.n_error_samples, + ) # The fit runs in double precision; taking the model back down to the data's # keeps dividing by it from promoting the fractional spectra. model_dtype = real_dtype(stokes_data.complex_pol_arr.dtype) @@ -1203,7 +1262,7 @@ def create_fractional_spectra( return FractionalSpectra( stokes_data=fractional_stokes_data, - fit_result=fit_result, + fit_result=kept_fit, no_nan_idx=no_nan_idx, ) @@ -1577,6 +1636,14 @@ def apply_weight_type( return np.where(_match_channel_mask(channel_mask, weight_arr), 0.0, weight_arr) +def stokes_i_template( + freq_arr_hz: NDArray[np.float64], alpha: float +) -> NDArray[np.float64]: + """Field-wide Stokes I spectrum `nu**alpha`, 1 at the band's log-centre.""" + centre = np.exp(np.nanmean(np.log(freq_arr_hz))) + return np.asarray((freq_arr_hz / centre) ** alpha, dtype=np.float64) + + def weighted_lam_sq_0( weight_arr: NDArray[np.float64], lambda_sq_arr_m2: NDArray[np.float64] ) -> float: diff --git a/tests/test_rmsynth.py b/tests/test_rmsynth.py index 227b690..27cccac 100644 --- a/tests/test_rmsynth.py +++ b/tests/test_rmsynth.py @@ -398,14 +398,18 @@ def test_stokes_i_terms_describe_the_fitted_model(): terms = results.stokes_i_terms assert terms["term_name"].to_list() == ["flux", "alpha"] assert terms["fit_function"].to_list() == ["log", "log"] + # ref_freq_hz is the FDF's own reference frequency. + term_ref_freq = float(terms["ref_freq_hz"][0]) + fdf_ref_freq = float(results.fdf_parameters["ref_freq_hz"][0]) + assert np.isclose(term_ref_freq, fdf_ref_freq, rtol=1e-6) # A pure power law, so the fitted flux and index are the input ones. - np.testing.assert_allclose(terms["term_value"].to_numpy(), [flux, alpha], rtol=1e-4) + expected_flux = flux * (term_ref_freq / ref_freq_hz) ** alpha + np.testing.assert_allclose( + terms["term_value"].to_numpy(), [expected_flux, alpha], rtol=1e-4 + ) assert (terms["term_error"].to_numpy() >= 0).all() - # ref_freq_hz is the FDF's own reference frequency, and the terms rebuild the - # model the fractional spectra were divided by. - term_ref_freq = float(terms["ref_freq_hz"][0]) - assert np.isclose(term_ref_freq, ref_freq_hz, rtol=1e-6) + # The terms rebuild the model the fractional spectra were divided by. rebuilt = power_law(len(terms) - 1)( freq_arr_hz / term_ref_freq, *terms["term_value"].to_numpy() ) diff --git a/tests/test_stokes_i_weighting.py b/tests/test_stokes_i_weighting.py new file mode 100644 index 0000000..ce5b2d1 --- /dev/null +++ b/tests/test_stokes_i_weighting.py @@ -0,0 +1,351 @@ +"""Tests for the inverse-variance weights that follow the Stokes I division.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import dask.array as da +import numpy as np +import pytest +import rm_lite.tools_3d.rmclean as rmclean3d_mod +import rm_lite.utils.fitting as fitting_mod +from dask.base import compute +from numpy.typing import NDArray +from rm_lite.tools_1d.rmsynth import run_rmsynth +from rm_lite.tools_3d.rmsynth import RMSynth3DResults, rmsynth_3d +from rm_lite.utils.fitting import StokesIFitOptions, field_spectral_index +from rm_lite.utils.synthesis import ( + FDFOptions, + StokesIWeighting, + freq_to_lambda2, + lambda2_to_freq, + stokes_i_template, +) + +if TYPE_CHECKING: + from collections.abc import Callable + +FREQ_ARR_HZ = (np.arange(744, 1032, 3) * 1e6).astype(np.float64) +WIDE_FREQ_ARR_HZ = (np.arange(800, 1800, 8) * 1e6).astype(np.float64) +MODES: list[StokesIWeighting | None] = [None, "global", "per_pixel"] + + +def fading_model() -> NDArray[np.float64]: + """Stokes I that falls to 2% of its peak at the top of the band.""" + t = (FREQ_ARR_HZ - FREQ_ARR_HZ[0]) / (FREQ_ARR_HZ[-1] - FREQ_ARR_HZ[0]) + return np.asarray(1 - 0.98 * t**6, dtype=np.float64) + + +def thin_source( + model: NDArray[np.float64], + frac_pol: float, + sigma: float, + rm_radm2: float = 30.0, + seed: int = 1, + freq_arr_hz: NDArray[np.float64] = FREQ_ARR_HZ, +) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Faraday-thin Q/U following `model`, with Gaussian noise of `sigma`.""" + rng = np.random.default_rng(seed) + angle = (2 * rm_radm2 * freq_to_lambda2(freq_arr_hz))[:, np.newaxis, np.newaxis] + stokes_q = frac_pol * model * np.cos(angle) + rng.normal(0, sigma, model.shape) + stokes_u = frac_pol * model * np.sin(angle) + rng.normal(0, sigma, model.shape) + return stokes_q, stokes_u + + +def synth_with_model( + model: NDArray[np.float64], + stokes_q: NDArray[np.float64], + stokes_u: NDArray[np.float64], + sigma: float, + d_phi_radm2: float = 2.0, + freq_arr_hz: NDArray[np.float64] = FREQ_ARR_HZ, + **kwargs: Any, +) -> RMSynth3DResults: + """3D synthesis with a supplied Stokes I model and 1/sigma^2 noise weights.""" + chunks = (-1, 20, 20) + return rmsynth_3d( + da.from_array(stokes_q, chunks=chunks), + da.from_array(stokes_u, chunks=chunks), + freq_arr_hz, + weight_arr=np.full(freq_arr_hz.size, 1 / sigma**2), + stokes_i_model=da.from_array(model, chunks=chunks), + d_phi_radm2=d_phi_radm2, + phi_max_radm2=400.0, + **kwargs, + ) + + +def broadcast(spectrum: NDArray[np.float64], ny: int, nx: int) -> NDArray[np.float64]: + """One spectrum in every pixel of an (n_freq, ny, nx) cube.""" + return np.broadcast_to(spectrum[:, None, None], (spectrum.size, ny, nx)).copy() + + +@pytest.mark.parametrize( + "kwargs", + [ + {"stokes_i_weighting": "bogus"}, + {"stokes_i_weighting": "per_pixel", "lam_sq_0_m2": "per_pixel"}, + {"stokes_i_weight_alpha": np.nan}, + {"stokes_i_weight_alpha": "bogus"}, + ], +) +def test_options_reject_bad_values(kwargs: dict[str, Any]) -> None: + """Unknown modes, a non-finite alpha and per-pixel weights with a per-pixel reference raise.""" + with pytest.raises(ValueError): + FDFOptions(n_samples=10.0, **kwargs) + + +@pytest.mark.parametrize("mode", MODES) +def test_theoretical_noise_matches_monte_carlo(mode: StokesIWeighting | None) -> None: + """The reported FDF noise is the scatter of the FDF over noise-only pixels.""" + model = broadcast(fading_model(), 40, 40) + sigma = 0.02 + stokes_q, stokes_u = thin_source(model, frac_pol=0.0, sigma=sigma) + synth = synth_with_model(model, stokes_q, stokes_u, sigma, stokes_i_weighting=mode) + + fdf = synth.fdf_dirty_cube.compute() + measured = np.std(fdf.real) + theory = np.median(np.asarray(synth.theoretical_noise.fdf_q_noise)) + assert measured == pytest.approx(theory, rel=0.05) + + +def test_weights_raise_snr_where_the_model_falls() -> None: + """Following the model gives more SNR, and PI stays unbiased.""" + model = broadcast(fading_model(), 40, 40) + sigma, frac_pol = 0.02, 0.1 + stokes_q, stokes_u = thin_source(model, frac_pol=frac_pol, sigma=sigma) + + snr = {} + for mode in MODES: + synth = synth_with_model( + model, stokes_q, stokes_u, sigma, stokes_i_weighting=mode + ) + truth = frac_pol * np.asarray(synth.stokes_i_ref_flux_map) + peak_pi = np.abs(synth.fdf_dirty_cube.compute()).max(axis=0) + snr[mode] = np.median(truth / np.asarray(synth.theoretical_noise.fdf_q_noise)) + if mode is not None: + assert np.median(peak_pi / truth) == pytest.approx(1.0, abs=0.02) + assert snr[None] < snr["global"] < snr["per_pixel"] + assert snr["per_pixel"] > 3 * snr[None] + + +def test_global_weights_set_lam_sq_0() -> None: + """The reference lambda^2 is the weighted mean of lambda^2 under template^2/sigma^2.""" + model = broadcast(FREQ_ARR_HZ ** -1.5, 2, 2) + sigma = np.linspace(0.01, 0.03, FREQ_ARR_HZ.size) + stokes_q, stokes_u = thin_source(model, frac_pol=0.1, sigma=0.0) + chunks = (-1, 2, 2) + synth = rmsynth_3d( + da.from_array(stokes_q, chunks=chunks), + da.from_array(stokes_u, chunks=chunks), + FREQ_ARR_HZ, + weight_arr=1 / sigma**2, + stokes_i_model=da.from_array(model, chunks=chunks), + d_phi_radm2=2.0, + stokes_i_weight_alpha=-1.5, + ) + weight = stokes_i_template(FREQ_ARR_HZ, -1.5) ** 2 / sigma**2 + lambda_sq = freq_to_lambda2(FREQ_ARR_HZ) + assert synth.lam_sq_0_m2 == pytest.approx(np.sum(weight * lambda_sq) / weight.sum()) + assert synth.stokes_i_weight_alpha == -1.5 + assert synth.stokes_i_weighting == "global" + + +def two_spectral_indices( + freq_arr_hz: NDArray[np.float64] = FREQ_ARR_HZ, +) -> NDArray[np.float64]: + """Two pixels, one flat-ish and one steep, as a (n_freq, 1, 2) model cube.""" + x = freq_arr_hz / 1e9 + return np.stack([x**-0.5, x**-2.5], axis=-1)[:, np.newaxis, :] + + +def test_global_mode_gives_every_pixel_one_rmsf_and_fdf_shape() -> None: + """Sources of different alpha get the same RMSF and the same FDF shape.""" + model = two_spectral_indices(WIDE_FREQ_ARR_HZ) + stokes_q, stokes_u = thin_source( + model, frac_pol=0.1, sigma=0.0, freq_arr_hz=WIDE_FREQ_ARR_HZ + ) + synth = synth_with_model( + model, + stokes_q, + stokes_u, + 1.0, + freq_arr_hz=WIDE_FREQ_ARR_HZ, + stokes_i_weight_alpha=-1.5, + per_pixel_rmsf=True, + ) + rmsf = np.asarray(synth.rmsf_cube.compute()) if synth.rmsf_cube is not None else None + assert rmsf is not None + np.testing.assert_allclose(rmsf[:, 0, 0], rmsf[:, 0, 1], atol=1e-10) + + amp = np.abs(synth.fdf_dirty_cube.compute()) + shape = amp / amp.max(axis=0) + np.testing.assert_allclose(shape[:, 0, 0], shape[:, 0, 1], atol=1e-10) + + +def test_per_pixel_mode_gives_each_pixel_its_own_rmsf() -> None: + """Per-pixel weights make the RMSF follow alpha, so the per-pixel RMSF is forced.""" + model = two_spectral_indices(WIDE_FREQ_ARR_HZ) + stokes_q, stokes_u = thin_source( + model, frac_pol=0.1, sigma=0.0, freq_arr_hz=WIDE_FREQ_ARR_HZ + ) + synth = synth_with_model( + model, + stokes_q, + stokes_u, + 1.0, + d_phi_radm2=0.2, + freq_arr_hz=WIDE_FREQ_ARR_HZ, + stokes_i_weighting="per_pixel", + ) + assert synth.rmsf_cube is not None + assert synth.per_pixel_rmsf is not None + rmsf = np.abs(np.asarray(synth.rmsf_cube.compute())) + phi = synth.phi_double_arr_radm2 + widths = [np.ptp(phi[rmsf[:, 0, i] >= 0.5]) for i in range(2)] + assert widths[1] > widths[0] * 1.04 + + +@pytest.mark.parametrize("mode", [None, "global"]) +def test_fallback_keeps_pi_continuous_across_the_snr_cut( + mode: StokesIWeighting | None, +) -> None: + """Pixels below the cut report the same fractional polarisation as those above.""" + # Wide and steep, so a flat fallback is well off the value at the reference. + freq_arr_hz = WIDE_FREQ_ARR_HZ + ny, nx = 1, 12 + frac_pol, sigma_i, alpha = 0.1, 1e-3, -4.0 + amplitude = np.geomspace(0.05, 1.0, nx)[np.newaxis, :] + stokes_i = (freq_arr_hz / 1.2e9)[:, None, None] ** alpha * amplitude[None] + stokes_i = stokes_i / np.median(stokes_i[:, 0, -1]) * 0.02 + stokes_q, stokes_u = thin_source( + stokes_i, frac_pol=frac_pol, sigma=0.0, freq_arr_hz=freq_arr_hz + ) + chunks = (-1, ny, nx) + synth = rmsynth_3d( + da.from_array(stokes_q, chunks=chunks), + da.from_array(stokes_u, chunks=chunks), + freq_arr_hz, + weight_arr=np.full(freq_arr_hz.size, 1 / 1e-3**2), + stokes_i=da.from_array(stokes_i, chunks=chunks), + stokes_i_error=np.full(freq_arr_hz.size, sigma_i), + fit_order=1, + stokes_i_snr_cut=50.0, + d_phi_radm2=0.5, + phi_max_radm2=100.0, + stokes_i_weighting=mode, + stokes_i_weight_alpha=alpha, + ) + fitted = np.isfinite(np.asarray(synth.stokes_i_model_order_map))[0] + assert fitted.any() and (~fitted).any(), "the cut does not split the row" + + ref_hz = float(lambda2_to_freq(synth.lam_sq_0_m2)) + i_at_ref = np.array( + [np.interp(ref_hz, freq_arr_hz, stokes_i[:, 0, i]) for i in range(nx)] + ) + peak_pi = np.abs(synth.fdf_dirty_cube.compute()).max(axis=0)[0] + frac = peak_pi / i_at_ref + np.testing.assert_allclose(frac[fitted], frac_pol, rtol=1e-3) + if mode is None: + assert np.all(np.abs(frac[~fitted] / frac_pol - 1) > 0.1) + else: + np.testing.assert_allclose(frac[~fitted], frac_pol, rtol=1e-3) + + +def test_field_spectral_index_recovers_alpha() -> None: + """The mean spectrum of a field of one power law gives back its index.""" + rng = np.random.default_rng(3) + amplitude = rng.uniform(0.5, 2.0, size=(1, 10, 10)) + stokes_i = FREQ_ARR_HZ[:, None, None] ** -1.2 * amplitude + stokes_i = stokes_i / stokes_i.mean() + rng.normal(0, 0.01, stokes_i.shape) + stokes_i[5:20, 0, 0] = np.nan + alpha = field_spectral_index( + da.from_array(stokes_i, chunks=(-1, 5, 5)), + np.full(FREQ_ARR_HZ.size, 0.01), + FREQ_ARR_HZ, + StokesIFitOptions(), + ) + assert alpha == pytest.approx(-1.2, abs=0.02) + + +@pytest.mark.parametrize("mode", ["global", "per_pixel"]) +def test_1d_matches_3d_under_variance_weighting(mode: StokesIWeighting) -> None: + """The 1D tool weights a spectrum as the 3D tool weights its pixel.""" + model = two_spectral_indices() + sigma = 0.01 + stokes_q, stokes_u = thin_source(model, frac_pol=0.2, sigma=sigma, seed=4) + synth = synth_with_model( + model, + stokes_q, + stokes_u, + sigma, + stokes_i_weighting=mode, + stokes_i_weight_alpha=-1.5, + ) + fdf_cube = synth.fdf_dirty_cube.compute() + for i in range(2): + ref = run_rmsynth( + freq_arr_hz=FREQ_ARR_HZ, + complex_pol_arr=stokes_q[:, 0, i] + 1j * stokes_u[:, 0, i], + complex_pol_error=np.full(FREQ_ARR_HZ.size, sigma * (1 + 1j)), + stokes_i_model_arr=model[:, 0, i], + stokes_i_model_error=np.zeros(FREQ_ARR_HZ.size), + d_phi_radm2=2.0, + phi_max_radm2=400.0, + stokes_i_weighting=mode, + stokes_i_weight_alpha=-1.5, + ) + ref_fdf = ref.fdf_arrs["fdf_dirty_complex_arr"].to_numpy().astype(complex) + np.testing.assert_allclose(fdf_cube[:, 0, i], ref_fdf, rtol=1e-6, atol=1e-10) + + +@pytest.mark.filterwarnings("ignore: All channels masked") +def test_per_pixel_mode_fits_once_per_chunk( + monkeypatch: pytest.MonkeyPatch, chunked: Callable[..., da.Array] +) -> None: + """The model feeds the weights, the FDF and CLEAN's RMSFs without a refit.""" + calls = {"n": 0} + original = fitting_mod._fit_stokes_i_block + + def counting(*args: Any, **kwargs: Any) -> Any: + calls["n"] += 1 + return original(*args, **kwargs) + + monkeypatch.setattr(fitting_mod, "_fit_stokes_i_block", counting) + rng = np.random.default_rng(5) + stokes_i = rng.uniform(1.0, 3.0, size=(1, 6, 8)) * ( + (FREQ_ARR_HZ / np.median(FREQ_ARR_HZ))[:, None, None] ** -0.8 + ) + stokes_q, stokes_u = thin_source(stokes_i, frac_pol=0.6, sigma=0.0) + synth = rmsynth_3d( + chunked(stokes_q, 3, 4), + chunked(stokes_u, 3, 4), + FREQ_ARR_HZ, + stokes_i=chunked(stokes_i, 3, 4), + stokes_i_error=np.full(FREQ_ARR_HZ.size, 1e-3), + phi_max_radm2=20.0, + d_phi_radm2=2.0, + stokes_i_weighting="per_pixel", + ) + assert synth.per_pixel_rmsf is not None + clean = rmclean3d_mod.run_rmclean( + synth.fdf_dirty_cube, + synth.rmsf_arr, + synth.phi_arr_radm2, + synth.phi_double_arr_radm2, + synth.fwhm_rmsf_radm2, + mask=1e-3, + threshold=1e-3, + per_pixel_rmsf=synth.per_pixel_rmsf, + ) + n_chunks = synth.fdf_dirty_cube.numblocks[1] * synth.fdf_dirty_cube.numblocks[2] + # `map_blocks` probes the block function once while the graph is built. + calls["n"] = 0 + compute( + clean.clean_fdf_cube, + clean.mom0_map, + synth.fdf_dirty_cube, + synth.theoretical_noise.fdf_error_noise, + scheduler="synchronous", + ) + assert calls["n"] == n_chunks diff --git a/tests/test_tools_3d_stokes_i.py b/tests/test_tools_3d_stokes_i.py index 3318370..a28b949 100644 --- a/tests/test_tools_3d_stokes_i.py +++ b/tests/test_tools_3d_stokes_i.py @@ -958,7 +958,7 @@ def artefact_synth(feature_width: float, **kwargs: Any) -> RMSynth3DResults: def test_artefact_spectrum_does_not_blow_up_the_fdf() -> None: """The bug this floor is for: a runaway fit took a 2.5 mJy signal to 1e30.""" - synth = artefact_synth(0.006) + synth = artefact_synth(0.006, stokes_i_weighting=None) fdf = np.asarray(synth.fdf_dirty_cube.compute()) model = np.asarray(require(synth.stokes_i_model_cube).compute()) From bd2e57883f1141851d2b090e528136162a537129 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 08:32:28 +0000 Subject: [PATCH 2/8] Lint the weighting tests Match each option error by its message, and split the SNR-cut assertion so a failure says which side of the cut is empty. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- tests/test_stokes_i_weighting.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/tests/test_stokes_i_weighting.py b/tests/test_stokes_i_weighting.py index ce5b2d1..ef6a6bf 100644 --- a/tests/test_stokes_i_weighting.py +++ b/tests/test_stokes_i_weighting.py @@ -81,17 +81,20 @@ def broadcast(spectrum: NDArray[np.float64], ny: int, nx: int) -> NDArray[np.flo @pytest.mark.parametrize( - "kwargs", + ("kwargs", "match"), [ - {"stokes_i_weighting": "bogus"}, - {"stokes_i_weighting": "per_pixel", "lam_sq_0_m2": "per_pixel"}, - {"stokes_i_weight_alpha": np.nan}, - {"stokes_i_weight_alpha": "bogus"}, + ({"stokes_i_weighting": "bogus"}, "stokes_i_weighting must be one of"), + ( + {"stokes_i_weighting": "per_pixel", "lam_sq_0_m2": "per_pixel"}, + "cannot be combined", + ), + ({"stokes_i_weight_alpha": np.nan}, "stokes_i_weight_alpha must be"), + ({"stokes_i_weight_alpha": "bogus"}, "stokes_i_weight_alpha must be"), ], ) -def test_options_reject_bad_values(kwargs: dict[str, Any]) -> None: +def test_options_reject_bad_values(kwargs: dict[str, Any], match: str) -> None: """Unknown modes, a non-finite alpha and per-pixel weights with a per-pixel reference raise.""" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match=match): FDFOptions(n_samples=10.0, **kwargs) @@ -131,7 +134,7 @@ def test_weights_raise_snr_where_the_model_falls() -> None: def test_global_weights_set_lam_sq_0() -> None: """The reference lambda^2 is the weighted mean of lambda^2 under template^2/sigma^2.""" - model = broadcast(FREQ_ARR_HZ ** -1.5, 2, 2) + model = broadcast(FREQ_ARR_HZ**-1.5, 2, 2) sigma = np.linspace(0.01, 0.03, FREQ_ARR_HZ.size) stokes_q, stokes_u = thin_source(model, frac_pol=0.1, sigma=0.0) chunks = (-1, 2, 2) @@ -174,7 +177,9 @@ def test_global_mode_gives_every_pixel_one_rmsf_and_fdf_shape() -> None: stokes_i_weight_alpha=-1.5, per_pixel_rmsf=True, ) - rmsf = np.asarray(synth.rmsf_cube.compute()) if synth.rmsf_cube is not None else None + rmsf = ( + np.asarray(synth.rmsf_cube.compute()) if synth.rmsf_cube is not None else None + ) assert rmsf is not None np.testing.assert_allclose(rmsf[:, 0, 0], rmsf[:, 0, 1], atol=1e-10) @@ -237,7 +242,8 @@ def test_fallback_keeps_pi_continuous_across_the_snr_cut( stokes_i_weight_alpha=alpha, ) fitted = np.isfinite(np.asarray(synth.stokes_i_model_order_map))[0] - assert fitted.any() and (~fitted).any(), "the cut does not split the row" + assert fitted.any(), "no pixel is above the cut" + assert (~fitted).any(), "no pixel is below the cut" ref_hz = float(lambda2_to_freq(synth.lam_sq_0_m2)) i_at_ref = np.array( From 50102cdae3883105aade7ef6a46b8f29a50350bd Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 11:45:20 +0000 Subject: [PATCH 3/8] Show the Stokes I weighting results, not just check them Each part of the new notebook section now draws what it tests: the channel weights and the peak PI each mode recovers, the RMSF each mode gives two spectral indices, and fractional polarisation across the SNR cut. The section ends with the numbers in one table. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- docs/examples/rmsynth_3d_stokes_i.ipynb | 172 +++++++++++++++++------- 1 file changed, 125 insertions(+), 47 deletions(-) diff --git a/docs/examples/rmsynth_3d_stokes_i.ipynb b/docs/examples/rmsynth_3d_stokes_i.ipynb index f27d9d3..c6d27c7 100644 --- a/docs/examples/rmsynth_3d_stokes_i.ipynb +++ b/docs/examples/rmsynth_3d_stokes_i.ipynb @@ -691,7 +691,7 @@ "\n", "It changes only the noise-based weight types (`\"variance\"`, `\"natural\"`, `\"uniform_lsq\"`, `\"briggs\"`). In both modes $\\lambda^2_0$ from `lam_sq_0_m2=\"auto\"`, and so the Stokes I reference frequency, is the weighted mean of $\\lambda^2$ under the template weights $T_j^2/\\sigma_j^2$.\n", "\n", - "Below is a Faraday-thin source with 10% fractional polarisation, in every pixel of a 40 x 40 cube, over an 800 to 1800 MHz band. Its Stokes I falls to 2% of its peak at the top of the band. The model is supplied, so the only difference between the runs is the weights." + "Below is a Faraday-thin source with 10% fractional polarisation, in every pixel of a 40 x 40 cube, over an 800 to 1800 MHz band. Its Stokes I falls to 2% of its peak at the top of the band. The model is supplied, so the only difference between the runs is the weights. The figure shows the model, the weight each mode gives each channel, and the peak PI each mode recovers in every pixel." ] }, { @@ -701,12 +701,15 @@ "metadata": {}, "outputs": [], "source": [ + "import polars as pl\n", "from rm_lite.utils.synthesis import lambda2_to_freq\n", "\n", "wide_hz = np.linspace(800e6, 1800e6, 64)\n", "wide_lsq = freq_to_lambda2(wide_hz)\n", "wchunks = (-1, 20, 20)\n", "sigma_w, frac_w = 0.02, 0.1\n", + "modes = (None, \"global\", \"per_pixel\")\n", + "mode_colours = {None: \"#2a78d6\", \"global\": \"#eb6834\", \"per_pixel\": \"#1baf7a\"}\n", "\n", "\n", "def thin_cube(model_cube, sigma, rm_radm2=30.0):\n", @@ -733,21 +736,58 @@ "\n", "\n", "t = (wide_hz - wide_hz[0]) / (wide_hz[-1] - wide_hz[0])\n", - "fading_i = np.broadcast_to((1 - 0.98 * t**6)[:, None, None], (wide_hz.size, 40, 40))\n", + "fading_spectrum = 1 - 0.98 * t**6\n", + "fading_i = np.broadcast_to(fading_spectrum[:, None, None], (wide_hz.size, 40, 40))\n", "q_w, u_w = thin_cube(fading_i, sigma_w)\n", "\n", - "snr = {}\n", - "for mode in (None, \"global\", \"per_pixel\"):\n", + "snr, pi_ratio, channel_weight = {}, {}, {}\n", + "for mode in modes:\n", " res = synth_wide(fading_i, q_w, u_w, stokes_i_weighting=mode)\n", " truth = frac_w * res.stokes_i_ref_flux_map.compute()\n", " pi = np.abs(res.fdf_dirty_cube.compute()).max(axis=0)\n", + " pi_ratio[mode] = (pi / truth).ravel()\n", " snr[mode] = float(np.median(truth / np.asarray(res.theoretical_noise.fdf_q_noise)))\n", - " print(\n", - " f\"{mode!s:>9}: median PI / truth {np.median(pi / truth):.3f}, SNR {snr[mode]:.1f}\"\n", + " # The weights each mode gives a channel, from the equations above.\n", + " if mode is None:\n", + " weight = np.ones_like(wide_hz)\n", + " elif mode == \"global\":\n", + " weight = (wide_hz / 1e9) ** (2 * res.stokes_i_weight_alpha)\n", + " else:\n", + " weight = fading_spectrum**2\n", + " channel_weight[mode] = weight / weight.max()\n", + "\n", + "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 4))\n", + "ax1.plot(wide_hz / 1e9, fading_spectrum, color=\"0.3\")\n", + "ax1.set(\n", + " xlabel=\"Frequency / GHz\",\n", + " ylabel=\"Stokes I / peak\",\n", + " title=\"Stokes I model in every pixel\",\n", + ")\n", + "for mode in modes:\n", + " ax2.plot(\n", + " wide_hz / 1e9, channel_weight[mode], color=mode_colours[mode], label=str(mode)\n", + " )\n", + " ax3.hist(\n", + " pi_ratio[mode],\n", + " bins=np.linspace(0.5, 1.8, 60),\n", + " histtype=\"step\",\n", + " lw=2,\n", + " color=mode_colours[mode],\n", + " label=f\"{mode}: SNR {snr[mode]:.1f}\",\n", " )\n", - " if mode is not None:\n", - " assert abs(np.median(pi / truth) - 1) < 0.02\n", + "ax2.set(\n", + " xlabel=\"Frequency / GHz\", ylabel=\"Weight / largest weight\", title=\"Channel weights\"\n", + ")\n", + "ax2.legend(title=\"stokes_i_weighting\")\n", + "ax3.axvline(1.0, color=\"0.5\", ls=\":\")\n", + "ax3.set(\n", + " xlabel=\"Peak PI / true PI\", ylabel=\"Pixels\", title=\"Peak PI over 1600 noisy pixels\"\n", + ")\n", + "ax3.legend(title=\"stokes_i_weighting\")\n", + "fig.tight_layout()\n", "\n", + "for mode in (\"global\", \"per_pixel\"):\n", + " assert abs(np.median(pi_ratio[mode]) - 1) < 0.02\n", "assert snr[None] < snr[\"global\"] < snr[\"per_pixel\"]" ] }, @@ -784,7 +824,8 @@ "\n", "\n", "rmsf_width = {}\n", - "for mode in (None, \"global\", \"per_pixel\"):\n", + "fig, axes = plt.subplots(1, 3, figsize=(16, 4), sharey=True)\n", + "for ax, mode in zip(axes, modes, strict=True):\n", " res = synth_wide(\n", " two_alpha,\n", " q_two,\n", @@ -794,13 +835,25 @@ " per_pixel_rmsf=True,\n", " )\n", " rmsf = res.rmsf_cube.compute()\n", - " rmsf_width[mode] = [\n", - " half_power_width(rmsf[:, 0, i], res.phi_double_arr_radm2) for i in range(2)\n", - " ]\n", - " print(\n", - " f\"{mode!s:>9}: RMSF FWHM {rmsf_width[mode][0]:.1f} (alpha -0.5) and \"\n", - " f\"{rmsf_width[mode][1]:.1f} (alpha -2.5) rad/m^2\"\n", + " phi = res.phi_double_arr_radm2\n", + " rmsf_width[mode] = [half_power_width(rmsf[:, 0, i], phi) for i in range(2)]\n", + " for i, (alpha, style) in enumerate(((-0.5, \"-\"), (-2.5, \"--\"))):\n", + " ax.plot(\n", + " phi,\n", + " np.abs(rmsf[:, 0, i]),\n", + " color=mode_colours[mode],\n", + " ls=style,\n", + " lw=2,\n", + " label=f\"alpha {alpha}: FWHM {rmsf_width[mode][i]:.1f}\",\n", + " )\n", + " ax.set(\n", + " xlim=(-120, 120),\n", + " xlabel=r\"$\\phi$ / (rad m$^{-2}$)\",\n", + " title=f\"stokes_i_weighting={mode}\",\n", " )\n", + " ax.legend(loc=\"upper right\")\n", + "axes[0].set(ylabel=\"|RMSF|\")\n", + "fig.tight_layout()\n", "\n", "assert np.isclose(rmsf_width[None][0], rmsf_width[None][1])\n", "assert np.isclose(rmsf_width[\"global\"][0], rmsf_width[\"global\"][1])\n", @@ -816,7 +869,7 @@ "\n", "A pixel below `stokes_i_snr_cut` has no fit. With the weighting on, it is divided by the template $T(\\nu)/T(\\nu_0)$ instead of a flat model. Its PI is then the PI at the reference frequency, like its fitted neighbours, to the extent that its spectrum follows the template. With `stokes_i_weighting=None` it keeps the flat model and reports a band-averaged PI instead. For a steep spectrum the two differ, so the PI map steps where the cut falls.\n", "\n", - "A row of pixels with the same $\\alpha = -3$ spectrum and a brightness that falls through the cut:" + "A row of pixels with the same $\\alpha = -3$ spectrum and a brightness that falls through the cut. The shaded pixels are below it:" ] }, { @@ -830,7 +883,7 @@ "row_i = 0.02 * (wide_hz / 1.2e9)[:, None, None] ** -3.0 * amp\n", "q_row, u_row = thin_cube(row_i, 0.0)\n", "frac_row = {}\n", - "for mode in (None, \"global\"):\n", + "for mode in modes:\n", " res = rmsynth_3d(\n", " da.from_array(q_row, chunks=(-1, 1, 12)),\n", " da.from_array(u_row, chunks=(-1, 1, 12)),\n", @@ -848,16 +901,48 @@ " ref_hz = float(lambda2_to_freq(res.lam_sq_0_m2))\n", " i_ref = np.array([np.interp(ref_hz, wide_hz, row_i[:, 0, i]) for i in range(12)])\n", " frac_row[mode] = np.abs(res.fdf_dirty_cube.compute()).max(axis=0)[0] / i_ref\n", - " print(\n", - " f\"{mode!s:>6}: fractional pol, fitted pixels {np.round(frac_row[mode][fitted_row], 3)}\"\n", - " )\n", - " print(\n", - " f\"{'':>6} fractional pol, below the cut {np.round(frac_row[mode][~fitted_row], 3)}\"\n", + "\n", + "brightness = amp.ravel()\n", + "cut = np.sqrt(brightness[~fitted_row].max() * brightness[fitted_row].min())\n", + "fig, ax = plt.subplots(figsize=(8, 4))\n", + "ax.axvspan(brightness.min() / 1.2, cut, color=\"0.93\", zorder=0)\n", + "ax.text(\n", + " cut / 1.08,\n", + " 0.097,\n", + " \"below the SNR cut:\\nno fit, fallback model\",\n", + " ha=\"right\",\n", + " color=\"0.3\",\n", + ")\n", + "ax.axhline(frac_w, color=\"0.5\", ls=\":\", label=\"true fractional polarisation\")\n", + "# global and per_pixel give the same values here, so global is drawn hollow.\n", + "for mode, marker, size, face in zip(\n", + " modes, (\"o\", \"s\", \"^\"), (8, 12, 7), (\"full\", \"none\", \"full\"), strict=True\n", + "):\n", + " ax.plot(\n", + " brightness,\n", + " frac_row[mode],\n", + " marker=marker,\n", + " ms=size,\n", + " mfc=mode_colours[mode] if face == \"full\" else \"none\",\n", + " mew=2,\n", + " lw=1.5,\n", + " color=mode_colours[mode],\n", + " label=str(mode),\n", " )\n", + "ax.set(\n", + " xscale=\"log\",\n", + " ylim=(0.095, 0.113),\n", + " xlabel=\"Stokes I brightness / brightest pixel\",\n", + " ylabel=r\"Peak PI / $I(\\nu_0)$\",\n", + " title=r\"A row of $\\alpha = -3$ pixels\",\n", + ")\n", + "ax.legend(title=\"stokes_i_weighting\", loc=\"upper right\")\n", + "fig.tight_layout()\n", "\n", "assert fitted_row.any()\n", "assert (~fitted_row).any()\n", - "np.testing.assert_allclose(frac_row[\"global\"], frac_w, rtol=1e-3)\n", + "for mode in (\"global\", \"per_pixel\"):\n", + " np.testing.assert_allclose(frac_row[mode], frac_w, rtol=1e-3)\n", "assert np.all(np.abs(frac_row[None][~fitted_row] / frac_w - 1) > 0.05)" ] }, @@ -868,7 +953,7 @@ "source": [ "#### Choosing a mode\n", "\n", - "The numbers from the cells above:" + "The numbers from the cells above, in one table:" ] }, { @@ -878,29 +963,22 @@ "metadata": {}, "outputs": [], "source": [ - "rows = [\n", - " (\n", - " \"SNR, Stokes I falling at the band top\",\n", - " snr[None],\n", - " snr[\"global\"],\n", - " snr[\"per_pixel\"],\n", - " ),\n", - " (\n", - " \"RMSF FWHM, alpha -0.5 (rad/m^2)\",\n", - " rmsf_width[None][0],\n", - " rmsf_width[\"global\"][0],\n", - " rmsf_width[\"per_pixel\"][0],\n", - " ),\n", - " (\n", - " \"RMSF FWHM, alpha -2.5 (rad/m^2)\",\n", - " rmsf_width[None][1],\n", - " rmsf_width[\"global\"][1],\n", - " rmsf_width[\"per_pixel\"][1],\n", - " ),\n", - "]\n", - "print(f\"{'':40s}{'None':>10}{'global':>10}{'per_pixel':>10}\")\n", - "for label, *values in rows:\n", - " print(f\"{label:40s}\" + \"\".join(f\"{v:10.1f}\" for v in values))" + "summary = pl.DataFrame(\n", + " {\n", + " \"stokes_i_weighting\": [str(mode) for mode in modes],\n", + " \"PI SNR, fading Stokes I\": [round(snr[mode], 1) for mode in modes],\n", + " \"median peak PI / true PI\": [\n", + " round(float(np.median(pi_ratio[mode])), 3) for mode in modes\n", + " ],\n", + " \"RMSF FWHM, alpha -0.5\": [round(rmsf_width[mode][0], 1) for mode in modes],\n", + " \"RMSF FWHM, alpha -2.5\": [round(rmsf_width[mode][1], 1) for mode in modes],\n", + " \"PI / true PI below the cut\": [\n", + " round(float(np.mean(frac_row[mode][~fitted_row] / frac_w)), 3)\n", + " for mode in modes\n", + " ],\n", + " }\n", + ")\n", + "summary" ] }, { From 17097f2fdcedb34c15e13f354857838518e30ff9 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 12:15:02 +0000 Subject: [PATCH 4/8] Show the Stokes I weighting on a mock field like SB56289 Replace the line plots with maps of one mock source on a RACS low, mid and high band whose spectrum steepens and curves down across it: PI SNR and RMSF width for each option, and PI across the Stokes I SNR cut. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- docs/examples/rmsynth_3d_stokes_i.ipynb | 343 +++++++++--------------- 1 file changed, 130 insertions(+), 213 deletions(-) diff --git a/docs/examples/rmsynth_3d_stokes_i.ipynb b/docs/examples/rmsynth_3d_stokes_i.ipynb index c6d27c7..eab1acb 100644 --- a/docs/examples/rmsynth_3d_stokes_i.ipynb +++ b/docs/examples/rmsynth_3d_stokes_i.ipynb @@ -691,7 +691,7 @@ "\n", "It changes only the noise-based weight types (`\"variance\"`, `\"natural\"`, `\"uniform_lsq\"`, `\"briggs\"`). In both modes $\\lambda^2_0$ from `lam_sq_0_m2=\"auto\"`, and so the Stokes I reference frequency, is the weighted mean of $\\lambda^2$ under the template weights $T_j^2/\\sigma_j^2$.\n", "\n", - "Below is a Faraday-thin source with 10% fractional polarisation, in every pixel of a 40 x 40 cube, over an 800 to 1800 MHz band. Its Stokes I falls to 2% of its peak at the top of the band. The model is supplied, so the only difference between the runs is the weights. The figure shows the model, the weight each mode gives each channel, and the peak PI each mode recovers in every pixel." + "The mock below shows what each option does to the maps. The earlier sections use a single RACS-low band and pure power laws, where the Stokes I model never gets small, so the three options give almost the same maps there. This one copies a RACS low, mid and high combination: 0.80 to 1.09, 1.30 to 1.44 and 1.52 to 1.80 GHz. It has one diffuse source with 30% fractional polarisation, whose spectrum steepens from left to right and curves down at the top of the band, so on the right the high band is close to the noise. Stokes I is fitted at `fit_order=2` with the default SNR cut of 5." ] }, { @@ -704,91 +704,109 @@ "import polars as pl\n", "from rm_lite.utils.synthesis import lambda2_to_freq\n", "\n", - "wide_hz = np.linspace(800e6, 1800e6, 64)\n", - "wide_lsq = freq_to_lambda2(wide_hz)\n", - "wchunks = (-1, 20, 20)\n", - "sigma_w, frac_w = 0.02, 0.1\n", - "modes = (None, \"global\", \"per_pixel\")\n", - "mode_colours = {None: \"#2a78d6\", \"global\": \"#eb6834\", \"per_pixel\": \"#1baf7a\"}\n", + "mock_hz = (\n", + " np.concatenate(\n", + " [np.arange(800, 1089, 8), np.arange(1296, 1441, 8), np.arange(1520, 1801, 8)]\n", + " )\n", + " * 1e6\n", + ")\n", + "mock_lsq = freq_to_lambda2(mock_hz)\n", + "n_mock, sigma_mock, frac_mock = 48, 1e-3, 0.3\n", + "yy, xx = np.mgrid[0:n_mock, 0:n_mock]\n", + "i_1ghz = 30e-3 * np.exp(\n", + " -((xx - n_mock / 2) ** 2 + (yy - n_mock / 2) ** 2) / (2 * 12.0**2)\n", + ")\n", + "alpha_mock = -0.6 - 1.8 * xx / (n_mock - 1)\n", + "beta_mock = -8.0 * xx / (n_mock - 1)\n", "\n", "\n", - "def thin_cube(model_cube, sigma, rm_radm2=30.0):\n", - " # Faraday-thin Q/U that follow the Stokes I model, plus noise.\n", - " angle = (2 * rm_radm2 * wide_lsq)[:, None, None]\n", - " noise = rng.normal(0, sigma, (2, *model_cube.shape))\n", - " return (\n", - " frac_w * model_cube * np.cos(angle) + noise[0],\n", - " frac_w * model_cube * np.sin(angle) + noise[1],\n", - " )\n", + "def mock_flux(freq_hz):\n", + " # The input Stokes I at freq_hz, a curved power law.\n", + " log_x = np.log10(np.asarray(freq_hz) / 1e9)\n", + " return i_1ghz * 10 ** (alpha_mock * log_x + beta_mock * log_x**2)\n", + "\n", + "\n", + "mock_i = np.stack([mock_flux(f) for f in mock_hz])\n", + "angle = 2 * (20.0 + 40.0 * yy / (n_mock - 1))[None] * mock_lsq[:, None, None]\n", + "noise = rng.normal(0, sigma_mock, (3, *mock_i.shape))\n", + "mock_q = frac_mock * mock_i * np.cos(angle) + noise[0]\n", + "mock_u = frac_mock * mock_i * np.sin(angle) + noise[1]\n", + "mock_i_obs = mock_i + noise[2]\n", + "mock_chunks = (-1, 24, 24)\n", "\n", "\n", - "def synth_wide(model_cube, q, u, d_phi_radm2=0.5, **kwargs):\n", + "def synth_mock(**kwargs):\n", " return rmsynth_3d(\n", - " da.from_array(q, chunks=wchunks),\n", - " da.from_array(u, chunks=wchunks),\n", - " wide_hz,\n", - " weight_arr=np.full(wide_hz.size, 1 / sigma_w**2),\n", - " stokes_i_model=da.from_array(model_cube, chunks=wchunks),\n", + " da.from_array(mock_q, chunks=mock_chunks),\n", + " da.from_array(mock_u, chunks=mock_chunks),\n", + " mock_hz,\n", + " weight_arr=np.full(mock_hz.size, 1 / sigma_mock**2),\n", + " stokes_i=da.from_array(mock_i_obs, chunks=mock_chunks),\n", + " stokes_i_error=np.full(mock_hz.size, sigma_mock),\n", + " fit_order=2,\n", " phi_max_radm2=300.0,\n", - " d_phi_radm2=d_phi_radm2,\n", + " d_phi_radm2=1.0,\n", " **kwargs,\n", " )\n", "\n", "\n", - "t = (wide_hz - wide_hz[0]) / (wide_hz[-1] - wide_hz[0])\n", - "fading_spectrum = 1 - 0.98 * t**6\n", - "fading_i = np.broadcast_to(fading_spectrum[:, None, None], (wide_hz.size, 40, 40))\n", - "q_w, u_w = thin_cube(fading_i, sigma_w)\n", + "def compare(result):\n", + " # Peak PI, the true PI at this run's reference frequency, and fitted pixels.\n", + " fdf, order, fdf_noise = da.compute(\n", + " result.fdf_dirty_cube,\n", + " result.stokes_i_model_order_map,\n", + " result.theoretical_noise.fdf_error_noise,\n", + " )\n", + " true_pi = frac_mock * mock_flux(lambda2_to_freq(result.lam_sq_0_m2))\n", + " peak_pi = np.abs(fdf).max(axis=0)\n", + " return peak_pi / true_pi, true_pi / np.asarray(fdf_noise), np.isfinite(order)\n", + "\n", "\n", - "snr, pi_ratio, channel_weight = {}, {}, {}\n", + "modes = (None, \"global\", \"per_pixel\")\n", + "on_source = i_1ghz > 5 * sigma_mock / np.sqrt(mock_hz.size)\n", + "left, right = on_source & (xx < 16), on_source & (xx > 32)\n", + "pi_snr, rmsf_fwhm, pi_error = {}, {}, {}\n", "for mode in modes:\n", - " res = synth_wide(fading_i, q_w, u_w, stokes_i_weighting=mode)\n", - " truth = frac_w * res.stokes_i_ref_flux_map.compute()\n", - " pi = np.abs(res.fdf_dirty_cube.compute()).max(axis=0)\n", - " pi_ratio[mode] = (pi / truth).ravel()\n", - " snr[mode] = float(np.median(truth / np.asarray(res.theoretical_noise.fdf_q_noise)))\n", - " # The weights each mode gives a channel, from the equations above.\n", - " if mode is None:\n", - " weight = np.ones_like(wide_hz)\n", - " elif mode == \"global\":\n", - " weight = (wide_hz / 1e9) ** (2 * res.stokes_i_weight_alpha)\n", - " else:\n", - " weight = fading_spectrum**2\n", - " channel_weight[mode] = weight / weight.max()\n", - "\n", - "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 4))\n", - "ax1.plot(wide_hz / 1e9, fading_spectrum, color=\"0.3\")\n", - "ax1.set(\n", - " xlabel=\"Frequency / GHz\",\n", - " ylabel=\"Stokes I / peak\",\n", - " title=\"Stokes I model in every pixel\",\n", + " result = synth_mock(stokes_i_weighting=mode, per_pixel_rmsf=True)\n", + " ratio, pi_snr[mode], _ = compare(result)\n", + " rmsf = np.abs(result.rmsf_cube.compute())\n", + " d_phi = np.diff(result.phi_double_arr_radm2)[0]\n", + " rmsf_fwhm[mode] = (rmsf >= 0.5 * rmsf.max(axis=0)).sum(axis=0) * d_phi\n", + " pi_error[mode] = np.abs(ratio - 1)\n", + "\n", + "fig, axes = plt.subplots(3, 3, figsize=(15, 13), constrained_layout=True)\n", + "inputs = (\n", + " (i_1ghz * 1e3, \"Stokes I at 1 GHz / mJy\", \"viridis\"),\n", + " (alpha_mock, r\"Spectral index $\\alpha$\", \"Spectral\"),\n", + " (beta_mock, r\"Curvature $\\beta$\", \"coolwarm\"),\n", ")\n", - "for mode in modes:\n", - " ax2.plot(\n", - " wide_hz / 1e9, channel_weight[mode], color=mode_colours[mode], label=str(mode)\n", + "for ax, (data, title, cmap) in zip(axes[0], inputs, strict=True):\n", + " im = ax.imshow(np.where(on_source, data, np.nan), origin=\"lower\", cmap=cmap)\n", + " fig.colorbar(im, ax=ax, shrink=0.8)\n", + " ax.set(title=f\"Input: {title}\")\n", + "for ax_snr, ax_rmsf, mode in zip(axes[1], axes[2], modes, strict=True):\n", + " im_snr = ax_snr.imshow(\n", + " np.where(on_source, pi_snr[mode], np.nan), origin=\"lower\", vmin=0, vmax=60\n", " )\n", - " ax3.hist(\n", - " pi_ratio[mode],\n", - " bins=np.linspace(0.5, 1.8, 60),\n", - " histtype=\"step\",\n", - " lw=2,\n", - " color=mode_colours[mode],\n", - " label=f\"{mode}: SNR {snr[mode]:.1f}\",\n", + " ax_snr.set(title=f\"stokes_i_weighting={mode}: PI SNR\")\n", + " im_rmsf = ax_rmsf.imshow(\n", + " np.where(on_source, rmsf_fwhm[mode], np.nan),\n", + " origin=\"lower\",\n", + " vmin=30,\n", + " vmax=48,\n", + " cmap=\"magma\",\n", " )\n", - "ax2.set(\n", - " xlabel=\"Frequency / GHz\", ylabel=\"Weight / largest weight\", title=\"Channel weights\"\n", - ")\n", - "ax2.legend(title=\"stokes_i_weighting\")\n", - "ax3.axvline(1.0, color=\"0.5\", ls=\":\")\n", - "ax3.set(\n", - " xlabel=\"Peak PI / true PI\", ylabel=\"Pixels\", title=\"Peak PI over 1600 noisy pixels\"\n", - ")\n", - "ax3.legend(title=\"stokes_i_weighting\")\n", - "fig.tight_layout()\n", + " ax_rmsf.set(title=f\"stokes_i_weighting={mode}: RMSF FWHM\")\n", + "fig.colorbar(im_snr, ax=axes[1], shrink=0.8, label=\"True PI / reported noise\")\n", + "fig.colorbar(im_rmsf, ax=axes[2], shrink=0.8, label=r\"FWHM / (rad m$^{-2}$)\")\n", "\n", "for mode in (\"global\", \"per_pixel\"):\n", - " assert abs(np.median(pi_ratio[mode]) - 1) < 0.02\n", - "assert snr[None] < snr[\"global\"] < snr[\"per_pixel\"]" + " assert np.median(pi_snr[mode][right]) > 2 * np.median(pi_snr[None][right])\n", + "for mode in (None, \"global\"):\n", + " assert np.ptp(rmsf_fwhm[mode][on_source]) == 0\n", + "assert np.median(rmsf_fwhm[\"per_pixel\"][right]) > np.median(\n", + " rmsf_fwhm[\"per_pixel\"][left]\n", + ")" ] }, { @@ -796,13 +814,17 @@ "id": "26", "metadata": {}, "source": [ - "Both modes give an unbiased PI and a higher SNR than $1/\\sigma_j^2$. The per-pixel weights are the inverse variance of each pixel's own corrected data, so they give the highest SNR that pixel can have. Global mode gets less here, because a power law cannot follow a spectrum that turns down this sharply.\n", + "The middle row is the SNR of the PI each option recovers, and the bottom row is the RMSF width in each pixel.\n", + "\n", + "- `None`: on the left, where the spectrum is close to flat, the SNR matches the other options. On the right the model falls to a few percent of its reference flux in the high band, Q and U are divided by it, and those channels still get full weight, so the SNR falls to a third. The source looks lopsided although its fractional polarisation is the same everywhere.\n", + "- `\"global\"`: the SNR follows the Stokes I brightness across the source, and every pixel has the same RMSF. Here that RMSF is wider than with `None`, because the template gives the high band less weight. The width follows the template, so on another band or field it can be narrower.\n", + "- `\"per_pixel\"`: the highest SNR, but the RMSF width follows the spectral index of each pixel, from about 31 rad/m$^2$ on the left to 49 on the right, and it scatters with the noise in each pixel's fit. A second moment or any other width measured against the RMSF would do the same.\n", "\n", - "#### The RMSF\n", + "#### Pixels below the Stokes I SNR cut\n", "\n", - "The weights also set the RMSF. In global mode every pixel has the same weights, so every pixel has the same RMSF and the same effective $\\lambda^2$, whatever its spectral index. In per-pixel mode the weights follow each pixel's Stokes I, so the RMSF and the effective $\\lambda^2$ change with the spectral index, and any Faraday-complexity measure made against the RMSF, such as the second moment of the FDF, is partly set by the Stokes I spectrum.\n", + "A pixel below `stokes_i_snr_cut` has no fit. With `stokes_i_weighting=None` it is not corrected, so its peak PI is a band-averaged PI, not the PI at the reference frequency. With either weighting it is divided by the field-wide power law instead, so it reports PI at the reference frequency, as far as its spectrum follows that power law.\n", "\n", - "Two noiseless Faraday-thin sources, with $\\alpha = -0.5$ and $\\alpha = -2.5$. The global template is fitted to their mean Stokes I:" + "The default cut of 5 leaves almost no polarised signal below it, so here it is raised to 60 to put the edge of the fitted region (black contour) on the source. The colour is peak PI over the true PI, and white is correct." ] }, { @@ -812,52 +834,28 @@ "metadata": {}, "outputs": [], "source": [ - "two_alpha = np.stack([(wide_hz / 1e9) ** -0.5, (wide_hz / 1e9) ** -2.5], axis=-1)[\n", - " :, None, :\n", - "]\n", - "q_two, u_two = thin_cube(two_alpha, 0.0)\n", - "\n", - "\n", - "def half_power_width(rmsf, phi):\n", - " # Width of the main lobe at half its peak.\n", - " return float(np.ptp(phi[np.abs(rmsf) >= 0.5 * np.abs(rmsf).max()]))\n", - "\n", - "\n", - "rmsf_width = {}\n", - "fig, axes = plt.subplots(1, 3, figsize=(16, 4), sharey=True)\n", + "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)\n", + "below_cut = {}\n", "for ax, mode in zip(axes, modes, strict=True):\n", - " res = synth_wide(\n", - " two_alpha,\n", - " q_two,\n", - " u_two,\n", - " d_phi_radm2=0.2,\n", - " stokes_i_weighting=mode,\n", - " per_pixel_rmsf=True,\n", + " ratio, snr_cut_map, fitted = compare(\n", + " synth_mock(stokes_i_weighting=mode, stokes_i_snr_cut=60.0)\n", " )\n", - " rmsf = res.rmsf_cube.compute()\n", - " phi = res.phi_double_arr_radm2\n", - " rmsf_width[mode] = [half_power_width(rmsf[:, 0, i], phi) for i in range(2)]\n", - " for i, (alpha, style) in enumerate(((-0.5, \"-\"), (-2.5, \"--\"))):\n", - " ax.plot(\n", - " phi,\n", - " np.abs(rmsf[:, 0, i]),\n", - " color=mode_colours[mode],\n", - " ls=style,\n", - " lw=2,\n", - " label=f\"alpha {alpha}: FWHM {rmsf_width[mode][i]:.1f}\",\n", - " )\n", - " ax.set(\n", - " xlim=(-120, 120),\n", - " xlabel=r\"$\\phi$ / (rad m$^{-2}$)\",\n", - " title=f\"stokes_i_weighting={mode}\",\n", + " shown = snr_cut_map > 5\n", + " below_cut[mode] = float(np.median(ratio[shown & ~fitted]))\n", + " im = ax.imshow(\n", + " np.where(shown, ratio, np.nan),\n", + " origin=\"lower\",\n", + " vmin=0.7,\n", + " vmax=1.3,\n", + " cmap=\"RdBu_r\",\n", " )\n", - " ax.legend(loc=\"upper right\")\n", - "axes[0].set(ylabel=\"|RMSF|\")\n", - "fig.tight_layout()\n", + " ax.contour(fitted, levels=[0.5], colors=\"k\", linewidths=1.5)\n", + " ax.set(title=f\"stokes_i_weighting={mode}: peak PI / true PI\")\n", + "fig.colorbar(im, ax=axes, shrink=0.8)\n", "\n", - "assert np.isclose(rmsf_width[None][0], rmsf_width[None][1])\n", - "assert np.isclose(rmsf_width[\"global\"][0], rmsf_width[\"global\"][1])\n", - "assert rmsf_width[\"per_pixel\"][1] > 1.04 * rmsf_width[\"per_pixel\"][0]" + "assert abs(below_cut[None] - 1) > 0.05\n", + "for mode in (\"global\", \"per_pixel\"):\n", + " assert abs(below_cut[mode] - 1) < 0.03" ] }, { @@ -865,117 +863,36 @@ "id": "28", "metadata": {}, "source": [ - "#### Pixels below the Stokes I SNR cut\n", - "\n", - "A pixel below `stokes_i_snr_cut` has no fit. With the weighting on, it is divided by the template $T(\\nu)/T(\\nu_0)$ instead of a flat model. Its PI is then the PI at the reference frequency, like its fitted neighbours, to the extent that its spectrum follows the template. With `stokes_i_weighting=None` it keeps the flat model and reports a band-averaged PI instead. For a steep spectrum the two differ, so the PI map steps where the cut falls.\n", - "\n", - "A row of pixels with the same $\\alpha = -3$ spectrum and a brightness that falls through the cut. The shaded pixels are below it:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "29", - "metadata": {}, - "outputs": [], - "source": [ - "amp = np.geomspace(0.05, 1.0, 12)[None, None, :]\n", - "row_i = 0.02 * (wide_hz / 1.2e9)[:, None, None] ** -3.0 * amp\n", - "q_row, u_row = thin_cube(row_i, 0.0)\n", - "frac_row = {}\n", - "for mode in modes:\n", - " res = rmsynth_3d(\n", - " da.from_array(q_row, chunks=(-1, 1, 12)),\n", - " da.from_array(u_row, chunks=(-1, 1, 12)),\n", - " wide_hz,\n", - " weight_arr=np.full(wide_hz.size, 1 / sigma_w**2),\n", - " stokes_i=da.from_array(row_i, chunks=(-1, 1, 12)),\n", - " stokes_i_error=np.full(wide_hz.size, 1e-3),\n", - " fit_order=1,\n", - " stokes_i_snr_cut=50.0,\n", - " phi_max_radm2=100.0,\n", - " d_phi_radm2=0.5,\n", - " stokes_i_weighting=mode,\n", - " )\n", - " fitted_row = np.isfinite(res.stokes_i_model_order_map.compute())[0]\n", - " ref_hz = float(lambda2_to_freq(res.lam_sq_0_m2))\n", - " i_ref = np.array([np.interp(ref_hz, wide_hz, row_i[:, 0, i]) for i in range(12)])\n", - " frac_row[mode] = np.abs(res.fdf_dirty_cube.compute()).max(axis=0)[0] / i_ref\n", - "\n", - "brightness = amp.ravel()\n", - "cut = np.sqrt(brightness[~fitted_row].max() * brightness[fitted_row].min())\n", - "fig, ax = plt.subplots(figsize=(8, 4))\n", - "ax.axvspan(brightness.min() / 1.2, cut, color=\"0.93\", zorder=0)\n", - "ax.text(\n", - " cut / 1.08,\n", - " 0.097,\n", - " \"below the SNR cut:\\nno fit, fallback model\",\n", - " ha=\"right\",\n", - " color=\"0.3\",\n", - ")\n", - "ax.axhline(frac_w, color=\"0.5\", ls=\":\", label=\"true fractional polarisation\")\n", - "# global and per_pixel give the same values here, so global is drawn hollow.\n", - "for mode, marker, size, face in zip(\n", - " modes, (\"o\", \"s\", \"^\"), (8, 12, 7), (\"full\", \"none\", \"full\"), strict=True\n", - "):\n", - " ax.plot(\n", - " brightness,\n", - " frac_row[mode],\n", - " marker=marker,\n", - " ms=size,\n", - " mfc=mode_colours[mode] if face == \"full\" else \"none\",\n", - " mew=2,\n", - " lw=1.5,\n", - " color=mode_colours[mode],\n", - " label=str(mode),\n", - " )\n", - "ax.set(\n", - " xscale=\"log\",\n", - " ylim=(0.095, 0.113),\n", - " xlabel=\"Stokes I brightness / brightest pixel\",\n", - " ylabel=r\"Peak PI / $I(\\nu_0)$\",\n", - " title=r\"A row of $\\alpha = -3$ pixels\",\n", - ")\n", - "ax.legend(title=\"stokes_i_weighting\", loc=\"upper right\")\n", - "fig.tight_layout()\n", + "- `None`: below the cut the PI is biased low, so the map steps at the contour.\n", + "- `\"global\"` and `\"per_pixel\"`: the step is gone on average. Where a pixel's spectrum differs from the field-wide power law the fallback is still off: on the left, where the spectrum is much flatter than the template, PI below the cut comes out high.\n", "\n", - "assert fitted_row.any()\n", - "assert (~fitted_row).any()\n", - "for mode in (\"global\", \"per_pixel\"):\n", - " np.testing.assert_allclose(frac_row[mode], frac_w, rtol=1e-3)\n", - "assert np.all(np.abs(frac_row[None][~fitted_row] / frac_w - 1) > 0.05)" - ] - }, - { - "cell_type": "markdown", - "id": "30", - "metadata": {}, - "source": [ "#### Choosing a mode\n", "\n", - "The numbers from the cells above, in one table:" + "The numbers from the maps above, over the source:" ] }, { "cell_type": "code", "execution_count": null, - "id": "31", + "id": "29", "metadata": {}, "outputs": [], "source": [ "summary = pl.DataFrame(\n", " {\n", " \"stokes_i_weighting\": [str(mode) for mode in modes],\n", - " \"PI SNR, fading Stokes I\": [round(snr[mode], 1) for mode in modes],\n", - " \"median peak PI / true PI\": [\n", - " round(float(np.median(pi_ratio[mode])), 3) for mode in modes\n", + " \"PI SNR, left half\": [\n", + " round(float(np.median(pi_snr[m][left])), 1) for m in modes\n", " ],\n", - " \"RMSF FWHM, alpha -0.5\": [round(rmsf_width[mode][0], 1) for mode in modes],\n", - " \"RMSF FWHM, alpha -2.5\": [round(rmsf_width[mode][1], 1) for mode in modes],\n", - " \"PI / true PI below the cut\": [\n", - " round(float(np.mean(frac_row[mode][~fitted_row] / frac_w)), 3)\n", - " for mode in modes\n", + " \"PI SNR, right half\": [\n", + " round(float(np.median(pi_snr[m][right])), 1) for m in modes\n", " ],\n", + " \"PI error, right half\": [\n", + " round(float(np.median(pi_error[m][right])), 3) for m in modes\n", + " ],\n", + " \"RMSF FWHM, left\": [float(np.median(rmsf_fwhm[m][left])) for m in modes],\n", + " \"RMSF FWHM, right\": [float(np.median(rmsf_fwhm[m][right])) for m in modes],\n", + " \"PI / true PI below the cut\": [round(below_cut[m], 3) for m in modes],\n", " }\n", ")\n", "summary" @@ -983,17 +900,17 @@ }, { "cell_type": "markdown", - "id": "32", + "id": "30", "metadata": {}, "source": [ - "- `\"global\"`: one RMSF and one effective frequency for the whole map, so PI, RM and Faraday-complexity maps compare pixel to pixel. A pixel whose spectrum is far from the template has less SNR than it could, and a pixel whose model falls far below the template in some channels is only partly protected from the noise there. The RMSF is set by the template, so it differs from the $1/\\sigma_j^2$ RMSF. Whether it is wider or narrower depends on the band and the noise: here it is narrower, and on a band with gaps it can be wider. With `\"auto\"` the template comes from each field, so pass a number to give several fields one RMSF.\n", + "- `\"global\"`: one RMSF and one effective frequency for the whole map, so PI, RM and Faraday-complexity maps compare pixel to pixel. A pixel whose spectrum is far from the template has less SNR than it could, a pixel whose model falls far below the template in some channels is only partly protected from the noise there, and below the SNR cut the fallback is only right where a pixel's spectrum follows the template. The RMSF is set by the template, so it differs from the $1/\\sigma_j^2$ RMSF. Whether it is wider or narrower depends on the band and the noise: here it is narrower, and on a band with gaps it can be wider. With `\"auto\"` the template comes from each field, so pass a number to give several fields one RMSF.\n", "- `\"per_pixel\"`: the highest SNR in every pixel. The RMSF and the effective frequency change with the spectral index, and it cannot be used with `lam_sq_0_m2=\"per_pixel\"`.\n", "- `None`: the weights before this option existed. The FDF noise is unbounded where a model falls towards zero." ] }, { "cell_type": "markdown", - "id": "33", + "id": "31", "metadata": {}, "source": [ "`rmsynth_3d_from_fits` takes the same Stokes I options as file paths (`stokes_i_file`, `stokes_i_error_file`, or `stokes_i_model_file`). See the [3D RM-CLEAN](rmclean_3d.ipynb) page to deconvolve the FDF cube." From 9a35581a68a846a61442f98431d1ac38b5edf5a4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 12:22:10 +0000 Subject: [PATCH 5/8] Split the mock figures, and give the mock two compact sources Label the mock input on its own figure with two example spectra, then show recovered PI SNR and RMSF width as separate figures. Use two compact sources on empty sky, as in the earlier mock: one nearly flat, one steep and curving down at the top of the band. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- docs/examples/rmsynth_3d_stokes_i.ipynb | 150 +++++++++++++++++------- 1 file changed, 109 insertions(+), 41 deletions(-) diff --git a/docs/examples/rmsynth_3d_stokes_i.ipynb b/docs/examples/rmsynth_3d_stokes_i.ipynb index eab1acb..1686e3d 100644 --- a/docs/examples/rmsynth_3d_stokes_i.ipynb +++ b/docs/examples/rmsynth_3d_stokes_i.ipynb @@ -691,7 +691,7 @@ "\n", "It changes only the noise-based weight types (`\"variance\"`, `\"natural\"`, `\"uniform_lsq\"`, `\"briggs\"`). In both modes $\\lambda^2_0$ from `lam_sq_0_m2=\"auto\"`, and so the Stokes I reference frequency, is the weighted mean of $\\lambda^2$ under the template weights $T_j^2/\\sigma_j^2$.\n", "\n", - "The mock below shows what each option does to the maps. The earlier sections use a single RACS-low band and pure power laws, where the Stokes I model never gets small, so the three options give almost the same maps there. This one copies a RACS low, mid and high combination: 0.80 to 1.09, 1.30 to 1.44 and 1.52 to 1.80 GHz. It has one diffuse source with 30% fractional polarisation, whose spectrum steepens from left to right and curves down at the top of the band, so on the right the high band is close to the noise. Stokes I is fitted at `fit_order=2` with the default SNR cut of 5." + "The mock below shows what each option does to the maps. The earlier sections use a single RACS-low band and pure power laws, where the Stokes I model never gets small, so the three options give almost the same maps there. This one copies a RACS low, mid and high combination: 0.80 to 1.09, 1.30 to 1.44 and 1.52 to 1.80 GHz. It has two compact sources, both 30% polarised. The left one has a nearly flat spectrum. The right one is steep and curves down at the top of the band, so its high band is close to the noise. Stokes I is fitted at `fit_order=2` with the default SNR cut of 5." ] }, { @@ -713,8 +713,15 @@ "mock_lsq = freq_to_lambda2(mock_hz)\n", "n_mock, sigma_mock, frac_mock = 48, 1e-3, 0.3\n", "yy, xx = np.mgrid[0:n_mock, 0:n_mock]\n", - "i_1ghz = 30e-3 * np.exp(\n", - " -((xx - n_mock / 2) ** 2 + (yy - n_mock / 2) ** 2) / (2 * 12.0**2)\n", + "# Two compact sources on empty sky, as in the mock above: one on the left\n", + "# with a nearly flat spectrum, one on the right that is steep and curves down.\n", + "left_pixel, right_pixel = (\n", + " (int(0.7 * n_mock), int(0.25 * n_mock)),\n", + " (int(0.35 * n_mock), int(0.75 * n_mock)),\n", + ")\n", + "i_1ghz = sum(\n", + " gaussian(np.hypot(xx - x0, yy - y0), amplitude=30e-3, mean=0.0, fwhm=8.0)\n", + " for y0, x0 in (left_pixel, right_pixel)\n", ")\n", "alpha_mock = -0.6 - 1.8 * xx / (n_mock - 1)\n", "beta_mock = -8.0 * xx / (n_mock - 1)\n", @@ -762,9 +769,70 @@ " return peak_pi / true_pi, true_pi / np.asarray(fdf_noise), np.isfinite(order)\n", "\n", "\n", - "modes = (None, \"global\", \"per_pixel\")\n", "on_source = i_1ghz > 5 * sigma_mock / np.sqrt(mock_hz.size)\n", - "left, right = on_source & (xx < 16), on_source & (xx > 32)\n", + "\n", + "fig, axes = plt.subplots(1, 4, figsize=(18, 4), constrained_layout=True)\n", + "maps = (\n", + " (i_1ghz * 1e3, \"Stokes I at 1 GHz / mJy\", \"viridis\"),\n", + " (frac_mock * i_1ghz * 1e3, \"Polarised intensity at 1 GHz / mJy\", \"viridis\"),\n", + " (alpha_mock, r\"Stokes I spectral index $\\alpha$\", \"Spectral\"),\n", + " (beta_mock, r\"Stokes I curvature $\\beta$\", \"coolwarm\"),\n", + ")\n", + "for ax, (data, title, cmap) in zip(axes, maps, strict=True):\n", + " im = ax.imshow(np.where(on_source, data, np.nan), origin=\"lower\", cmap=cmap)\n", + " fig.colorbar(im, ax=ax, shrink=0.8)\n", + " ax.set(title=title)\n", + "for (y, x), marker in zip((left_pixel, right_pixel), (\"o\", \"s\"), strict=True):\n", + " axes[0].plot(x, y, marker, color=\"w\", mec=\"k\", ms=9)\n", + "fig.suptitle(\"Mock input, the same for every run below\")\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(13, 4), sharey=True, constrained_layout=True)\n", + "for ax, (y, x), name in zip(\n", + " axes,\n", + " (left_pixel, right_pixel),\n", + " (\"left source (circle)\", \"right source (square)\"),\n", + " strict=True,\n", + "):\n", + " ax.plot(\n", + " mock_hz / 1e9,\n", + " mock_i_obs[:, y, x] * 1e3,\n", + " \".\",\n", + " color=\"0.6\",\n", + " label=\"Stokes I, observed\",\n", + " )\n", + " ax.plot(\n", + " mock_hz / 1e9, mock_i[:, y, x] * 1e3, \"-\", color=\"k\", label=\"Stokes I, input\"\n", + " )\n", + " ax.plot(\n", + " mock_hz / 1e9,\n", + " np.hypot(mock_q, mock_u)[:, y, x] * 1e3,\n", + " \".\",\n", + " color=\"#2a78d6\",\n", + " label=\"PI, observed\",\n", + " )\n", + " ax.axhline(sigma_mock * 1e3, color=\"0.4\", ls=\":\", label=\"noise per channel\")\n", + " ax.set(xlabel=\"Frequency / GHz\", title=f\"Spectra of the {name} pixel\")\n", + " ax.legend()\n", + "axes[0].set(ylabel=\"Flux density / mJy\")" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "The spectra show the problem. The left source's Stokes I is close to flat and far above the noise. The right source's steepens and curves down, so by the top of the band Stokes I and PI are both near the noise, and a Stokes I model fitted there can fall far below the reference flux." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "modes = (None, \"global\", \"per_pixel\")\n", + "left, right = on_source & (xx < n_mock / 2), on_source & (xx >= n_mock / 2)\n", "pi_snr, rmsf_fwhm, pi_error = {}, {}, {}\n", "for mode in modes:\n", " result = synth_mock(stokes_i_weighting=mode, per_pixel_rmsf=True)\n", @@ -774,34 +842,30 @@ " rmsf_fwhm[mode] = (rmsf >= 0.5 * rmsf.max(axis=0)).sum(axis=0) * d_phi\n", " pi_error[mode] = np.abs(ratio - 1)\n", "\n", - "fig, axes = plt.subplots(3, 3, figsize=(15, 13), constrained_layout=True)\n", - "inputs = (\n", - " (i_1ghz * 1e3, \"Stokes I at 1 GHz / mJy\", \"viridis\"),\n", - " (alpha_mock, r\"Spectral index $\\alpha$\", \"Spectral\"),\n", - " (beta_mock, r\"Curvature $\\beta$\", \"coolwarm\"),\n", - ")\n", - "for ax, (data, title, cmap) in zip(axes[0], inputs, strict=True):\n", - " im = ax.imshow(np.where(on_source, data, np.nan), origin=\"lower\", cmap=cmap)\n", - " fig.colorbar(im, ax=ax, shrink=0.8)\n", - " ax.set(title=f\"Input: {title}\")\n", - "for ax_snr, ax_rmsf, mode in zip(axes[1], axes[2], modes, strict=True):\n", - " im_snr = ax_snr.imshow(\n", + "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)\n", + "for ax, mode in zip(axes, modes, strict=True):\n", + " im = ax.imshow(\n", " np.where(on_source, pi_snr[mode], np.nan), origin=\"lower\", vmin=0, vmax=60\n", " )\n", - " ax_snr.set(title=f\"stokes_i_weighting={mode}: PI SNR\")\n", - " im_rmsf = ax_rmsf.imshow(\n", + " ax.set(title=f\"stokes_i_weighting={mode}\")\n", + "fig.colorbar(im, ax=axes, shrink=0.8, label=\"True PI / reported noise\")\n", + "fig.suptitle(\"Recovered polarised intensity: SNR of the peak PI\")\n", + "\n", + "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)\n", + "for ax, mode in zip(axes, modes, strict=True):\n", + " im = ax.imshow(\n", " np.where(on_source, rmsf_fwhm[mode], np.nan),\n", " origin=\"lower\",\n", " vmin=30,\n", " vmax=48,\n", " cmap=\"magma\",\n", " )\n", - " ax_rmsf.set(title=f\"stokes_i_weighting={mode}: RMSF FWHM\")\n", - "fig.colorbar(im_snr, ax=axes[1], shrink=0.8, label=\"True PI / reported noise\")\n", - "fig.colorbar(im_rmsf, ax=axes[2], shrink=0.8, label=r\"FWHM / (rad m$^{-2}$)\")\n", + " ax.set(title=f\"stokes_i_weighting={mode}\")\n", + "fig.colorbar(im, ax=axes, shrink=0.8, label=r\"FWHM / (rad m$^{-2}$)\")\n", + "fig.suptitle(\"RMSF width in each pixel\")\n", "\n", "for mode in (\"global\", \"per_pixel\"):\n", - " assert np.median(pi_snr[mode][right]) > 2 * np.median(pi_snr[None][right])\n", + " assert np.median(pi_snr[mode][right]) > np.median(pi_snr[None][right])\n", "for mode in (None, \"global\"):\n", " assert np.ptp(rmsf_fwhm[mode][on_source]) == 0\n", "assert np.median(rmsf_fwhm[\"per_pixel\"][right]) > np.median(\n", @@ -811,14 +875,14 @@ }, { "cell_type": "markdown", - "id": "26", + "id": "28", "metadata": {}, "source": [ - "The middle row is the SNR of the PI each option recovers, and the bottom row is the RMSF width in each pixel.\n", + "Each option is run on this same input. The first figure below is the SNR of the peak polarised intensity each option recovers, the second is the RMSF width in each pixel.\n", "\n", - "- `None`: on the left, where the spectrum is close to flat, the SNR matches the other options. On the right the model falls to a few percent of its reference flux in the high band, Q and U are divided by it, and those channels still get full weight, so the SNR falls to a third. The source looks lopsided although its fractional polarisation is the same everywhere.\n", - "- `\"global\"`: the SNR follows the Stokes I brightness across the source, and every pixel has the same RMSF. Here that RMSF is wider than with `None`, because the template gives the high band less weight. The width follows the template, so on another band or field it can be narrower.\n", - "- `\"per_pixel\"`: the highest SNR, but the RMSF width follows the spectral index of each pixel, from about 31 rad/m$^2$ on the left to 49 on the right, and it scatters with the noise in each pixel's fit. A second moment or any other width measured against the RMSF would do the same.\n", + "- `None`: on the left source, whose spectrum is close to flat, the SNR matches the other options. On the right source the model falls to a few percent of its reference flux in the high band, Q and U are divided by it, and those channels still get full weight, so its SNR falls although both sources are equally bright and equally polarised.\n", + "- `\"global\"`: both sources come out at a similar SNR, and every pixel has the same RMSF. Here that RMSF is wider than with `None`, because the template gives the high band less weight. The width follows the template, so on another band or field it can be narrower.\n", + "- `\"per_pixel\"`: the highest SNR, but the RMSF width follows the spectral index of each pixel, so the two sources get different RMSFs, and it scatters with the noise in each pixel's fit. A second moment or any other width measured against the RMSF would do the same.\n", "\n", "#### Pixels below the Stokes I SNR cut\n", "\n", @@ -830,7 +894,7 @@ { "cell_type": "code", "execution_count": null, - "id": "27", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -851,16 +915,18 @@ " )\n", " ax.contour(fitted, levels=[0.5], colors=\"k\", linewidths=1.5)\n", " ax.set(title=f\"stokes_i_weighting={mode}: peak PI / true PI\")\n", - "fig.colorbar(im, ax=axes, shrink=0.8)\n", + "fig.colorbar(im, ax=axes, shrink=0.8, label=\"Peak PI / true PI\")\n", + "fig.suptitle(\n", + " \"Polarised intensity across the Stokes I SNR cut (black contour), same input\"\n", + ")\n", "\n", - "assert abs(below_cut[None] - 1) > 0.05\n", "for mode in (\"global\", \"per_pixel\"):\n", - " assert abs(below_cut[mode] - 1) < 0.03" + " assert abs(below_cut[mode] - 1) < abs(below_cut[None] - 1)" ] }, { "cell_type": "markdown", - "id": "28", + "id": "30", "metadata": {}, "source": [ "- `None`: below the cut the PI is biased low, so the map steps at the contour.\n", @@ -874,24 +940,26 @@ { "cell_type": "code", "execution_count": null, - "id": "29", + "id": "31", "metadata": {}, "outputs": [], "source": [ "summary = pl.DataFrame(\n", " {\n", " \"stokes_i_weighting\": [str(mode) for mode in modes],\n", - " \"PI SNR, left half\": [\n", + " \"PI SNR, left source\": [\n", " round(float(np.median(pi_snr[m][left])), 1) for m in modes\n", " ],\n", - " \"PI SNR, right half\": [\n", + " \"PI SNR, right source\": [\n", " round(float(np.median(pi_snr[m][right])), 1) for m in modes\n", " ],\n", - " \"PI error, right half\": [\n", + " \"PI error, right source\": [\n", " round(float(np.median(pi_error[m][right])), 3) for m in modes\n", " ],\n", - " \"RMSF FWHM, left\": [float(np.median(rmsf_fwhm[m][left])) for m in modes],\n", - " \"RMSF FWHM, right\": [float(np.median(rmsf_fwhm[m][right])) for m in modes],\n", + " \"RMSF FWHM, left source\": [float(np.median(rmsf_fwhm[m][left])) for m in modes],\n", + " \"RMSF FWHM, right source\": [\n", + " float(np.median(rmsf_fwhm[m][right])) for m in modes\n", + " ],\n", " \"PI / true PI below the cut\": [round(below_cut[m], 3) for m in modes],\n", " }\n", ")\n", @@ -900,7 +968,7 @@ }, { "cell_type": "markdown", - "id": "30", + "id": "32", "metadata": {}, "source": [ "- `\"global\"`: one RMSF and one effective frequency for the whole map, so PI, RM and Faraday-complexity maps compare pixel to pixel. A pixel whose spectrum is far from the template has less SNR than it could, a pixel whose model falls far below the template in some channels is only partly protected from the noise there, and below the SNR cut the fallback is only right where a pixel's spectrum follows the template. The RMSF is set by the template, so it differs from the $1/\\sigma_j^2$ RMSF. Whether it is wider or narrower depends on the band and the noise: here it is narrower, and on a band with gaps it can be wider. With `\"auto\"` the template comes from each field, so pass a number to give several fields one RMSF.\n", @@ -910,7 +978,7 @@ }, { "cell_type": "markdown", - "id": "31", + "id": "33", "metadata": {}, "source": [ "`rmsynth_3d_from_fits` takes the same Stokes I options as file paths (`stokes_i_file`, `stokes_i_error_file`, or `stokes_i_model_file`). See the [3D RM-CLEAN](rmclean_3d.ipynb) page to deconvolve the FDF cube." From 724feb21357de873cbd759e1b6886e4a2c78f685 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 12:30:33 +0000 Subject: [PATCH 6/8] Make the noise-based weight types a type, not a constant WeightType is now built from a NoiseWeightType literal, so the list lives in one place and callers test membership with get_args. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- rm_lite/tools_1d/rmsynth.py | 6 +++--- rm_lite/tools_3d/rmsynth.py | 8 ++++---- rm_lite/utils/synthesis.py | 17 +++++------------ 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/rm_lite/tools_1d/rmsynth.py b/rm_lite/tools_1d/rmsynth.py index 893e25c..b3a13f7 100644 --- a/rm_lite/tools_1d/rmsynth.py +++ b/rm_lite/tools_1d/rmsynth.py @@ -4,7 +4,7 @@ import time from dataclasses import replace -from typing import Literal, NamedTuple +from typing import Literal, NamedTuple, get_args import dask.array as da import numpy as np @@ -22,9 +22,9 @@ ) from rm_lite.utils.logging import logger from rm_lite.utils.synthesis import ( - NOISE_WEIGHT_TYPES, FDFOptions, LamSq0Mode, + NoiseWeightType, StokesData, StokesIWeighting, WeightType, @@ -294,7 +294,7 @@ def _run_rmsynth( weighting = ( fdf_options.stokes_i_weighting - if not ignore_stokes_i and fdf_options.weight_type in NOISE_WEIGHT_TYPES + if not ignore_stokes_i and fdf_options.weight_type in get_args(NoiseWeightType) else None ) noise_error = stokes_data.complex_pol_error diff --git a/rm_lite/tools_3d/rmsynth.py b/rm_lite/tools_3d/rmsynth.py index a23eec5..0550b1a 100644 --- a/rm_lite/tools_3d/rmsynth.py +++ b/rm_lite/tools_3d/rmsynth.py @@ -6,7 +6,7 @@ from dataclasses import replace from functools import partial from pathlib import Path -from typing import Any, Literal, NamedTuple, cast +from typing import Any, Literal, NamedTuple, cast, get_args import dask.array as da import numpy as np @@ -43,9 +43,9 @@ ) from rm_lite.utils.logging import logger, quiet_logs from rm_lite.utils.synthesis import ( - NOISE_WEIGHT_TYPES, FDFOptions, LamSq0Mode, + NoiseWeightType, RMSynthParams, StokesIWeighting, TheoreticalNoise, @@ -854,7 +854,7 @@ def rmsynth_3d( stokes_i_source = stokes_i_model if stokes_i_model is not None else stokes_i weighting = ( fdf_options.stokes_i_weighting - if stokes_i_source is not None and weight_type in NOISE_WEIGHT_TYPES + if stokes_i_source is not None and weight_type in get_args(NoiseWeightType) else None ) noise_weight_arr = weight_arr @@ -1412,7 +1412,7 @@ def rmsynth_3d_from_fits( # Noise-based types use 1/sigma^2 as their base (uniform_lsq/briggs then apply # the geometric lambda^2 factor); per-channel `uniform` deliberately ignores noise. - if weight_arr is None and weight_type in NOISE_WEIGHT_TYPES: + if weight_arr is None and weight_type in get_args(NoiseWeightType): weight_arr = get_weight_arr_from_fits( stokes_q_file, stokes_u_file, diff --git a/rm_lite/utils/synthesis.py b/rm_lite/utils/synthesis.py index 82e044e..5e172eb 100644 --- a/rm_lite/utils/synthesis.py +++ b/rm_lite/utils/synthesis.py @@ -140,9 +140,11 @@ class TheoreticalNoise(NamedTuple): """Amplitude units of an FDF: RM-synthesis output is per RMSF, a CLEAN component model is already integrated (each component is a flux).""" -WeightType: TypeAlias = Literal[ - "variance", "natural", "uniform", "uniform_lsq", "briggs" -] +NoiseWeightType: TypeAlias = Literal["variance", "natural", "uniform_lsq", "briggs"] +""" Weight types built on the Q/U noise, so the only ones the Stokes I +weighting can change. """ + +WeightType: TypeAlias = Literal[NoiseWeightType, "uniform"] """ RM-synthesis weighting: `variance`/`natural` (1/sigma^2, equivalent), `uniform` (equal per channel), `uniform_lsq` (equal per lambda^2 interval, narrows the RMSF), `briggs` (robust interpolation between natural and @@ -153,15 +155,6 @@ class TheoreticalNoise(NamedTuple): every pixel's weights by one field-wide spectrum, `per_pixel` by each pixel's own model. """ -NOISE_WEIGHT_TYPES: tuple[WeightType, ...] = ( - "variance", - "natural", - "uniform_lsq", - "briggs", -) -""" Weight types built on the Q/U noise, so the only ones the Stokes I -weighting can change. """ - @dataclass(frozen=True, kw_only=True, slots=True) class FDFOptions: From eec22bbc5f30f30d1cdb50c9e09e9c12e901c2ba Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 14:03:01 +0000 Subject: [PATCH 7/8] Blank pixels with no usable Stokes I model Below the Stokes I SNR cut, or where a fitted model is rejected, the 3D tool used to divide by a flat (or template) model, so those pixels kept a band-averaged PI next to fitted pixels reporting PI at the reference frequency. The maps stepped at the edge of the fitted region. Such pixels now get a NaN model, so their FDF and every map made from it are NaN, as in RM-Tools. Run without stokes_i for PI at every pixel. This removes fallback_alpha and fallback_model. Blank FDFs are now NaN in every moment map, not mom0 = 0, the 3D RM-CLEAN model is NaN where the clean FDF is, and the theoretical noise of a pixel with no weight is NaN without a divide warning. The 1D tool is unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- docs/examples/rmsynth_3d_stokes_i.ipynb | 131 ++++++---------- rm_lite/tools_1d/rmsynth.py | 5 +- rm_lite/tools_3d/rmclean.py | 2 + rm_lite/tools_3d/rmsynth.py | 27 ++-- rm_lite/utils/fitting.py | 52 ++---- rm_lite/utils/synthesis.py | 54 +++---- tests/test_stokes_i_weighting.py | 43 +++-- tests/test_tools_3d_dask.py | 16 +- tests/test_tools_3d_stokes_i.py | 200 ++++++++++-------------- 9 files changed, 206 insertions(+), 324 deletions(-) diff --git a/docs/examples/rmsynth_3d_stokes_i.ipynb b/docs/examples/rmsynth_3d_stokes_i.ipynb index 1686e3d..fc438a0 100644 --- a/docs/examples/rmsynth_3d_stokes_i.ipynb +++ b/docs/examples/rmsynth_3d_stokes_i.ipynb @@ -318,7 +318,7 @@ "source": [ "### Masking low-SNR pixels\n", "\n", - "Fitting a Stokes I model to a noise pixel is pointless and can inject fake spectral structure. `rmsynth_3d` fits only pixels whose frequency-averaged Stokes I SNR, `median(I) * sqrt(N) / median(noise)`, is at least `stokes_i_snr_cut` (default 5). Below the cut, or if a fit fails to converge, the pixel gets a flat model, so it has no spectral correction and its FDF is the plain Q/U FDF. This is with `weight_type=\"uniform\"`, as on this page. With a noise-based weight type the pixel gets a power law instead (see [Weights after the Stokes I division](#Weights-after-the-Stokes-I-division)). The FDF cube is never blanked, but `stokes_i_ref_flux_map` is NaN on those pixels, along with alpha, order and the term cubes: the flat model is a mean of noise, so reporting it as a flux would give a meaningless (often negative) fractional polarisation.\n", + "Fitting a Stokes I model to a noise pixel is pointless and can inject fake spectral structure. `rmsynth_3d` fits only pixels whose frequency-averaged Stokes I SNR, `median(I) * sqrt(N) / median(noise)`, is at least `stokes_i_snr_cut` (default 5). A pixel below the cut has no Stokes I model, so it is blanked: its model, its FDF and every map made from them are NaN, as in RM-Tools. Leaving it uncorrected instead would put a band-averaged PI next to fitted pixels that report the PI at the reference frequency, and the map would step at the edge of the fitted region. For polarisation at every pixel, run RM-synthesis again without `stokes_i` (see [Blanked Stokes I](#Blanked-Stokes-I)).\n", "\n", "Medians, not the mean and rms, so one bad channel cannot decide whether a pixel is fitted: an rms over the errors turns one inflated error channel into a near-zero SNR, and a mean over the fluxes lets one negative spike make it negative. Either quietly flattens a good pixel.\n", "\n", @@ -369,9 +369,9 @@ "source": [ "### Blanked Stokes I\n", "\n", - "The fallback has two steps, and neither blanks a pixel. As above, this is the fallback with `weight_type=\"uniform\"`. A pixel the fitter cannot use, below the cut, not converged, or with a model that runs away, gets a flat model at its mean Stokes I, which cancels out of the FDF. A pixel with no mean to divide by either, negative or with Stokes I blank there, gets no correction at all: RM-synthesis on the raw Q/U.\n", + "A pixel is blanked whenever it has no usable Stokes I model: below the cut, with no Stokes I error to measure the SNR against, with Stokes I blank, or with a model that cannot divide Q/U (see the next section). The second case happens at a mosaic edge, where the linmos weight goes to zero, so the Stokes I error rm-lite derives from it is infinite.\n", "\n", - "The first step covers more than the cut, which is why it does not need an error of its own. At a mosaic edge the linmos weight goes to zero, so the Stokes I error rm-lite derives from it is infinite and the cut has nothing to measure against, while Q and U still hold real data. A peak PI map has to have a value at those pixels, so the Stokes I pass hands back a usable model whatever it was given." + "Q and U can still hold real data at those pixels. Run `rmsynth_3d` without `stokes_i` to get an FDF at every pixel. That FDF has no Stokes I correction anywhere, so keep it as a separate set of maps rather than using it to fill the gaps in the corrected ones." ] }, { @@ -390,38 +390,37 @@ "\n", "q_da = da.from_array(stokes_q_obs, chunks=chunks)\n", "u_da = da.from_array(stokes_u_obs, chunks=chunks)\n", - "fallback_kwargs = {\n", + "blank_kwargs = {\n", " \"freq_arr_hz\": freq_hz,\n", " \"phi_max_radm2\": 150.0,\n", " \"d_phi_radm2\": 1.0,\n", " \"weight_type\": \"uniform\",\n", "}\n", - "result_fallback = rmsynth_3d(\n", + "result_blanked = rmsynth_3d(\n", " q_da,\n", " u_da,\n", " stokes_i=da.from_array(no_i, chunks=chunks),\n", " stokes_i_error=da.from_array(no_i_err, chunks=chunks),\n", " stokes_i_snr_cut=5.0,\n", - " **fallback_kwargs,\n", + " **blank_kwargs,\n", ")\n", - "uncorrected = rmsynth_3d(q_da, u_da, **fallback_kwargs)\n", + "uncorrected = rmsynth_3d(q_da, u_da, **blank_kwargs)\n", "\n", - "model_fallback, fdf_fallback, fdf_plain = da.compute(\n", - " result_fallback.stokes_i_model_cube,\n", - " result_fallback.fdf_dirty_cube,\n", + "model_blanked, fdf_blanked, fdf_plain = da.compute(\n", + " result_blanked.stokes_i_model_cube,\n", + " result_blanked.fdf_dirty_cube,\n", " uncorrected.fdf_dirty_cube,\n", ")\n", - "# Never blanked, and each fallback pixel gives back the plain Q/U FDF.\n", - "assert np.isfinite(model_fallback).all()\n", - "assert np.isfinite(fdf_fallback).all()\n", + "# Blanked in the corrected run, present in the uncorrected one.\n", "for i in (ib, ib + 1, ib + 2):\n", - " np.testing.assert_allclose(model_fallback[:, jb, i], model_fallback[0, jb, i])\n", - " np.testing.assert_allclose(fdf_fallback[:, jb, i], fdf_plain[:, jb, i], atol=1e-6)\n", - "# Blank or negative Stokes I: no correction at all, i.e. a model of exactly 1.\n", - "np.testing.assert_allclose(model_fallback[:, jb, ib + 1 : ib + 3], 1.0)\n", + " assert np.isnan(model_blanked[:, jb, i]).all()\n", + " assert np.isnan(fdf_blanked[:, jb, i]).all()\n", + " assert np.isfinite(fdf_plain[:, jb, i]).all()\n", + "# The pixels next to them keep their fit.\n", + "assert np.isfinite(fdf_blanked[:, jb, ib - 1]).all()\n", "print(\n", - " f\"peak PI at the blanked-error pixel: {np.abs(fdf_fallback[:, jb, ib]).max():.4f} \"\n", - " f\"(plain Q/U {np.abs(fdf_plain[:, jb, ib]).max():.4f})\"\n", + " f\"peak PI at the blanked-error pixel: corrected NaN, \"\n", + " f\"uncorrected {np.abs(fdf_plain[:, jb, ib]).max():.4f}\"\n", ")" ] }, @@ -434,8 +433,7 @@ "\n", "An SNR cut keeps the fitter off empty sky, but it does nothing about a pixel that is bright and badly behaved e.g. a deconvolution artefact, with a Stokes I spectrum that passes through zero inside the band. Stokes Q/U are divided by a model fitted to that spectrum, and the FDF is rescaled by the model's reference flux. So the correction amplifies them by `model(ref_freq) / min(model)`, with nothing bounding it.\n", "\n", - "`stokes_i_model_floor_sigma` (default 0.01) rejects a model that dips below that many sigma of the pixel's band-averaged Stokes I noise (the same noise `stokes_i_snr_cut` measures against).\n", - "\n" + "`stokes_i_model_floor_sigma` (default 0.01) rejects a model that dips below that many sigma of the pixel's band-averaged Stokes I noise (the same noise `stokes_i_snr_cut` measures against), and blanks the pixel, as below the cut." ] }, { @@ -466,8 +464,7 @@ "}\n", "floored = rmsynth_3d(q_da, u_da, **artefact_kwargs)\n", "unfloored = rmsynth_3d(q_da, u_da, stokes_i_model_floor_sigma=0.0, **artefact_kwargs)\n", - "# No Stokes I at all: what the flat fallback should reduce to, since dividing by\n", - "# the mean and rescaling by it cancel exactly.\n", + "# No Stokes I at all, to compare the unfloored run against.\n", "uncorrected = rmsynth_3d(\n", " q_da,\n", " u_da,\n", @@ -477,19 +474,22 @@ " weight_type=\"uniform\",\n", ")\n", "\n", - "peak, runaway, plain = (\n", + "floored_fdf = floored.fdf_dirty_cube[:, jb, ib].compute()\n", + "runaway, plain = (\n", " np.abs(result.fdf_dirty_cube[:, jb, ib].compute()).max()\n", - " for result in (floored, unfloored, uncorrected)\n", + " for result in (unfloored, uncorrected)\n", ")\n", - "print(f\"peak PI: uncorrected {plain:.4f}, floored {peak:.4f}, unfloored {runaway:.3e}\")\n", - "assert runaway > 1e3 * peak\n", - "assert np.isclose(peak, plain, rtol=1e-5)\n", + "print(f\"peak PI: uncorrected {plain:.4f}, unfloored {runaway:.3e}, floored NaN\")\n", + "assert np.isnan(floored_fdf).all()\n", + "assert runaway > 1e3 * plain\n", "\n", - "# What the floor lets through reads as lost SNR, not as signal.\n", + "# What the floor would let through reads as lost SNR, not as signal.\n", "noise = np.asarray(unfloored.theoretical_noise.fdf_error_noise)[jb, ib]\n", - "floored_noise = np.asarray(floored.theoretical_noise.fdf_error_noise)[jb, ib]\n", - "print(f\"peak SNR: floored {peak / floored_noise:.1f}, unfloored {runaway / noise:.1f}\")\n", - "assert runaway / noise < peak / floored_noise" + "plain_noise = np.asarray(uncorrected.theoretical_noise.fdf_error_noise)[jb, ib]\n", + "print(\n", + " f\"peak SNR: uncorrected {plain / plain_noise:.1f}, unfloored {runaway / noise:.1f}\"\n", + ")\n", + "assert runaway / noise < plain / plain_noise" ] }, { @@ -832,20 +832,22 @@ "outputs": [], "source": [ "modes = (None, \"global\", \"per_pixel\")\n", - "left, right = on_source & (xx < n_mock / 2), on_source & (xx >= n_mock / 2)\n", "pi_snr, rmsf_fwhm, pi_error = {}, {}, {}\n", "for mode in modes:\n", " result = synth_mock(stokes_i_weighting=mode, per_pixel_rmsf=True)\n", - " ratio, pi_snr[mode], _ = compare(result)\n", + " ratio, pi_snr[mode], fitted = compare(result)\n", " rmsf = np.abs(result.rmsf_cube.compute())\n", " d_phi = np.diff(result.phi_double_arr_radm2)[0]\n", " rmsf_fwhm[mode] = (rmsf >= 0.5 * rmsf.max(axis=0)).sum(axis=0) * d_phi\n", " pi_error[mode] = np.abs(ratio - 1)\n", + "# Faint source edges are below the Stokes I cut, so blank in every run.\n", + "shown = on_source & fitted\n", + "left, right = shown & (xx < n_mock / 2), shown & (xx >= n_mock / 2)\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)\n", "for ax, mode in zip(axes, modes, strict=True):\n", " im = ax.imshow(\n", - " np.where(on_source, pi_snr[mode], np.nan), origin=\"lower\", vmin=0, vmax=60\n", + " np.where(shown, pi_snr[mode], np.nan), origin=\"lower\", vmin=0, vmax=60\n", " )\n", " ax.set(title=f\"stokes_i_weighting={mode}\")\n", "fig.colorbar(im, ax=axes, shrink=0.8, label=\"True PI / reported noise\")\n", @@ -854,7 +856,7 @@ "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)\n", "for ax, mode in zip(axes, modes, strict=True):\n", " im = ax.imshow(\n", - " np.where(on_source, rmsf_fwhm[mode], np.nan),\n", + " np.where(shown, rmsf_fwhm[mode], np.nan),\n", " origin=\"lower\",\n", " vmin=30,\n", " vmax=48,\n", @@ -867,7 +869,7 @@ "for mode in (\"global\", \"per_pixel\"):\n", " assert np.median(pi_snr[mode][right]) > np.median(pi_snr[None][right])\n", "for mode in (None, \"global\"):\n", - " assert np.ptp(rmsf_fwhm[mode][on_source]) == 0\n", + " assert np.ptp(rmsf_fwhm[mode][shown]) == 0\n", "assert np.median(rmsf_fwhm[\"per_pixel\"][right]) > np.median(\n", " rmsf_fwhm[\"per_pixel\"][left]\n", ")" @@ -882,56 +884,14 @@ "\n", "- `None`: on the left source, whose spectrum is close to flat, the SNR matches the other options. On the right source the model falls to a few percent of its reference flux in the high band, Q and U are divided by it, and those channels still get full weight, so its SNR falls although both sources are equally bright and equally polarised.\n", "- `\"global\"`: both sources come out at a similar SNR, and every pixel has the same RMSF. Here that RMSF is wider than with `None`, because the template gives the high band less weight. The width follows the template, so on another band or field it can be narrower.\n", - "- `\"per_pixel\"`: the highest SNR, but the RMSF width follows the spectral index of each pixel, so the two sources get different RMSFs, and it scatters with the noise in each pixel's fit. A second moment or any other width measured against the RMSF would do the same.\n", - "\n", - "#### Pixels below the Stokes I SNR cut\n", - "\n", - "A pixel below `stokes_i_snr_cut` has no fit. With `stokes_i_weighting=None` it is not corrected, so its peak PI is a band-averaged PI, not the PI at the reference frequency. With either weighting it is divided by the field-wide power law instead, so it reports PI at the reference frequency, as far as its spectrum follows that power law.\n", - "\n", - "The default cut of 5 leaves almost no polarised signal below it, so here it is raised to 60 to put the edge of the fitted region (black contour) on the source. The colour is peak PI over the true PI, and white is correct." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "29", - "metadata": {}, - "outputs": [], - "source": [ - "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)\n", - "below_cut = {}\n", - "for ax, mode in zip(axes, modes, strict=True):\n", - " ratio, snr_cut_map, fitted = compare(\n", - " synth_mock(stokes_i_weighting=mode, stokes_i_snr_cut=60.0)\n", - " )\n", - " shown = snr_cut_map > 5\n", - " below_cut[mode] = float(np.median(ratio[shown & ~fitted]))\n", - " im = ax.imshow(\n", - " np.where(shown, ratio, np.nan),\n", - " origin=\"lower\",\n", - " vmin=0.7,\n", - " vmax=1.3,\n", - " cmap=\"RdBu_r\",\n", - " )\n", - " ax.contour(fitted, levels=[0.5], colors=\"k\", linewidths=1.5)\n", - " ax.set(title=f\"stokes_i_weighting={mode}: peak PI / true PI\")\n", - "fig.colorbar(im, ax=axes, shrink=0.8, label=\"Peak PI / true PI\")\n", - "fig.suptitle(\n", - " \"Polarised intensity across the Stokes I SNR cut (black contour), same input\"\n", - ")\n", - "\n", - "for mode in (\"global\", \"per_pixel\"):\n", - " assert abs(below_cut[mode] - 1) < abs(below_cut[None] - 1)" + "- `\"per_pixel\"`: the highest SNR, but the RMSF width follows the spectral index of each pixel, so the two sources get different RMSFs, and it scatters with the noise in each pixel's fit. A second moment or any other width measured against the RMSF would do the same." ] }, { "cell_type": "markdown", - "id": "30", + "id": "29", "metadata": {}, "source": [ - "- `None`: below the cut the PI is biased low, so the map steps at the contour.\n", - "- `\"global\"` and `\"per_pixel\"`: the step is gone on average. Where a pixel's spectrum differs from the field-wide power law the fallback is still off: on the left, where the spectrum is much flatter than the template, PI below the cut comes out high.\n", - "\n", "#### Choosing a mode\n", "\n", "The numbers from the maps above, over the source:" @@ -940,7 +900,7 @@ { "cell_type": "code", "execution_count": null, - "id": "31", + "id": "30", "metadata": {}, "outputs": [], "source": [ @@ -960,7 +920,6 @@ " \"RMSF FWHM, right source\": [\n", " float(np.median(rmsf_fwhm[m][right])) for m in modes\n", " ],\n", - " \"PI / true PI below the cut\": [round(below_cut[m], 3) for m in modes],\n", " }\n", ")\n", "summary" @@ -968,17 +927,17 @@ }, { "cell_type": "markdown", - "id": "32", + "id": "31", "metadata": {}, "source": [ - "- `\"global\"`: one RMSF and one effective frequency for the whole map, so PI, RM and Faraday-complexity maps compare pixel to pixel. A pixel whose spectrum is far from the template has less SNR than it could, a pixel whose model falls far below the template in some channels is only partly protected from the noise there, and below the SNR cut the fallback is only right where a pixel's spectrum follows the template. The RMSF is set by the template, so it differs from the $1/\\sigma_j^2$ RMSF. Whether it is wider or narrower depends on the band and the noise: here it is narrower, and on a band with gaps it can be wider. With `\"auto\"` the template comes from each field, so pass a number to give several fields one RMSF.\n", + "- `\"global\"`: one RMSF and one effective frequency for the whole map, so PI, RM and Faraday-complexity maps compare pixel to pixel. A pixel whose spectrum is far from the template has less SNR than it could, and a pixel whose model falls far below the template in some channels is only partly protected from the noise there. The RMSF is set by the template, so it differs from the $1/\\sigma_j^2$ RMSF. Whether it is wider or narrower depends on the band and the noise: here it is wider. With `\"auto\"` the template comes from each field, so pass a number to give several fields one RMSF.\n", "- `\"per_pixel\"`: the highest SNR in every pixel. The RMSF and the effective frequency change with the spectral index, and it cannot be used with `lam_sq_0_m2=\"per_pixel\"`.\n", "- `None`: the weights before this option existed. The FDF noise is unbounded where a model falls towards zero." ] }, { "cell_type": "markdown", - "id": "33", + "id": "32", "metadata": {}, "source": [ "`rmsynth_3d_from_fits` takes the same Stokes I options as file paths (`stokes_i_file`, `stokes_i_error_file`, or `stokes_i_model_file`). See the [3D RM-CLEAN](rmclean_3d.ipynb) page to deconvolve the FDF cube." diff --git a/rm_lite/tools_1d/rmsynth.py b/rm_lite/tools_1d/rmsynth.py index b3a13f7..b9929cb 100644 --- a/rm_lite/tools_1d/rmsynth.py +++ b/rm_lite/tools_1d/rmsynth.py @@ -3,7 +3,6 @@ from __future__ import annotations import time -from dataclasses import replace from typing import Literal, NamedTuple, get_args import dask.array as da @@ -240,8 +239,7 @@ def run_rmsynth( stokes_i_weight_alpha=stokes_i_weight_alpha, ) # snr_cut=None: the 1D fractional fit has only one spectrum, so an SNR cut - # would just silently drop fractional polarisation rather than fall back to - # a flat per-pixel model as it does in 3D. + # would just silently drop its fractional polarisation. fit_options = StokesIFitOptions( fit_order=fit_order, fit_function=fit_function, @@ -304,7 +302,6 @@ def _run_rmsynth( weight_error = noise_error / stokes_i_template( stokes_data.freq_arr_hz, weight_alpha ) - fit_options = replace(fit_options, fallback_alpha=weight_alpha) rmsynth_params = compute_rmsynth_params( freq_arr_hz=stokes_data.freq_arr_hz, diff --git a/rm_lite/tools_3d/rmclean.py b/rm_lite/tools_3d/rmclean.py index 9234ad6..e6e251d 100644 --- a/rm_lite/tools_3d/rmclean.py +++ b/rm_lite/tools_3d/rmclean.py @@ -599,6 +599,8 @@ def run_rmclean( multiscale_options=multiscale_options, log_level=log_level, ) + # A blank pixel has no components rather than zero of them. + model = da.where(da.isnan(clean), np.nan, model) maps = faraday_maps( clean, diff --git a/rm_lite/tools_3d/rmsynth.py b/rm_lite/tools_3d/rmsynth.py index 0550b1a..4bb1c4a 100644 --- a/rm_lite/tools_3d/rmsynth.py +++ b/rm_lite/tools_3d/rmsynth.py @@ -3,7 +3,6 @@ from __future__ import annotations import logging -from dataclasses import replace from functools import partial from pathlib import Path from typing import Any, Literal, NamedTuple, cast, get_args @@ -127,7 +126,7 @@ class RMSynth3DResults(NamedTuple): """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 + was not fitted (below the SNR cut or an unusable model). None unless a Stokes I cube was fitted (a supplied model has no fitted order).""" stokes_i_coeff_cube: da.Array | None = None """Fitted Stokes I model terms, shape (n_coeff, ny, nx) with @@ -168,8 +167,8 @@ class RMSynth3DResults(NamedTuple): stokes_i_weighting: StokesIWeighting | None = None """How the weights followed the Stokes I division; None if they did not.""" stokes_i_weight_alpha: float | None = None - """Spectral index of the field-wide power law the weights and the fallback - used; None if the weights did not follow the Stokes I division.""" + """Spectral index of the field-wide power law the weights used; None if the + weights did not follow the Stokes I division.""" def _compute_global_params( @@ -763,14 +762,16 @@ def rmsynth_3d( 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. + Stokes I SNR a pixel is not fitted, and its FDF is NaN rather than + left uncorrected beside corrected pixels. None fits every pixel. + Fit path only. Needs a Stokes I error to measure SNR against, so raises unless one of `stokes_i_error` / `estimate_stokes_i_noise` is given. Defaults to 5.0. stokes_i_model_floor_sigma (float, optional): Reject a fitted model dipping this many sigma below the pixel's band-averaged Stokes I - noise, falling back to a flat one. 0 disables. Defaults to 0.01. + noise, blanking the pixel as below the SNR cut. 0 disables. + Defaults to 0.01. stokes_i_robust_loss (RobustLoss, optional): Downweight channels far from the Stokes I model, so one bad channel cannot drag the fit. "cauchy" (default), "soft_l1" or "huber"; "linear" is plain least squares. @@ -784,10 +785,8 @@ def rmsynth_3d( field-wide power law for every pixel: one RMSF and one reference frequency for the whole map. "per_pixel" uses each pixel's own model: the most sensitive, but the RMSF then varies with the - spectral index. Either way pixels without a kept fit fall back to the - field-wide power law rather than a flat model. None keeps 1/sigma^2 - and the flat fallback. No effect without Stokes I or with - weight_type="uniform". Defaults to "global". + spectral index. None keeps 1/sigma^2. No effect without Stokes I or + with weight_type="uniform". Defaults to "global". stokes_i_weight_alpha (float | "auto", optional): Spectral index of the field-wide power law, or "auto" to fit it to the Stokes I (or model) averaged over the field. Pin it to share one RMSF between fields. @@ -872,7 +871,6 @@ def rmsynth_3d( weight_arr = _scale_channels( weight_arr, stokes_i_template(freq_arr_hz, weight_alpha) ** 2 ) - fit_options = replace(fit_options, fallback_alpha=weight_alpha) if weighting == "per_pixel" and not per_pixel_rmsf: logger.info( "stokes_i_weighting='per_pixel' weights each pixel by its own Stokes I " @@ -1045,11 +1043,6 @@ def rmsynth_3d( real_dtype(fdf_dtype) ) - if ref_flux_map is not None and order_map is not None: - # The rescale needs the flat fallback value, but on unfitted pixels - # it is a mean of noise, so don't report it as a flux. - ref_flux_map = da.where(da.isfinite(order_map), ref_flux_map, np.nan) - if per_pixel_ref: # Synthesised at the cube's reference, then moved to each pixel's own. # Exact, and a phase ramp rather than a transform per pixel. diff --git a/rm_lite/utils/fitting.py b/rm_lite/utils/fitting.py index 62dc779..5b9ae0f 100644 --- a/rm_lite/utils/fitting.py +++ b/rm_lite/utils/fitting.py @@ -60,9 +60,8 @@ class StokesIFitOptions: """How far, in sigma, before a channel is downweighted. Flat from 1 to 10""" model_floor_sigma: float = 0.01 """Reject a model dipping this many sigma below the band-averaged Stokes I - noise, falling back to a flat one (see `model_noise_floor`); 0 disables""" - fallback_alpha: float | None = None - """Spectral index of the model used where no fit is kept; None keeps a flat one""" + noise (see `model_noise_floor`): 3D blanks the pixel, 1D leaves Q/U + uncorrected. 0 disables""" def __post_init__(self) -> None: if self.fit_function not in ("log", "linear"): @@ -725,7 +724,7 @@ def fit_stokes_i_model( Masks channels that cannot be fitted, then returns None if too few remain (`< abs(options.fit_order) + 2`) or, when `options.snr_cut` is given, the - frequency-averaged SNR is below it, letting the caller impose a flat model. A + frequency-averaged SNR is below it, so the caller leaves the pixel out. A fit that cannot converge does not raise: `static_fit` falls back to a flat (mean) model. Right at the minimum channel count the fit is real but its AIC is inf (see `aic_lsq`), so a negative `fit_order` settles on a lower order @@ -941,12 +940,8 @@ class PixelFit(NamedTuple): """y pixel""" x: int """x pixel""" - i_spec: NDArray[np.float64] - """The pixel's Stokes I spectrum (unmasked), for the flat-model fallback.""" e_spec: NDArray[np.float64] """The pixel's error spectrum, for the noise floor.""" - good: NDArray[np.bool_] - """Finite-channel mask, for the flat-model fallback.""" fit: FitResult | None """The fit, or None if the pixel was skipped (too few channels / low SNR).""" @@ -971,7 +966,6 @@ def _iter_pixel_fits( for x in range(cx): i_spec = i_block[:, y, x] e_spec = _pixel_stokes_i_error(err_block, err_1d, n_freq, y, x) - good = np.isfinite(i_spec) fit = fit_stokes_i_model( freq_arr_hz=freq_arr_hz, ref_freq_hz=ref_freq_for_pixel(ref_freq_hz, y, x), @@ -979,7 +973,7 @@ def _iter_pixel_fits( stokes_i_error_arr=e_spec, options=fit_options, ) - yield PixelFit(y, x, i_spec, e_spec, good, fit) + yield PixelFit(y, x, e_spec, fit) class BlockPlanes(NamedTuple): @@ -1067,23 +1061,6 @@ def _write_error_planes( out[planes.alpha_error, y, x] = abs(a_high - a_low) -def fallback_model( - freq_arr_hz: NDArray[np.float64], - ref_freq_hz: float, - mean_flux: float, - alpha: float | None, -) -> NDArray[np.float64]: - """Model to divide by where no fit is kept. - - Flat at the mean without `alpha`, so Q/U get no correction. With it, a power - law of that index, 1 at the reference frequency: the amplitude cancels once - the FDF is rescaled, so it needs no flux. - """ - if alpha is None: - return np.full_like(freq_arr_hz, flat_model_value(mean_flux)) - return np.asarray((freq_arr_hz / ref_freq_hz) ** alpha, dtype=np.float64) - - RefFreqHz: TypeAlias = float | NDArray[np.float64] """A reference frequency in Hz: one for the whole image, or one per pixel.""" @@ -1121,12 +1098,9 @@ def _fit_stokes_i_block( optional (see `_pixel_stokes_i_error`). A pixel that was not fitted (too few finite channels or SNR below `fit_options.snr_cut`) or whose model is unusable (non-finite anywhere, or below `fit_options.model_floor_sigma` - times the pixel's band-averaged noise, see `model_is_usable`) falls back to - `fallback_model`: a power law of `fit_options.fallback_alpha` when that is - set, else a flat model at its mean Stokes I, and one whose mean cannot divide - either (negative, or no finite Stokes I at all) to a flat 1.0, leaving Q/U - uncorrected. Either way the model is finite, so no pixel of the FDF is - blanked by the Stokes I pass alone; alpha, order, terms and errors stay NaN. + times the pixel's band-averaged noise, see `model_is_usable`) keeps a NaN + model, so its Q/U are blanked rather than left uncorrected beside corrected + neighbours. """ i_block = arrays[0] # A per-pixel reference arrives as a (cy, cx) block after the data; @@ -1144,17 +1118,12 @@ def _fit_stokes_i_block( # The 1D fitter logs per fit and per failure. At cube scale that floods, so # quiet it to at least ERROR whatever the caller's log_level. with quiet_logs(max(log_level, logging.ERROR)): - for y, x, i_spec, e_spec, good, fit in _iter_pixel_fits( + for y, x, e_spec, fit in _iter_pixel_fits( i_block, err_block, err_1d, freq_arr_hz, ref_freq_hz, fit_options ): - pixel_ref_hz = ref_freq_for_pixel(ref_freq_hz, y, x) - mean_flux = float(np.mean(i_spec[good])) if good.any() else np.nan - fallback = fallback_model( - freq_arr_hz, pixel_ref_hz, mean_flux, fit_options.fallback_alpha - ) if fit is None: - out[planes.model, y, x] = fallback continue + pixel_ref_hz = ref_freq_for_pixel(ref_freq_hz, y, x) model = fit.stokes_i_model_func( freq_arr_hz / pixel_ref_hz, *np.asarray(fit.popt) ) @@ -1163,7 +1132,6 @@ def _fit_stokes_i_block( # flagged and Q/U did not, so all of it has to be usable. if not model_is_usable(model, noise_floor): n_rejected += 1 - out[planes.model, y, x] = fallback continue _write_model_planes( out, y, x, planes, fit, model, freq_arr_hz, pixel_ref_hz @@ -1182,7 +1150,7 @@ def _fit_stokes_i_block( if n_rejected: logger.warning( f"{n_rejected} of {cy * cx} pixels in this chunk fitted an unusable " - "Stokes I model and fell back to `fallback_model` (see " + "Stokes I model and were blanked (see " "`rm_lite.utils.fitting.model_is_usable`). Expect this on pixels with " "no real Stokes I signal, i.e. when `stokes_i_snr_cut` is None, and " "on artefacts, whose Stokes I passes through zero." diff --git a/rm_lite/utils/synthesis.py b/rm_lite/utils/synthesis.py index 5e172eb..0811c5c 100644 --- a/rm_lite/utils/synthesis.py +++ b/rm_lite/utils/synthesis.py @@ -40,7 +40,6 @@ FitResult, StokesIFitOptions, check_snr_cut_has_error, - fallback_model, fit_fdf, fit_rmsf, fit_sampled_peak, @@ -429,8 +428,9 @@ def calc_faraday_moments( Returns: FaradayMoments: mom0 (flux units), mom1 (rad/m^2), mom2 (dispersion, rad/m^2), the polarised intensity and angle at the - reference lambda^2, and an error for each. Spectra with no valid - amplitude have mom0 = 0 and mom1 = mom2 = NaN. + reference lambda^2, and an error for each. Spectra with no + amplitude above the threshold have mom0 = 0 and mom1 = mom2 = NaN; + blank (all-NaN) spectra are NaN throughout. Errors propagate `fdf_error` to first order, scaled by `fwhm_rmsf_radm2 / delta_phi` since the noise is correlated over one RMSF @@ -445,6 +445,7 @@ def calc_faraday_moments( raise ValueError(msg) phi_arr_radm2 = validate_phi_arr(complex_fdf_arr, phi_arr_radm2, axis) + blank = ~np.any(np.isfinite(complex_fdf_arr), axis=axis) if debias: if auto_threshold_sigma is not None: @@ -584,7 +585,7 @@ def calc_faraday_moments( np.where(pi_lam_sq_0 > 0, pi_lam_sq_0, np.nan), pi_lam_sq_0_error ) - return FaradayMoments( + moments = FaradayMoments( mom0=mom0, mom0_debias=mom0_debias, mom0_error=mom0_error, @@ -598,6 +599,8 @@ def calc_faraday_moments( pa_lam_sq_0=pa_lam_sq_0, pa_lam_sq_0_error=pa_lam_sq_0_error, ) + # A blank spectrum has no FDF at all, unlike one the threshold emptied. + return FaradayMoments(*(np.where(blank, np.nan, m) for m in moments)) def debias_polarised_intensity( @@ -1191,35 +1194,21 @@ def create_fractional_spectra( noise_floor = model_noise_floor( stokes_data.stokes_i_error_arr[no_nan_idx], fit_options.model_floor_sigma ) - kept_fit: FitResult | None = fit_result if not model_is_usable(model_good, noise_floor): logger.warning( "The fitted Stokes I model cannot safely divide Q/U (see " - "`rm_lite.utils.fitting.model_is_usable`); falling back to " - "`rm_lite.utils.fitting.fallback_model`." + "`rm_lite.utils.fitting.model_is_usable`); falling back to a flat " + "model at the mean Stokes I, so Q/U get no spectral correction." ) - kept_fit = None - if fit_options.fallback_alpha is None: - kept_fit = flat_fit_result( - flat_model_value(float(np.mean(i_good))), - len(np.asarray(fit_result.popt)) - 1, - fit_options.fit_function, - ) - - if kept_fit is None: - stokes_i_model_arr = fallback_model( - stokes_data.freq_arr_hz, - ref_freq_hz, - float(np.mean(i_good)), - fit_options.fallback_alpha, - ) - stokes_i_model_error = np.zeros_like(stokes_i_model_arr) - else: - stokes_i_model_arr, stokes_i_model_error = sample_model_error( - kept_fit, - stokes_data.freq_arr_hz / ref_freq_hz, - fit_options.n_error_samples, + fit_result = flat_fit_result( + flat_model_value(float(np.mean(i_good))), + len(np.asarray(fit_result.popt)) - 1, + fit_options.fit_function, ) + + stokes_i_model_arr, stokes_i_model_error = sample_model_error( + fit_result, stokes_data.freq_arr_hz / ref_freq_hz, fit_options.n_error_samples + ) # The fit runs in double precision; taking the model back down to the data's # keeps dividing by it from promoting the fractional spectra. model_dtype = real_dtype(stokes_data.complex_pol_arr.dtype) @@ -1255,7 +1244,7 @@ def create_fractional_spectra( return FractionalSpectra( stokes_data=fractional_stokes_data, - fit_result=kept_fit, + fit_result=fit_result, no_nan_idx=no_nan_idx, ) @@ -1303,10 +1292,15 @@ def compute_theoretical_noise( """ weight_arr = zero_nonfinite(weight_arr) complex_pol_error_flagged = zero_nonfinite(complex_pol_error) + # A blank pixel has no weight; divide it by 1 and blank it after, since a + # complex divide by NaN warns and `np.errstate` never reaches a dask worker. + weight_sum = np.sum(weight_arr, axis=0) + has_weight = weight_sum > 0 fdf_complex_noise = np.sqrt( np.nansum(weight_arr**2 * complex_pol_error_flagged**2, axis=0) - / (np.sum(weight_arr, axis=0)) ** 2 + / np.where(has_weight, weight_sum, 1.0) ** 2 ) + fdf_complex_noise = np.where(has_weight, fdf_complex_noise, np.nan + 1j * np.nan) fdf_error_noise = (fdf_complex_noise.real + fdf_complex_noise.imag) / 2 return TheoreticalNoise( diff --git a/tests/test_stokes_i_weighting.py b/tests/test_stokes_i_weighting.py index ef6a6bf..90935c8 100644 --- a/tests/test_stokes_i_weighting.py +++ b/tests/test_stokes_i_weighting.py @@ -211,15 +211,14 @@ def test_per_pixel_mode_gives_each_pixel_its_own_rmsf() -> None: assert widths[1] > widths[0] * 1.04 -@pytest.mark.parametrize("mode", [None, "global"]) -def test_fallback_keeps_pi_continuous_across_the_snr_cut( +@pytest.mark.parametrize("mode", MODES) +def test_pixels_below_the_snr_cut_are_blank_in_every_map( mode: StokesIWeighting | None, ) -> None: - """Pixels below the cut report the same fractional polarisation as those above.""" - # Wide and steep, so a flat fallback is well off the value at the reference. + """Only corrected pixels reach the maps, and they report the right fraction.""" freq_arr_hz = WIDE_FREQ_ARR_HZ ny, nx = 1, 12 - frac_pol, sigma_i, alpha = 0.1, 1e-3, -4.0 + frac_pol, sigma_i, alpha = 0.1, 1e-3, -2.0 amplitude = np.geomspace(0.05, 1.0, nx)[np.newaxis, :] stokes_i = (freq_arr_hz / 1.2e9)[:, None, None] ** alpha * amplitude[None] stokes_i = stokes_i / np.median(stokes_i[:, 0, -1]) * 0.02 @@ -239,7 +238,7 @@ def test_fallback_keeps_pi_continuous_across_the_snr_cut( d_phi_radm2=0.5, phi_max_radm2=100.0, stokes_i_weighting=mode, - stokes_i_weight_alpha=alpha, + per_pixel_rmsf=True, ) fitted = np.isfinite(np.asarray(synth.stokes_i_model_order_map))[0] assert fitted.any(), "no pixel is above the cut" @@ -250,12 +249,32 @@ def test_fallback_keeps_pi_continuous_across_the_snr_cut( [np.interp(ref_hz, freq_arr_hz, stokes_i[:, 0, i]) for i in range(nx)] ) peak_pi = np.abs(synth.fdf_dirty_cube.compute()).max(axis=0)[0] - frac = peak_pi / i_at_ref - np.testing.assert_allclose(frac[fitted], frac_pol, rtol=1e-3) - if mode is None: - assert np.all(np.abs(frac[~fitted] / frac_pol - 1) > 0.1) - else: - np.testing.assert_allclose(frac[~fitted], frac_pol, rtol=1e-3) + np.testing.assert_allclose(peak_pi[fitted] / i_at_ref[fitted], frac_pol, rtol=1e-3) + + clean = rmclean3d_mod.run_rmclean_from_synth(synth, max_iter=10) + maps = { + f"{label}.{name}": value + for label, cube in ( + ("dirty", synth.fdf_dirty_cube), + ("clean", clean.clean_fdf_cube), + ("model", clean.model_fdf_cube), + ) + for name, value in rmclean3d_mod.faraday_maps( + cube, + phi_arr_radm2=synth.phi_arr_radm2, + fwhm_rmsf_radm2=synth.fwhm_rmsf_radm2, + fdf_units="integrated" if label == "model" else "per_rmsf", + lam_sq_0_m2=synth.lam_sq_0_m2, + lambda_sq_arr_m2=synth.lambda_sq_arr_m2, + fdf_noise=synth.theoretical_noise.fdf_error_noise, + moment_threshold=None, + ).items() + } + (computed,) = compute(maps) + finite = [ + k for k, v in computed.items() if np.isfinite(np.asarray(v)[0, ~fitted]).any() + ] + assert not finite def test_field_spectral_index_recovers_alpha() -> None: diff --git a/tests/test_tools_3d_dask.py b/tests/test_tools_3d_dask.py index c608247..9986984 100644 --- a/tests/test_tools_3d_dask.py +++ b/tests/test_tools_3d_dask.py @@ -1509,7 +1509,7 @@ def test_rmsynth_3d_from_fits_on_a_dummy_stokes_axis( def test_weight_as_stokes_i_error_is_quiet_outside_the_beam(tmp_path): - """A zeroed linmos weight costs the fit, not the FDF, and warns nothing.""" + """A zeroed linmos weight blanks the pixels outside the beam, and warns nothing.""" rng = np.random.default_rng(2025) freq_arr_hz = (np.arange(744, 1032, 6) * 1e6).astype(np.float64) ny = nx = 12 @@ -1576,19 +1576,11 @@ def test_weight_as_stokes_i_error_is_quiet_outside_the_beam(tmp_path): assert blanked == [] # An infinite error leaves the fit nothing to weight or cut against, so - # those pixels report no alpha. Q and U still have data, so the FDF keeps - # them: uncorrected, which is what the flat fallback reduces to. + # those pixels report no alpha and, with no Stokes I correction, no FDF. assert np.all(np.isnan(computed[~inside])) np.testing.assert_allclose(computed[inside], alpha, atol=1e-5) - assert np.all(np.isfinite(fdf_cube)) - plain = rmsynth_3d_from_fits( - tmp_path / "q.fits", - tmp_path / "u.fits", - d_phi_radm2=D_PHI_RADM2, - phi_max_radm2=150.0, - ).fdf_dirty_cube.compute() - # rtol: the flat model divides and rescales in the cube's float32. - np.testing.assert_allclose(fdf_cube[:, ~inside], plain[:, ~inside], rtol=1e-5) + assert np.all(np.isnan(fdf_cube[:, ~inside])) + assert np.all(np.isfinite(fdf_cube[:, inside])) def test_channel_noise_from_channel_chunks_matches_whole_cube( diff --git a/tests/test_tools_3d_stokes_i.py b/tests/test_tools_3d_stokes_i.py index a28b949..fa5ac8f 100644 --- a/tests/test_tools_3d_stokes_i.py +++ b/tests/test_tools_3d_stokes_i.py @@ -23,6 +23,7 @@ pad_coefficients, polynomial, power_law, + stokes_i_snr, ) from rm_lite.utils.synthesis import calc_faraday_peaks, freq_to_lambda2 from scipy import optimize @@ -391,8 +392,23 @@ def cube_with_faint_pixels(faint: list[tuple[int, int]], noise: float = 1e-3): return stokes_q, stokes_u, stokes_i_obs, err, freq_arr_hz -def test_stokes_i_snr_cut_falls_back_to_flat_model(chunked: Callable[..., da.Array]): - """Pixels below the SNR cut fall back to a flat model, with alpha masked.""" +def assert_only_blanked( + result: RMSynth3DResults, blanked: list[tuple[int, int]] +) -> None: + """The Stokes I model and FDF are NaN on `blanked` and finite elsewhere.""" + model = require(result.stokes_i_model_cube).compute() + fdf = result.fdf_dirty_cube.compute() + mask = np.zeros(model.shape[1:], dtype=bool) + for j, i in blanked: + mask[j, i] = True + assert np.isnan(model[:, mask]).all() + assert np.isnan(fdf[:, mask]).all() + assert np.isfinite(model[:, ~mask]).all() + assert np.isfinite(fdf[:, ~mask]).all() + + +def test_stokes_i_snr_cut_blanks_faint_pixels(chunked: Callable[..., da.Array]): + """Pixels below the SNR cut are blanked, not left uncorrected.""" faint = [(0, 0), (2, 3)] q, u, i_obs, err, freq = cube_with_faint_pixels(faint) common: dict[str, Any] = { @@ -409,23 +425,18 @@ def test_stokes_i_snr_cut_falls_back_to_flat_model(chunked: Callable[..., da.Arr stokes_i_snr_cut=5.0, **common, ) - # The uncorrected reference FDF (no Stokes I at all). - raw = rmsynth_3d(chunked(q), chunked(u), freq, **common) - model = require(result.stokes_i_model_cube).compute() alpha = require(result.stokes_i_alpha_map).compute() order = require(result.stokes_i_model_order_map).compute() fdf = result.fdf_dirty_cube.compute() - fdf_raw = raw.fdf_dirty_cube.compute() faint_mask = np.zeros(model.shape[1:], dtype=bool) for j, i in faint: faint_mask[j, i] = True - # Flat model: finite and constant across frequency. - assert np.isfinite(model[:, j, i]).all() - np.testing.assert_allclose(model[:, j, i], model[0, j, i], rtol=1e-10) - # No spectral correction -> FDF matches the uncorrected Q/U FDF. - np.testing.assert_allclose(fdf[:, j, i], fdf_raw[:, j, i], atol=1e-8) + assert np.isnan(model[:, faint_mask]).all() + assert np.isnan(fdf[:, faint_mask]).all() + assert np.isfinite(model[:, ~faint_mask]).all() + assert np.isfinite(fdf[:, ~faint_mask]).all() # Masked (unfitted) pixels have NaN alpha/order; fitted pixels are finite. assert np.isnan(alpha[faint_mask]).all() @@ -436,12 +447,10 @@ def test_stokes_i_snr_cut_falls_back_to_flat_model(chunked: Callable[..., da.Arr assert np.isfinite(fitted_order).all() assert (fitted_order >= 0).all() np.testing.assert_array_equal(fitted_order, np.round(fitted_order)) - # The model cube is never blanked, even for masked pixels. - assert np.isfinite(model).all() -def test_blank_stokes_i_error_keeps_the_fdf(chunked: Callable[..., da.Array]): - """A pixel with no usable Stokes I error falls back, it does not blank. +def test_blank_stokes_i_error_blanks_the_fdf(chunked: Callable[..., da.Array]): + """A pixel with no usable Stokes I error has no SNR, so it is blanked. A linmos weight cube is zero outside the primary beam, so the error rm-lite derives from it is inf there while Q/U still hold real data. @@ -463,22 +472,11 @@ def test_blank_stokes_i_error_keeps_the_fdf(chunked: Callable[..., da.Array]): stokes_i_snr_cut=5.0, **common, ) - raw = rmsynth_3d(chunked(q), chunked(u), freq, **common) + assert_only_blanked(result, [(1, 1), (2, 0)]) - model = require(result.stokes_i_model_cube).compute() - fdf = result.fdf_dirty_cube.compute() - fdf_raw = raw.fdf_dirty_cube.compute() - assert np.isfinite(model).all() - assert np.isfinite(fdf).all() - for j, i in ((1, 1), (2, 0)): - # Fallback #1: flat at the pixel's mean Stokes I, which cancels out of - # the FDF, leaving the uncorrected Q/U one. - np.testing.assert_allclose(model[:, j, i], model[0, j, i], rtol=1e-10) - np.testing.assert_allclose(fdf[:, j, i], fdf_raw[:, j, i], atol=1e-8) - -def test_blank_stokes_i_keeps_the_fdf(chunked: Callable[..., da.Array]): - """A pixel with no Stokes I at all, or a negative mean, still gets an FDF.""" +def test_blank_stokes_i_blanks_the_fdf(chunked: Callable[..., da.Array]): + """A pixel with no Stokes I at all, or a negative one, is blanked.""" q, u, i_obs, err, freq = cube_with_faint_pixels([]) i_obs[:, 0, 2] = np.nan i_obs[:, 1, 3] *= -1.0 @@ -496,17 +494,7 @@ def test_blank_stokes_i_keeps_the_fdf(chunked: Callable[..., da.Array]): stokes_i_snr_cut=5.0, **common, ) - raw = rmsynth_3d(chunked(q), chunked(u), freq, **common) - - model = require(result.stokes_i_model_cube).compute() - fdf = result.fdf_dirty_cube.compute() - fdf_raw = raw.fdf_dirty_cube.compute() - assert np.isfinite(model).all() - assert np.isfinite(fdf).all() - for j, i in ((0, 2), (1, 3)): - # Fallback #2: no mean to divide by, so no Stokes I correction at all. - np.testing.assert_allclose(model[:, j, i], 1.0, rtol=1e-10) - np.testing.assert_allclose(fdf[:, j, i], fdf_raw[:, j, i], atol=1e-8) + assert_only_blanked(result, [(0, 2), (1, 3)]) def test_model_must_be_usable_outside_the_fitted_channels( @@ -534,12 +522,8 @@ def test_model_must_be_usable_outside_the_fitted_channels( weight_type="uniform", phi_max_radm2=200.0, ) - model = require(result.stokes_i_model_cube).compute()[:, 1, 2] - assert model_is_usable(model, model_noise_floor(err[:, 1, 2], 0.01)) - # Rejected, so the flat fallback: constant, and alpha/order stay NaN. - np.testing.assert_allclose(model, model[0], rtol=1e-10) + assert_only_blanked(result, [(1, 2)]) assert np.isnan(require(result.stokes_i_model_order_map).compute()[1, 2]) - assert np.isfinite(result.fdf_dirty_cube.compute()).all() def runaway_stokes_i_pixel() -> tuple[ @@ -593,8 +577,8 @@ def test_runaway_fit_is_rejected_before_it_overflows_float32() -> None: assert not model_is_usable(model, floor) -def test_runaway_fit_keeps_the_pixel(chunked: Callable[..., da.Array]) -> None: - """The runaway pixel falls back instead of blanking, end to end.""" +def test_runaway_fit_blanks_only_its_pixel(chunked: Callable[..., da.Array]) -> None: + """The runaway pixel is blanked, end to end, and its neighbour is kept.""" freq_arr_hz, i_spec, err, _ = runaway_stokes_i_pixel() n_freq = freq_arr_hz.size lambda_sq_arr_m2 = freq_to_lambda2(freq_arr_hz) @@ -628,37 +612,34 @@ def test_runaway_fit_keeps_the_pixel(chunked: Callable[..., da.Array]) -> None: phi_max_radm2=200.0, weight_type="uniform", ) - model = require(result.stokes_i_model_cube).compute() + assert_only_blanked(result, [(0, 1)]) ref_flux = require(result.stokes_i_ref_flux_map).compute() + assert np.isnan(ref_flux[0, 1]) fdf = result.fdf_dirty_cube.compute() - # The model stays inside float32, so the reference flux is a number and the - # FDF is not multiplied by NaN. - assert np.isfinite(model).all() - with np.errstate(over="ignore"): - assert np.isfinite(np.asarray(model, dtype=np.float32)).all() - assert np.isfinite(fdf).all() peaks = calc_faraday_peaks(fdf, result.phi_arr_radm2, result.fwhm_rmsf_radm2) - assert np.isfinite(peaks.peak_pi).all() - # Rejected, so the flat fallback: constant, and no reported reference flux. - np.testing.assert_allclose(model[:, 0, 1], model[0, 0, 1], rtol=1e-10) - assert np.isnan(ref_flux[0, 1]) + assert np.isfinite(peaks.peak_pi[0, 0]) def test_stokes_i_snr_cut_zero_fits_all_pixels(chunked: Callable[..., da.Array]): """A cut of 0 disables the SNR gate, so even faint pixels are fitted.""" - faint = [(0, 0)] - q, u, i_obs, err, freq = cube_with_faint_pixels(faint) - result = rmsynth_3d( - chunked(q), - chunked(u), - freq, - stokes_i=chunked(i_obs), - stokes_i_error=chunked(err), - stokes_i_snr_cut=0.0, - d_phi_radm2=D_PHI_RADM2, - weight_type="uniform", - ) - assert np.isfinite(require(result.stokes_i_model_cube).compute()).all() + q, u, i_obs, err, freq = cube_with_faint_pixels([(0, 0)], noise=6e-4) + assert stokes_i_snr(i_obs[:, 0, 0], err[:, 0, 0]) < 5.0 + + def order_at_faint_pixel(cut: float) -> float: + result = rmsynth_3d( + chunked(q), + chunked(u), + freq, + stokes_i=chunked(i_obs), + stokes_i_error=chunked(err), + stokes_i_snr_cut=cut, + d_phi_radm2=D_PHI_RADM2, + weight_type="uniform", + ) + return float(require(result.stokes_i_model_order_map).compute()[0, 0]) + + assert np.isnan(order_at_faint_pixel(5.0)) + assert np.isfinite(order_at_faint_pixel(0.0)) def test_fit_stokes_i_model_flat_fallback_on_failure(monkeypatch): @@ -890,13 +871,11 @@ def test_model_is_usable_rejects_models_that_cannot_divide( assert not model_is_usable(model) -def test_unusable_model_takes_the_flat_fallback() -> None: - """A rejected pixel gets a flat model at its mean I, so no correction.""" +def test_unusable_model_blanks_the_pixel() -> None: + """A rejected pixel is blanked rather than left uncorrected.""" rng = np.random.default_rng(20260823) n_freq, ny, nx = 125, 4, 6 freq = np.arange(800e6, 1800e6, 8e6)[:n_freq] - # Just enough offset that every pixel mean stays positive, so the flat - # fallback is well defined, while the fits themselves still run away. stokes_i = rng.normal(0, 0.05, (n_freq, ny, nx)) + 0.01 q = rng.normal(0, 0.02, (n_freq, ny, nx)) u = rng.normal(0, 0.02, (n_freq, ny, nx)) @@ -911,16 +890,11 @@ def test_unusable_model_takes_the_flat_fallback() -> None: d_phi_radm2=D_PHI_RADM2, weight_type="uniform", ) - model = np.asarray(require(result.stokes_i_model_cube).compute()) - fdf = np.asarray(result.fdf_dirty_cube.compute()) - assert np.all(np.isfinite(model)) - assert np.all(model > 0) - assert np.all(np.isfinite(fdf)) - - mean_i = stokes_i.mean(axis=0) - flat = np.isclose(model.max(axis=0), model.min(axis=0)) - assert flat.any(), "no pixel took the fallback, so it is not under test" - assert np.allclose(model[:, flat], mean_i[flat]) + order = np.asarray(require(result.stokes_i_model_order_map).compute()) + rejected = np.isnan(order) + assert rejected.any(), "no fit was rejected, so it is not under test" + assert rejected.sum() < rejected.size, "no fit was kept to compare against" + assert_only_blanked(result, list(zip(*np.nonzero(rejected), strict=True))) def artefact_cube( @@ -956,14 +930,10 @@ def artefact_synth(feature_width: float, **kwargs: Any) -> RMSynth3DResults: ) -def test_artefact_spectrum_does_not_blow_up_the_fdf() -> None: +def test_artefact_spectrum_is_blanked_not_blown_up() -> None: """The bug this floor is for: a runaway fit took a 2.5 mJy signal to 1e30.""" synth = artefact_synth(0.006, stokes_i_weighting=None) - fdf = np.asarray(synth.fdf_dirty_cube.compute()) - model = np.asarray(require(synth.stokes_i_model_cube).compute()) - - assert np.isclose(model.max(), model.min()), "fit was kept, so no floor applied" - assert np.abs(fdf).max() == pytest.approx(0.0025, rel=0.1) + assert np.isnan(np.asarray(synth.fdf_dirty_cube.compute())).all() def test_unfloored_artefact_fit_carries_its_amplification_into_the_noise() -> None: @@ -973,11 +943,19 @@ def test_unfloored_artefact_fit_carries_its_amplification_into_the_noise() -> No noise = np.asarray(synth.theoretical_noise.fdf_error_noise).item() assert peak / 0.0025 > 1e3, "not the runaway regime this is testing" - floored = artefact_synth(0.006) - floored_peak = np.abs(np.asarray(floored.fdf_dirty_cube.compute())).max() - floored_noise = np.asarray(floored.theoretical_noise.fdf_error_noise).item() + freq, stokes_q, stokes_u, _ = artefact_cube(0.006) + raw = rmsynth_3d( + da.from_array(stokes_q, chunks=stokes_q.shape), + da.from_array(stokes_u, chunks=stokes_u.shape), + freq, + weight_arr=np.full(freq.size, 1.0 / 1e-4**2), + d_phi_radm2=D_PHI_RADM2, + phi_max_radm2=1000.0, + ) + raw_peak = np.abs(np.asarray(raw.fdf_dirty_cube.compute())).max() + raw_noise = float(cast("float", raw.theoretical_noise.fdf_error_noise)) # The noise outruns the peak, which averages the amplified channels down. - assert peak / noise < floored_peak / floored_noise + assert peak / noise < raw_peak / raw_noise def test_flat_stokes_i_leaves_the_theoretical_noise_alone() -> None: @@ -1009,8 +987,8 @@ def test_flat_stokes_i_leaves_the_theoretical_noise_alone() -> None: ) -def test_pixel_with_no_finite_channels_gets_no_correction() -> None: - """A fully blanked Stokes I pixel goes uncorrected, and leaves its neighbours alone.""" +def test_pixel_with_no_finite_channels_is_blanked() -> None: + """A fully blanked Stokes I pixel is blanked, and leaves its neighbours alone.""" rng = np.random.default_rng(20260824) n_freq, ny, nx = 64, 2, 3 freq = np.linspace(800e6, 1800e6, n_freq) @@ -1030,15 +1008,9 @@ def test_pixel_with_no_finite_channels_gets_no_correction() -> None: d_phi_radm2=D_PHI_RADM2, weight_type="uniform", ) - model = np.asarray(require(result.stokes_i_model_cube).compute()) + assert_only_blanked(result, [(0, 0)]) alpha = np.asarray(require(result.stokes_i_alpha_map).compute()) - - # No mean to divide by, so a model of 1: Q/U pass through untouched. - np.testing.assert_allclose(model[:, 0, 0], 1.0) assert np.isnan(alpha[0, 0]) - # Every other pixel is untouched by its blank neighbour. - assert np.isfinite(model[:, 0, 1:]).all() - assert np.isfinite(model[:, 1, :]).all() assert np.isfinite(alpha[0, 1:]).all() assert np.isfinite(alpha[1, :]).all() @@ -1451,7 +1423,7 @@ def test_flat_model_value(mean_flux: float, expected: float) -> None: def test_unfitted_pixels_report_no_reference_flux( chunked: Callable[..., da.Array], ) -> None: - """A pixel below the SNR cut gets no flux to divide by.""" + """A pixel below the SNR cut reports no flux.""" faint = [(0, 0), (2, 3)] q, u, i_obs, err, freq = cube_with_faint_pixels(faint) common: dict[str, Any] = { @@ -1468,11 +1440,7 @@ def test_unfitted_pixels_report_no_reference_flux( stokes_i_snr_cut=5.0, **common, ) - raw = rmsynth_3d(chunked(q), chunked(u), freq, **common) - ref_flux = np.asarray(require(result.stokes_i_ref_flux_map).compute()) - fdf = result.fdf_dirty_cube.compute() - fdf_raw = raw.fdf_dirty_cube.compute() faint_mask = np.zeros(ref_flux.shape, dtype=bool) for j, i in faint: @@ -1480,14 +1448,9 @@ def test_unfitted_pixels_report_no_reference_flux( assert np.isnan(ref_flux[faint_mask]).all() assert np.isfinite(ref_flux[~faint_mask]).all() assert (ref_flux[~faint_mask] > 0).all() - # The rescale still used the flat value, so the FDF is uncorrected, not NaN. - for j, i in faint: - np.testing.assert_allclose(fdf[:, j, i], fdf_raw[:, j, i], atol=1e-8) -def test_negative_stokes_i_reports_no_flux_and_a_finite_fdf( - chunked: Callable[..., da.Array], -) -> None: +def test_negative_stokes_i_is_blanked(chunked: Callable[..., da.Array]) -> None: """A negative pixel must not export its mean as a flux.""" cube = make_stokes_i_cube(ny=2, nx=2, alpha=-0.8, noise=0.01) stokes_i = cube.stokes_i.copy() @@ -1503,11 +1466,6 @@ def test_negative_stokes_i_reports_no_flux_and_a_finite_fdf( weight_type="uniform", ) ref_flux = np.asarray(require(result.stokes_i_ref_flux_map).compute()) - model = np.asarray(require(result.stokes_i_model_cube).compute()) - fdf = result.fdf_dirty_cube.compute() - assert np.isnan(ref_flux[0, 0]) assert np.isfinite(ref_flux[1:]).all() - # Nothing divides by zero or a negative, so the FDF stays finite everywhere. - assert (model > 0).all() - assert np.isfinite(fdf).all() + assert_only_blanked(result, [(0, 0)]) From dfabc2577d0671e934d162f2262ac18221376c71 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 28 Sep 2026 14:19:45 +0000 Subject: [PATCH 8/8] Leave the RM-CLEAN model at zero on blank pixels A blank pixel has no CLEAN components, and zero says that. The dirty and clean maps stay NaN there. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F6ajgfHAviQKM4UXZ3gNi7 --- rm_lite/tools_3d/rmclean.py | 2 -- tests/test_stokes_i_weighting.py | 6 +++++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/rm_lite/tools_3d/rmclean.py b/rm_lite/tools_3d/rmclean.py index e6e251d..9234ad6 100644 --- a/rm_lite/tools_3d/rmclean.py +++ b/rm_lite/tools_3d/rmclean.py @@ -599,8 +599,6 @@ def run_rmclean( multiscale_options=multiscale_options, log_level=log_level, ) - # A blank pixel has no components rather than zero of them. - model = da.where(da.isnan(clean), np.nan, model) maps = faraday_maps( clean, diff --git a/tests/test_stokes_i_weighting.py b/tests/test_stokes_i_weighting.py index 90935c8..140b0e0 100644 --- a/tests/test_stokes_i_weighting.py +++ b/tests/test_stokes_i_weighting.py @@ -271,10 +271,14 @@ def test_pixels_below_the_snr_cut_are_blank_in_every_map( ).items() } (computed,) = compute(maps) + # The CLEAN model of a blank pixel has no components, so it is zero, not NaN. finite = [ - k for k, v in computed.items() if np.isfinite(np.asarray(v)[0, ~fitted]).any() + k + for k, v in computed.items() + if not k.startswith("model.") and np.isfinite(np.asarray(v)[0, ~fitted]).any() ] assert not finite + assert not np.asarray(clean.model_fdf_cube.compute())[:, 0, ~fitted].any() def test_field_spectral_index_recovers_alpha() -> None: