{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum likelihood mapmaking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import maria\n",
    "from maria.io import fetch\n",
    "\n",
    "input_map = maria.map.load(fetch(\"maps/cluster2.fits\"), nu=150e9)\n",
    "input_map.data *= 5e1 \n",
    "\n",
    "input_map[..., 256:-256, 256:-256].plot(cmap=\"cmb\")\n",
    "print(input_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from maria import Planner\n",
    "\n",
    "planner = Planner(target=input_map, site=\"cerro_toco\", constraints={\"el\": (70, 90)})\n",
    "plans = planner.generate_plans(total_duration=900, \n",
    "                               max_chunk_duration=900, \n",
    "                               sample_rate=50,\n",
    "                               scan_options={\"radius\": input_map.width.deg / 2})\n",
    "\n",
    "plans[0].plot()\n",
    "print(plans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import maria\n",
    "from maria.instrument import Band\n",
    "\n",
    "f150 = Band(\n",
    "    center=150e9, \n",
    "    width=30e9, \n",
    "    NET_RJ=30e-6, \n",
    "    knee=1e1, \n",
    "    gain_error=2e-2)\n",
    "\n",
    "array = {\"field_of_view\": 0.1, \n",
    "         \"beam_spacing\": 1.2,\n",
    "         \"primary_size\": 50, \n",
    "         \"shape\": \"circle\",\n",
    "         \"bands\": [f150]}\n",
    "\n",
    "instrument = maria.get_instrument(array=array)\n",
    "\n",
    "print(instrument)\n",
    "instrument.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = maria.Simulation(\n",
    "    instrument,\n",
    "    plans=plans,\n",
    "    site=\"cerro_toco\",\n",
    "    map=input_map,\n",
    "    atmosphere=\"2d\",\n",
    "    atmosphere_kwargs={\"weather\": {\"pwv\": 1.0}},\n",
    ")\n",
    "\n",
    "print(sim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tods = sim.run()\n",
    "tods[0].plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from maria.mappers import MaximumLikelihoodMapper\n",
    "\n",
    "ml_mapper = MaximumLikelihoodMapper(tods=tods, \n",
    "                                    width=0.8*input_map.width.deg,\n",
    "                                    height=0.8*input_map.height.deg,\n",
    "                                    units=\"Jy/pixel\")\n",
    "print(f\"{ml_mapper.loss() = }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initial map is a \"guess\" constructing by heavily filtering the input TODs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ml_mapper.map)\n",
    "ml_mapper.map.plot(cmap=\"cmb\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To fit the map we run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_mapper.fit(epochs=4, steps_per_epoch=32, lr=1e-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which gives us an improved map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_mapper.map.plot(cmap=\"cmb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "that will improve more as it continues to fit."
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}