{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3D RM-synthesis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import tempfile\n",
    "from pathlib import Path\n",
    "\n",
    "import astropy.units as u\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import zarr\n",
    "from astropy.io import fits\n",
    "from astropy.visualization import quantity_support\n",
    "\n",
    "plt.rcParams[\"figure.dpi\"] = 150\n",
    "\n",
    "_ = quantity_support()\n",
    "rng = np.random.default_rng(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll simulate a small Stokes Q/U cutout cube with a per-pixel Faraday rotation measure, write it to FITS with a dummy Stokes axis (as in a typical ASKAP/EMU cutout cube, e.g. `image.restored.q.<field>.fits`), and read it back lazily with `rm_lite.utils.dask_io`.\n",
    "\n",
    "As with the 1D example, we simulate RACS-all frequency coverage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now build a small image with two compact polarised sources, one bottom-right and one top-left, each a 2D Gaussian blob. We use `rm_lite.utils.fitting.gaussian` on a radial distance grid, on an otherwise unpolarised background. Keeping most of the field source-free matters below: the robust per-channel noise estimator needs mostly-empty sky to lock onto, as it would on a real image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.fitting import gaussian\n",
    "from rm_lite.utils.synthesis import faraday_simple_spectrum\n",
    "\n",
    "ny, nx = 64, 64\n",
    "y_grid, x_grid = np.mgrid[0:ny, 0:nx]\n",
    "blob_y, blob_x = ny * 0.3, nx * 0.7\n",
    "blob2_y, blob2_x = ny * 0.7, nx * 0.3\n",
    "rm_map = (\n",
    "    80.0 * (x_grid / nx - 0.5) * 2\n",
    "    + np.exp(-((x_grid - blob_x) ** 2 + (y_grid - blob_y) ** 2) / (2 * 4.0**2))\n",
    "    - np.exp(-((x_grid - blob2_x) ** 2 + (y_grid - blob2_y) ** 2) / (2 * 4.0**2))\n",
    ")\n",
    "\n",
    "radius_grid = np.hypot(x_grid - blob_x, y_grid - blob_y)\n",
    "radius2_grid = np.hypot(x_grid - blob2_x, y_grid - blob2_y)\n",
    "frac_pol_map = gaussian(radius_grid, amplitude=0.6, mean=0.0, fwhm=6.0) + gaussian(\n",
    "    radius2_grid, amplitude=0.6, mean=0.0, fwhm=6.0\n",
    ")\n",
    "\n",
    "psi0_deg = 30.0\n",
    "rms_noise = 0.03\n",
    "\n",
    "stokes_q = np.empty((freq_hz.size, ny, nx), dtype=np.float32)\n",
    "stokes_u = np.empty((freq_hz.size, ny, nx), dtype=np.float32)\n",
    "for j in range(ny):\n",
    "    for i in range(nx):\n",
    "        complex_spectrum = faraday_simple_spectrum(\n",
    "            freq_hz,\n",
    "            frac_pol=frac_pol_map[j, i],\n",
    "            psi0_deg=psi0_deg,\n",
    "            rm_radm2=rm_map[j, i],\n",
    "        )\n",
    "        stokes_q[:, j, i] = complex_spectrum.real + rng.normal(\n",
    "            0, rms_noise, freq_hz.size\n",
    "        )\n",
    "        stokes_u[:, j, i] = complex_spectrum.imag + rng.normal(\n",
    "            0, rms_noise, freq_hz.size\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    "im1 = ax1.imshow(rm_map, origin=\"lower\", cmap=\"RdBu_r\", vmin=-100, vmax=100)\n",
    "fig.colorbar(im1, ax=ax1, label=f\"RM / ({u.rad / u.m**2:latex_inline})\")\n",
    "ax1.set(title=\"Input (true) RM map\")\n",
    "im2 = ax2.imshow(frac_pol_map, origin=\"lower\")\n",
    "fig.colorbar(im2, ax=ax2, label=\"Fractional polarisation\")\n",
    "ax2.set(title=\"Input (true) fractional polarisation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now write Stokes Q and U to separate FITS cubes, each with a degenerate leading Stokes axis (`NAXIS4 = 1`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmpdir = Path(tempfile.mkdtemp())\n",
    "\n",
    "\n",
    "def write_stokes_fits(path: Path, data: np.ndarray, stokes: int) -> None:\n",
    "    header = fits.Header()\n",
    "    header[\"CTYPE1\"] = \"RA---SIN\"\n",
    "    header[\"CTYPE2\"] = \"DEC--SIN\"\n",
    "    header[\"CTYPE3\"] = \"FREQ\"\n",
    "    header[\"CRVAL3\"] = freq_hz[0]\n",
    "    header[\"CDELT3\"] = float(np.diff(freq_hz)[0])\n",
    "    header[\"CRPIX3\"] = 1.0\n",
    "    header[\"CUNIT3\"] = \"Hz\"\n",
    "    header[\"CTYPE4\"] = \"STOKES\"\n",
    "    header[\"CRVAL4\"] = stokes\n",
    "    header[\"CDELT4\"] = 1.0\n",
    "    header[\"CRPIX4\"] = 1.0\n",
    "    fits.PrimaryHDU(data=data[np.newaxis, ...], header=header).writeto(\n",
    "        path, overwrite=True\n",
    "    )\n",
    "\n",
    "\n",
    "stokes_q_fits = tmpdir / \"cutout.q.fits\"\n",
    "stokes_u_fits = tmpdir / \"cutout.u.fits\"\n",
    "write_stokes_fits(stokes_q_fits, stokes_q, stokes=2)\n",
    "write_stokes_fits(stokes_u_fits, stokes_u, stokes=3)\n",
    "\n",
    "fits.getdata(stokes_q_fits).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`read_fits_cube_dask` reads the cube one spatial block at a time. Each block reopens the FITS file with `memmap=True` and slices out just that block's rows and columns, wrapped in a `dask.delayed` call and assembled into a single chunked `dask.array`. Actual disk reads are deferred until a block is computed. The frequency axis is always kept whole in a chunk, since RM-synthesis needs every channel per pixel, and the spatial chunk size comes from a target per-chunk memory footprint. Here it is set artificially small so multiple chunks show up in this small demo cube.\n",
    "\n",
    "In practice, peak memory during processing is a few times `target_chunk_mb` (roughly 2 to 4x, from buffer copies made when a block is materialised and again when it is written out), not a hard cap. It scales with chunk size, not cube size, so a `target_chunk_mb` set well below your available memory stays safe however large the on-disk cube is. Total RM-synthesis wall-clock time is set by the cube's total size, not how it is chunked, so smaller chunks cost little in speed. When in doubt, prefer a smaller `target_chunk_mb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.dask_io import read_fits_cube_dask\n",
    "\n",
    "q_dask, q_header = read_fits_cube_dask(\n",
    "    stokes_q_fits, target_chunk_mb=32 * 1024 / 1024**2\n",
    ")\n",
    "u_dask, _ = read_fits_cube_dask(stokes_u_fits, target_chunk_mb=32 * 1024 / 1024**2)\n",
    "q_dask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before running RM-synthesis, estimate a robust per-channel noise directly from the cube with `estimate_channel_noise_mad`. It takes the MAD-based standard deviation over every pixel in each channel plane, combined across Q and U the same way `compute_rmsynth_params` combines a per-channel error spectrum. This feeds `weight_arr` below and, in the RM-CLEAN notebook, the auto-mask/threshold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.dask_io import estimate_channel_noise_mad\n",
    "\n",
    "channel_noise = estimate_channel_noise_mad(q_dask, u_dask)\n",
    "weight_arr = 1.0 / channel_noise**2\n",
    "\n",
    "print(f\"true per-channel noise:      {rms_noise:.4f}\")\n",
    "print(f\"estimated per-channel noise: {channel_noise.mean():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run 3D RM-synthesis with `rmsynth_3d`. This applies the same `rmsynth_nufft`/`get_rmsf_nufft` math used by the 1D and 2D tools across the whole cube via `dask.array.map_blocks`, one call per spatial chunk.\n",
    "\n",
    "We pin `phi_max_radm2`/`d_phi_radm2` to a small, coarse Faraday depth grid (just enough planes to cover the simulated RM range) so this demo runs quickly. A real search would use finer sampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.tools_3d.rmsynth import rmsynth_3d\n",
    "\n",
    "help(rmsynth_3d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = rmsynth_3d(\n",
    "    q_dask,\n",
    "    u_dask,\n",
    "    freq_hz,\n",
    "    weight_arr=weight_arr,\n",
    "    phi_max_radm2=150.0,\n",
    "    d_phi_radm2=10.0,\n",
    ")\n",
    "result.fdf_dirty_cube"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`rmsynth_3d` takes dask arrays, for callers who already have Q/U loaded (e.g. from zarr). When Q/U are plain FITS files on disk, `rmsynth_3d_from_fits` folds the `read_fits_cube_dask` calls above into `rmsynth_3d` itself, so the two files-to-lazy-cube reads don't need to be written out by hand each time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.tools_3d.rmsynth import rmsynth_3d_from_fits\n",
    "\n",
    "result_from_fits = rmsynth_3d_from_fits(\n",
    "    stokes_q_fits,\n",
    "    stokes_u_fits,\n",
    "    weight_arr=weight_arr,\n",
    "    phi_max_radm2=150.0,\n",
    "    d_phi_radm2=10.0,\n",
    "    target_chunk_mb=32 * 1024 / 1024**2,\n",
    ")\n",
    "np.allclose(result_from_fits.fdf_dirty_cube.compute(), result.fdf_dirty_cube.compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`fdf_dirty_cube` and `rmsf_cube` are still lazy; nothing has been computed yet. Let's materialise the dirty FDF cube and look at the peak polarised intensity and recovered RM maps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_dirty_cube = result.fdf_dirty_cube.compute()\n",
    "peak_pi_map = np.max(np.abs(fdf_dirty_cube), axis=0)\n",
    "peak_idx_map = np.argmax(np.abs(fdf_dirty_cube), axis=0)\n",
    "peak_rm_map = result.phi_arr_radm2[peak_idx_map]\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    "im1 = ax1.imshow(peak_pi_map, origin=\"lower\")\n",
    "fig.colorbar(im1, ax=ax1, label=\"Polarised intensity\")\n",
    "ax1.set(title=\"Peak polarised intensity\")\n",
    "im2 = ax2.imshow(peak_rm_map, origin=\"lower\", cmap=\"RdBu_r\", vmin=-100, vmax=100)\n",
    "fig.colorbar(im2, ax=ax2, label=f\"RM / ({u.rad / u.m**2:latex_inline})\")\n",
    "ax2.set(title=\"Peak RM (recovered)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Serialisation: zarr vs FITS\n",
    "\n",
    "The dirty FDF and RMSF cubes are complex-valued, which matters for how they get written out:\n",
    "\n",
    "- **zarr** supports complex128 natively, and `dask.array.to_zarr` writes chunk by chunk without ever materialising the full cube in memory. The write scales with chunk size, not cube size, same as the computation itself (with the same buffer-copy overhead noted above, so budget a few times `target_chunk_mb`, not exactly `target_chunk_mb`).\n",
    "- **FITS** has no native complex dtype, and `astropy.io.fits` needs a full in-memory array to write. Writing an FDF cube to FITS means computing the whole cube first, then splitting it into real/imaginary (or real/imaginary/polarised-intensity) HDUs. This is the convention classic RM-Tools uses for its `FDF_real_dirty.fits`/`FDF_im_dirty.fits`/`FDF_tot_dirty.fits` outputs.\n",
    "\n",
    "zarr is the better fit for cube-sized outputs; FITS remains useful for interoperability with tools that expect it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.dask_io import write_zarr_group\n",
    "\n",
    "zarr_store = tmpdir / \"rmsynth3d.zarr\"\n",
    "write_zarr_group(\n",
    "    zarr_store,\n",
    "    {\"fdf_dirty\": result.fdf_dirty_cube, \"rmsf\": result.rmsf_cube},\n",
    ")\n",
    "\n",
    "group = zarr.open(zarr_store)\n",
    "group[\"fdf_dirty\"].shape, group[\"fdf_dirty\"].dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fits.PrimaryHDU(fdf_dirty_cube.real.astype(np.float32)).writeto(\n",
    "    tmpdir / \"FDF_real_dirty.fits\", overwrite=True\n",
    ")\n",
    "fits.PrimaryHDU(fdf_dirty_cube.imag.astype(np.float32)).writeto(\n",
    "    tmpdir / \"FDF_im_dirty.fits\", overwrite=True\n",
    ")\n",
    "fits.PrimaryHDU(np.abs(fdf_dirty_cube).astype(np.float32)).writeto(\n",
    "    tmpdir / \"FDF_tot_dirty.fits\", overwrite=True\n",
    ")\n",
    "\n",
    "roundtrip = fits.getdata(tmpdir / \"FDF_tot_dirty.fits\")\n",
    "np.allclose(roundtrip, np.abs(fdf_dirty_cube), atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Divide out the Stokes I spectral index first on the [3D RM-synthesis with Stokes I correction](rmsynth_3d_stokes_i.ipynb) page, and see the [3D RM-CLEAN](rmclean_3d.ipynb) page to deconvolve these dirty FDF/RMSF 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": 2
}