{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 3D RM-synthesis with Stokes I correction\n",
    "\n",
    "Builds on the [3D RM-synthesis](rmsynth_3d.ipynb) notebook and adds the optional per-pixel Stokes I correction.\n",
    "\n",
    "A Stokes I spectral index reweights the polarised signal across $\\lambda^2$. That changes the FDF's effective RMSF, so it no longer matches the `rmsf_cube` the tool reports (built from the channel weights alone). RM-CLEAN and the moments assume `rmsf_cube`, so the mismatch biases them. The fix is the same as in 1D: divide by a Stokes I model to get fractional spectra $q = Q/I$, $u = U/I$, run RM-synthesis, then rescale the FDF back to flux by the Stokes I at the reference frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import astropy.units as u\n",
    "import dask.array as da\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from astropy.visualization import quantity_support\n",
    "\n",
    "plt.rcParams[\"figure.dpi\"] = 150\n",
    "_ = quantity_support()\n",
    "rng = np.random.default_rng(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "Simulate a small cube with RACS-all frequency coverage. The intrinsic fractional polarisation is flat (the same at every pixel), so a correct pipeline recovers a flat fractional-polarisation map whatever the Stokes I does. The RM varies across the field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.fitting import gaussian\n",
    "from rm_lite.utils.synthesis import faraday_simple_spectrum\n",
    "\n",
    "bw_low = 288\n",
    "freqs = np.linspace(943.5 - bw_low / 2, 943.5 + bw_low / 2, 36) * u.MHz\n",
    "freq_hz = freqs.to(u.Hz).value\n",
    "\n",
    "ny, nx = 48, 48\n",
    "y_grid, x_grid = np.mgrid[0:ny, 0:nx]\n",
    "blob_y, blob_x = ny * 0.35, nx * 0.65\n",
    "blob2_y, blob2_x = ny * 0.7, nx * 0.3\n",
    "\n",
    "rm_map = 60.0 * (x_grid / nx - 0.5) * 2 + np.exp(\n",
    "    -((x_grid - blob_x) ** 2 + (y_grid - blob_y) ** 2) / (2 * 4.0**2)\n",
    ")\n",
    "psi0_deg = 30.0\n",
    "\n",
    "frac_pol_true = 0.1  # intrinsic fractional polarisation: a FLAT map\n",
    "q_frac = np.empty((freq_hz.size, ny, nx), dtype=np.float32)\n",
    "u_frac = np.empty((freq_hz.size, ny, nx), dtype=np.float32)\n",
    "for j in range(ny):\n",
    "    for i in range(nx):\n",
    "        cs = faraday_simple_spectrum(\n",
    "            freq_hz, frac_pol=frac_pol_true, psi0_deg=psi0_deg, rm_radm2=rm_map[j, i]\n",
    "        )\n",
    "        q_frac[:, j, i] = cs.real\n",
    "        u_frac[:, j, i] = cs.imag"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "Stokes I: two Gaussian blobs on empty sky, with a spectral index that varies across the field. None of this is Faraday structure, but it distorts an uncorrected FDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_freq_hz = float(np.median(freq_hz))\n",
    "radius = np.hypot(x_grid - blob_x, y_grid - blob_y)\n",
    "radius2 = np.hypot(x_grid - blob2_x, y_grid - blob2_y)\n",
    "i_ref_map = (\n",
    "    0.0\n",
    "    + gaussian(radius, amplitude=1.0, mean=0.0, fwhm=8.0)\n",
    "    + gaussian(radius2, amplitude=0.6, mean=0.0, fwhm=8.0)\n",
    ")\n",
    "alpha_map_true = -1.0 + 0.6 * (x_grid / nx - 0.5) * 2  # ~ -1.6 .. -0.4\n",
    "stokes_i = (\n",
    "    i_ref_map[None] * (freq_hz / ref_freq_hz)[:, None, None] ** alpha_map_true[None]\n",
    ").astype(np.float32)\n",
    "\n",
    "rms_noise = 0.02\n",
    "stokes_q_obs = q_frac * stokes_i + rng.normal(0, rms_noise, stokes_i.shape).astype(\n",
    "    np.float32\n",
    ")\n",
    "stokes_u_obs = u_frac * stokes_i + rng.normal(0, rms_noise, stokes_i.shape).astype(\n",
    "    np.float32\n",
    ")\n",
    "stokes_i_obs = stokes_i + rng.normal(0, rms_noise, stokes_i.shape).astype(np.float32)\n",
    "\n",
    "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(14, 4))\n",
    "im1 = ax1.imshow(i_ref_map, origin=\"lower\")\n",
    "fig.colorbar(im1, ax=ax1, label=\"Stokes I\")\n",
    "ax1.set(title=\"Input Stokes I (ref freq)\")\n",
    "im2 = ax2.imshow(alpha_map_true, origin=\"lower\", cmap=\"Spectral\", vmin=-2, vmax=+2)\n",
    "fig.colorbar(im2, ax=ax2, label=r\"$\\alpha$\")\n",
    "ax2.set(title=\"Input spectral index\")\n",
    "im3 = ax3.imshow(rm_map, origin=\"lower\", cmap=\"RdBu_r\", vmin=-80, vmax=80)\n",
    "fig.colorbar(im3, ax=ax3, label=f\"RM / ({u.rad / u.m**2:latex_inline})\")\n",
    "ax3.set(title=\"Input RM\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "Pass the Stokes I cube as `stokes_i`. `rmsynth_3d` fits a model per pixel (the same fitter as 1D), forms the fractional spectra, runs RM-synthesis, and rescales back to flux. The fit is quiet at cube scale. Pass a ready-made `stokes_i_model` cube to skip fitting.\n",
    "\n",
    "`fit_order=-2` iterates orders 0 to 2 and keeps the best by AIC. `stokes_i_snr_cut=5` fits only pixels above SNR 5 (see below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.tools_3d.rmsynth import rmsynth_3d\n",
    "\n",
    "chunks = (-1, 16, 16)\n",
    "result_i = rmsynth_3d(\n",
    "    da.from_array(stokes_q_obs, chunks=chunks),\n",
    "    da.from_array(stokes_u_obs, chunks=chunks),\n",
    "    freq_hz,\n",
    "    stokes_i=da.from_array(stokes_i_obs, chunks=chunks),\n",
    "    estimate_stokes_i_noise=True,  # per-channel noise from the I cube for the fit\n",
    "    fit_function=\"log\",  # power-law fit\n",
    "    fit_order=-2,  # iterate orders 0..2, keep the best by AIC\n",
    "    stokes_i_snr_cut=5.0,  # only fit pixels above SNR 5\n",
    "    phi_max_radm2=150.0,\n",
    "    d_phi_radm2=1.0,\n",
    "    weight_type=\"uniform\",\n",
    ")\n",
    "model_cube = result_i.stokes_i_model_cube.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "`rmsynth_3d` also returns 2D Stokes I maps, like the moment maps: the model at the reference frequency (`stokes_i_ref_flux_map`), the fitted spectral index (`stokes_i_alpha_map`), and the fitted polynomial order per pixel (`stokes_i_model_order_map`). With `fit_order=-2` the order is chosen per pixel by AIC, so the order map shows how much spectral structure each pixel justified. The fit tracks the observed Stokes I at a bright pixel, and the recovered maps match the inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "jb, ib = int(blob_y), int(blob_x)  # a bright pixel\n",
    "ref_flux_map = result_i.stokes_i_ref_flux_map.compute()\n",
    "alpha_map = result_i.stokes_i_alpha_map.compute()\n",
    "order_map = result_i.stokes_i_model_order_map.compute()\n",
    "\n",
    "fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(18, 4))\n",
    "ax1.plot(freq_hz / 1e9, stokes_i_obs[:, jb, ib], \".\", ms=4, label=\"observed I\")\n",
    "ax1.plot(freq_hz / 1e9, model_cube[:, jb, ib], \"-\", label=\"fitted model\")\n",
    "ax1.set(xlabel=\"Frequency / GHz\", ylabel=\"Stokes I\", title=\"Per-pixel Stokes I fit\")\n",
    "ax1.legend()\n",
    "im2 = ax2.imshow(ref_flux_map, origin=\"lower\")\n",
    "fig.colorbar(im2, ax=ax2, label=\"Stokes I\")\n",
    "ax2.set(title=\"Recovered Stokes I (ref freq)\")\n",
    "im3 = ax3.imshow(alpha_map, origin=\"lower\", cmap=\"Spectral\", vmin=-2, vmax=+2)\n",
    "fig.colorbar(im3, ax=ax3, label=r\"$\\alpha = d\\,\\ln I / d\\,\\ln \\nu$\")\n",
    "ax3.set(title=\"Recovered spectral index\")\n",
    "im4 = ax4.imshow(order_map, origin=\"lower\", cmap=\"viridis\")\n",
    "fig.colorbar(im4, ax=ax4, label=\"fitted order\")\n",
    "ax4.set(title=\"Fitted model order\")\n",
    "\n",
    "# Self-check: order is an integer >= 0 on fitted pixels, and NaN exactly where\n",
    "# alpha is NaN (the unfitted sky) -- same fitted-pixel set for both maps.\n",
    "fitted = np.isfinite(order_map)\n",
    "assert (order_map[fitted] >= 0).all()\n",
    "np.testing.assert_array_equal(order_map[fitted], np.round(order_map[fitted]))\n",
    "np.testing.assert_array_equal(fitted, np.isfinite(alpha_map))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### Recovering the flat fractional-polarisation map\n",
    "\n",
    "The corrected FDF is in flux units, so its peak follows the Stokes I brightness and says nothing on its own about fractional polarisation. Divide the peak by the reference Stokes I map (`stokes_i_ref_flux_map`) to get the intrinsic fractional polarisation: flat at `frac_pol_true`, as put in, whatever the Stokes I brightness and index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "snr_map = np.mean(stokes_i_obs, axis=0) * np.sqrt(freq_hz.size) / rms_noise\n",
    "source = snr_map >= 5.0  # where the Stokes I fit was applied\n",
    "\n",
    "fdf_i = result_i.fdf_dirty_cube.compute()\n",
    "peak_pi_corrected = np.max(np.abs(fdf_i), axis=0)\n",
    "# Fractional pol is only meaningful where there is Stokes I signal.\n",
    "frac_pol_recovered = np.where(source, peak_pi_corrected / ref_flux_map, np.nan)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    "im1 = ax1.imshow(peak_pi_corrected, origin=\"lower\")\n",
    "fig.colorbar(im1, ax=ax1, label=\"Polarised intensity\")\n",
    "ax1.set(title=\"Peak PI (absolute, follows Stokes I)\")\n",
    "im2 = ax2.imshow(frac_pol_recovered, origin=\"lower\", vmin=0, vmax=2 * frac_pol_true)\n",
    "fig.colorbar(im2, ax=ax2, label=\"Fractional polarisation\")\n",
    "ax2.set(title=f\"Recovered fractional pol (flat ~ {frac_pol_true})\")\n",
    "print(\n",
    "    f\"recovered fractional pol over source: \"\n",
    "    f\"median={np.nanmedian(frac_pol_recovered):.3f}, input={frac_pol_true}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "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, `mean(I) * sqrt(N) / rms(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 at its mean: no spectral correction, so its FDF is the plain Q/U FDF. Nothing is blanked.\n",
    "\n",
    "Most of this field is empty sky. With the cut, only the sources are fitted. Without it (`stokes_i_snr_cut=None`) every noise pixel is fitted too, and the spectral-index map fills with junk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_nomask = rmsynth_3d(\n",
    "    da.from_array(stokes_q_obs, chunks=chunks),\n",
    "    da.from_array(stokes_u_obs, chunks=chunks),\n",
    "    freq_hz,\n",
    "    stokes_i=da.from_array(stokes_i_obs, chunks=chunks),\n",
    "    estimate_stokes_i_noise=True,\n",
    "    fit_function=\"linear\",\n",
    "    fit_order=-2,\n",
    "    stokes_i_snr_cut=None,  # no mask: fit every pixel\n",
    "    phi_max_radm2=150.0,\n",
    "    d_phi_radm2=1.0,\n",
    "    weight_type=\"uniform\",\n",
    ")\n",
    "alpha_nomask = result_nomask.stokes_i_alpha_map.compute()\n",
    "\n",
    "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(14, 4))\n",
    "im1 = ax1.imshow(snr_map, origin=\"lower\")\n",
    "fig.colorbar(im1, ax=ax1, label=\"Stokes I SNR\")\n",
    "ax1.contour(snr_map, levels=[5.0], colors=\"w\", linewidths=1)\n",
    "ax1.set(title=\"Stokes I SNR (contour at cut = 5)\")\n",
    "im2 = ax2.imshow(alpha_map, origin=\"lower\", cmap=\"Spectral\", vmin=-2, vmax=2)\n",
    "fig.colorbar(im2, ax=ax2, label=r\"$\\alpha$\")\n",
    "ax2.set(title=\"alpha, SNR cut = 5 (masked)\")\n",
    "im3 = ax3.imshow(alpha_nomask, origin=\"lower\", cmap=\"Spectral\", vmin=-2, vmax=2)\n",
    "fig.colorbar(im3, ax=ax3, label=r\"$\\alpha$\")\n",
    "ax3.set(title=\"alpha, no cut (fit all)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Optional per-pixel model and spectral-index error\n",
    "\n",
    "Pass `compute_model_error=True` for a per-pixel 1-sigma Stokes I model error cube *and* a matching spectral-index error map (`stokes_i_alpha_error_map`), both from a single Monte-Carlo over each pixel's fit covariance. The alpha error is the spread of the per-realisation spectral index at the reference frequency.\n",
    "\n",
    "One thing to know: the model cube and the error outputs come out of a single dask task, so the Monte-Carlo cost is attached to the model cube. The FDF divides by the model cube, so **computing the FDF also runs the Monte-Carlo** once the flag is on, even if you never look at the error outputs. So:\n",
    "\n",
    "- Leave `compute_model_error=False` (default) unless you actually want the error, and no Monte-Carlo runs.\n",
    "- When you do want it, compute the error outputs **together** with the model and FDF in one pass (one `dask.compute(...)` or `write_zarr_group`) so the shared fit runs once for everything."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_err = rmsynth_3d(\n",
    "    da.from_array(stokes_q_obs, chunks=chunks),\n",
    "    da.from_array(stokes_u_obs, chunks=chunks),\n",
    "    freq_hz,\n",
    "    stokes_i=da.from_array(stokes_i_obs, chunks=chunks),\n",
    "    estimate_stokes_i_noise=True,\n",
    "    fit_function=\"log\",\n",
    "    fit_order=-3,\n",
    "    stokes_i_snr_cut=5.0,\n",
    "    compute_model_error=True,  # error rides along in the same fit pass\n",
    "    n_error_samples=200,\n",
    "    phi_max_radm2=150.0,\n",
    "    d_phi_radm2=1.0,\n",
    "    weight_type=\"uniform\",\n",
    ")\n",
    "\n",
    "# Compute everything you need in ONE pass so the shared per-pixel fit runs once.\n",
    "# (Computing the FDF alone would still run the Monte-Carlo, since it divides by\n",
    "# the model cube and model+error share a task -- so ask for them together.)\n",
    "fdf_err, model_err, model_error, alpha_err_val, alpha_err_map, order_err_map = (\n",
    "    da.compute(\n",
    "        result_err.fdf_dirty_cube,\n",
    "        result_err.stokes_i_model_cube,\n",
    "        result_err.stokes_i_model_error_cube,\n",
    "        result_err.stokes_i_alpha_map,\n",
    "        result_err.stokes_i_alpha_error_map,\n",
    "        result_err.stokes_i_model_order_map,\n",
    "    )\n",
    ")\n",
    "\n",
    "jb, ib = int(blob_y), int(blob_x)  # a bright pixel\n",
    "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 4))\n",
    "ax1.plot(\n",
    "    freq_hz / 1e9, stokes_i_obs[:, jb, ib], \".\", ms=4, color=\"0.5\", label=\"observed I\"\n",
    ")\n",
    "ax1.plot(freq_hz / 1e9, model_err[:, jb, ib], \"-\", label=\"model\")\n",
    "ax1.fill_between(\n",
    "    freq_hz / 1e9,\n",
    "    model_err[:, jb, ib] - model_error[:, jb, ib],\n",
    "    model_err[:, jb, ib] + model_error[:, jb, ib],\n",
    "    alpha=0.3,\n",
    "    label=r\"$\\pm 1\\sigma$\",\n",
    ")\n",
    "ax1.set(xlabel=\"Frequency / GHz\", ylabel=\"Stokes I\", title=\"Model with 1-sigma error\")\n",
    "ax1.legend()\n",
    "im2 = ax2.imshow(alpha_err_map, origin=\"lower\", cmap=\"magma\")\n",
    "fig.colorbar(im2, ax=ax2, label=r\"$\\sigma_\\alpha$\")\n",
    "ax2.set(title=\"Spectral-index error (fitted pixels)\")\n",
    "im3 = ax3.imshow(order_err_map, origin=\"lower\", cmap=\"viridis\")\n",
    "fig.colorbar(im3, ax=ax3, label=\"fitted order\")\n",
    "ax3.set(title=\"Fitted model order (AIC)\")\n",
    "\n",
    "# Self-check: the error map is finite and positive on fitted pixels (NaN on the\n",
    "# masked sky), and the recovered index at the bright pixel sits within ~3 sigma\n",
    "# of the input, i.e. the reported error is a sane 1-sigma.\n",
    "fitted = np.isfinite(alpha_err_map)\n",
    "assert (alpha_err_map[fitted] > 0).all()\n",
    "assert abs(alpha_err_val[jb, ib] - alpha_map_true[jb, ib]) < 3 * alpha_err_map[jb, ib]\n",
    "# Order map shares the same fitted-pixel set and holds integer orders there.\n",
    "np.testing.assert_array_equal(np.isfinite(order_err_map), fitted)\n",
    "np.testing.assert_array_equal(order_err_map[fitted], np.round(order_err_map[fitted]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "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 these cubes."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rm-lite (3.11.13)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}