{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1D RM-synthesis"
   ]
  },
  {
   "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 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 generate some synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we'll simulate RACS-all frequency coverage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we make a Faraday simple spectrum with a single RM component. We will use the following parameters:"
   ]
  },
  {
   "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_simple_spectrum(\n",
    "    freq_arr_hz=freq_hz,\n",
    "    frac_pol=frac_pol,\n",
    "    psi0_deg=psi0_deg,\n",
    "    rm_radm2=rm_radm2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(\n",
    "    freq_hz, np.real(complex_data_noiseless), \".\", label=\"Stokes Q\", color=\"tab:red\"\n",
    ")\n",
    "ax.plot(\n",
    "    freq_hz, np.imag(complex_data_noiseless), \".\", label=\"Stokes U\", color=\"tab:blue\"\n",
    ")\n",
    "ax.legend()\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\nu$ / {u.Hz:latex_inline}\",\n",
    "    ylabel=\"Flux density\",\n",
    "    title=\"Stokes Q and U\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can run RM-synthesis by calling `rmsynth.run_rmsynth`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(rmsynth.run_rmsynth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = rmsynth.run_rmsynth(\n",
    "    freq_arr_hz=freq_hz,\n",
    "    complex_pol_arr=complex_data_noiseless,\n",
    "    complex_pol_error=np.zeros_like(complex_data_noiseless),\n",
    "    do_fit_rmsf=True,\n",
    "    n_samples=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output values are Polars dataframes that can be inspected easily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_arrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmsf_arrs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we provided no Stokes $I$ data, the stokes I model will just be unity with 0 error. The `flag_arr` array tells us which channels were not used in RM-synthesis or model fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stokes_i_arrs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also easily visualise the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "x1, x2, y1, y2 = 95, 105, 0.45, 0.55  # subregion of the original image\n",
    "axins = ax.inset_axes(\n",
    "    (0.9, 0.6, 0.4, 0.4), xlim=(x1, x2), ylim=(y1, y2), xticklabels=[], yticklabels=[]\n",
    ")\n",
    "for _ax in [ax, axins]:\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",
    "    # xlim=[50, 150],\n",
    ")\n",
    "ax.indicate_inset_zoom(axins, edgecolor=\"black\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phi2_arr_radm2 = rmsf_arrs[\"phi2_arr_radm2\"].to_numpy()\n",
    "rmsf_arr = rmsf_arrs[\"rmsf_complex_arr\"].to_numpy().astype(complex)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(\n",
    "    phi2_arr_radm2,\n",
    "    rmsf_arr.real,\n",
    "    color=\"tab:red\",\n",
    "    label=\"Stokes Q\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi2_arr_radm2,\n",
    "    rmsf_arr.imag,\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Stokes U\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi2_arr_radm2,\n",
    "    np.abs(rmsf_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Polarized intensity\",\n",
    ")\n",
    "ax.legend()\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"RMSF\",\n",
    "    title=\"RMSF\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets do a more complex example. We'll add noise and a Stokes $I$ spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.fitting import power_law"
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "Now we enable Stokes $I$ model fitting through providing the data, and enabling `fit_order`. If `fit_order<0` an iterative fit will be performed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = 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)\n",
    "    * (rms_noise + rms_noise * 1j),\n",
    "    stokes_i_arr=stokes_i_noisy,\n",
    "    stokes_i_error_arr=np.ones_like(stokes_i_noisy) * rms_noise,\n",
    "    do_fit_rmsf=True,\n",
    "    n_samples=100,\n",
    "    fit_order=-3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf_arrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stokes_i_arrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(freq_hz, stokes_i_noiseless, label=\"Input model\")\n",
    "ax.plot(freq_hz, stokes_i_noisy, \".\", label=\"Noisy data\")\n",
    "ax.plot(\n",
    "    stokes_i_arrs[\"freq_arr_hz\"],\n",
    "    stokes_i_arrs[\"stokes_i_model_arr\"],\n",
    "    \"k--\",\n",
    "    label=\"Fitted model\",\n",
    ")\n",
    "ax.fill_between(\n",
    "    stokes_i_arrs[\"freq_arr_hz\"],\n",
    "    stokes_i_arrs[\"stokes_i_model_arr\"] - stokes_i_arrs[\"stokes_i_model_error\"],\n",
    "    stokes_i_arrs[\"stokes_i_model_arr\"] + stokes_i_arrs[\"stokes_i_model_error\"],\n",
    "    alpha=0.3,\n",
    "    color=\"k\",\n",
    "    label=\"Fitted model error\",\n",
    ")\n",
    "ax.legend()\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\nu$ / {u.Hz:latex_inline}\",\n",
    "    ylabel=\"Flux density\",\n",
    "    title=\"Stokes I\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": [
    "phi2_arr_radm2 = rmsf_arrs[\"phi2_arr_radm2\"].to_numpy()\n",
    "rmsf_arr = rmsf_arrs[\"rmsf_complex_arr\"].to_numpy().astype(complex)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(\n",
    "    phi2_arr_radm2,\n",
    "    rmsf_arr.real,\n",
    "    color=\"tab:red\",\n",
    "    label=\"Stokes Q\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi2_arr_radm2,\n",
    "    rmsf_arr.imag,\n",
    "    color=\"tab:blue\",\n",
    "    label=\"Stokes U\",\n",
    ")\n",
    "ax.plot(\n",
    "    phi2_arr_radm2,\n",
    "    np.abs(rmsf_arr),\n",
    "    color=\"k\",\n",
    "    label=\"Polarized intensity\",\n",
    ")\n",
    "ax.legend()\n",
    "ax.set(\n",
    "    xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n",
    "    ylabel=\"RMSF\",\n",
    "    title=\"RMSF\",\n",
    ")"
   ]
  }
 ],
 "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
}