{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Customize PySM models\n",
    "\n",
    "There are 3 straightforward ways to customize PySM models.\n",
    "\n",
    "First we can still keep the same algorithms used to generate the components and just provide different templates/paramers, we can achieve this either [using configuration files](#use_a_configuration_file) or, for more flexibility, [work directly with the component Python object](#work_with_python_class).\n",
    "\n",
    "Instead if we need to implement different algorithms, we need to [implement a subclass of pysm.Model](#implement_a_subclass_of_pysm.Model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import pysm\n",
    "import pysm.units as u\n",
    "import healpy as hp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use a configuration file\n",
    "\n",
    "The PySM models are implemented as classes and are then configured with input parameters in [`presets.cfg`](https://github.com/healpy/pysm/blob/master/pysm/data/presets.cfg), the same facilities used to load that configuration file and create the component object can be leverage by creating our own `cfg` file, for example here we are using a scalar value for `map_mbb_index` instead of using the template map provided in the `d1` model.\n",
    "But we could as well replace the maps with our own custom maps.\n",
    "\n",
    "The `%%file` magic creates a file with the content of the cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting mydust.cfg\n"
     ]
    }
   ],
   "source": [
    "%%file mydust.cfg\n",
    "\n",
    "[dustfixedspec]\n",
    "class = \"ModifiedBlackBody\"\n",
    "map_I = \"pysm_2/dust_t_new.fits\"\n",
    "map_Q = \"pysm_2/dust_q_new.fits\"\n",
    "map_U = \"pysm_2/dust_u_new.fits\"\n",
    "unit_I = \"uK_RJ\"\n",
    "unit_Q = \"uK_RJ\"\n",
    "unit_U = \"uK_RJ\"\n",
    "map_mbb_index = 1.5\n",
    "map_mbb_temperature = \"pysm_2/dust_temp.fits\"\n",
    "unit_mbb_temperature = \"K\"\n",
    "freq_ref_I = \"545 GHz\"\n",
    "freq_ref_P = \"353 GHz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "sky_mydust = pysm.Sky(nside=128, component_config=\"mydust.cfg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "sky_d1 = pysm.Sky(nside=128, preset_strings=[\"d1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_100GHz_mydust = sky_mydust.get_emission(100 * u.GHz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_100GHz_d1 = sky_d1.get_emission(100 * u.GHz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 249272 characters]\n",
      "text/plain": [
       "<Figure size 612x388.8 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 253128 characters]\n",
      "text/plain": [
       "<Figure size 612x388.8 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "hp.mollview(map_100GHz_mydust[0], min=0, max=100, title=\"I map\", unit=map_100GHz_d1.unit)\n",
    "hp.mollview(map_100GHz_d1[0], min=0, max=100, title=\"I map\", unit=map_100GHz_d1.unit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 110592 characters]\n",
      "text/plain": [
       "<Figure size 612x388.8 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "hp.mollview(map_100GHz_mydust[0] - map_100GHz_d1[0], min=-100, max=100, title=\"I map\", unit=map_100GHz_d1.unit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Work with the component class\n",
    "\n",
    "We have more flexibility if we use the component class itself, we can create an instance and we can provide different input parameters and that is equivalent of using a configuration file.\n",
    "\n",
    "But we can also create the class and then modify in place the class attributes, see an example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "dust = pysm.ModifiedBlackBody(\n",
    "    nside = 128,\n",
    "    map_I = \"pysm_2/dust_t_new.fits\",\n",
    "    map_Q = \"pysm_2/dust_q_new.fits\",\n",
    "    map_U = \"pysm_2/dust_u_new.fits\",\n",
    "    unit_I = \"uK_RJ\",\n",
    "    unit_Q = \"uK_RJ\",\n",
    "    unit_U = \"uK_RJ\",\n",
    "    map_mbb_index = \"pysm_2/dust_beta.fits\",\n",
    "    map_mbb_temperature = \"pysm_2/dust_temp.fits\",\n",
    "    unit_mbb_temperature = \"K\",\n",
    "    freq_ref_I = \"545 GHz\",\n",
    "    freq_ref_P = \"353 GHz\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "dust.mbb_temperature *= .8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_100GHz_mydust = dust.get_emission(100 * u.GHz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "sky_d1 = pysm.Sky(nside = 128, preset_strings=[\"d1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_100GHz_d1 = sky_d1.get_emission(100 * u.GHz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 258404 characters]\n",
      "text/plain": [
       "<Figure size 612x388.8 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 253344 characters]\n",
      "text/plain": [
       "<Figure size 612x388.8 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "hp.mollview(map_100GHz_mydust[0], min=0, max=100, title=\"I map mydust\", unit=map_100GHz_d1.unit)\n",
    "hp.mollview(map_100GHz_d1[0], min=0, max=100, title=\"I map d1\", unit=map_100GHz_d1.unit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 122896 characters]\n",
      "text/plain": [
       "<Figure size 612x388.8 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "hp.mollview(map_100GHz_mydust[0] - map_100GHz_d1[0], min=-100, max=100, title=\"I map\", unit=map_100GHz_d1.unit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implement a subclass of pysm.Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to implement your own model, subclass `pysm.Model` and implement the `get_emission` method which gets an array of frequency `freqs` and a bandpass `weights` and performs the computation and bandpass integration and returns the output in `uK_RJ`.\n",
    "\n",
    "If we want to reuse some of the algorithms implemented in one of the PySM classes, it is possible to directly subclass one of those instead of starting from `pysm.Model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " This function evaluates the component model at a either\n",
      "        a single frequency, an array of frequencies, or over a bandpass.\n",
      "\n",
      "        Parameters\n",
      "        ----------\n",
      "        freqs: float\n",
      "            Frequency at which the model should be evaluated, assumed to be\n",
      "            given in GHz.\n",
      "\n",
      "        Returns\n",
      "        -------\n",
      "        ndarray\n",
      "            Set of maps at the given frequency or frequencies. This will have\n",
      "            shape (nfreq, 3, npix).\n",
      "        \n"
     ]
    }
   ],
   "source": [
    "print(pysm.models.ModifiedBlackBody.get_emission.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Monopole(pysm.Model):\n",
    "    def __init__(self, nside, value, map_dist=None):\n",
    "        super().__init__(nside=nside, map_dist=map_dist)\n",
    "        self.value = value << u.uK_RJ\n",
    "    def get_emission(self, freqs, weights=None):\n",
    "        return np.ones((3, hp.nside2npix(self.nside))) * self.value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "monopole = Monopole(nside=128, value=1*u.mK_RJ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$[[1000,~1000,~1000,~\\dots,~1000,~1000,~1000],~\n",
       " [1000,~1000,~1000,~\\dots,~1000,~1000,~1000],~\n",
       " [1000,~1000,~1000,~\\dots,~1000,~1000,~1000]] \\; \\mathrm{\\mu K_{{RJ}}}$"
      ],
      "text/plain": [
       "<Quantity [[1000., 1000., 1000., ..., 1000., 1000., 1000.],\n",
       "           [1000., 1000., 1000., ..., 1000., 1000., 1000.],\n",
       "           [1000., 1000., 1000., ..., 1000., 1000., 1000.]] uK_RJ>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "monopole.get_emission(10 * u.GHz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "sky = pysm.Sky(nside=128, preset_strings=[\"d1\"], component_objects=[monopole])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$[1001.2895,~1001.3444,~1001.556] \\; \\mathrm{\\mu K_{{RJ}}}$"
      ],
      "text/plain": [
       "<Quantity [1001.28953135, 1001.34435296, 1001.55604243] uK_RJ>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sky.get_emission(100 * u.GHz)[0, :3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For performance better use `numba` to make use of all cores on the machine and execute at compiled code speed, for example:\n",
    "\n",
    "https://github.com/healpy/pysm/blob/master/pysm/models/power_law.py"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PySM 3",
   "language": "python",
   "name": "pysm3"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
