{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 1D Multiscale RM-CLEAN\n",
    "\n",
    "*NOTE:* This is still under development and experimentation. Consider this mode experimental.\n",
    "\n",
    "Standard RM-CLEAN models the Faraday dispersion function (FDF) as a sum of delta functions convolved with the RMSF. That is ideal for Faraday-thin (point) sources, but a Faraday-thick source is genuinely extended in Faraday depth, and a pile of deltas is a poor model of it: the components form a spiky comb that overshoots the true envelope, even though the restored FDF looks fine once the beam smooths it over.\n",
    "\n",
    "Multiscale RM-CLEAN is the RM-domain analogue of image-plane multiscale CLEAN (Cornwell 2008; Offringa & Smirnov 2017). Each component is a tapered kernel at one of several Faraday-depth *scales* (in units of the RMSF FWHM), not just a delta. What it buys you is a better *component model* of thick structure, and on thick sources it also tends to need fewer components: single-scale must build a deep comb of deltas to represent the width, while multiscale traces it with a handful of wide kernels. Every sub-minor iteration still places one component at a fixed gain, so the honest cost counter is `n_sub_minor_iter`.\n",
    "\n",
    "Scale selection is width-gated: a point-like peak is cleaned exactly like single-scale CLEAN (point safety), while a genuinely wide peak engages the extended scales via matched-filter scoring. Turn the whole thing on with `multiscale=True`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## How a scale gets chosen\n",
    "\n",
    "At every step CLEAN has to decide which scale to place the next component on. The obvious score for \"how well does scale $s$ explain the current residual\" is the peak of the residual $R$ convolved with that scale's kernel $K_s$,\n",
    "\n",
    "$$\\mathrm{peak}_s = \\max_\\phi \\, \\lvert R * K_s \\rvert .$$\n",
    "\n",
    "The problem is that a wider kernel sums over more of the FDF, and the noise in an FDF is strongly correlated (the RMSF has broad wings), so a wide kernel piles up correlated noise and its raw peak is inflated even on pure noise. Comparing raw peaks would always favour the widest scale.\n",
    "\n",
    "### Why not the Offringa & Smirnov (2017) bias\n",
    "\n",
    "O&S17 solve this for image-plane multiscale CLEAN with a fixed *scale bias*: multiply each score by a factor that geometrically down-weights larger scales,\n",
    "\n",
    "$$\\mathrm{score}_s = \\mathrm{peak}_s \\cdot \\beta^{-1-\\log_2(s/s_0)},$$\n",
    "\n",
    "with one tuning constant $\\beta$ (WSClean uses 0.6). That is fine when the noise is white and roughly the same on every scale, as it is in an image. It does not describe the RM domain: here the per-scale noise penalty is set by the RMSF shape, not by a clean geometric law, so no single $\\beta$ can match it. On the deep scale grids used here a fixed bias mis-weights the scales and pulls point sources onto oversized kernels, which blows up their recovered flux. We tried an exact re-implementation; it is not point-safe here, which is why rm-lite does not use it.\n",
    "\n",
    "### Matched-filter scoring instead\n",
    "\n",
    "Rather than guess a penalty, we normalise each score by the noise the kernel actually lets through. A few definitions make this precise.\n",
    "\n",
    "**Correlated FDF noise.** Independent noise in the measured $Q$ and $U$ channels does not stay independent after RM synthesis: every point of the FDF is a weighted sum of the same channels, so neighbouring Faraday depths share noise. The shape of that correlation is exactly the RMSF, so the RMSF plays the role of the noise *covariance* of the FDF. This is the root of the problem above: a wide kernel averages over many correlated cells, and correlated noise does not average down the way independent noise would.\n",
    "\n",
    "**Scale kernel $K_s$.** The tapered kernel for scale $s$ (in RMSF FWHM units), sum-normalised so that scale 0 is a delta and cleaning on it is ordinary single-scale CLEAN. The smoothed residual $R * K_s$ is what the score looks at, and $\\max_\\phi \\lvert R * K_s\\rvert$ is the numerator above.\n",
    "\n",
    "**Per-scale noise level $\\sigma_s$.** Because the covariance is the RMSF, the standard deviation of $\\lvert R * K_s\\rvert$ on a noise-only spectrum has a closed form,\n",
    "\n",
    "$$\\sigma_s = \\sqrt{\\frac{(K_s * K_s * \\mathrm{RMSF})(0)}{\\mathrm{RMSF}(0)}}, \\qquad \\sigma_0 = 1 .$$\n",
    "\n",
    "Reading it left to right: $*$ is convolution; $K_s * K_s$ applies the kernel twice (once for the kernel in the score, once because the noise itself is pushed through the same filter); convolving that with the RMSF propagates the noise covariance through the filter; and $(\\,0\\,)$ evaluates the result at zero lag (the RMSF peak channel), which is the variance the filter passes. Dividing by $\\mathrm{RMSF}(0)$ measures everything in units of the ordinary single-scale noise, which is what fixes $\\sigma_0 = 1$. So $\\sigma_s$ is simply \"how many times noisier than a single-scale clean is scale $s$\".\n",
    "\n",
    "**Matched filter.** A matched filter correlates the data with the template you are looking for and normalises by the noise that template picks up. Here the template is the scale kernel, so dividing the kernel response by its own noise level $\\sigma_s$ is exactly that:\n",
    "\n",
    "$$\\mathrm{score}_s = \\frac{\\max_\\phi \\lvert R * K_s \\rvert}{\\sigma_s}.$$\n",
    "\n",
    "The numerator measures how strongly scale $s$ responds to real structure in the residual; the denominator measures how strongly it responds to noise alone. Their ratio is a detection signal-to-noise, directly comparable across scales with no constant to hand-tune. A scale wins only when it detects structure above its *own* noise floor, not merely because it is wide.\n",
    "\n",
    "### Width gating keeps points safe\n",
    "\n",
    "A raw signal-to-noise argmax is still slightly scale-degenerate under correlated noise, so on its own it rarely engages a wide scale in a useful way. We gate it on the width of the feature being cleaned. An extended scale is engaged only when both:\n",
    "\n",
    "1. the residual peak is genuinely wider than the beam: its half-max width exceeds $1.2\\times$ the measured RMSF (dirty-beam) width, where a point source sits at $1.0\\times$; and\n",
    "2. the best extended score is at least $0.85\\times$ the delta-scale score, so structure that happens to be wide but is poorly matched to any wide kernel does not trigger it.\n",
    "\n",
    "If either test fails the step falls back to the signal-to-noise score, and among scales within a small margin (`multiscale_selection_margin`, default 0.08) of the best it takes the *smallest*. The result is point safety: a thin source is cleaned exactly like single-scale CLEAN, and wide scales activate only on structure that is really resolved.\n",
    "\n",
    "## How masking works\n",
    "\n",
    "Masking happens in stages, and none of them is a global cut on the raw FDF.\n",
    "\n",
    "First a **single-scale pass** cleans the spectrum on the delta scale alone. Wherever it places components defines the *source region*: a tight footprint around real emission. A noise-only spectrum leaves this empty, so multiscale CLEAN then does nothing, which is exactly what we want.\n",
    "\n",
    "Each scale then gets **its own mask** through an opening filter. The dirty FDF is convolved with the scale kernel, morphologically opened (erode then dilate by the scale width) so only structure at least that wide survives, thresholded at the flux-correct level $\\mathrm{mask}\\times$(kernel peak response), and finally confined to the source region dilated by the scale width. That dilation is what stops a wide kernel leaking onto the RMSF sidelobe islands that sit off to the side of a bright source. A scale whose mask comes out empty is switched off for this spectrum.\n",
    "\n",
    "Cleaning is then two-phase. Phase 1 cleans down to the mask level and records where each scale actually placed flux. Phase 2 deep-cleans to the final threshold, but restricts each scale to its own phase-1 footprint, so nothing smears fresh flux across the window late in the clean."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import logging\n",
    "\n",
    "import astropy.units as u\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from astropy.visualization import quantity_support\n",
    "from rm_lite.tools_1d import rmclean, rmsynth\n",
    "from rm_lite.utils.fitting import unit_centred_gaussian\n",
    "from rm_lite.utils.logging import logger\n",
    "from rm_lite.utils.simulate import build_model_fdf, delta, gauss, simulate_fdf\n",
    "from rm_lite.utils.synthesis import (\n",
    "    calc_faraday_moments,\n",
    "    freq_to_lambda2,\n",
    ")\n",
    "\n",
    "plt.rcParams[\"figure.dpi\"] = 150\n",
    "_ = quantity_support()\n",
    "logger.setLevel(logging.WARNING)\n",
    "rng = np.random.default_rng(42)\n",
    "\n",
    "# A wide fractional bandwidth (300-1800 MHz) so extended Faraday structure\n",
    "# survives depolarisation over several RMSF widths.\n",
    "freqs = np.linspace(300, 1800, 300) * u.MHz\n",
    "freq_hz = freqs.to(u.Hz).value\n",
    "sigma_chan = 0.02  # per-channel noise in Q and U\n",
    "\n",
    "\n",
    "def run_synth(complex_pol_arr):\n",
    "    return rmsynth.run_rmsynth(\n",
    "        freq_arr_hz=freq_hz,\n",
    "        complex_pol_arr=complex_pol_arr,\n",
    "        complex_pol_error=np.ones_like(complex_pol_arr)\n",
    "        * (sigma_chan + 1j * sigma_chan),\n",
    "        n_samples=50,\n",
    "        phi_max_radm2=100,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Example 1: Faraday-thin source (point safety)\n",
    "\n",
    "A thin source at RM = +3.3 RMSF FWHM. Multiscale must reproduce single-scale CLEAN here: same flux, same number of component-placement steps, all components on the delta scale at the source RM. The `simulate_fdf` helper forward-models the source through channel space, so the noise in the FDF is the physically correct correlated field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "thin_spec = delta(center_fwhm=3.3, amp=1.0)\n",
    "sim_thin = simulate_fdf(thin_spec, freq_hz, rng=rng, sigma=sigma_chan)\n",
    "synth_thin = run_synth(sim_thin.complex_pol_arr)\n",
    "\n",
    "fwhm = float(synth_thin.fdf_parameters[\"fwhm_rmsf_radm2\"][0])\n",
    "rm_true = 3.3 * fwhm * u.rad / u.m**2\n",
    "print(f\"RMSF FWHM = {fwhm:.2f} rad/m2, source RM = {rm_true:.1f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "help(rmclean.run_rmclean_from_synth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "single_thin = rmclean.run_rmclean_from_synth(synth_thin, auto_mask=7, auto_threshold=1)\n",
    "multi_thin = rmclean.run_rmclean_from_synth(\n",
    "    synth_thin, auto_mask=7, auto_threshold=1, multiscale=True\n",
    ")\n",
    "\n",
    "\n",
    "def unpack(res):\n",
    "    arrs = res.fdf_arrs\n",
    "    params = res.clean_parameters\n",
    "    return {\n",
    "        \"phi\": arrs[\"phi_arr_radm2\"].to_numpy() * u.rad / u.m**2,\n",
    "        \"dirty\": arrs[\"fdf_dirty_complex_arr\"].to_numpy().astype(complex),\n",
    "        \"clean\": arrs[\"fdf_clean_complex_arr\"].to_numpy().astype(complex),\n",
    "        \"model\": arrs[\"fdf_model_complex_arr\"].to_numpy().astype(complex),\n",
    "        \"resid\": arrs[\"fdf_residual_complex_arr\"].to_numpy().astype(complex),\n",
    "        \"mask\": float(params[\"mask\"][0]),\n",
    "        \"threshold\": float(params[\"threshold\"][0]),\n",
    "        \"n_iter\": int(params[\"n_iter\"][0]),\n",
    "        \"n_sub_minor_iter\": int(params[\"n_sub_minor_iter\"][0]),\n",
    "    }\n",
    "\n",
    "\n",
    "s = unpack(single_thin)\n",
    "m = unpack(multi_thin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(m[\"phi\"], np.abs(m[\"dirty\"]), color=\"k\", alpha=0.4, label=\"dirty\")\n",
    "ax.plot(\n",
    "    m[\"phi\"],\n",
    "    np.abs(m[\"clean\"] - m[\"resid\"]),\n",
    "    color=\"k\",\n",
    "    alpha=1,\n",
    "    label=\"convolved model\",\n",
    ")\n",
    "ax.step(m[\"phi\"], np.abs(m[\"model\"]), color=\"tab:red\", where=\"mid\", label=\"model\")\n",
    "ax.plot(\n",
    "    m[\"phi\"], np.abs(m[\"resid\"]), color=\"tab:blue\", alpha=1, zorder=10, label=\"residual\"\n",
    ")\n",
    "ax.axhline(m[\"threshold\"], color=\"k\", ls=\":\", label=\"threshold\", lw=1)\n",
    "ax.axhline(m[\"mask\"], color=\"k\", ls=\"--\", label=\"mask\", lw=1)\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"polarised intensity\",\n",
    "    title=\"Thin source, multiscale CLEAN\",\n",
    "    xlim=(rm_true.value - 15 * fwhm, rm_true.value + 15 * fwhm),\n",
    "    yscale=\"log\",\n",
    "    ylim=(1e-5, 2e0),\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "The component model is a handful of deltas at the source RM, exactly as single-scale would place them. The assertions below check the guarantees rather than describing them: the model flux matches single-scale, the components sit at the source RM, and the component-placement count is identical.\n",
    "\n",
    "One counter subtlety: `n_iter` means different things in the two modes. For single-scale it counts components; for multiscale it counts *minor cycles* (scale re-selections), each of which can place many components. `n_sub_minor_iter` is the component count in both modes, so that is the number to compare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Components sit at the source RM.\n",
    "support = np.abs(m[\"model\"]) > 0\n",
    "assert support.any()\n",
    "assert np.all(np.abs(m[\"phi\"][support] - rm_true) < 1.5 * fwhm * u.rad / u.m**2)\n",
    "\n",
    "# Flux parity with single-scale.\n",
    "m0_single = float(calc_faraday_moments(np.abs(s[\"clean\"]), s[\"phi\"].value, fwhm).mom0)\n",
    "m0_multi = float(calc_faraday_moments(np.abs(m[\"clean\"]), m[\"phi\"].value, fwhm).mom0)\n",
    "assert np.isclose(m0_single, m0_multi, rtol=1e-3)\n",
    "\n",
    "# Identical work: same number of component placements.\n",
    "assert m[\"n_sub_minor_iter\"] == s[\"n_sub_minor_iter\"]\n",
    "\n",
    "print(f\"flux: single {m0_single:.3f}, multi {m0_multi:.3f}\")\n",
    "print(\n",
    "    f\"components placed: single {s['n_sub_minor_iter']}, multi {m['n_sub_minor_iter']}\"\n",
    ")\n",
    "print(f\"multiscale minor cycles (scale re-selections): {m['n_iter']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Example 2: Faraday-thick source\n",
    "\n",
    "An external-dispersion Gaussian with sigma = 1.5 RMSF FWHM. On this band such a source is still polarised over much of the band, so its width is genuinely recoverable.\n",
    "\n",
    "This is where the two algorithms differ. Both recover the flux and both leave a comparable residual, but the *component models* are very different: single-scale builds a comb of deltas that overshoots the true envelope, while multiscale builds the source from wide kernels that trace it. The four panels show the input model and the clean models, bare and convolved with the clean beam. The beam hides the difference, which is exactly why the model matters: any physical interpretation of Faraday thickness uses the model, not the restored FDF.\n",
    "\n",
    "Because single-scale has to represent the width with a deep comb of deltas, it also places far *more* components than multiscale, which traces the source with a few wide kernels. So on a thick source multiscale wins on both counts, model shape and component count, checked below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "thick_spec = gauss(1.5, amp=1.0)\n",
    "sim_thick = simulate_fdf(thick_spec, freq_hz, rng=rng, sigma=sigma_chan)\n",
    "synth_thick = run_synth(sim_thick.complex_pol_arr)\n",
    "\n",
    "single_thick = rmclean.run_rmclean_from_synth(\n",
    "    synth_thick, auto_mask=10, auto_threshold=1\n",
    ")\n",
    "multi_thick = rmclean.run_rmclean_from_synth(\n",
    "    synth_thick,\n",
    "    auto_mask=10,\n",
    "    auto_threshold=1,\n",
    "    multiscale=True,\n",
    ")\n",
    "s = unpack(single_thick)\n",
    "m = unpack(multi_thick)\n",
    "\n",
    "print(\n",
    "    f\"components placed: single {s['n_sub_minor_iter']}, multi {m['n_sub_minor_iter']}\"\n",
    ")\n",
    "print(f\"multiscale minor cycles (scale re-selections): {m['n_iter']}\")\n",
    "\n",
    "phi = m[\"phi\"].value\n",
    "truth_model = build_model_fdf(thick_spec, phi, fwhm)\n",
    "\n",
    "# Unit-peak clean beam, matching how rm_lite restores the model.\n",
    "beam = unit_centred_gaussian(phi - phi[phi.size // 2], fwhm=fwhm)\n",
    "\n",
    "\n",
    "def conv_beam(arr):\n",
    "    return np.convolve(arr.real, beam, mode=\"same\") + 1j * np.convolve(\n",
    "        arr.imag, beam, mode=\"same\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 2, sharex=\"col\", sharey=\"row\", figsize=(12, 6))\n",
    "\n",
    "\n",
    "axs[0, 0].plot(\n",
    "    phi, np.abs(s[\"clean\"] - s[\"resid\"]), color=\"k\", alpha=1, label=\"convolved model\"\n",
    ")\n",
    "axs[0, 0].plot(\n",
    "    phi, np.abs(s[\"resid\"]), color=\"tab:blue\", alpha=1, zorder=10, label=\"residual\"\n",
    ")\n",
    "axs[0, 0].step(\n",
    "    phi, np.abs(s[\"model\"]), color=\"tab:red\", where=\"mid\", label=\"model\", lw=0.5\n",
    ")\n",
    "\n",
    "axs[0, 0].axhline(m[\"threshold\"], color=\"k\", ls=\":\", label=\"threshold\", lw=1)\n",
    "axs[0, 0].axhline(m[\"mask\"], color=\"k\", ls=\"--\", label=\"mask\", lw=1)\n",
    "axs[0, 0].legend()\n",
    "axs[0, 0].set_title(\"Delta scale\")\n",
    "axs[0, 0].set(ylabel=\"polarised intensity\")\n",
    "\n",
    "\n",
    "axs[0, 1].plot(phi, np.abs(m[\"clean\"] - m[\"resid\"]), color=\"k\", alpha=1)\n",
    "axs[0, 1].plot(phi, np.abs(m[\"resid\"]), color=\"tab:blue\", alpha=1, zorder=10)\n",
    "axs[0, 1].step(\n",
    "    phi, np.abs(m[\"model\"]), color=\"tab:red\", where=\"mid\", label=\"Single-scale\", lw=0.5\n",
    ")\n",
    "axs[0, 1].axhline(m[\"threshold\"], color=\"k\", ls=\":\", label=\"Threshold\", lw=1)\n",
    "axs[0, 1].axhline(m[\"mask\"], color=\"k\", ls=\"--\", label=\"Mask\", lw=1)\n",
    "axs[0, 1].set_title(\"Multi-scale\")\n",
    "\n",
    "axs[1, 0].plot(phi, np.abs(s[\"resid\"]), color=\"tab:blue\", zorder=10)\n",
    "axs[1, 0].step(\n",
    "    phi, np.abs(s[\"model\"]), color=\"tab:red\", where=\"mid\", label=\"Single-scale\", lw=0.5\n",
    ")\n",
    "axs[1, 0].axhline(m[\"threshold\"], color=\"k\", ls=\":\", label=\"Threshold\", lw=1)\n",
    "axs[1, 0].axhline(m[\"mask\"], color=\"k\", ls=\"--\", label=\"Mask\", lw=1)\n",
    "axs[1, 0].set(\n",
    "    ylabel=\"polarised intensity\", xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\"\n",
    ")\n",
    "\n",
    "axs[1, 1].plot(phi, np.abs(m[\"resid\"]), color=\"tab:blue\", zorder=10)\n",
    "axs[1, 1].step(\n",
    "    phi, np.abs(m[\"model\"]), color=\"tab:red\", where=\"mid\", label=\"Single-scale\", lw=0.5\n",
    ")\n",
    "axs[1, 1].axhline(m[\"threshold\"], color=\"k\", ls=\":\", label=\"Threshold\", lw=1)\n",
    "axs[1, 1].axhline(m[\"mask\"], color=\"k\", ls=\"--\", label=\"Mask\", lw=1)\n",
    "axs[1, 1].set(\n",
    "    yscale=\"log\", ylim=(1e-4, 2e-1), xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\"\n",
    ")\n",
    "\n",
    "fig.suptitle(\"Faraday thick source\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_shape_error(model, truth):\n",
    "    \"\"\"Scale-free rms of |model| against best-fit-amplitude |truth|.\"\"\"\n",
    "    a, t = np.abs(model), np.abs(truth)\n",
    "    amp = float((a * t).sum() / (t * t).sum())\n",
    "    return float(np.sqrt(np.mean((a - amp * t) ** 2)) / (amp * t.max()))\n",
    "\n",
    "\n",
    "err_single = model_shape_error(s[\"model\"], truth_model)\n",
    "err_multi = model_shape_error(m[\"model\"], truth_model)\n",
    "print(\n",
    "    f\"model shape error: single {err_single:.3f}, multi {err_multi:.3f} \"\n",
    "    f\"(ratio {err_multi / err_single:.2f})\"\n",
    ")\n",
    "\n",
    "# The multiscale model is much closer to the true envelope.\n",
    "assert err_multi < 0.5 * err_single\n",
    "\n",
    "# At no cost in flux ...\n",
    "m0_single = float(calc_faraday_moments(np.abs(s[\"clean\"]), phi, fwhm).mom0)\n",
    "m0_multi = float(calc_faraday_moments(np.abs(m[\"clean\"]), phi, fwhm).mom0)\n",
    "assert 0.9 < m0_multi / m0_single < 1.2\n",
    "\n",
    "# ... and with fewer components: single-scale needs a deep comb of deltas to\n",
    "# represent the width, while multiscale traces it with a handful of wide kernels.\n",
    "steps_ratio = m[\"n_sub_minor_iter\"] / s[\"n_sub_minor_iter\"]\n",
    "assert steps_ratio < 1.0\n",
    "print(f\"flux: single {m0_single:.3f}, multi {m0_multi:.3f}\")\n",
    "print(\n",
    "    f\"components placed: single {s['n_sub_minor_iter']}, multi {m['n_sub_minor_iter']} \"\n",
    "    f\"(ratio {steps_ratio:.2f})\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Example 3: narrow bands degenerate to single-scale\n",
    "\n",
    "The scales are in units of the RMSF FWHM, and the largest recoverable Faraday scale is set by the shortest wavelength in the band: `phi_max_scale = pi / lambda_sq_min`. On a narrow band the RMSF is so wide that no extended scale fits below that ceiling, so the automatic grid collapses to `[0.0]` (with a warning) and multiscale CLEAN is exactly single-scale CLEAN. This is safe, and it is also physics: a source wider than the ceiling is depolarised out of the band, so there is nothing for a wide kernel to recover."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.clean import MultiscaleOptions, default_scales\n",
    "from rm_lite.utils.synthesis import get_fwhm_rmsf, make_phi_arr\n",
    "\n",
    "narrow_hz = np.linspace(744e6, 1032e6, 288)  # RACS-low band\n",
    "lam2 = freq_to_lambda2(narrow_hz)\n",
    "narrow_fwhm = float(get_fwhm_rmsf(lam2).fwhm_rmsf_radm2)\n",
    "phi_max_scale = float(np.pi / lam2.min())\n",
    "phi_narrow = make_phi_arr(500.0, narrow_fwhm / 10)\n",
    "\n",
    "grid = default_scales(\n",
    "    phi_narrow, narrow_fwhm, MultiscaleOptions(), phi_max_scale_radm2=phi_max_scale\n",
    ")\n",
    "print(f\"RMSF FWHM = {narrow_fwhm:.1f} rad/m2\")\n",
    "print(\n",
    "    f\"largest recoverable scale = {phi_max_scale:.1f} rad/m2 \"\n",
    "    f\"= {phi_max_scale / narrow_fwhm:.2f} FWHM\"\n",
    ")\n",
    "print(f\"auto scale grid (FWHM units): {grid}\")\n",
    "\n",
    "# No extended scale fits below the ceiling: the grid is delta-only.\n",
    "assert grid.tolist() == [0.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And the clean itself is identical to single-scale, at identical cost.\n",
    "sim_narrow = simulate_fdf(thin_spec, narrow_hz, rng=rng, sigma=sigma_chan)\n",
    "synth_narrow = rmsynth.run_rmsynth(\n",
    "    freq_arr_hz=narrow_hz,\n",
    "    complex_pol_arr=sim_narrow.complex_pol_arr,\n",
    "    complex_pol_error=np.ones_like(sim_narrow.complex_pol_arr)\n",
    "    * (sigma_chan + 1j * sigma_chan),\n",
    "    n_samples=50,\n",
    ")\n",
    "single_n = unpack(\n",
    "    rmclean.run_rmclean_from_synth(synth_narrow, auto_mask=7, auto_threshold=1)\n",
    ")\n",
    "multi_n = unpack(\n",
    "    rmclean.run_rmclean_from_synth(\n",
    "        synth_narrow, auto_mask=7, auto_threshold=1, multiscale=True\n",
    "    )\n",
    ")\n",
    "assert np.allclose(single_n[\"clean\"], multi_n[\"clean\"])\n",
    "assert single_n[\"n_sub_minor_iter\"] == multi_n[\"n_sub_minor_iter\"]\n",
    "print(\n",
    "    \"narrowband multiscale == single-scale, \"\n",
    "    f\"{multi_n['n_sub_minor_iter']} components each\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Setting scales manually\n",
    "\n",
    "`multiscale_scales` overrides the automatic grid. The legitimate use is on a band that *does* support extended structure, when you want a leaner grid than the automatic one. Forcing extended scales onto a narrow band cannot beat the depolarisation ceiling: what little a wide kernel sees still fits at beam width, so the width-gated selection cleans it on the delta scale anyway.\n",
    "\n",
    "One caution: a bigger scale is not a better scale. Under the correlated noise of the FDF the matched-filter score of a very wide kernel never competes with the delta scale, so forcing a grid like `[0, 8]` just disengages multiscale. The scale the automatic run actually engages on this source is 3, so the leaner grid worth forcing is `[0, 3]`. The model win survives with it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_forced = unpack(\n",
    "    rmclean.run_rmclean_from_synth(\n",
    "        synth_thick,\n",
    "        auto_mask=10,\n",
    "        auto_threshold=1,\n",
    "        multiscale=True,\n",
    "        multiscale_scales=np.array([0.0, 3.0]),\n",
    "    )\n",
    ")\n",
    "err_forced = model_shape_error(multi_forced[\"model\"], truth_model)\n",
    "print(\n",
    "    f\"model shape error with forced grid [0, 3]: {err_forced:.3f} \"\n",
    "    f\"(single-scale {err_single:.3f})\"\n",
    ")\n",
    "assert err_forced < 0.7 * err_single\n",
    "\n",
    "# Flux is still consistent with single-scale.\n",
    "m0_forced = float(calc_faraday_moments(np.abs(multi_forced[\"clean\"]), phi, fwhm).mom0)\n",
    "assert 0.9 < m0_forced / m0_single < 1.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(\n",
    "    multi_forced[\"phi\"],\n",
    "    np.abs(multi_forced[\"dirty\"]),\n",
    "    color=\"k\",\n",
    "    alpha=0.4,\n",
    "    label=\"dirty\",\n",
    ")\n",
    "ax.plot(\n",
    "    multi_forced[\"phi\"],\n",
    "    np.abs(multi_forced[\"clean\"] - multi_forced[\"resid\"]),\n",
    "    color=\"k\",\n",
    "    alpha=1,\n",
    "    label=\"convolved model\",\n",
    ")\n",
    "ax.step(\n",
    "    multi_forced[\"phi\"],\n",
    "    np.abs(multi_forced[\"model\"]),\n",
    "    color=\"tab:red\",\n",
    "    where=\"mid\",\n",
    "    label=\"model\",\n",
    ")\n",
    "ax.plot(\n",
    "    multi_forced[\"phi\"],\n",
    "    np.abs(multi_forced[\"resid\"]),\n",
    "    color=\"tab:blue\",\n",
    "    alpha=1,\n",
    "    zorder=10,\n",
    "    label=\"residual\",\n",
    ")\n",
    "ax.axhline(multi_forced[\"threshold\"], color=\"k\", ls=\":\", label=\"threshold\", lw=1)\n",
    "ax.axhline(multi_forced[\"mask\"], color=\"k\", ls=\"--\", label=\"mask\", lw=1)\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"polarised intensity\",\n",
    "    title=\"Fixed scales [0, 3]\",\n",
    "    yscale=\"log\",\n",
    "    ylim=(1e-4, 2e0),\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- `multiscale=True` is safe to leave on: thin sources are cleaned exactly like single-scale CLEAN, and narrow bands degenerate to single-scale by construction.\n",
    "- The win is the *component model* of Faraday-thick structure; on thick sources multiscale also places far fewer components than single-scale, which needs a deep delta comb. Compare costs with `n_sub_minor_iter`; `n_iter` counts scale re-selections in multiscale mode.\n",
    "- The largest recoverable Faraday scale is `pi / lambda_sq_min`. Wider structure is depolarised out of the band and no algorithm can recover it.\n",
    "- `multiscale_scales` forces a custom grid (in RMSF FWHM units); it cannot beat the depolarisation ceiling, and very wide scales simply disengage."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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
}