02_atmosphere.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Atmosphere\n",
    "> Load and scale the maps of the noise from the atmosphere\n",
    "- toc: True\n",
    "\n",
    "* We have 8 independent realizations\n",
    "* They scale with the number of telescopes, i.e. it doesn't make any difference if there are 1 or 3 channels in the same telescope, they do not help beat the atmosphere\n",
    "* They were simulated for 10 days at 100% efficiency at Chile, 46.29% (SAT) and 37.23% (LAT) at Pole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp atmosphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "import healpy as hp\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import logging as log\n",
    "\n",
    "from s4_design_sim_tool.core import (\n",
    "    get_telescope,\n",
    "    get_observing_efficiency,\n",
    "    base_folder,\n",
    "    simulations_observing_efficiency,\n",
    "    mapmaking_naming,\n",
    "    read_instrument_model,\n",
    ")\n",
    "\n",
    "from s4_design_sim_tool.noise import get_thinfp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import toml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = toml.load(\"s4_design.toml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log.basicConfig(level=log.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exports\n",
    "\n",
    "\n",
    "def get_telecope_years(config, site, channel):\n",
    "    \"\"\"Compute the number of telescope/years in the CMB-S4 configuration\n",
    "\n",
    "    config_telescopes : dict\n",
    "        CMB-S4 telescopes configuration,\n",
    "        generally loaded from a TOML file\n",
    "    site : str\n",
    "        'Pole' or 'Chile', case doesn't matter\n",
    "    channel : str\n",
    "        Channel tag, e.g. 'MFHS1'\n",
    "    \"\"\"\n",
    "    telescope_years = 0\n",
    "    for telescope_name, telescope_config in config[\"telescopes\"][\n",
    "        get_telescope(channel)\n",
    "    ].items():\n",
    "        if telescope_config[\"site\"].lower() == site.lower():\n",
    "            has_band = telescope_config.get(channel[:-1], 0) > 0\n",
    "            telescope_years += has_band * telescope_config.get(\n",
    "                \"years\", config[\"experiment\"][\"total_experiment_length_years\"]\n",
    "            )\n",
    "    return telescope_years"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s4 = read_instrument_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tail -30 s4_design.toml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute telescope/years for the reference design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for site in [\"Pole\", \"Chile\"]:\n",
    "    for channel in s4:\n",
    "        telescope_years = get_telecope_years(config, site, channel[\"band\"])\n",
    "        print(site, channel[\"band\"], telescope_years)\n",
    "        telescope = get_telescope(channel[\"band\"])\n",
    "        if site == \"Chile\":\n",
    "            if telescope == \"SAT\":\n",
    "                assert telescope_years == 0, \"All SAT at Pole\"\n",
    "            elif channel[\"band\"].startswith(\"ULFL\"):\n",
    "                assert telescope_years == 0, \"No ULFL in Chile\"              \n",
    "            else:\n",
    "                if \"P\" in channel[\"band\"]:\n",
    "                    assert telescope_years == 0, \"Pole only LAT\"\n",
    "                else:                   \n",
    "                    assert telescope_years == 14, \"2 LAT each band\"\n",
    "        if site == \"Pole\":\n",
    "            if telescope == \"SAT\":\n",
    "                assert telescope_years == 14, \"2 SAT telescopes each band\"\n",
    "            else:\n",
    "                if \"P\" in channel[\"band\"]:\n",
    "                    assert telescope_years == 7, \"1 LAT with all 4 bands\"\n",
    "                else:\n",
    "                    assert telescope_years == 0, \"Chile only LAT\"\n",
    "    print(30 * \"=\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exports\n",
    "\n",
    "\n",
    "def load_atmosphere(config, site, channel, realization=0, raw=False):\n",
    "    \"\"\"Load foreground maps for a channel\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : dict\n",
    "        CMB-S4 configuration,\n",
    "        generally loaded from a TOML file\n",
    "    site : str\n",
    "        'Pole' or 'Chile', case doesn't matter\n",
    "    channel : str\n",
    "        Channel tag, e.g. 'MFHS1'\n",
    "    realization : int\n",
    "        Choose one of the available 8 realizations\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    output_map : numpy array\n",
    "        Output map with all emissions combined, uses nan for missing pixels\n",
    "    \"\"\"\n",
    "\n",
    "    telescope = get_telescope(channel)\n",
    "    channel_noP = channel.replace(\"P\", \"\")\n",
    "\n",
    "    map_filename = (\n",
    "        Path(f\"{realization:08d}\")\n",
    "        / f\"{site.lower()}_atmosphere_{telescope}_{channel_noP}_{mapmaking_naming[telescope]}\"\n",
    "    )\n",
    "    log.info(f\"Reading {map_filename}\")\n",
    "    atmosphere_map = hp.read_map(\n",
    "        Path(base_folder) / map_filename, (0, 1, 2), dtype=None, verbose=False\n",
    "    )\n",
    "    if raw:\n",
    "        atmosphere_map[atmosphere_map == 0] = hp.UNSEEN\n",
    "        return atmosphere_map\n",
    "\n",
    "    atmosphere_map[atmosphere_map == hp.UNSEEN] = np.nan\n",
    "    atmosphere_map[atmosphere_map == 0] = np.nan\n",
    "\n",
    "    # input map is 10 days at 100% efficiency\n",
    "    atmosphere_map *= np.sqrt(\n",
    "        10\n",
    "        * simulations_observing_efficiency[site.lower()][channel]\n",
    "        / (\n",
    "            365.25\n",
    "            * get_observing_efficiency(\n",
    "                config[\"experiment\"][\"observing_efficiency\"], site, telescope, channel\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "    atmosphere_map /= np.sqrt(get_telecope_years(config, site, channel))\n",
    "\n",
    "    atmosphere_map[0] *= config[\"experiment\"].get(\"atmosphere_scaling_T\", 1)\n",
    "    atmosphere_map[1:] *= config[\"experiment\"].get(\"atmosphere_scaling_P\", 1)\n",
    "    atmosphere_map[1:] /= np.sqrt(get_thinfp(channel))\n",
    "\n",
    "    return atmosphere_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run on a channel and plot results\n",
    "\n",
    "Available atmosphere maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filenames = !ls $base_folder/00000000/*atmo*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path\n",
    "for f in map(os.path.basename, filenames):\n",
    "    print(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "channel = \"HFS1\"\n",
    "site = \"Pole\"\n",
    "telescope = get_telescope(channel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_map = hp.read_map(\n",
    "    base_folder + \"/00000000/\" + \\\n",
    "    f\"{site.lower()}_atmosphere_{telescope}_{channel}_telescope_all_time_all_filtered.fits.gz\"\n",
    ", (0,1,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_map[input_map == 0] = hp.UNSEEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map = load_atmosphere(config, site, channel, realization=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.nanmin(output_map), np.nanmax(output_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map[np.isnan(output_map)] = hp.UNSEEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare mask\n",
    "np.testing.assert_allclose(input_map == hp.UNSEEN, output_map == hp.UNSEEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observing_efficiency = get_observing_efficiency(\n",
    "                config[\"experiment\"][\"observing_efficiency\"], site, telescope, channel\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_map[input_map == hp.UNSEEN] = np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expected_map = input_map * np.sqrt(10 * simulations_observing_efficiency[site.lower()][channel] / (7 * 365.25 * observing_efficiency) / 2)\n",
    "expected_map[1:] /= np.sqrt(get_thinfp(channel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(expected_map[np.logical_not(np.isnan(expected_map))],\n",
    "    output_map[output_map != hp.UNSEEN], rtol=1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(output_map[0], min=-1e-5, max=1e-5, unit=\"K\", title=\"Atmosphere I\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are no systematics in the half-wave plate or bandpass mismatch, so almost all the atmosphere signal is rejected in polarization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(output_map[1], min=-1e-9, max=1e-9, unit=\"K\", title=\"Atmosphere Q\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(output_map[2], min=-1e-9, max=1e-9, unit=\"K\", title=\"Atmosphere U\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test scaling\n",
    "\n",
    "We want to be able to freely scale the atmosphere noise differently for Temperature and for Polarization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config[\"experiment\"][\"atmosphere_scaling_T\"] = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config[\"experiment\"][\"atmosphere_scaling_P\"] = 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map = load_atmosphere(config, site, channel, realization=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map[np.isnan(output_map)] = hp.UNSEEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(config[\"experiment\"][\"atmosphere_scaling_T\"] * expected_map[0, np.logical_not(np.isnan(expected_map[0]))],\n",
    "    output_map[0][output_map[0] != hp.UNSEEN], rtol=1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(config[\"experiment\"][\"atmosphere_scaling_P\"]  * expected_map[1:][np.logical_not(np.isnan(expected_map[1:]))],\n",
    "    output_map[1:][output_map[1:] != hp.UNSEEN], rtol=1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "condajupynersc",
   "language": "python",
   "name": "condajupynersc"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}