{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verify synchrotron templates\n"
   ],
   "id": "2c5d5abfde25"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a15c079-f15d-42fb-b66e-77843711d11a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\n",
    "    \"OMP_NUM_THREADS\"\n",
    "] = \"200\"  # for jupyter.nersc.gov otherwise the notebook only uses 2 cores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a40afd7-f8b9-45e1-a618-4a0f1053b925",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pysm3\n",
    "import healpy as hp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12da229e-d783-4800-9378-4b9929d1e955",
   "metadata": {},
   "outputs": [],
   "source": [
    "cd /global/cfs/cdirs/cmb/www/pysm-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4602125-d810-41d5-9bf9-050cd33803ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "cd synch/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8581f5b7-ac8f-4c52-9c02-4c729c5f7b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d0b279b-d58f-4767-8707-908141255b0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "name = \"template\"\n",
    "#name = \"curvature\"\n",
    "#name = \"beta\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bce9499-e85e-4c35-afe9-9bc7db5191de",
   "metadata": {},
   "outputs": [],
   "source": [
    "pol = (0,1,2) if name == \"template\" else (0,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72bec66b-d27d-443d-af34-8a19ef70691c",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_date = \"_2023.02.25\" if name in [\"template\"] else \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd29706-7098-4b55-97ca-94b31080ce7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "nsides = [2048, 4096, 8192]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c636ba1-7459-4f96-9089-1440e6e34987",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = {}\n",
    "for nside in nsides:\n",
    "    m[nside] = hp.read_map(f\"synch_{name}_nside{nside}{template_date}.fits\", pol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "953817f1-39bf-49f9-8947-b3d503514d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "cl = {}\n",
    "for nside in nsides:\n",
    "    cl[nside] = hp.anafast(m[nside], lmax=int(min(2.5*nside,2*8192)))\n",
    "    if cl[nside].ndim == 1:\n",
    "        cl[nside] = cl[nside].reshape((1, -1))\n",
    "    print(nside, cl[nside].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4893273-385c-4ef0-aa13-4feaedd159f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36d90f03-08f7-4dc2-adac-0275d2156858",
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir = Path(\"raw\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "971ec2f6-94dc-4f0b-ab6a-cceb80f9f1ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "if name == \"template\":\n",
    "    largescale_filename = datadir / \"synch_largescale_template_logpoltens_alm_lmax128_2023.02.24.fits.gz\"\n",
    "elif name == \"curvature\":\n",
    "    largescale_filename = datadir / \"synch_curvature_alm_nside8192_lmax16384.fits\"\n",
    "else:\n",
    "    largescale_filename = datadir / f\"synch_largescale_{name}_alm_nside512_lmax768.fits.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5fdb5bf-7cb1-4751-a634-c6027ebdad07",
   "metadata": {},
   "outputs": [],
   "source": [
    "alm_large_scale = hp.read_alm(\n",
    "    largescale_filename,\n",
    "    hdu=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98b597d-1ce2-4ef9-a9cc-1e5f1ffa8fff",
   "metadata": {},
   "outputs": [],
   "source": [
    "if name == \"template\":\n",
    "    smallscale_filename = datadir / \"synch_small_scales_cl_lmax16384_2023.02.24.fits.gz\"\n",
    "else:\n",
    "    smallscale_filename = datadir / f\"synch_small_scales_{name}_cl_lmax16384.fits.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2003da95-0d7b-49bc-a96e-ec370dacc1bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "if name != \"curvature\":\n",
    "    cl_small_scale = hp.read_cl(\n",
    "        smallscale_filename\n",
    "    )\n",
    "    cl_small_scale = cl_small_scale.reshape((len(pol) if name != \"template\" else 4, -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61b405a9-6aeb-415b-9657-8aca19e37007",
   "metadata": {},
   "outputs": [],
   "source": [
    "pol_label = \"TEB\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cb6aae8-8fa8-460d-a0ab-9915b26e6aae",
   "metadata": {},
   "outputs": [],
   "source": [
    "for p in pol:\n",
    "    plt.figure(figsize=(12,6))\n",
    "    for nside in reversed(nsides):\n",
    "        plt.loglog(cl[nside][p], label=f\"map at Nside={nside}\")\n",
    "    if name != \"template\":\n",
    "        plt.loglog(hp.alm2cl(alm_large_scale.astype(complex).reshape((len(pol), -1)))[p], \"--\", alpha=.5, label=\"large scale\")\n",
    "        if name != \"curvature\":\n",
    "            plt.loglog(cl_small_scale[p], \"--\", label=\"small scales\")\n",
    "    plt.title(f\"{name} maps spectra comparison {pol_label[p]}\")\n",
    "    plt.axvline([3150], linestyle=\"--\", color=\"black\", label=\"ell = 3150\", alpha=.5)\n",
    "    plt.legend()\n",
    "\n",
    "    plt.grid();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c31eb0c-201a-4360-82f9-5f5f8e03beac",
   "metadata": {},
   "outputs": [],
   "source": [
    "modulation_alm = hp.read_alm(datadir / \"synch_amplitude_modulation_alms_lmax768.fits.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfafe7e1-6445-4611-8678-dcb33c7a62f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "modulation_map = hp.alm2map(modulation_alm.astype(complex), nside=2048)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36d1bd9b-4378-45f1-a95f-4fa6621eb3e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(modulation_map, title=f\"Modulation map, mean: {modulation_map.mean():.2f}\", unit=\"dimensionless\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cd55ef4-0bc2-4457-9cdd-d75e96a2e9a1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}