{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate GNILC spectral index and temperature maps\n"
   ],
   "id": "e97c58170ef0"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import healpy as hp\n",
    "from pathlib import Path"
   ],
   "id": "ddb21c064382"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = Path(\"production-data\") / \"dust_gnilc\""
   ],
   "id": "c15f23bedbdd"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir = output_dir / \"raw\""
   ],
   "id": "a036e0a6971e"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "output_nside = 2048\n",
    "name = \"beta\""
   ],
   "id": "39faa21c6697"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_lmax = int(min(2.5 * output_nside, 8192 * 2))"
   ],
   "id": "6d7075b923e6"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Galactic mask"
   ],
   "id": "bef0d3e66863"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "galactic_mask = (\n",
    "    hp.ud_grade(\n",
    "        hp.read_map(datadir / \"HFI_Mask_GalPlane-apo2_2048_R2.00_GAL080_noapo.fits.gz\"),\n",
    "        output_nside,\n",
    "    )\n",
    "    == 1\n",
    ")"
   ],
   "id": "a5bf06594b37"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Large scales"
   ],
   "id": "52d94f38731f"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm_large_scale = hp.read_alm(\n",
    "    datadir / f\"gnilc_dust_largescale_template_{name}_alm_nside2048_lmax1024.fits.gz\",\n",
    "    hdu=1,\n",
    ")"
   ],
   "id": "8295f556f137"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_large_scale = hp.alm2map(alm_large_scale.astype(np.complex128), nside=output_nside)"
   ],
   "id": "2b90c5ac912d"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small scales modulation"
   ],
   "id": "750fe2db2e8d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modulate_alm = hp.read_alm(\n",
    "    datadir / f\"gnilc_dust_temperature_modulation_alms_lmax768.fits.gz\"\n",
    ").astype(np.complex128)"
   ],
   "id": "a9e76ec26fcf"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small scales"
   ],
   "id": "694048261d40"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_small_scale = hp.read_cl(\n",
    "    datadir / f\"gnilc_dust_small_scales_{name}_cl_lmax16384_2023.06.06.fits.gz\"\n",
    ")"
   ],
   "id": "f4a85c3e6d8a"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(cl_small_scale)"
   ],
   "id": "1032856b4af4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_small_scale"
   ],
   "id": "d40b1d13091d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.loglog(cl_small_scale)"
   ],
   "id": "d10e511355e7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synalm_lmax = 8192 * 2  # for reproducibility\n",
    "# synalm_lmax = 1024\n",
    "seed = 777 if name == \"beta\" else 888\n",
    "np.random.seed(seed)\n",
    "\n",
    "alm_small_scale = hp.synalm(\n",
    "    cl_small_scale,\n",
    "    lmax=synalm_lmax,\n",
    "    new=True,\n",
    ")\n",
    "\n",
    "alm_small_scale = hp.almxfl(alm_small_scale, np.ones(int(2.5 * output_nside+1)))\n",
    "map_small_scale = hp.alm2map(alm_small_scale, nside=output_nside)\n",
    "assert np.isnan(map_small_scale).sum() == 0"
   ],
   "id": "fffb67421d42"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_small_scale = hp.alm2map(alm_small_scale, nside=output_nside)"
   ],
   "id": "ba3e951de89d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(map_small_scale)"
   ],
   "id": "c9c60541edaa"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_small_scale *= hp.alm2map(modulate_alm, output_nside)"
   ],
   "id": "d86c70d3b65c"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(map_small_scale)"
   ],
   "id": "4db22476050a"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(galactic_mask)"
   ],
   "id": "c86e22874a32"
  },
  {
   "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": "3154be7648a2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_total = map_large_scale + map_small_scale"
   ],
   "id": "50047ddc8d41"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "galplane_fix_mask = hp.read_map(datadir / \"gnilc_dust_galplane.fits.gz\", 3)"
   ],
   "id": "adbcf1a1f5ec"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ext_number = {\"beta\":0, \"Td\":1}"
   ],
   "id": "ef1560f6e44c"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "galplane_fix = hp.read_map(datadir / \"gnilc_dust_beta_Td_galplane.fits.gz\", ext_number[name])"
   ],
   "id": "268d83a93f92"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_total *= hp.ud_grade(galplane_fix_mask, output_nside)\n",
    "map_total += hp.ud_grade(galplane_fix * (1 - galplane_fix_mask), output_nside)"
   ],
   "id": "430a14f1d39b"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(map_total - map_large_scale - map_small_scale, title=\"Effect of galactic plane fix\")"
   ],
   "id": "51565cb88355"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.write_map(\n",
    "    output_dir / f\"gnilc_dust_{name}_nside{output_nside}_2023.06.06.fits\",\n",
    "    map_total,\n",
    "    coord=\"G\",\n",
    "    column_units=\"\" if name == \"beta\" else \"K\",\n",
    "    extra_header = [(\"lmax\", output_lmax), (\"ref_freq\", \"353 GHz\")],\n",
    "    dtype=np.float32,\n",
    "    overwrite=True,\n",
    ")"
   ],
   "id": "b7943f217895"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(map_large_scale, title=\"Large scale\")\n",
    "hp.mollview(map_small_scale, title=\"Small scale\")\n",
    "hp.mollview(map_total, title=\"Total\")"
   ],
   "id": "0fae37f8b4bf"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [],
   "id": "ed26d346cd08"
  }
 ],
 "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
}