diff --git a/docs/examples/rmsynth_3d_stokes_i.ipynb b/docs/examples/rmsynth_3d_stokes_i.ipynb index c69ce33..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. 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. 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" ] }, { @@ -678,6 +678,267 @@ "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", + "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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "import polars as pl\n", + "from rm_lite.utils.synthesis import lambda2_to_freq\n", + "\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", + "# 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", + "\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_mock(**kwargs):\n", + " return rmsynth_3d(\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=1.0,\n", + " **kwargs,\n", + " )\n", + "\n", + "\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", + "on_source = i_1ghz > 5 * sigma_mock / np.sqrt(mock_hz.size)\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", + "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], 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(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", + "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(shown, rmsf_fwhm[mode], np.nan),\n", + " origin=\"lower\",\n", + " vmin=30,\n", + " vmax=48,\n", + " cmap=\"magma\",\n", + " )\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]) > np.median(pi_snr[None][right])\n", + "for mode in (None, \"global\"):\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", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "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 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." + ] + }, + { + "cell_type": "markdown", + "id": "29", + "metadata": {}, + "source": [ + "#### Choosing a mode\n", + "\n", + "The numbers from the maps above, over the source:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "summary = pl.DataFrame(\n", + " {\n", + " \"stokes_i_weighting\": [str(mode) for mode in modes],\n", + " \"PI SNR, left source\": [\n", + " round(float(np.median(pi_snr[m][left])), 1) for m in modes\n", + " ],\n", + " \"PI SNR, right source\": [\n", + " round(float(np.median(pi_snr[m][right])), 1) for m in modes\n", + " ],\n", + " \"PI error, right source\": [\n", + " round(float(np.median(pi_error[m][right])), 3) for m in modes\n", + " ],\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", + " }\n", + ")\n", + "summary" + ] + }, + { + "cell_type": "markdown", + "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, 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": "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 073ed26..b9929cb 100644 --- a/rm_lite/tools_1d/rmsynth.py +++ b/rm_lite/tools_1d/rmsynth.py @@ -3,8 +3,9 @@ from __future__ import annotations import time -from typing import Literal, NamedTuple +from typing import Literal, NamedTuple, get_args +import dask.array as da import numpy as np import polars as pl from numpy.typing import NDArray @@ -16,13 +17,17 @@ StokesIFitOptions, coefficient_errors, coefficient_names, + field_spectral_index, ) from rm_lite.utils.logging import logger from rm_lite.utils.synthesis import ( FDFOptions, LamSq0Mode, + NoiseWeightType, StokesData, + StokesIWeighting, WeightType, + apply_weight_type, compute_rmsynth_params, compute_theoretical_noise, create_fractional_spectra, @@ -32,6 +37,7 @@ get_rmsf_nufft, lambda2_to_freq, rmsynth_nufft, + stokes_i_template, ) @@ -114,6 +120,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 +159,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 +197,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,10 +235,11 @@ 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 - # 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, @@ -254,10 +290,23 @@ 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 get_args(NoiseWeightType) + 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 + ) + 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 +325,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 +391,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 +399,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..4bb1c4a 100644 --- a/rm_lite/tools_3d/rmsynth.py +++ b/rm_lite/tools_3d/rmsynth.py @@ -5,7 +5,7 @@ import logging 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 @@ -36,6 +36,7 @@ StokesIFitOptions, alpha_from_model_block, coefficient_names, + field_spectral_index, fit_stokes_cube, ref_flux_from_block, ) @@ -43,7 +44,9 @@ from rm_lite.utils.synthesis import ( FDFOptions, LamSq0Mode, + NoiseWeightType, RMSynthParams, + StokesIWeighting, TheoreticalNoise, WeightType, apply_weight_type, @@ -58,6 +61,7 @@ lambda2_to_freq, make_double_phi_arr, rmsynth_nufft, + stokes_i_template, ) @@ -122,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 @@ -160,6 +164,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 used; None if the + weights did not follow the Stokes I division.""" def _compute_global_params( @@ -204,13 +213,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 +654,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 +714,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, @@ -705,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. @@ -720,6 +779,18 @@ 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. 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. + 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 +834,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 +850,34 @@ 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 get_args(NoiseWeightType) + 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 + ) + 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 +1001,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, @@ -931,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. @@ -1004,6 +1111,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 +1260,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 +1312,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 +1405,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 get_args(NoiseWeightType): weight_arr = get_weight_arr_from_fits( stokes_q_file, stokes_u_file, @@ -1350,6 +1458,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..5b9ae0f 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 @@ -59,7 +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""" + 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"): @@ -722,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 @@ -847,6 +849,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, @@ -887,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).""" @@ -917,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), @@ -925,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): @@ -1013,17 +1061,6 @@ 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, - 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) - - RefFreqHz: TypeAlias = float | NDArray[np.float64] """A reference frequency in Hz: one for the whole image, or one per pixel.""" @@ -1061,11 +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 - 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; @@ -1083,12 +1118,10 @@ 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 ): - mean_flux = float(np.mean(i_spec[good])) if good.any() else np.nan if fit is None: - _write_flat_model(out, y, x, planes, mean_flux) continue pixel_ref_hz = ref_freq_for_pixel(ref_freq_hz, y, x) model = fit.stokes_i_model_func( @@ -1099,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 - _write_flat_model(out, y, x, planes, mean_flux) continue _write_model_planes( out, y, x, planes, fit, model, freq_arr_hz, pixel_ref_hz @@ -1118,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 a flat one (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 63039be..0811c5c 100644 --- a/rm_lite/utils/synthesis.py +++ b/rm_lite/utils/synthesis.py @@ -139,14 +139,21 @@ 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 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. """ + @dataclass(frozen=True, kw_only=True, slots=True) class FDFOptions: @@ -170,6 +177,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 +204,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) @@ -391,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 @@ -407,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: @@ -546,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, @@ -560,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( @@ -1251,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( @@ -1577,6 +1623,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..140b0e0 --- /dev/null +++ b/tests/test_stokes_i_weighting.py @@ -0,0 +1,380 @@ +"""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", "match"), + [ + ({"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], match: str) -> None: + """Unknown modes, a non-finite alpha and per-pixel weights with a per-pixel reference raise.""" + with pytest.raises(ValueError, match=match): + 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", MODES) +def test_pixels_below_the_snr_cut_are_blank_in_every_map( + mode: StokesIWeighting | None, +) -> None: + """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, -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 + 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, + 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" + 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( + [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] + 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) + # 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 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: + """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_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 3318370..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) - 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) + synth = artefact_synth(0.006, stokes_i_weighting=None) + 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)])