{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use `get_rm_from_skycoord` and compare with values from RMextract"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example script shows how to get ionospheric rotation measures for a give source,\n",
    "observer location and array of times. This script uses the default ionospheric settings,\n",
    "which are good for most purposes."
   ]
  },
  {
   "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.coordinates import EarthLocation, SkyCoord\n",
    "from astropy.time import Time\n",
    "from astropy.visualization import quantity_support, time_support\n",
    "from spinifex import get_rm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Required to load local data for example - not needed for normal use\n",
    "from importlib import resources\n",
    "\n",
    "with resources.as_file(resources.files(\"spinifex.data.tests\")) as datapath:\n",
    "    spinifex_data = datapath\n",
    "###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = Time(\"2020-01-08T01:13:00\") + np.arange(10) * 1 * u.hr\n",
    "# create source object\n",
    "source = SkyCoord(ra=350.85 * u.deg, dec=58.815 * u.deg)\n",
    "# create Earth Location\n",
    "lon = 6.367 * u.deg\n",
    "lat = 52.833 * u.deg\n",
    "dwingeloo = EarthLocation(lon=lon, lat=lat, height=0 * u.km)\n",
    "# get rotation measures\n",
    "rm = get_rm.get_rm_from_skycoord(\n",
    "    loc=dwingeloo,\n",
    "    times=times,\n",
    "    source=source,\n",
    "    # We set these options to use the data packaged with spinifex\n",
    "    # Unsetting them will cause the function to download the data from the internet\n",
    "    output_directory=spinifex_data,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the values from the same server, source and time from RMextract, install RMextract and run the following lines:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "have_rmextract = False\n",
    "\n",
    "if have_rmextract:\n",
    "    import RMextract\n",
    "\n",
    "    starttime = times[0].mjd * 24 * 3600  # MJD in seconds (casacore definition)\n",
    "    endtime = times[-1].mjd * 24 * 3600\n",
    "    statpos = [dwingeloo.get_itrs().cartesian.xyz.to(u.m).value]\n",
    "    pointing = [source.ra.rad, source.dec.rad]  # 3C196 Ra, Dec in radians\n",
    "    rm_dict = RMextract.getRM.getRM(\n",
    "        ionexPath=\"./ionex_data/\",\n",
    "        radec=pointing,\n",
    "        timestep=3600,\n",
    "        timerange=[starttime, endtime],\n",
    "        stat_positions=statpos,\n",
    "        earth_rot=1,\n",
    "        prefix=\"uqrg\",\n",
    "        server=\"http://chapman.upc.es\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we directly take the values we got from the above code, to avoid installing RMextract"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rm_dict = {}\n",
    "rm_dict[\"RM\"] = {}\n",
    "rm_dict[\"RM\"][\"st1\"] = np.array(\n",
    "    [\n",
    "        [\n",
    "            *[0.10242756, 0.08804738, 0.07332737, 0.08251355],\n",
    "            *[0.07816817, 0.08137425, 0.11095578, 0.15794444],\n",
    "            *[0.26828872, 0.36951645, 0.44763038, 0.49045776],\n",
    "        ]\n",
    "    ]\n",
    ").T\n",
    "rm_dict[\"times\"] = np.array(\n",
    "    [\n",
    "        *[5.08515918e09, 5.08516278e09, 5.08516638e09, 5.08516998e09],\n",
    "        *[5.08517358e09, 5.08517718e09, 5.08518078e09, 5.08518438e09],\n",
    "        *[5.08518798e09, 5.08519158e09, 5.08519518e09, 5.08519878e09],\n",
    "    ]\n",
    ")\n",
    "rm_dict[\"azimuth\"] = {}\n",
    "rm_dict[\"azimuth\"][\"st1\"] = np.array(\n",
    "    [\n",
    "        *[-0.5065002, -0.37364739, -0.23410709, -0.09012949],\n",
    "        *[0.05561559, 0.20026593, 0.34108016, 0.475697],\n",
    "        *[0.60222255, 0.7190818, 0.82458167, 0.9159993],\n",
    "    ]\n",
    ")\n",
    "rm_dict[\"elev\"] = {}\n",
    "rm_dict[\"elev\"][\"st1\"] = np.array(\n",
    "    [\n",
    "        *[0.52492472, 0.45731976, 0.40983145, 0.38421338],\n",
    "        *[0.38147449, 0.40172543, 0.44416086, 0.50718446],\n",
    "        *[0.58862976, 0.68600301, 0.79668206, 0.91802653],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rm_rme = rm_dict[\"RM\"][\"st1\"]\n",
    "times_rme = Time(rm_dict[\"times\"] / (3600 * 24), format=\"mjd\")\n",
    "az_rme = np.degrees(rm_dict[\"azimuth\"][\"st1\"])\n",
    "elev_rme = np.degrees(rm_dict[\"elev\"][\"st1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print to screen\n",
    "rotation_measures = rm.rm\n",
    "rm_times = rm.times\n",
    "elevations = rm.elevation\n",
    "azimuths = rm.azimuth\n",
    "print(\"time      RM (rad/lambda^2)      azimuth (deg)      elevation (deg)\")\n",
    "for myrm, tm, az, el in zip(rotation_measures, rm_times, azimuths, elevations):\n",
    "    print(f\"{tm.isot} {myrm} {az:3.2f} {el:2.2f}\")\n",
    "print(\"Values from RMextract:\")\n",
    "for myrm, tm, az, el in zip(rm_rme, times_rme, az_rme, elev_rme):\n",
    "    print(f\"{tm.isot} {myrm} {az:3.2f} {el:2.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with time_support(), quantity_support():\n",
    "    fig, ax = plt.subplots()\n",
    "    im = ax.errorbar(rm.times.datetime, rm.rm, rm.rm_error, fmt=\".\", label=\"spinifex\")\n",
    "    im2 = ax.plot(times_rme.datetime, rm_rme, \"o-\", label=\"RMextract\")\n",
    "    ax.set_ylabel(\"RM / rad/m$^2$\")\n",
    "    ax.legend()\n",
    "    plt.xticks(rotation=20)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}