Pixell_simulations.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yPF5lzLN9DJB"
      },
      "source": [
        "# Simulations with pixell\n",
        "\n",
        "*Written by the ACT Collaboration*\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/simonsobs/pixell/notebooks/pixell_simulations.ipynb)\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "This notebook, and the accompanying notebooks included in this set, are designed to help users who are new to working with [`pixell`](https://github.com/simonsobs/pixell/) get started with the package. As a set these notebooks will guide users through examples of how to read in and display maps, how to perform spherical harmonic transform and calculate simple spectra, how to transform the maps and how to study point sources in the maps.\n",
        "\n",
        "The `pixell` library allows users to load,\n",
        "manipulate and analyze maps stored in rectangular pixelization. It is\n",
        "mainly targeted for use with maps of the sky (e.g. CMB intensity and polarization maps, stacks of 21 cm intensity maps, binned galaxy positions or shear) in cylindrical projection.\n",
        "\n",
        "In this notebook we will use pixell functions to simulate millimeter-wave sky maps, including different components such as CMB, CMB lensing, Extragalactic Point Sources and Galaxy Clusters, and simplistic noise.\n",
        "\n",
        "- [x] `enmap.randmap` (flat-sky Gaussian simulations, CMB and noise)\n",
        "\n",
        "- [x] `curvedsky.randmap` (curved-sky Gaussian simulations, CMB and noise)\n",
        "\n",
        "- [x] `lensing.lens_map` (flat-sky lensing operation)\n",
        "\n",
        "- [x] `lensing.lens_map_curved` (curved-sky lensing operation)\n",
        "\n",
        "- [x] `pointsrcs.sim_objects` (injecting point sources)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kDFKT4bz89Ra"
      },
      "outputs": [],
      "source": [
        "# Install neccesary packages\n",
        "!pip install pixell camb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s3We41ht94xX"
      },
      "outputs": [],
      "source": [
        "# Import packages\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from pixell import powspec\n",
        "from pixell import utils\n",
        "from pixell import enmap\n",
        "from pixell import curvedsky\n",
        "from pixell import lensing\n",
        "from pixell import pointsrcs\n",
        "from pixell import enplot"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RFGNWqMC9yb5"
      },
      "source": [
        "**Making a map from scratch**\n",
        "\n",
        "For details on how to manipulate maps in pixell take a look at [\"Pixell_map_manipulation.ipynb\"](https://github.com/simonsobs/pixell_tutorials/blob/master/Pixell_map_manipulation.ipynb). First lets make a zeros map on a certain patch of the sky, let's use the region covering right ascension from 15 to 35 degrees and declination from -10 to -2 (you can try any patch). Here we need to define the vertices of a rectangle from left bottom to top right, these are RA0,DEC0 = 35,-10 and RA1,DEC1 = 15,-2 (think about the sky as going from right ascension 180 to -180 and declinations -90 to +90. We will use 0.5 arcminutes pixels in plate carrée projection, this means all pixels cover 0.5 arcminutes in latitude and longitude"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NX2A3eIk-y0J"
      },
      "outputs": [],
      "source": [
        "              # DEC0, RA0  DEC1 RA1\n",
        "box = np.array([[-10, 35], [-2, 15]]) * utils.degree\n",
        "shape, wcs = enmap.geometry(pos=box, res=0.5*utils.arcmin, proj='car')\n",
        "# we will use only the temperature component or parameter stokes I\n",
        "omap = enmap.zeros(shape, wcs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-d8ndl4btd96"
      },
      "outputs": [],
      "source": [
        "# we can use enplot and see what we have\n",
        "enplot.pshow(omap, range=500, colorbar=True, downgrade=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G9UeyFDdtipD"
      },
      "source": [
        "looks like an empty map in the region that we wanted. We can also check the pixel size, since all pixels cover 0.5 in width and height we should have 8x60/0.5 = 960 pixels in height and 20x60/0.5 = 2400 in width\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DTXdCdZ2tnLZ"
      },
      "outputs": [],
      "source": [
        "omap.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6PyaTjLxPhOZ"
      },
      "source": [
        "This means that dimension = (number of stokes parameters, declination pixels, right ascension pixels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_vJ5OYvI9yjp"
      },
      "source": [
        "**Gaussian realization of CMB**\n",
        "\n",
        "There are two ways of simulating things, using flat sky approximation or curved sky. Let's start with flat sky, there is a handy function from `enmap` called `rand_map`, which draws a random realization from a power spectrum. To call this function we need to pass the following arguments:\n",
        "\n",
        "`shape`: shape of the map that we want to simulate, in our case omap.shape\n",
        "\n",
        "`wcs`: world coordinates of the map, in our case omap.wcs\n",
        "\n",
        "`cov`: power spectrum that we want to include, more on this later\n",
        "\n",
        "`scalar = False`: If we have polarization components, treat them as polarization\n",
        "\n",
        "`seed = None`: The random seed that we want to use, if None if will generate one\n",
        "\n",
        "`pixel_units = False`: The input power spectrum uses steradians\n",
        "\n",
        "`iau = False`: The default polarization convention in the CMB community (sad)\n",
        "\n",
        "`spin = [0,2]`: Apply a EB -> QU rotation\n",
        "\n",
        "The three first argument are mandatory, but you may need to modify the others depending on your application (see the `pixell` documentation)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7JsQibuUPyii"
      },
      "source": [
        "We need a power spectrum, specifically we want to simulate the CMB, for that we will use the Planck TT PS (you can try any other cosmology).\n",
        "\n",
        "First we download the file in the format L TT EE BB TE\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hZB9rpVyRDT8"
      },
      "outputs": [],
      "source": [
        "!wget -O planck_lensedCls.dat https://raw.githubusercontent.com/simonsobs/nemo/main/nemo/data/planck_lensedCls.dat"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UvkiTJlPRHGb"
      },
      "source": [
        "We can plot the file to be sure if it contains the typical CMB power spectrum in the following way\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LmpnSCHxRTK6"
      },
      "outputs": [],
      "source": [
        "fileps = \"planck_lensedCls.dat\"\n",
        "L, TT, EE, BB, TE = np.loadtxt(fileps, usecols=(0, 1, 2, 3, 4), unpack=True)\n",
        "plt.loglog(L, TT)\n",
        "# plt.yscale(\"log\")\n",
        "plt.xlabel(r\"$\\ell$\")\n",
        "plt.ylabel(r\"$\\ell (\\ell+1) C_{\\ell}^{TT}/(2 \\pi)$\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6aGoL8QvGIfU"
      },
      "outputs": [],
      "source": [
        "TT.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UsSrZQqRtzhu"
      },
      "source": [
        "We can then pass this Planck TT power spectrum to the pixell function `powspec.read_spectrum`, so that the array is reshaped in the form that Pixell expects, and the 2$\\pi/(\\ell (\\ell+1))$ factors are applied."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UEIWlvj4Tvxh"
      },
      "outputs": [],
      "source": [
        "ps = powspec.read_spectrum(fileps, scale=True, expand=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hswCbdCvFeJP"
      },
      "outputs": [],
      "source": [
        "ell = np.arange(len(ps[0]))\n",
        "plt.semilogx(L, TE)\n",
        "plt.semilogx(ell, ps[3] * (ell*(ell+1)) / (2*np.pi))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XtGDnuEpvsZg"
      },
      "source": [
        "The pixell function `enmap.rand_map` will then take in the shape and wcs that we defined above, as well as the power spectrum from the cell above to make the corresponding map."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-dIh1pIIRvoZ"
      },
      "outputs": [],
      "source": [
        "# T only\n",
        "cmbmap = enmap.rand_map(omap.shape, omap.wcs, cov=ps[0])\n",
        "enplot.pshow(cmbmap, range=500, colorbar=True, downgrade=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A7xJ4kx9UN2g"
      },
      "source": [
        "If you run the previous code multiple times, you will see a different version each time, that's because each random realization is different but it is sampled from the same probability distribution (given by the power spectrum).\n",
        "\n",
        "We can also make this polarized. We also need a 3-dimensional map to hold the simulation. We also need to change the `cov` parameter so that it includes all polarization components. We have a TT, EE, BB, and TE spectrum. We need to put this into a `(3, 3, nl)` array, rather than just an `(nl,)` array:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4nWl2s4hj7L0"
      },
      "outputs": [],
      "source": [
        "nl = ps.shape[-1]\n",
        "l = np.arange(nl)\n",
        "\n",
        "ps_square = np.zeros((3, 3, nl))\n",
        "print(ps_square.shape)\n",
        "\n",
        "ps_square[0, 0] = ps[0]\n",
        "ps_square[0, 1] = ps[3]\n",
        "ps_square[1, 0] = ps[3]\n",
        "ps_square[1, 1] = ps[1]\n",
        "ps_square[2, 2] = ps[2]\n",
        "\n",
        "plt.semilogy(l, l*(l+1)/2/np.pi*ps_square[0, 0], label='00')\n",
        "plt.semilogy(l, l*(l+1)/2/np.pi*ps_square[1, 1], label='11')\n",
        "plt.semilogy(l, l*(l+1)/2/np.pi*ps_square[2, 2], label='22')\n",
        "plt.ylim(1e-3, 1e4)\n",
        "plt.legend()\n",
        "plt.show()\n",
        "\n",
        "plt.plot(l, l*(l+1)/2/np.pi*ps_square[0, 1], label='01')\n",
        "plt.legend()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nM_pj4yLmgFH"
      },
      "source": [
        "Then we provide this square array as the `cov` parameter instead:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XinX36L4miXy"
      },
      "outputs": [],
      "source": [
        "# T, Q, U\n",
        "cmbmap_pol = enmap.rand_map((3, *shape), wcs, cov=ps_square)\n",
        "for i in range(3):\n",
        "  enplot.pshow(cmbmap_pol[i], range=[500, 25, 25][i], colorbar=True, downgrade=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "46fOKO6QUQh2"
      },
      "source": [
        "Now this is using flat sky approximation, so it uses the Fourier transform of the patch. If we want to account for the curvature of the sky, which is important for big patches we need to use the `curvedsky` module, in that module we can also find a `rand_map` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "klCE_gs5USqX"
      },
      "outputs": [],
      "source": [
        "# full sky goes from ra,dec +180, -90 to ra, dec -180, +90\n",
        "# we can use bigger pixels such as 4 arcmin to avoid using too much memory\n",
        "                #DEC0, RA0  DEC1 RA1\n",
        "shape, wcs = enmap.fullsky_geometry(res=4*utils.arcmin, proj='car', variant='CC')\n",
        "# we will use only the temperature component or parameter stokes I\n",
        "omap = enmap.zeros((1, *shape), wcs)\n",
        "# checking the map with a plot\n",
        "enplot.pshow(omap, range=500, downgrade=4, ticks=15, colorbar=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F3tSf_0rv5l_"
      },
      "source": [
        "This `curvedsky.rand_map` takes slightly different parameters than `enmap.rand_map`, in order they are\n",
        "\n",
        "`shape`: shape of the map that we want to simulate, in our case omap.shape\n",
        "\n",
        "`wcs`: world coordinates of the map, in our case omap.wcs\n",
        "\n",
        "`ps`: power spectrum that we want to include. Can take an `(nl,)` shaped ps or a `(ncomp, ncomp, nl)` ps.\n",
        "\n",
        "`lmax=None`: maximum multipole range to sample\n",
        "\n",
        "`dtype=np.float32`: type of the resulting array\n",
        "\n",
        "`seed=None`: number of the seed, if None it will generate one\n",
        "\n",
        "`spin=[0,2]`: type of spin, in we have a IQU map it is [0,2], if we give as an input a single power spectrum it will take that into consideration to generate only one map\n",
        "\n",
        "`method=\"auto\"`: the default \"auto\" is fine\n",
        "\n",
        "`verbose=False`: gives text information about the process\n",
        "\n",
        "Note: it is typically sufficient to use single-precision maps (`dtype=np.float32`) which would speed up computation and use less memory. The default for this function is double precision (`dtype=np.float64`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0Bc0gGwaUjVv"
      },
      "outputs": [],
      "source": [
        "# now simulating the CMB taking into account the curvature of the sky\n",
        "# this may take some time since it uses spherical harmonics\n",
        "\n",
        "# T only\n",
        "cmbmap = curvedsky.rand_map(omap.shape, omap.wcs, ps=ps[0], lmax=3000, dtype=np.float32, verbose=True)\n",
        "enplot.pshow(cmbmap, range=500, downgrade=4, ticks=15, colorbar=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jQpjNRplUZT6"
      },
      "source": [
        "Ok, it looks interesting. Notice the distortion at the poles; is this what you expected? To put it in context, consider this image of the earth in the same CAR projection:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HXP3TMUKcyjz"
      },
      "source": [
        "![](https://upload.wikimedia.org/wikipedia/commons/thumb/8/83/Equirectangular_projection_SW.jpg/1920px-Equirectangular_projection_SW.jpg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0NQ3gf4GyFR0"
      },
      "source": [
        "This makes sense now. Each pixel in the top row of pixels (and bottom row) are in fact the same physical point -- the pole! Using `enmap.rand_map` in this case might make a more uniform looking picture but it would be physically wrong -- we can't use the flat sky approximation when our range of declinations is large!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tshbkq-GUVNY"
      },
      "outputs": [],
      "source": [
        "# now simulating the CMB in this patch in the flat sky approximation\n",
        "# NOTE: THIS IS WRONG\n",
        "cmbmap = enmap.rand_map(omap.shape, omap.wcs, cov=ps[0])\n",
        "enplot.pshow(cmbmap, range=500, colorbar=True, downgrade=4, ticks=15)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7f70XuQq1jmN"
      },
      "source": [
        "There are no distortions at the poles -- this is *wrong*. We need the distortions!\n",
        "\n",
        "Just like the flatsky case, we can draw a polarized simulation. We again need a map to hold the 3 polarization components, and need to give it a power spectrum that puts the components in a square:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c3TSGZwepku0"
      },
      "outputs": [],
      "source": [
        "# T, Q, U\n",
        "cmbmap_pol = curvedsky.rand_map((3, *shape), wcs, ps=ps_square, lmax=3000, dtype=np.float32, verbose=True)\n",
        "for i in range(3):\n",
        "  enplot.pshow(cmbmap_pol[i], range=[500, 25, 25][i], downgrade=4, ticks=15, colorbar=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FIgq7BMw9ysb"
      },
      "source": [
        "**Lensing**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y2TT6w_8xldi"
      },
      "source": [
        "To make flat-sky lensing simulations, we need unlensed CMB power spectra, as well as a lensing potential power spectrum. We obtain these from CAMB by running the following cells:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "12A-gpoJ4QyW"
      },
      "outputs": [],
      "source": [
        "import camb\n",
        "\n",
        "pars = camb.set_params(H0=67.5, ombh2=0.022, omch2=0.122, mnu=0.06, omk=0, tau=0.06,\n",
        "                       As=2e-9, ns=0.965, halofit_version='mead', lmax=4000)\n",
        "pars.set_for_lmax(4000, lens_potential_accuracy=4)\n",
        "\n",
        "# calculate results for these parameters\n",
        "results = camb.get_results(pars)\n",
        "powers = results.get_cmb_power_spectra(pars, CMB_unit='muK', raw_cl=True)\n",
        "TT, EE, BB, TE = powers[\"unlensed_scalar\"].T\n",
        "ell = np.arange(len(TT))\n",
        "\n",
        "# for simplicity assume lensing is uncorrelated with the CMB power spectra, but\n",
        "# we could include correlations too as CAMB calculates them\n",
        "PP = powers[\"lens_potential\"][:, 0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ey-PxhdX1sGT"
      },
      "source": [
        "Before we create the lensed simulations, we need to package the CMB and lensing potential power spectra in a format that pixell will understand:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x1VNVLpyObKx"
      },
      "outputs": [],
      "source": [
        "ps_init = np.vstack((ell, TT, EE, TE, PP))\n",
        "ps_init = np.atleast_2d(ps_init)\n",
        "ps_init = powspec.expand_inds(np.array(ps_init[0],dtype=int), ps_init[1:])\n",
        "ps_cmb = ps_init[:3]\n",
        "ps_cmb = powspec.sym_expand(ps_cmb, scheme=\"diag\", ncomp=3)\n",
        "ps_lens = ps_init[3]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ueUvhz_19A2"
      },
      "source": [
        "We are then ready to use the pixell function `enmap.rand_map` to make a $\\phi$ map from the lensing potential power spectrum, as well as unlensed CMB maps from the CMB power spectra.\n",
        "\n",
        "Let's be careful to go back to the small patch of sky for which we used `enmap.rand_map` in the beginning of this notebook, where the flat-sky limit is approximately valid:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UFF6_hOB17n4"
      },
      "outputs": [],
      "source": [
        "              # DEC0, RA0  DEC1 RA1\n",
        "box = np.array([[-10, 35], [-2, 15]]) * utils.degree\n",
        "shape, wcs = enmap.geometry(pos=box, res=0.5*utils.arcmin, proj='car')\n",
        "\n",
        "phi_map = enmap.rand_map(shape, wcs, cov=ps_lens)\n",
        "cmbmap = enmap.rand_map((3, *shape), wcs, cov=ps_cmb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PTeynEA72dDM"
      },
      "source": [
        "The pixell function `lensing.lens_map` will take in the CMB map we created above, as well as the gradient of the $\\phi$ map and make a lensed version of the CMB maps. It has the following arguments:\n",
        "\n",
        "`imap`: the unlensed map with shape `(..., ny, nx)` to be lensed\n",
        "\n",
        "`grad_phi`: the map of the gradient of the lensing potential with shape `(2, ny, nx)`, obtained from `enmap.grad_phi`\n",
        "\n",
        "`order=3`: related to the pixel-space interpolation (fine to keep)\n",
        "\n",
        "`mode=\"spline\"`: same as above\n",
        "\n",
        "`border=\"cyclic\"`: same as above\n",
        "\n",
        "`trans=False`: if True, perform adjoint of lensing (not delensing)\n",
        "\n",
        "`deriv=False`: whether to return derivatives of the interpolation\n",
        "\n",
        "`h=1e-7`: finite difference size for the derivative\n",
        "\n",
        "For standard lensing, we can accept all the defaults and just pass our unlensed map and lensing realization:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2-XwNNvI2at7"
      },
      "outputs": [],
      "source": [
        "grad_phi = enmap.grad(phi_map)\n",
        "lensedcmb = lensing.lens_map(cmbmap, grad_phi, order=3, mode=\"spline\", border=\"cyclic\", trans=False, deriv=False, h=1e-7)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "blYntsOt2v46"
      },
      "source": [
        "Let's take a look at the lensed CMB maps we just created, followed by the corresponding unlensed CMB map:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3wEccFDQxOit"
      },
      "outputs": [],
      "source": [
        "enplot.pshow(lensedcmb[0], colorbar=True, downgrade=2, range=300)\n",
        "enplot.pshow(cmbmap[0], colorbar=True, downgrade=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zyTrdUeOe_rv"
      },
      "source": [
        "The effect of lensing is made more obvious by looking at the difference between the output and input (unlensed) CMB maps:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tEojjqfve-JE"
      },
      "outputs": [],
      "source": [
        "enplot.pshow(lensedcmb[0] - cmbmap[0], colorbar=True, downgrade=2, range=50)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jJ1yHWKa21Iy"
      },
      "source": [
        "We can then compare the power spectra of these lensed maps to the lensed power spectra from CAMB. To calculate the power spectra of the lensed maps we created, we will first apodize the maps, and use the pixell functions `enmap.map2harm` and enmap.lbin to calculate the power spectra. If you need a refresher of this procedure, please refer to the notebook [\"Fourier Operations with pixell\"](https://github.com/simonsobs/pixell_tutorials/blob/master/Pixell_fourier_space_operations.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VBX8FYSb4H7E"
      },
      "source": [
        "Apodization is done as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kLXytVYRo46N"
      },
      "outputs": [],
      "source": [
        "apod_pix = 200 # number of pixels at the edge to apodize\n",
        "taper = enmap.apod(cmbmap*0.0 + 1.0, apod_pix)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J5X1wU3mo9OZ"
      },
      "outputs": [],
      "source": [
        "for i in range(3):\n",
        "  enplot.pshow((taper * lensedcmb)[i], downgrade=2, colorbar=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xRHS9Zq64Lb2"
      },
      "source": [
        "We can then compute the 2D FFT with `enmap.map2harm` and use `enmap.lbin` to bin the power spectra into 1D."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ek0NoO8tNhj_"
      },
      "outputs": [],
      "source": [
        "harm_lensed = enmap.map2harm(taper * lensedcmb, normalize=\"phys\")\n",
        "w2 = np.mean(taper**2)\n",
        "power = (harm_lensed[:, None] * np.conj(harm_lensed)).real/w2\n",
        "cl, ls1 = power.lbin(bsize=40)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P4oKvpPf4z0R"
      },
      "source": [
        "Here, we also load the CAMB lensed power spectra to compare to the lensed power spectra from the simulation we created in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u9SE_LlSk_qR"
      },
      "outputs": [],
      "source": [
        "powers_lensed_camb =  powers[\"lensed_scalar\"].T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hxHxVjwfdcJm"
      },
      "outputs": [],
      "source": [
        "dll_fac = (ls1*(ls1+1))/(2*np.pi)\n",
        "camb_dll_fac = (ell*(ell+1))/(2*np.pi)\n",
        "\n",
        "fig, ax = plt.subplots(2, 2)\n",
        "\n",
        "ax[0,0].plot(ls1,cl[0, 0]*dll_fac,label=\"This notebook\")\n",
        "ax[0,0].plot(powers_lensed_camb[0]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[0,0].set_xlim(20, 4000)\n",
        "\n",
        "ax[0,1].plot(ls1,cl[1, 1]*dll_fac,label=\"This notebook\")\n",
        "ax[0,1].plot(powers_lensed_camb[1]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[0,1].set_xlim(20, 4000)\n",
        "\n",
        "ax[1,0].plot(ls1,cl[2, 2]*dll_fac,label=\"This notebook\")\n",
        "ax[1,0].plot(powers_lensed_camb[2]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[1,0].set_xlim(20, 4000)\n",
        "\n",
        "ax[1,1].plot(ls1,cl[0, 1]*dll_fac,label=\"This notebook\")\n",
        "ax[1,1].plot(powers_lensed_camb[3]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[1,1].set_xlim(20, 4000)\n",
        "\n",
        "ax[0,0].set_xlabel(\"$\\ell$\")\n",
        "ax[0,0].set_ylabel(\"TT\")\n",
        "ax[0,1].set_xlabel(\"$\\ell$\")\n",
        "ax[0,1].set_ylabel(\"EE\")\n",
        "ax[1,0].set_xlabel(\"$\\ell$\")\n",
        "ax[1,0].set_ylabel(\"BB\")\n",
        "ax[1,1].set_xlabel(\"$\\ell$\")\n",
        "ax[1,1].set_ylabel(\"TE\")\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HUYzOURssu-1"
      },
      "source": [
        "Nice! There are clearly some lensing-induced B-modes, but they don't have quite the right shape. This is likely due to the fact that we've done everything on the flat sky and/or haven't properly accounted for the mode-coupling due to the sky mask.\n",
        "\n",
        "Fortunately, just as for the Gaussian realization from the CMB, `pixell` also supports full-sky lensing operations that accounts for spherical geometry. Because we can work on the full-sky, we won't need to apodize the map; however, this will not apply to a realistic simulation for ACT or SO data, which only observes a fraction of the sky.\n",
        "\n",
        "As for the Gaussian realization, the curved-sky lensing operation `lensing.lens_map_curved` has new arguments:\n",
        "\n",
        "`shape`: with `wcs`, the geometry of the output map\n",
        "\n",
        "`wcs`: with shape, the geometry of the output map\n",
        "\n",
        "`phi_alm`: the spherical harmonic realization of the lensing potential realization\n",
        "\n",
        "`cmb_alm`: the spherical harmonics of the unlensed CMB realization\n",
        "\n",
        "`phi_ainfo=None`: describes the spherical harmonic ordering convention for `phi` (you will know if you need to adjust this)\n",
        "\n",
        "`maplmax=None`: deprecated\n",
        "\n",
        "`dtype=np.float64`: output precision, needs to be double for accurate calculations (may crash otherwise)\n",
        "\n",
        "`spin=[0,2]`: like in `curvedsky.rand_map`\n",
        "\n",
        "`output=\"l\"`: returns only the lensed CMB map. `lu` returns the lensed and unlensed CMB. More options for lensing maps in the [function documentation](https://github.com/simonsobs/pixell/blob/a61ca1a3ae99e4c8627c86a0fe2401d4bfe67d66/pixell/lensing.py#L134)\n",
        "\n",
        "`geodesic=True`: slower but accurate, `False` for faster but less accurate\n",
        "\n",
        "`verbose=False`: print helpful messages\n",
        "\n",
        "`delta_theta=None`: can leave this as-is\n",
        "\n",
        "In the below, besides supplying the required arguments, the only other important argument we change is to also return the unlensed CMB so we can compare:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PpV-vQ77wt6T"
      },
      "outputs": [],
      "source": [
        "# first get a full-sky realization of the phi map and the polarized, unlensed cmb\n",
        "# use low-resolution to reduce memory and runtime\n",
        "shape, wcs = enmap.fullsky_geometry(res=4*utils.arcmin, proj='car')\n",
        "\n",
        "# pixell needs the maps in spherical harmonics\n",
        "phi_alm = curvedsky.rand_alm(ps=ps_lens, lmax=2700)\n",
        "cmb_alm = curvedsky.rand_alm(ps=ps_cmb, lmax=2700)\n",
        "\n",
        "# get the lensed cmb map\n",
        "lensedcmb, cmbmap = lensing.lens_map_curved((3, *shape), wcs, phi_alm, cmb_alm, output=\"lu\", verbose=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MbTUbm2M023R"
      },
      "source": [
        "As before, let's take a look at the lensed CMB map, the unlensed CMB map, and their difference:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G-gBkklV09An"
      },
      "outputs": [],
      "source": [
        "enplot.pshow(lensedcmb[0], colorbar=True, downgrade=4, ticks=15)\n",
        "enplot.pshow(cmbmap[0], colorbar=True, downgrade=4, ticks=15)\n",
        "enplot.pshow(lensedcmb[0] - cmbmap[0], colorbar=True, downgrade=4, ticks=15)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qt44MhuCL2WK"
      },
      "source": [
        "Now we can take the full-sky power spectrum without a mask to see if the lensing produced the correct lensed CMB power spectrum. A nice benefit of this test compared to the flat-sky test is that we have more sky area and so the measured power spectrum is also less noisy:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "36bGziOGL1rR"
      },
      "outputs": [],
      "source": [
        "lensed_alm = curvedsky.map2alm(lensedcmb, lmax=2700)\n",
        "cl = curvedsky.alm2cl(lensed_alm[:, None], lensed_alm)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K_M_UVJ2MY2P"
      },
      "outputs": [],
      "source": [
        "ell = np.arange(cl.shape[-1])\n",
        "dll_fac = (ell*(ell+1))/(2*np.pi)\n",
        "\n",
        "camb_ell = np.arange(powers_lensed_camb[0].shape[-1])\n",
        "camb_dll_fac = (camb_ell*(camb_ell+1))/(2*np.pi)\n",
        "\n",
        "fig, ax = plt.subplots(2, 2, figsize=(10, 6))\n",
        "\n",
        "ax[0,0].plot(cl[0, 0]*dll_fac,label=\"This notebook\")\n",
        "ax[0,0].plot(powers_lensed_camb[0]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[0,0].set_xlim(20, 2700)\n",
        "\n",
        "ax[0,1].plot(cl[1, 1]*dll_fac,label=\"This notebook\")\n",
        "ax[0,1].plot(powers_lensed_camb[1]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[0,1].set_xlim(20, 2700)\n",
        "\n",
        "ax[1,0].plot(cl[2, 2]*dll_fac,label=\"This notebook\")\n",
        "ax[1,0].plot(powers_lensed_camb[2]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[1,0].set_xlim(20, 2700)\n",
        "\n",
        "ax[1,1].plot(cl[0, 1]*dll_fac,label=\"This notebook\")\n",
        "ax[1,1].plot(powers_lensed_camb[3]*camb_dll_fac,label=\"CAMB lensed\")\n",
        "ax[1,1].set_xlim(20, 2700)\n",
        "\n",
        "ax[0,0].set_xlabel(\"$\\ell$\")\n",
        "ax[0,0].set_ylabel(\"TT\")\n",
        "ax[0,1].set_xlabel(\"$\\ell$\")\n",
        "ax[0,1].set_ylabel(\"EE\")\n",
        "ax[1,0].set_xlabel(\"$\\ell$\")\n",
        "ax[1,0].set_ylabel(\"BB\")\n",
        "ax[1,1].set_xlabel(\"$\\ell$\")\n",
        "ax[1,1].set_ylabel(\"TE\")\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7hBJxBOkNny0"
      },
      "source": [
        "Hooray, that worked well -- look at how much better the BB spectrum matches! Evidently the flat-sky lensing example was indeed biased either by its flat-sky treatment and/or simplistic handling of the mask when measuring the power spectrum. However, we should keep in mind that the lensing is not accurate out to the bandlimit -- we need to build-in some buffer to our analysis by lensing to a high bandlimit but only using larger scales (in this case, e.g., <2000)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AzJ6umw3QylK"
      },
      "source": [
        "**Point Sources**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vMfdDrc6Uz-3"
      },
      "source": [
        "In CMB maps we typically have extragalactic point sources emission and galaxy clusters through the Sunyaev Zel'dovich effect. The former are bright spots in the map and the latter it is a decrement for frequencies below 220 GHz and increment above 220 GHz, the effect it null approximately at 220 GHz.\n",
        "\n",
        "To simulate point sources we need a profile, since these galaxies are usually unresolved points in the sky their profile is the telescope beam. For a 6-meter telescope the full width at half maximum of the beam is roughly 1.4 arcminutes at 150 GHz, we can simplify the profile as Gaussian.\n",
        "\n",
        "Considering that a unit normalized Gaussian centered at zero has the form $\\exp{\\left( -\\frac{r^2}{2\\sigma^2} \\right)}$ and the full width at half maximum is $\\mathrm{FWHM} = 2 \\sqrt{2 \\ln(2)} \\sigma$, we have the following functional form for the beam: $\\exp{\\left(-\\frac{4\\ln(2)r^2}{\\mathrm{FWHM}^2}\\right)}$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NvSpC9LmVHGm"
      },
      "outputs": [],
      "source": [
        "# Generating a Gaussian beam\n",
        "FWHM = 1.4\n",
        "r = np.linspace(0, 60, 1000) # 60 arcminutes or 1 degree\n",
        "B = np.exp(-(4*np.log(2)*r**2) / (FWHM**2))\n",
        "plt.plot(r, B)\n",
        "plt.xlim(0, 5)\n",
        "plt.ylabel(\"Beam amplitude\")\n",
        "plt.xlabel(\"Radius (arcmin)\")\n",
        "# we need to pass the radius to radians\n",
        "r *= 1/60 * np.pi / 180"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_nvc372KVKd-"
      },
      "source": [
        "Now we can generate sources, for that we will use the `pointsrcs` module, in specific the function `sim_objects`, this function requires the following:\n",
        "\n",
        "`shape`: Shape of the map in which we will inject sources, in our case omap.shape\n",
        "\n",
        "`wcs`: World coordinates of the map in which we will inject sources in our case omap.wcs\n",
        "\n",
        "`poss`: Position of the sources that we want to inject in the form [dec,ra] where dec and ra are arrays of floats with the declination and right ascension of the sources in radians\n",
        "\n",
        "`amp`: Amplitude of the sources that we want to inject, this is the peak value of the beam in temperature units as an array of floats\n",
        "\n",
        "`profile`: The radial profile of the sources, in case of point sources this profile is the beam, but for example for cluster it could be a more extended profile, the profile is given as [r,B] where r is the radius in radians and B is the profile amplitude\n",
        "\n",
        "`vmin=None`: The lowest value to simulate in map units, by default it takes 1e-3*amp.\n",
        "\n",
        "`pixwin=False`: If we want to apply the pixel window function after simulating the objects.\n",
        "\n",
        "A nice feature of this function is that it automatically handles the \"distorted\" geometry near the poles, so that radially-symmetric objects will appear horizontally stretched (as they should). This won't be evident in the small cutout below but it is important for a full-sky simulation:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F9-AaBFoVMsW"
      },
      "outputs": [],
      "source": [
        "             # DEC0, RA0  DEC1 RA1\n",
        "box = np.array([[-10, 35],[-2, 15]]) * utils.degree\n",
        "shape, wcs = enmap.geometry(pos=box, res=0.5 * utils.arcmin, proj='car')\n",
        "# we will use only the temperature component or parameter stokes I\n",
        "\n",
        "# we will generate 100 sources\n",
        "nsrc = 100\n",
        "# we choose a logspace between 100 and 10000\n",
        "amp = np.logspace(2.0, 4.0, nsrc)\n",
        "# the position are random values inside omap\n",
        "dec = np.random.uniform(-10, -2, nsrc) * np.pi / 180\n",
        "ra = np.random.uniform(15, 35, nsrc)  *np.pi / 180\n",
        "\n",
        "# we generate the sourcemap here\n",
        "srcmap = pointsrcs.sim_objects(shape, wcs, [dec, ra], amp, [r, B],\n",
        "                               min=np.min(amp)*1e-4)\n",
        "\n",
        "# plotting the result\n",
        "enplot.pshow(srcmap, range=500)\n",
        "\n",
        "# adding to a cmb map\n",
        "cmbmap = curvedsky.rand_map(shape, wcs, ps=ps[0])\n",
        "enplot.pshow(srcmap + cmbmap, range=500)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Simplistic Noise**\n",
        "\n",
        "Finally we consider simple, additive noise that you can add to any of the previous realizations of the underlying \"signal.\" Unlike the signal, which dies off quickly as a function of $\\ell$, the noise is not bandlimited. Realistic noise in the maps is quite complicated (see https://arxiv.org/abs/2303.04180), so here we just show a simplified example. In fact, we've already used all of the tools.\n",
        "\n",
        "We want to go back the case of using `enmap.rand_map` -- the flat-sky code -- to generate Gaussian realizations on the full-sky. This was unphysical for the signal, but for the noise, it is faster, sufficient, and in some ways better (again, the realization is not bandlimited).\n",
        "\n",
        "We define a realistic noise power spectrum and level, and can also exploit that noise is ~symmetric in E and B to just perform a \"scalar\" rather than spin-2 transform:"
      ],
      "metadata": {
        "id": "8mmWPSZwTsbn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WlPiL9zOiBtS"
      },
      "outputs": [],
      "source": [
        "# make a polarized noise ps function\n",
        "def T_and_P_noise_ps(ell, white_level=30, noise_ps_scaling=-4, T_knee=3000,\n",
        "                     T_cap=300, P_knee=300, P_cap=100):\n",
        "  \"\"\"Get the temperature and polarization noise power spectra evaluated at ell.\n",
        "  Follows this model:\n",
        "\n",
        "  PS(ell) = white_level**2 * ((ell/knee)**scaling + 1), ell > cap\n",
        "  PS(ell) = white_level**2 * ((cap/knee)**scaling + 1), ell <= cap\n",
        "\n",
        "  Parameters\n",
        "  ----------\n",
        "  ell : (...) np.ndarray\n",
        "    Angular scales.\n",
        "  white_level : scalar\n",
        "    Temperature white noise level in uK-arcmin. Polarization is this times\n",
        "    sqrt(2).\n",
        "  noise_ps_scaling : scalar\n",
        "    Power-law scaling of the low-ell noise power spectra.\n",
        "  T_knee : scalar\n",
        "    Ell-knee of temperature power spectrum.\n",
        "  T_cap : scalar\n",
        "    Minimum ell at which the spectrum is capped.\n",
        "  P_knee : scalar\n",
        "    Ell-knee of polarization power spectrum.\n",
        "  P_cap : scalar\n",
        "    Minimum ell at which the spectrum is capped.\n",
        "\n",
        "  Returns\n",
        "  -------\n",
        "  (3, ...) np.ndarray\n",
        "    The polarization noise power spectra in T, Q, U. Assumed diagonal over\n",
        "    polarization.\n",
        "  \"\"\"\n",
        "  T = np.zeros_like(ell)\n",
        "  mask = ell <= T_cap\n",
        "  T[mask] = white_level**2 * ((T_cap/T_knee)**noise_ps_scaling + 1)\n",
        "  T[~mask] = white_level**2 * ((ell[~mask]/T_knee)**noise_ps_scaling + 1)\n",
        "\n",
        "  P = np.zeros_like(ell)\n",
        "  mask = ell <= P_cap\n",
        "  P[mask] = 2 * white_level**2 * ((P_cap/P_knee)**noise_ps_scaling + 1)\n",
        "  P[~mask] = 2 * white_level**2 * ((ell[~mask]/P_knee)**noise_ps_scaling + 1)\n",
        "\n",
        "\n",
        "  # convert to steradians (note, not square radians!). the below lines first\n",
        "  # convert to square radians, then from square radians to steradians\n",
        "  T *= utils.arcmin**2 * (4*np.pi / (np.pi * 2*np.pi))\n",
        "  P *= utils.arcmin**2 * (4*np.pi / (np.pi * 2*np.pi))\n",
        "\n",
        "  # put into square shape and return\n",
        "  out = np.zeros((3, 3, *ell.shape), ell.dtype)\n",
        "  out[0, 0] = T\n",
        "  out[1, 1] = P\n",
        "  out[2, 2] = P\n",
        "\n",
        "  return out"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's see what the default looks like as a function of $\\ell$. This includes large-scale correlated noise from the atmosphere as would be observed by the SO LAT, but the SO SAT includes a half-wave plate and processing filters to mitigate this, so the noise spectra would need adjustment:"
      ],
      "metadata": {
        "id": "4649BFRFGl57"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ell = np.arange(10000, dtype=np.float64)\n",
        "noise_ps = T_and_P_noise_ps(ell)\n",
        "plt.loglog(ell, noise_ps[0, 0])\n",
        "plt.loglog(ell, noise_ps[1, 1])"
      ],
      "metadata": {
        "id": "idE-iUuHGWLA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's make a convenience function to draw noise realizations, since this will also need to take into account the area of each pixel to keep the white-noise level consistent:"
      ],
      "metadata": {
        "id": "p3frLSYIAny_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_noise_sim(shape, wcs, seed=None, **T_and_P_noise_ps_kwargs):\n",
        "  \"\"\"Draw a noise realization from the constructed noise powre spectrum\n",
        "  that also accounts for the smaller, and thus noisier, pixels near the\n",
        "  poles.\n",
        "\n",
        "  Parameters\n",
        "  ----------\n",
        "  shape : (ny, nx) tuple\n",
        "    The footprint of the map.\n",
        "  wcs : astropy.wcs.wcs.WCS\n",
        "    The geometry of the map. Assumes units of degrees (the pixell\n",
        "    default for wcs).\n",
        "  seed : int or list of int\n",
        "    Random seed.\n",
        "  T_and_P_noise_ps_kwargs : dict\n",
        "    Keyword arguments to be passed to T_and_P_noise_ps.\n",
        "\n",
        "  Returns\n",
        "  -------\n",
        "  (3, ny, nx) enmap.ndmap\n",
        "    Noise realization (polarized) drawn from T_and_P_noise_ps, with the\n",
        "    correct white-noise level, correlated noise shape, and corrected for\n",
        "    pixel areas.\n",
        "  \"\"\"\n",
        "  # get the noise ps. for enmap.rand_map to work, this needs to be\n",
        "  # evaluated at integer ells and cover all the angular scales in\n",
        "  # out 2d fourier space\n",
        "  ell = np.arange(0, np.ceil(enmap.modlmap(shape, wcs).max()) + 1)\n",
        "  noise_ps = T_and_P_noise_ps(ell, **T_and_P_noise_ps_kwargs)\n",
        "\n",
        "  # draw a noise realization\n",
        "  noise_sim = enmap.rand_map((3, *shape), wcs, cov=noise_ps, seed=seed,\n",
        "                             scalar=True)\n",
        "\n",
        "  # normalize by pixel area. we want this in terms of fraction of a\n",
        "  # \"flat-sky\" pixel\n",
        "  pixsize_steradians = enmap.pixsizemap(shape, wcs, broadcastable=True)\n",
        "  pix_area_deg = np.abs(np.prod(wcs.wcs.cdelt))\n",
        "  pix_area_steradians = pix_area_deg * utils.degree**2\n",
        "  frac_pixsize = pixsize_steradians / pix_area_steradians\n",
        "\n",
        "  return noise_sim / np.sqrt(frac_pixsize)"
      ],
      "metadata": {
        "id": "xHb547KoHEd2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's try it out on the full-sky:"
      ],
      "metadata": {
        "id": "zbaSqjRtKnTc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "shape, wcs = enmap.fullsky_geometry(res=4*utils.arcmin, proj='car')\n",
        "\n",
        "noise_sim = get_noise_sim(shape, wcs)\n",
        "\n",
        "for i in range(3):\n",
        "  enplot.pshow(noise_sim[i], downgrade=4, ticks=15, colorbar=True)"
      ],
      "metadata": {
        "id": "Sj0bsLPVKqBy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Put it all together\n",
        "\n",
        "Let's take all the above pieces and make a not-so-unrealistic simulation of data as would be observed by an actual CMB experiment:\n",
        "\n",
        "\n",
        "\n",
        "1.   We'll draw a full-sky unlensed T, Q, U CMB realization and a full-sky realization of the lensing potential and use it to \"lens\" the CMB\n",
        "2.   We'll add a new step to the above: convolve the resulting signal with\n",
        "the instrumental beam so that it is consistent with the point sources\n",
        "3.   We'll inject a sample of point sources with the same beam size. Assume they are not polarized\n",
        "4.   We'll add a T, Q, U noise realization to the signal\n",
        "\n",
        "This simulation omits quite a lot of real-world complexity in each of its parts --- we won't include any scan-synchronous signal, beam leakage or pixel window function, polarized or extended or \"SZ-like\" sources, or realistic noise --- but it is actually fairly representative.\n",
        "\n",
        "We'll wrap it in a big convenience function. Play around with all the inputs or add even more arguments to make it more tunable!\n",
        "\n",
        "(Exercise: An obvious way to streamline the below function would be to allow the user to pass pre-built signal or noise maps. That way someone could more efficiently add multiple noise realizations to the same signal realization, without rebuilding the same signal map each time)\n"
      ],
      "metadata": {
        "id": "oOnru2LQLR-3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import healpy as hp\n",
        "\n",
        "def get_signal_and_noise_sim(shape, wcs, lmax, ps_cmb, ps_lens, cmb_seed=None,\n",
        "                             phi_seed=None, src_pos_seed=None, noise_seed=None,\n",
        "                             beam_fwhm=1.4, nsrc=10000, white_level=30,\n",
        "                             noise_ps_scaling=-4, T_knee=3000, T_cap=300,\n",
        "                             P_knee=300, P_cap=100):\n",
        "  \"\"\"Draw a semi-realistic end-to-end sim including lensed CMB, temperature\n",
        "  point sources, instrumental beam, and instrumental and atmospheric noise.\n",
        "\n",
        "  Parameters\n",
        "  ----------\n",
        "  shape : (ny, nx) tuple\n",
        "    The footprint of the map.\n",
        "  wcs : astropy.wcs.wcs.WCS\n",
        "    The geometry of the map. Assumes units of degrees (the pixell\n",
        "    default for wcs).\n",
        "  lmax : int\n",
        "    maximum multipole for the signal\n",
        "  ps_cmb : (3, 3, nl) np.ndarray\n",
        "    CMB TEB power spectra with covariances\n",
        "  ps_lens : (nl,) np.ndarray\n",
        "    Lensing potential power spectrum\n",
        "  cmb_seed : int or list of int\n",
        "    Seed for unlensed CMB\n",
        "  phi_seed : int or list of int\n",
        "    Seed for lensing potential\n",
        "  src_pos_seed : int or list of int\n",
        "    Seed for point source locations\n",
        "  noise_seed : int or list of int\n",
        "    Seed for noise\n",
        "  beam_fwhm : scalar\n",
        "    Beam width in arcmin\n",
        "  nsrc : int\n",
        "    Number of point sources logarithmically distributed between\n",
        "    100 and 10000 uK\n",
        "  white_level : scalar\n",
        "    Temperature white noise level in uK-arcmin. Polarization is this times\n",
        "    sqrt(2).\n",
        "  noise_ps_scaling : scalar\n",
        "    Power-law scaling of the low-ell noise power spectra.\n",
        "  T_knee : scalar\n",
        "    Ell-knee of temperature power spectrum.\n",
        "  T_cap : scalar\n",
        "    Minimum ell at which the spectrum is capped.\n",
        "  P_knee : scalar\n",
        "    Ell-knee of polarization power spectrum.\n",
        "  P_cap : scalar\n",
        "    Minimum ell at which the spectrum is capped.\n",
        "\n",
        "  Returns\n",
        "  -------\n",
        "  (3, ny, nx) enmap.ndmap\n",
        "    The simulation\n",
        "  \"\"\"\n",
        "\n",
        "  # get lensed CMB. assume lensing indep. of CMB (not true)\n",
        "  phi_alm = curvedsky.rand_alm(ps=ps_lens, lmax=lmax, seed=phi_seed)\n",
        "  cmb_alm = curvedsky.rand_alm(ps=ps_cmb, lmax=lmax, seed=cmb_seed)\n",
        "  lensedcmb = lensing.lens_map_curved((3, *shape), wcs, phi_alm, cmb_alm,\n",
        "                                      output=\"l\", verbose=True)[0]\n",
        "\n",
        "  # convolve it with the beam. fwhm in radians\n",
        "  bl = hp.gauss_beam(beam_fwhm * utils.arcmin, lmax=lmax)\n",
        "  lensedcmb_alm = curvedsky.map2alm(lensedcmb, lmax=lmax)\n",
        "  curvedsky.almxfl(lensedcmb_alm, bl, out=lensedcmb_alm)\n",
        "  curvedsky.alm2map(lensedcmb_alm, lensedcmb)\n",
        "\n",
        "  # add srcs (already convolved with beam)\n",
        "  r = np.linspace(0, 60, 1000) * utils.arcmin # 60 arcminutes or 1 degree\n",
        "  B = np.exp(-4*np.log(2)*r**2 / (beam_fwhm*utils.arcmin)**2)\n",
        "  amp = np.logspace(2.0, 4.0, nsrc)\n",
        "\n",
        "  src_rng = np.random.default_rng(src_pos_seed)\n",
        "  dec = src_rng.uniform(-90, 90, nsrc) * utils.degree\n",
        "  ra = src_rng.uniform(-180, 180, nsrc)  * utils.degree\n",
        "  srcmap = pointsrcs.sim_objects(shape, wcs, [dec, ra], amp, [r, B],\n",
        "                                 vmin=np.min(amp)*1e-4)\n",
        "\n",
        "  # and noise\n",
        "  noisemap = get_noise_sim(shape, wcs, seed=noise_seed, white_level=white_level,\n",
        "                           noise_ps_scaling=noise_ps_scaling, T_knee=T_knee,\n",
        "                           T_cap=T_cap, P_knee=P_knee, P_cap=P_cap)\n",
        "\n",
        "  out = lensedcmb + noisemap\n",
        "  out[0] += srcmap\n",
        "  return out"
      ],
      "metadata": {
        "id": "d7YD88NvLU_B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "totalsim = get_signal_and_noise_sim(shape, wcs, 2700, ps_cmb, ps_lens)\n",
        "\n",
        "for i in range(3):\n",
        "  enplot.pshow(totalsim[i], downgrade=4, ticks=15, colorbar=True)"
      ],
      "metadata": {
        "id": "tA4kJ4eLTIFa"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}