{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Compute a transfer function\n",
    "\n",
    "Atmospheric emission can dominate the noise budget and must be removed from the time-ordered data before mapping. This filtering inevitably projects out correlated sky modes, suppressing signal at large angular scales. The spatial transfer function $T(u)$ provides a direct estimate of this scale-dependent signal loss. `maria` includes a built-in function to compute this, using the cross-spectrum between the output and the known input map:\n",
    "\n",
    "$$T(u) = \\frac{\\mathrm{Re}\\langle \\tilde{s}_{\\rm in}^*(u)\\,\\tilde{s}_{\\rm out}(u)\\rangle}{\\langle|\\tilde{s}_{\\rm in}(u)|^2\\rangle}$$\n",
    "\n",
    "Please note that the cross-spectrum is used instead of the power ratio $P_{\\rm out}/P_{\\rm in}$ because noise in the output is uncorrelated with the input and averages to zero, leaving an unbiased estimate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Input sky\n",
    "\n",
    "In this tutorial, we simulate a two-frequency observation of a galaxy cluster at 150 and 270 GHz and recover a map with the `BinMapper`. We load the same map at both frequencies — any difference in the recovered signal will come from the instrument and the filtering alone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import maria\n",
    "from maria.io import fetch\n",
    "\n",
    "map_filename = fetch(\"maps/cluster2.fits\")\n",
    "\n",
    "nu = np.array([150e9, 270e9])\n",
    "\n",
    "m1 = maria.map.load(filename=map_filename, nu=nu[0], width=20.00/60.00)\n",
    "m2 = maria.map.load(filename=map_filename, nu=nu[1], width=20.00/60.00)\n",
    "\n",
    "# Arbitrary signal boost to improve the SNR of the output map\n",
    "m1.data *= -5.0e3\n",
    "m2.data *= -5.0e3\n",
    "\n",
    "maps = maria.map.concatenate([m1, m2], dim=\"nu\")\n",
    "print(maps)\n",
    "maps.to(\"K_RJ\").plot(slices=dict(nu=[0, 1]), cmap=\"RdYlBu_r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13de8738",
   "metadata": {},
   "source": [
    "## Instrument and observation plan\n",
    "\n",
    "We use TolTEC on the LMT, keeping only the 150 and 270 GHz arrays and dropping the 220 GHz one for efficiency. Because the beam is diffraction-limited, the 270 GHz array resolves angular scales roughly 1.8× finer than the 150 GHz array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from maria.instrument import get_instrument_config, Instrument\n",
    "\n",
    "config = get_instrument_config(\"TolTEC\")\n",
    "config[\"arrays\"] = {k: config[\"arrays\"][k] for k in [\"array-1\", \"array-3\"]}\n",
    "instrument = Instrument.from_config(config)\n",
    "\n",
    "print(instrument)\n",
    "instrument.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from maria import Planner\n",
    "\n",
    "planner = Planner(\n",
    "    target=maps,\n",
    "    start_time=\"2024-01-01T22:00:00\",\n",
    "    site=\"llano_de_chajnantor\",\n",
    "    constraints={\"el\": (60, 90)},\n",
    ")\n",
    "\n",
    "plans = planner.generate_plans(\n",
    "    total_duration=0.1 * 3600,\n",
    "    max_chunk_duration=3600,\n",
    "    scan_options={\"radius\": 6.5 / 60.0, \"miss_factor\": 0.3},\n",
    ")\n",
    "\n",
    "print(plans)\n",
    "plans[0].plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Simulation\n",
    "\n",
    "Passing `map=maps` attaches the ground-truth sky to each output TOD so the mapper can propagate it to the output map automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = maria.Simulation(\n",
    "    instrument,\n",
    "    plans=plans,\n",
    "    site=\"llano_de_chajnantor\",\n",
    "    map=maps,\n",
    "    atmosphere=\"2d\",\n",
    ")\n",
    "\n",
    "print(sim)\n",
    "\n",
    "tods = sim.run()\n",
    "tods[0].plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Mapmaking\n",
    "\n",
    "Common-mode subtraction (`remove_modes`) and per-detector spline removal (`remove_spline`) suppress large-scale correlated signal. Both operations remove power at long time scales, which maps to large angular scales on the sky, i.e., where $T$ will fall below 1.\n",
    "\n",
    "> **Note:** `map_postprocessing` is intentionally left empty here. Any map-level post-processing (smoothing, filtering, etc.) applied after binning would itself suppress or modify signal in the output map, and its effect would be absorbed into the measured $T(u)$. To isolate the transfer function of the TOD filtering alone, always set `map_postprocessing={}` when computing transfer functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from maria.mappers import BinMapper\n",
    "\n",
    "mapper = BinMapper(\n",
    "    tods=tods,\n",
    "    units=\"uK_RJ\",\n",
    "    stokes=\"I\",\n",
    "    resolution=maps.resolution,\n",
    "    tod_preprocessing={\n",
    "        \"remove_modes\": {\"modes_to_remove\": 1},\n",
    "        \"remove_spline\": {\"knot_spacing\": 60, \"remove_el_gradient\": True},\n",
    "    },\n",
    "    map_postprocessing={},\n",
    ")\n",
    "\n",
    "output_map = mapper.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map.plot(slices=dict(nu=[0, 1]), cmap=\"RdYlBu_r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Transfer function\n",
    "\n",
    "To compute the transfer function, call `transfer_function()` on the output map. This will return a `TransferFunction` object containing the transfer function $T(u)$ and the corresponding spatial frequencies $u$. The transfer function can be plotted to visualize the scale-dependent signal loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "001f2d38",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf = output_map.transfer_function(window=True)\n",
    "print(tf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aae0aa07",
   "metadata": {},
   "source": [
    "To visualize the transfer function, it is possible to use the built-in `plot()` method. Solid curves show the measured $T(u)$, while dashed curves show the Gaussian beam per channel. The large-scale rolloff reflects the TOD filtering above, while the small-scale rolloff tracks the beam."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.plot(x_unit=\"arcmin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "Finally, individual channels can be selected with `slices=dict(nu=[...])`, consistent with how `slices` is used in `.plot()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.plot(slices=dict(nu=[0]), x_unit=\"arcmin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "windowing-options",
   "metadata": {},
   "source": [
    "## Apodization window\n",
    "\n",
    "Before the FFT, `transfer_function()` multiplies both maps by a 2D apodization window to suppress spectral leakage from sharp map edges. The `window` argument controls which window is applied:\n",
    "\n",
    "| Value | Shape | When to use |\n",
    "|-------|-------|-------------|\n",
    "| `\"hann\"` (default) | Full cosine taper to zero at both edges | Strongest leakage suppression; best when the field is much larger than the scales of interest, since it discards edge signal |\n",
    "| `\"tukey\"` | Cosine taper on outer `taper` fraction, central region = 1 | Mild apodization that preserves most of the map signal and substantially reduces edge leakage |\n",
    "| `np.ndarray` of shape `(ny, nx)` | User-supplied 2D window | Full control — e.g. to mask point sources or weight by the hit-count map |\n",
    "| `False` / `None` | No window (rectangular) | Maximum large-scale sensitivity; avoid unless the map has periodic or artificially zero-padded boundaries |\n",
    "\n",
    "The `taper` parameter (default 0.1) sets the fraction of each axis tapered by the Tukey roll-off. Increasing it toward 1 makes the Tukey window approach a Hann window.\n",
    "\n",
    "> **Tip:** for typical BinMapper output maps, where valid signal extends close to the boundary, `\"tukey\"` with a small `taper` is preferred over `\"hann\"`. The full Hann taper downweights map edges heavily, effectively shrinking the usable field and worsening large-scale mode recovery."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "windowing-comparison",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True)\n",
    "\n",
    "windows = [\n",
    "    (\"tukey\",  dict(window=\"tukey\", taper=0.1)),\n",
    "    (\"hann\",   dict(window=\"hann\")),\n",
    "    (\"none\",   dict(window=False)),\n",
    "]\n",
    "\n",
    "for i, (label, kwargs) in enumerate(windows):\n",
    "    tf_w = output_map.transfer_function(slices=dict(nu=[0]), **kwargs)\n",
    "    tf_w.plot(ax=ax, x_unit=\"arcmin\", add_beam=False)\n",
    "    ax.lines[-1].set_label(label)\n",
    "    ax.lines[-1].set_color(f\"C{i}\")\n",
    "\n",
    "ax.legend(frameon=False, fontsize=9)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "maria",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}