{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate synchrotron spectral index map\n"
   ],
   "id": "a80007f234a5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import healpy as hp\n",
    "from pathlib import Path"
   ],
   "id": "3a79602dbb73"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir=Path(\"data/\")\n",
    "output_dir = Path(\"production-data/synch\")\n",
    "output_dir_raw = output_dir / \"raw\""
   ],
   "id": "bc120a52e495"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "output_nside = 2048"
   ],
   "id": "1f455305d060"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_lmax = int(min(2.5 * output_nside, 8192 * 2))"
   ],
   "id": "8ed458499ed2"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Large scales"
   ],
   "id": "e59950616873"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm_large_scale = hp.read_alm(\n",
    "    output_dir_raw / f\"synch_largescale_beta_alm_nside512_lmax768.fits.gz\",\n",
    "    hdu=1,\n",
    ")"
   ],
   "id": "9190d17dea8d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_large_scale = hp.alm2map(alm_large_scale.astype(np.complex128), nside=output_nside)"
   ],
   "id": "76faf8f23019"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small scales modulation"
   ],
   "id": "5eff1698a440"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modulate_alm = hp.read_alm(\n",
    "    output_dir_raw / f\"synch_amplitude_modulation_alms_lmax768.fits.gz\"\n",
    ").astype(np.complex128)"
   ],
   "id": "4d387c248da0"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small scales"
   ],
   "id": "0f84fa682230"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_small_scale = hp.read_cl(\n",
    "    output_dir_raw / f\"synch_small_scales_beta_cl_lmax16384.fits.gz\"\n",
    ")"
   ],
   "id": "e61db6ef2447"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synalm_lmax = 8192 * 2  # for reproducibility\n",
    "# synalm_lmax = output_lmax\n",
    "seed = 444\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(output_lmax+1))\n",
    "map_small_scale = hp.alm2map(alm_small_scale, nside=output_nside)\n",
    "map_small_scale *= hp.alm2map(modulate_alm, output_nside)\n",
    "assert np.isnan(map_small_scale).sum() == 0"
   ],
   "id": "54b741e82519"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine scales\n",
    "\n",
    "* Combine small and large scale maps\n",
    "* Write output map"
   ],
   "id": "ff8479d7d554"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_map = map_large_scale + map_small_scale - 3.1"
   ],
   "id": "47dd6ec1299f"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.write_map(output_dir / f\"synch_beta_nside{output_nside}.fits\", output_map, dtype=np.float32, overwrite=True)"
   ],
   "id": "f4e63d9f1415"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pysm3.utils import add_metadata"
   ],
   "id": "e8e4c34b81c9"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "add_metadata([output_dir / f\"synch_beta_nside{output_nside}.fits\"], coord=\"G\", unit=\"\", ref_freq=\"23 GHz\", ell_max=str(output_lmax))"
   ],
   "id": "aa60cd252eac"
  },
  {
   "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(output_map, title=\"Total\")"
   ],
   "id": "dee9284c9425"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pycmb",
   "language": "python",
   "name": "pycmb"
  },
  "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}