{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate PySM 3 dust templates\n"
   ],
   "id": "2ed068588048"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import healpy as hp\n",
    "from pathlib import Path"
   ],
   "id": "48cfaae0621d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = Path(\"production-data\") / \"dust_gnilc\""
   ],
   "id": "47e6bd8f0b59"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir = output_dir / \"raw\""
   ],
   "id": "7b9eb1d70fb1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "output_nside = 2048"
   ],
   "id": "dccbdbc48b71"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_lmax = int(min(2.5 * output_nside, 8192 * 2))"
   ],
   "id": "950b1909e84c"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Large scales"
   ],
   "id": "53af8b33c36a"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm_log_pol_tens_large_scale = hp.read_alm(\n",
    "    datadir\n",
    "    / \"gnilc_dust_largescale_template_logpoltens_alm_nside2048_lmax1024_complex64.fits.gz\",\n",
    "    hdu=(1, 2, 3),\n",
    ")"
   ],
   "id": "4e48f1c05433"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_large_scale = hp.alm2map(\n",
    "    alm_log_pol_tens_large_scale.astype(np.complex128), nside=output_nside\n",
    ")"
   ],
   "id": "98eb63a9250e"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_large_scale"
   ],
   "id": "7908e4405531"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small scales modulation"
   ],
   "id": "8ce22a074df8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modulate_alm = {\n",
    "    k: hp.read_alm(datadir / f\"gnilc_dust_{k}_modulation_alms_lmax768.fits.gz\").astype(\n",
    "        np.complex128\n",
    "    )\n",
    "    for k in [\"temperature\", \"polarization\"]\n",
    "}"
   ],
   "id": "80e047e38757"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small scales"
   ],
   "id": "c15b62f4b866"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_small_scale = hp.read_cl(\n",
    "    datadir / \"gnilc_dust_small_scales_logpoltens_cl_lmax16384.fits.gz\"\n",
    ")"
   ],
   "id": "65b6f864c42e"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synalm_lmax = 8192 * 2  # it needs to be the same for all output nside\n",
    "# synalm_lmax = 1024\n",
    "np.random.seed(8192)\n",
    "\n",
    "alm_log_pol_tens_small_scale = hp.synalm(\n",
    "    list(cl_small_scale),\n",
    "    lmax=synalm_lmax,\n",
    "    new=True,\n",
    ")"
   ],
   "id": "86d65d0ddc44"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm_log_pol_tens_small_scale = [\n",
    "    hp.almxfl(each, np.ones(output_lmax))\n",
    "    for each in alm_log_pol_tens_small_scale\n",
    "]\n",
    "map_log_pol_tens_small_scale = hp.alm2map(\n",
    "    alm_log_pol_tens_small_scale, nside=output_nside\n",
    ")"
   ],
   "id": "b2c5877b295c"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_small_scale"
   ],
   "id": "b3330f0697f1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_small_scale[0] *= hp.alm2map(\n",
    "    modulate_alm[\"temperature\"], output_nside\n",
    ")\n",
    "map_log_pol_tens_small_scale[1:] *= hp.alm2map(\n",
    "    modulate_alm[\"polarization\"], output_nside\n",
    ")\n",
    "assert np.isnan(map_log_pol_tens_small_scale).sum() == 0"
   ],
   "id": "2cc72d118afb"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_small_scale"
   ],
   "id": "f75e306938de"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine scales\n",
    "\n",
    "* Combine small and large scale maps\n",
    "* Transform from logpoltens to IQU\n",
    "* Write output map"
   ],
   "id": "af33e7eec6de"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens = map_log_pol_tens_large_scale\n",
    "map_log_pol_tens += map_log_pol_tens_small_scale"
   ],
   "id": "ce59ea49c4a8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del map_log_pol_tens_small_scale"
   ],
   "id": "66a6eab0e160"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pysm3.utils import log_pol_tens_to_map"
   ],
   "id": "212e42f5fa92"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map = log_pol_tens_to_map(map_log_pol_tens)"
   ],
   "id": "16bfc43c96f1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map"
   ],
   "id": "5c9942bb12d8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del map_log_pol_tens"
   ],
   "id": "1bae29d4ad6f"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Galactic plane fix"
   ],
   "id": "58d9dd0b5071"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "galplane_fix = hp.read_map(datadir / \"gnilc_dust_galplane.fits.gz\", (0, 1, 2, 3))"
   ],
   "id": "5b5d84f4aa7c"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map *= hp.ud_grade(galplane_fix[3], output_nside)\n",
    "output_map += hp.ud_grade(galplane_fix[:3] * (1 - galplane_fix[3]), output_nside)"
   ],
   "id": "477f09d0394f"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map"
   ],
   "id": "80e079069071"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Color correction"
   ],
   "id": "e4f3f733db2e"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Planck 353 GHz color correction https://github.com/galsci/pysm/issues/99"
   ],
   "id": "787262707ace"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map *= 0.911"
   ],
   "id": "139dffbc9734"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map"
   ],
   "id": "0f64832c7fa0"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.write_map(\n",
    "    output_dir / f\"gnilc_dust_template_nside{output_nside}.fits\",\n",
    "    output_map,\n",
    "    dtype=np.float32,\n",
    "    overwrite=True,\n",
    "    column_units = \"uK_RJ\",\n",
    "    extra_header = [(\"lmax\", output_lmax), (\"ref_freq\", \"353 GHz\")]\n",
    ")"
   ],
   "id": "c9b2f02097aa"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "condanamaster2",
   "language": "python",
   "name": "condanamaster2"
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}