{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import healpy as hp\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = Path(\"production-data\") / \"dust_gnilc\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir = output_dir / \"raw\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "output_nside = 2048"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_lmax = int(min(2.5 * output_nside, 8192 * 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Large scales"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_large_scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Small scales modulation"
   ]
  },
  {
   "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",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Small scales"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_small_scale"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_log_pol_tens_small_scale"
   ]
  },
  {
   "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"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del map_log_pol_tens_small_scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pysm3.utils import log_pol_tens_to_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map = log_pol_tens_to_map(map_log_pol_tens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del map_log_pol_tens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Galactic plane fix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "galplane_fix = hp.read_map(datadir / \"gnilc_dust_galplane.fits.gz\", (0, 1, 2, 3))"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Color correction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Planck 353 GHz color correction https://github.com/galsci/pysm/issues/99"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map *= 0.911"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map"
   ]
  },
  {
   "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",
    ")"
   ]
  }
 ],
 "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
}