{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2D 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",
    "\n",
    "plt.rcParams[\"figure.dpi\"] = 150\n",
    "\n",
    "_ = quantity_support()\n",
    "rng = np.random.default_rng(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set up some time-dependent spectra. We'll vary the RM and fractional polarisation as a function of tim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "freqs = np.linspace(1.1, 3.1, 128) * u.GHz\n",
    "freq_hz = freqs.to(u.Hz).value\n",
    "n_times = 1024\n",
    "time_chan = np.arange(n_times)\n",
    "rm_time = np.sin(2 * np.pi * time_chan / n_times) * 100.0\n",
    "frac_pol_time = (-(np.linspace(-1, 1, n_times) ** 2) + 1) * 0.7\n",
    "# psi0_time = rng.uniform(0.0, 180.0, n_times)\n",
    "psi0_time = time_chan % 180"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(8, 6), sharex=True)\n",
    "ax1.plot(time_chan, rm_time)\n",
    "ax2.plot(time_chan, frac_pol_time)\n",
    "ax3.plot(\n",
    "    time_chan,\n",
    "    psi0_time,\n",
    ")\n",
    "ax1.set(\n",
    "    ylabel=f\"RM / ({u.rad / u.m**2:latex_inline})\",\n",
    "    title=\"Input data for RM synthesis\",\n",
    ")\n",
    "ax2.set(\n",
    "    ylabel=\"Fractional Polarisation\",\n",
    ")\n",
    "\n",
    "ax3.set(\n",
    "    xlabel=\"Time Channel\",\n",
    "    ylabel=\"Polaristion angle / deg\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll simulate the spectra and place in a 2D array"
   ]
  },
  {
   "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",
    "dynamic_spectrum = np.empty((len(freqs), n_times), dtype=np.complex128)\n",
    "total_dynamic_spectrum = np.empty((len(freqs), n_times), dtype=np.float64)\n",
    "\n",
    "\n",
    "for time_step, (rm_radm2, frac_pol, psi0_deg) in enumerate(\n",
    "    zip(rm_time, frac_pol_time, psi0_time)\n",
    "):\n",
    "    complex_data_noiseless = faraday_simple_spectrum(\n",
    "        freq_hz,\n",
    "        frac_pol=frac_pol,\n",
    "        psi0_deg=psi0_deg,\n",
    "        rm_radm2=rm_radm2,\n",
    "    )\n",
    "    stokes_i_flux = 1.0\n",
    "    spectral_index = -0.7\n",
    "    rms_noise = 0.5\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",
    "    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",
    "\n",
    "    dynamic_spectrum[:, time_step] = complex_data_noisy\n",
    "    total_dynamic_spectrum[:, time_step] = stokes_i_noisy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 2, figsize=(12, 8), sharex=True, sharey=True)\n",
    "ax1, ax2, ax3, ax4 = axs.flatten()\n",
    "\n",
    "im = ax1.imshow(\n",
    "    total_dynamic_spectrum,\n",
    "    aspect=\"auto\",\n",
    "    origin=\"lower\",\n",
    "    extent=(0, n_times, np.min(freqs), np.max(freqs)),\n",
    ")\n",
    "fig.colorbar(im, ax=ax1)\n",
    "ax1.set(ylabel=\"Frequency / GHz\", title=\"Stokes I\")\n",
    "\n",
    "im = ax2.imshow(\n",
    "    np.real(dynamic_spectrum),\n",
    "    aspect=\"auto\",\n",
    "    origin=\"lower\",\n",
    "    extent=(0, n_times, np.min(freqs), np.max(freqs)),\n",
    "    cmap=\"coolwarm\",\n",
    ")\n",
    "ax2.set(\n",
    "    title=\"Stokes Q\",\n",
    ")\n",
    "fig.colorbar(im, ax=ax2)\n",
    "\n",
    "im = ax3.imshow(\n",
    "    np.imag(dynamic_spectrum),\n",
    "    aspect=\"auto\",\n",
    "    origin=\"lower\",\n",
    "    extent=(0, n_times, np.min(freqs), np.max(freqs)),\n",
    "    cmap=\"coolwarm\",\n",
    ")\n",
    "ax3.set(\n",
    "    title=\"Stokes U\",\n",
    "    xlabel=\"Time step\",\n",
    "    ylabel=\"Frequency / GHz\",\n",
    ")\n",
    "fig.colorbar(im, ax=ax3)\n",
    "\n",
    "im = ax4.imshow(\n",
    "    np.abs(dynamic_spectrum),\n",
    "    aspect=\"auto\",\n",
    "    origin=\"lower\",\n",
    "    extent=(0, n_times, np.min(freqs), np.max(freqs)),\n",
    "    cmap=\"magma\",\n",
    ")\n",
    "fig.colorbar(im, ax=ax4)\n",
    "ax4.set(\n",
    "    xlabel=\"Time step\",\n",
    "    title=\"pI\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To do the RM synthesis, we'll use some of the utility functions directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rm_lite.utils.synthesis import make_phi_arr, rmsynth_nufft\n",
    "\n",
    "help(rmsynth_nufft)\n",
    "help(make_phi_arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phis = make_phi_arr(500, 0.1)\n",
    "lam_sq_0_m2 = float(np.mean(freq_to_lambda2(freq_hz)))\n",
    "\n",
    "fdf_spectrum = rmsynth_nufft(\n",
    "    complex_pol_arr=dynamic_spectrum,\n",
    "    lambda_sq_arr_m2=freq_to_lambda2(freq_hz),\n",
    "    phi_arr_radm2=phis,\n",
    "    weight_arr=np.ones_like(freq_hz),\n",
    "    lam_sq_0_m2=lam_sq_0_m2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.imshow(\n",
    "    np.abs(fdf_spectrum),\n",
    "    # aspect=\"auto\",\n",
    "    origin=\"lower\",\n",
    "    extent=(0, n_times, np.min(phis), np.max(phis)),\n",
    ")\n",
    "ax.set(\n",
    "    xlabel=\"Time step\",\n",
    "    ylabel=f\"Faraday depth / ({u.rad / u.m**2:latex_inline})\",\n",
    "    title=\"Dynamic spectrum\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's recover the PI and RM from the Faraday spectrum. Taking the mean will not perform well due to bandwidth depolarisation, but RM-sythnesis gives us the full-bandwidth sensitivity with a coherent sum. \n",
    "\n",
    "Note that at low SNR Ricean bias becomes significant. Further, our uncertainty in RM also goes up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peak_pi_spectrum = np.max(np.abs(fdf_spectrum), axis=0)\n",
    "\n",
    "max_pixels = np.argmax(np.abs(fdf_spectrum), axis=0)\n",
    "\n",
    "peak_rm_spectrum = phis[max_pixels]\n",
    "peak_q_spectrum = np.real(fdf_spectrum)[max_pixels, np.arange(fdf_spectrum.shape[1])]\n",
    "peak_u_spectrum = np.imag(fdf_spectrum)[max_pixels, np.arange(fdf_spectrum.shape[1])]\n",
    "peak_pa_spectrum = np.rad2deg(0.5 * np.arctan2(peak_u_spectrum, peak_q_spectrum)) % 180\n",
    "peak_pa_spectrum_detrot = (\n",
    "    np.rad2deg(np.deg2rad(peak_pa_spectrum) - (peak_rm_spectrum * lam_sq_0_m2)) % 180\n",
    ")\n",
    "\n",
    "\n",
    "fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1, sharex=True, figsize=(8, 12))\n",
    "ax1.plot(time_chan, peak_pi_spectrum, label=\"measured - RM synthesis\")\n",
    "ax1.plot(time_chan, frac_pol_time, label=\"input\")\n",
    "\n",
    "ax2.plot(\n",
    "    time_chan,\n",
    "    peak_q_spectrum,\n",
    "    c=\"tab:blue\",\n",
    "    label=\"Stokes Q - RM synthesis\",\n",
    ")\n",
    "ax2.plot(\n",
    "    time_chan,\n",
    "    peak_u_spectrum,\n",
    "    c=\"tab:red\",\n",
    "    label=\"Stokes U - RM synthesis\",\n",
    ")\n",
    "ax2.set(\n",
    "    ylabel=\"Stokes Q, U\",\n",
    ")\n",
    "ax2.legend()\n",
    "\n",
    "ax3.plot(\n",
    "    time_chan,\n",
    "    peak_pa_spectrum,\n",
    "    label=\"measured - RM synthesis\",\n",
    ")\n",
    "ax3.plot(\n",
    "    time_chan,\n",
    "    peak_pa_spectrum_detrot,\n",
    "    label=\"detrotated\",\n",
    ")\n",
    "ax3.plot(\n",
    "    time_chan,\n",
    "    psi0_time,\n",
    "    label=\"input\",\n",
    ")\n",
    "ax3.legend()\n",
    "ax3.set(\n",
    "    ylabel=\"Polarisation angle / deg\",\n",
    ")\n",
    "\n",
    "ax1.legend()\n",
    "\n",
    "ax4.plot(time_chan, peak_rm_spectrum, label=\"measured\")\n",
    "ax4.plot(time_chan, rm_time, label=\"input\")\n",
    "ax4.legend()\n",
    "\n",
    "ax4.set(\n",
    "    xlabel=\"Time step\",\n",
    "    ylabel=f\"RM / ({u.rad / u.m**2:latex_inline})\",\n",
    "    title=\"Peak RM spectrum\",\n",
    ")\n",
    "ax1.set(\n",
    "    ylabel=\"Peak polarized intensity\",\n",
    "    title=\"Peak polarized intensity spectrum\",\n",
    ")"
   ]
  },
  {
   "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
}