{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Polarized observations\n",
    "\n",
    "This tutorial covers working with polarized instrument and maps, and recovering polarized maps from observations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "We start with a normal instrument, and create two orthogonally polarized copies of each detector by setting ``polarized: True`` in the ``Array`` config:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import maria\n",
    "from maria.instrument import Band\n",
    "\n",
    "f090 = Band(\n",
    "    center=90e9,  # in Hz\n",
    "    width=20e9,  # in Hz\n",
    "    NET_RJ=40e-6,  # in K sqrt(s)\n",
    "    knee=1e0,    # in Hz\n",
    "    gain_error=5e-2)\n",
    "\n",
    "f150 = Band(\n",
    "    center=150e9, \n",
    "    width=30e9, \n",
    "    NET_RJ=60e-6, \n",
    "    knee=1e0, \n",
    "    gain_error=5e-2)\n",
    "\n",
    "array = {\"field_of_view\": 0.5, \n",
    "         \"shape\": \"circle\", \n",
    "         \"beam_spacing\": 1.5,\n",
    "         \"primary_size\": 10, \n",
    "         \"polarized\": True,\n",
    "         \"bands\": [f090, f150]}\n",
    "\n",
    "instrument = maria.get_instrument(array=array)\n",
    "instrument.arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "We can see the resulting polarization footprint in the instrument plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(instrument)\n",
    "instrument.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "Let's observe the use the Einstein map, which has a faint polarization signature underneath the unpolarized signal of Einstein's face. Remember that all maps are five dimensional (stokes, frequency, time, y, x); this map has four channels in the stokes dimensions (the I, Q, U, and V [Stokes parameters](https://en.wikipedia.org/wiki/Stokes_parameters)). We can plot all the channels by giving ``plot`` a shaped set of stokes parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from maria.io import fetch\n",
    "\n",
    "input_map = maria.map.load(fetch(\"maps/einstein.h5\"))\n",
    "input_map.plot(stokes=[[\"I\", \"Q\"], \n",
    "                       [\"U\", \"V\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plan = maria.Plan(\n",
    "    start_time=\"2024-08-06T03:00:00\",\n",
    "    scan_pattern=\"daisy\",\n",
    "    scan_options={\"radius\": 0.5, \"speed\": 0.1},  # in degrees\n",
    "    duration=900,  # in seconds\n",
    "    sample_rate=50,  # in Hz\n",
    "    scan_center=(0, -23),\n",
    "    frame=\"ra_dec\")\n",
    "\n",
    "print(plan)\n",
    "plan.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = maria.Simulation(\n",
    "    instrument,\n",
    "    plan=plan,\n",
    "    site=\"llano_de_chajnantor\",\n",
    "    atmosphere=\"2d\",\n",
    "    atmosphere_kwargs={\"weather\": {\"pwv\": 0.5}},\n",
    "    map=input_map\n",
    ")\n",
    "\n",
    "print(sim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tod = sim.run()\n",
    "\n",
    "print(tod)\n",
    "tod.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from maria.mappers import BinMapper\n",
    "\n",
    "mapper = BinMapper(\n",
    "    center=(0, -23),\n",
    "    stokes=\"IQU\",\n",
    "    frame=\"ra_dec\",\n",
    "    width=1.0,\n",
    "    height=1.0,\n",
    "    resolution=1.0 / 256,\n",
    "    tod_preprocessing={\n",
    "        \"window\": {\"name\": \"tukey\", \"kwargs\": {\"alpha\": 0.1}},\n",
    "        \"remove_spline\": {\"knot_spacing\": 30, \"remove_el_gradient\": True},\n",
    "        \"remove_modes\": {\"modes_to_remove\": [0]},\n",
    "    },\n",
    "    map_postprocessing={\n",
    "        \"gaussian_filter\": {\"sigma\": 1},\n",
    "        # \"median_filter\": {\"size\": 1},\n",
    "    },\n",
    "    units=\"mK_RJ\",\n",
    ")\n",
    "\n",
    "mapper.add_tods(tod)\n",
    "\n",
    "output_map = mapper.run()\n",
    "\n",
    "print(output_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "Note that we can't see any of the circular polarization, since our instrument isn't sensitive to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map.plot(stokes=[\"I\", \"Q\", \"U\"], nu_index=[[0], [1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}