{
 "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)"
   ]
  },
  {
   "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, strict=False)\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.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",
    "    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",
    "\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=float(np.mean(freq_to_lambda2(freq_hz))),\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. We'll compare to just taking the mean across frequency. Taking the mean will not perform well due to bandwidth depolarisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peak_pi_spectrum = np.max(np.abs(fdf_spectrum), axis=0)\n",
    "max_pixels = np.argmax(np.abs(fdf_spectrum), axis=0)\n",
    "\n",
    "peak_rm_spectrum = phis[max_pixels]\n",
    "\n",
    "\n",
    "## This is the *wrong* thing to do:\n",
    "stokes_q, stokes_u = np.real(dynamic_spectrum), np.imag(dynamic_spectrum)\n",
    "stokes_q_mean = np.mean(stokes_q, axis=0)\n",
    "stokes_u_mean = np.mean(stokes_u, axis=0)\n",
    "pi_mean = np.hypot(stokes_q_mean, stokes_u_mean)\n",
    "\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(8, 8))\n",
    "ax1.plot(time_chan, peak_pi_spectrum, label=\"measured - RM synthesis\")\n",
    "ax1.plot(time_chan, frac_pol_time, label=\"input\")\n",
    "ax1.plot(time_chan, pi_mean, label=\"'measured' - mean\")\n",
    "\n",
    "ax1.legend()\n",
    "\n",
    "ax2.plot(time_chan, peak_rm_spectrum, label=\"measured\")\n",
    "ax2.plot(time_chan, rm_time, label=\"input\")\n",
    "ax2.legend()\n",
    "\n",
    "ax2.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
}