{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1D RM-CLEAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\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 numpy.typing import NDArray\n",
    "from rm_lite.tools_1d import rmclean, rmsynth\n",
    "\n",
    "plt.rcParams[\"figure.dpi\"] = 150\n",
    "\n",
    "_ = quantity_support()\n",
    "rng = np.random.default_rng(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll start with the same data s the 1D RM-synth example. Keeping it simple, we'll use a Faraday simple source with a bit of noise.\n",
    "\n",
    "I'll skip the details here, see the RM-synth page for more info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.fitting import power_law\n",
    "from rm_lite.utils.synthesis import faraday_simple_spectrum, freq_to_lambda2\n",
    "\n",
    "\n",
    "def faraday_slab_spectrum(\n",
    "    freq_arr_hz: NDArray[np.float64],\n",
    "    frac_pol: float,\n",
    "    psi0_deg: float,\n",
    "    rm_radm2: float,\n",
    "    delta_rm_radm2: float,\n",
    ") -> NDArray[np.complex128]:\n",
    "    lambda_sq_arr_m2 = freq_to_lambda2(freq_arr_hz)\n",
    "\n",
    "    return (\n",
    "        frac_pol\n",
    "        * np.exp(2j * (np.deg2rad(psi0_deg) + rm_radm2 * lambda_sq_arr_m2))\n",
    "        * (\n",
    "            np.sin(delta_rm_radm2 * lambda_sq_arr_m2)\n",
    "            / (delta_rm_radm2 * lambda_sq_arr_m2)\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "def faraday_gaussian_spectrum(\n",
    "    freq_arr_hz: NDArray[np.float64],\n",
    "    frac_pol: float,\n",
    "    psi0_deg: float,\n",
    "    rm_radm2: float,\n",
    "    sigma_rm_radm2: float,\n",
    "):\n",
    "    lambda_sq_arr_m2 = freq_to_lambda2(freq_arr_hz)\n",
    "    rm_term = np.exp(2j * (np.deg2rad(psi0_deg) + rm_radm2 * lambda_sq_arr_m2))\n",
    "    depol_term = np.exp(-2.0 * sigma_rm_radm2**2 * lambda_sq_arr_m2**2)\n",
    "    return frac_pol * rm_term * depol_term\n",
    "\n",
    "\n",
    "bw_low = 288\n",
    "bw_mid = 144\n",
    "bw_high = 288\n",
    "low = np.linspace(943.5 - bw_low / 2, 943.5 + bw_low / 2, 36) * u.MHz\n",
    "mid = np.linspace(1367.5 - bw_mid / 2, 1367.5 + bw_mid / 2, 9) * u.MHz\n",
    "high = np.linspace(1655.5 - bw_high / 2, 1655.5 + bw_high / 2, 9) * u.MHz\n",
    "freqs = np.concatenate([low, mid, high])\n",
    "freq_hz = freqs.to(u.Hz).value\n",
    "\n",
    "delta_rm_radm2 = 30\n",
    "rm_radm2 = 100\n",
    "frac_pol = 0.7\n",
    "psi0_deg = 10\n",
    "\n",
    "complex_data_noiseless = faraday_simple_spectrum(\n",
    "    freq_arr_hz=freq_hz,\n",
    "    frac_pol=frac_pol,\n",
    "    psi0_deg=psi0_deg,\n",
    "    rm_radm2=rm_radm2,\n",
    ")\n",
    "stokes_i_flux = 2.0\n",
    "spectral_index = -0.7\n",
    "rms_noise = 0.1\n",
    "\n",
    "stokes_i_model = power_law(order=1)\n",
    "stokes_i_noiseless = stokes_i_model(\n",
    "    freq_hz / (np.mean(freq_hz)), stokes_i_flux, spectral_index\n",
    ")\n",
    "stokes_i_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n",
    "stokes_i_noisy = stokes_i_noiseless + stokes_i_noise\n",
    "\n",
    "\n",
    "stokes_q_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n",
    "stokes_u_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n",
    "complex_noise = stokes_q_noise + 1j * stokes_u_noise\n",
    "\n",
    "complex_flux = complex_data_noiseless * stokes_i_noiseless\n",
    "complex_data_noisy = complex_data_noiseless + complex_noise\n",
    "rm_syth_results = rmsynth.run_rmsynth(\n",
    "    freq_arr_hz=freq_hz,\n",
    "    complex_pol_arr=complex_data_noisy,\n",
    "    complex_pol_error=np.ones_like(complex_data_noiseless) * rms_noise,\n",
    "    do_fit_rmsf=True,\n",
    "    n_samples=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's remind ourselves of the what the dirty spectrum looks like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = rm_syth_results\n",
    "\n",
    "phi_arr_radm2 = fdf_arrs[\"phi_arr_radm2\"].to_numpy()\n",
    "fdf_dirty_arr = fdf_arrs[\"fdf_dirty_complex_arr\"].to_numpy().astype(complex)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    fdf_dirty_arr.real,\n",
    "    color=\"tab:red\",\n",
    "    label=\"Stokes Q\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    fdf_dirty_arr.imag,\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Stokes U\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_dirty_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Polarized intensity\",\n",
    ")\n",
    "\n",
    "ax.errorbar(\n",
    "    fdf_parameters[\"peak_rm_fit\"],\n",
    "    fdf_parameters[\"peak_pi_fit\"],\n",
    "    xerr=fdf_parameters[\"peak_rm_fit_error\"],\n",
    "    yerr=fdf_parameters[\"peak_pi_error\"],\n",
    "    fmt=\"o\",\n",
    "    lw=1,\n",
    "    color=\"red\",\n",
    "    mfc=\"none\",\n",
    "    label=\"Fitted peak\",\n",
    ")\n",
    "\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"Flux density\",\n",
    "    title=\"Dirty FDF\",\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the convenience function `rmclean.run_rmclean_from_synth` to directly take the outputs of 1D RM-synthesis for RM-CLEAN.\n",
    "\n",
    "Inspired by WSClean, `rm_lite` implements a deep cleaning mode with auto-masking and auto-thresholding based on the noise profile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(rmclean.run_rmclean_from_synth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmclean_results = rmclean.run_rmclean_from_synth(\n",
    "    rm_synth_1d_results=rm_syth_results, auto_mask=10, auto_threshold=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmclean_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the DataFrames\n",
    "clean_arrs = rmclean_results.fdf_arrs\n",
    "clean_fdf_params = rmclean_results.fdf_parameters\n",
    "clean_params = rmclean_results.clean_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract some arrays\n",
    "phi_arr_radm2 = clean_arrs[\"phi_arr_radm2\"].to_numpy()\n",
    "fdf_dirty_arr = clean_arrs[\"fdf_dirty_complex_arr\"].to_numpy().astype(complex)\n",
    "fdf_clean_arr = clean_arrs[\"fdf_clean_complex_arr\"].to_numpy().astype(complex)\n",
    "fdf_model_arr = clean_arrs[\"fdf_model_complex_arr\"].to_numpy().astype(complex)\n",
    "fdf_residual_arr = clean_arrs[\"fdf_residual_complex_arr\"].to_numpy().astype(complex)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at how the CLEAN went. Note that the deep clean allows us to model the entire simple component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_dirty_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Dirty FDF\",\n",
    "    alpha=0.5,\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_clean_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Clean FDF\",\n",
    ")\n",
    "ax.step(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_model_arr),\n",
    "    color=\"tab:red\",\n",
    "    label=\"Model FDF\",\n",
    "    where=\"mid\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_residual_arr),\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Residual FDF\",\n",
    "    alpha=0.5,\n",
    ")\n",
    "\n",
    "ax.errorbar(\n",
    "    clean_fdf_params[\"peak_rm_fit\"],\n",
    "    clean_fdf_params[\"peak_pi_fit\"],\n",
    "    xerr=clean_fdf_params[\"peak_rm_fit_error\"],\n",
    "    yerr=clean_fdf_params[\"peak_pi_error\"],\n",
    "    fmt=\"o\",\n",
    "    lw=1,\n",
    "    color=\"red\",\n",
    "    mfc=\"none\",\n",
    "    label=\"Fitted peak\",\n",
    ")\n",
    "\n",
    "ax.axhline(\n",
    "    y=clean_params[\"threshold\"][0],\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Threshold\",\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "ax.axhline(\n",
    "    y=clean_params[\"mask\"][0],\n",
    "    color=\"tab:orange\",\n",
    "    label=\"Mask\",\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"Flux density\",\n",
    "    title=\"Dirty FDF\",\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's take a look at the more complex example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_rm_radm2 = 30\n",
    "rm_radm2 = 100\n",
    "frac_pol = 0.5\n",
    "psi0_deg = 10\n",
    "complex_data_noiseless = faraday_slab_spectrum(\n",
    "    freq_arr_hz=freq_hz,\n",
    "    frac_pol=frac_pol,\n",
    "    psi0_deg=psi0_deg,\n",
    "    rm_radm2=rm_radm2,\n",
    "    delta_rm_radm2=delta_rm_radm2,\n",
    ")\n",
    "\n",
    "\n",
    "stokes_i_flux = 1.0\n",
    "spectral_index = -0.7\n",
    "rms_noise = 0.1\n",
    "\n",
    "\n",
    "stokes_i_model = power_law(order=1)\n",
    "stokes_i_noiseless = stokes_i_model(\n",
    "    freq_hz / (np.mean(freq_hz)), stokes_i_flux, spectral_index\n",
    ")\n",
    "stokes_i_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n",
    "stokes_i_noisy = stokes_i_noiseless + stokes_i_noise\n",
    "\n",
    "\n",
    "stokes_q_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n",
    "stokes_u_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n",
    "complex_noise = stokes_q_noise + 1j * stokes_u_noise\n",
    "\n",
    "complex_flux = complex_data_noiseless * stokes_i_noiseless\n",
    "complex_data_noisy = complex_data_noiseless + complex_noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rm_syth_results = rmsynth.run_rmsynth(\n",
    "    freq_arr_hz=freq_hz,\n",
    "    complex_pol_arr=complex_data_noisy,\n",
    "    complex_pol_error=np.ones_like(complex_data_noiseless) * rms_noise,\n",
    "    do_fit_rmsf=True,\n",
    "    n_samples=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = rm_syth_results\n",
    "\n",
    "phi_arr_radm2 = fdf_arrs[\"phi_arr_radm2\"].to_numpy()\n",
    "fdf_dirty_arr = fdf_arrs[\"fdf_dirty_complex_arr\"].to_numpy().astype(complex)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    fdf_dirty_arr.real,\n",
    "    color=\"tab:red\",\n",
    "    label=\"Stokes Q\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    fdf_dirty_arr.imag,\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Stokes U\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_dirty_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Polarized intensity\",\n",
    ")\n",
    "\n",
    "ax.errorbar(\n",
    "    fdf_parameters[\"peak_rm_fit\"],\n",
    "    fdf_parameters[\"peak_pi_fit\"],\n",
    "    xerr=fdf_parameters[\"peak_rm_fit_error\"],\n",
    "    yerr=fdf_parameters[\"peak_pi_error\"],\n",
    "    fmt=\"o\",\n",
    "    lw=1,\n",
    "    color=\"red\",\n",
    "    mfc=\"none\",\n",
    "    label=\"Fitted peak\",\n",
    ")\n",
    "\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"Flux density\",\n",
    "    title=\"Dirty FDF\",\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmclean_results = rmclean.run_rmclean_from_synth(\n",
    "    rm_synth_1d_results=rm_syth_results, auto_mask=10, auto_threshold=0.5\n",
    ")\n",
    "\n",
    "clean_arrs = rmclean_results.fdf_arrs\n",
    "clean_fdf_params = rmclean_results.fdf_parameters\n",
    "clean_params = rmclean_results.clean_parameters\n",
    "\n",
    "\n",
    "phi_arr_radm2 = clean_arrs[\"phi_arr_radm2\"].to_numpy()\n",
    "fdf_dirty_arr = clean_arrs[\"fdf_dirty_complex_arr\"].to_numpy().astype(complex)\n",
    "fdf_clean_arr = clean_arrs[\"fdf_clean_complex_arr\"].to_numpy().astype(complex)\n",
    "fdf_model_arr = clean_arrs[\"fdf_model_complex_arr\"].to_numpy().astype(complex)\n",
    "fdf_residual_arr = clean_arrs[\"fdf_residual_complex_arr\"].to_numpy().astype(complex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_dirty_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Dirty FDF\",\n",
    "    alpha=0.5,\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_clean_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Clean FDF\",\n",
    ")\n",
    "ax.step(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_model_arr),\n",
    "    color=\"tab:red\",\n",
    "    label=\"Model FDF\",\n",
    "    where=\"mid\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi_arr_radm2,\n",
    "    np.abs(fdf_residual_arr),\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Residual FDF\",\n",
    "    alpha=0.5,\n",
    ")\n",
    "\n",
    "ax.errorbar(\n",
    "    clean_fdf_params[\"peak_rm_fit\"],\n",
    "    clean_fdf_params[\"peak_pi_fit\"],\n",
    "    xerr=clean_fdf_params[\"peak_rm_fit_error\"],\n",
    "    yerr=clean_fdf_params[\"peak_pi_error\"],\n",
    "    fmt=\"o\",\n",
    "    lw=1,\n",
    "    color=\"red\",\n",
    "    mfc=\"none\",\n",
    "    label=\"Fitted peak\",\n",
    ")\n",
    "ax.axhline(\n",
    "    y=clean_params[\"threshold\"][0],\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Threshold\",\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "ax.axhline(\n",
    "    y=clean_params[\"mask\"][0],\n",
    "    color=\"tab:orange\",\n",
    "    label=\"Mask\",\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"Flux density\",\n",
    "    title=\"CLEANed FDF\",\n",
    ")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rm-lite",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}