CMB_School_Part_09.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jeffmcm1977/CMBAnalysis_SummerSchool/blob/master/CMB_School_Part_09.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mapmaking\n",
    "\n",
    "### Code from Sigurd Naess, addapted by Jeff McMahon\n",
    "\n",
    "\n",
    "CMB telescopes observe by scanning accross the sky and recording the response of each detector as a funciton of time as well as pointing information.  The detector response time series is commonly refered to as time ordered data or a TOD.  Each CMB field is usually observed many times over the course of one or more seasons yielding a large nunber of overlapping TODs.   At the pole the sky rotation is alligmed with the elevation axis of a telescope resulting in paralell scans in sky coordinates.  At midlattitude sites there is no such alignment and sky rotation can produce cross-linked scans.  Offline the TODs are processed to reconstruct the maps.   \n",
    "\n",
    "In this notebook you will (1) model TODs for a map that is obsered horizontally and vertically to create a cross-linked set of TODs, (2) add 1/f noise to create TODs similar to what would be output from a CMB instrument, and (3) you will use several mapmaking codes to reconstruct maps from these time streams.   The mapmaking codes include (i) a simple binned average that works well if the TODs have negliagble white noise (from the instrument (e.g., ddemodulated half-wave plate polarziation data) or by prewhitening the data with a filter), and (ii) an iterative maximum liklihood mapmaker that accounts for the noise correlations from the atmosphere.\n",
    "\n",
    "Excercises will guide you through the differences between cross-linked and non-crosslinked maps, and the impacts of filtering before and after mapmaking."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First we read in the necessary libraries from previous code modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -c \"import cmb_modules\" || ( \\\n",
    "    wget https:/jeffmcm1977/CMBAnalysis_SummerSchool/raw/master/cmb_school.tar.gz && \\\n",
    "    tar xzvf cmb_school.tar.gz \\\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib.mlab as mlab\n",
    "import matplotlib.pyplot as plt\n",
    "import astropy.io.fits as fits\n",
    "\n",
    "import constants as cs # the constants module\n",
    "import cmb_modules # the module of functions\n",
    "\n",
    "N = int(cs.N)\n",
    "c_min = cs.c_min\n",
    "c_max = cs.c_max\n",
    "X_width =cs.X_width\n",
    "Y_width = cs.Y_width\n",
    "beam_size_fwhp = cs.beam_size_fwhp\n",
    "\n",
    "pix_size = cs.pix_size\n",
    "print(\"pixel size: \", pix_size)\n",
    "\n",
    "Number_of_Sources  = cs.Number_of_Sources\n",
    "Amplitude_of_Sources = cs.Amplitude_of_Sources\n",
    "Number_of_Sources_EX = cs.Number_of_Sources_EX\n",
    "Amplitude_of_Sources_EX = cs.Amplitude_of_Sources_EX\n",
    "\n",
    "Number_of_SZ_Clusters  = cs.Number_of_SZ_Clusters\n",
    "Mean_Amplitude_of_SZ_Clusters = cs.Mean_Amplitude_of_SZ_Clusters\n",
    "SZ_beta = cs.SZ_beta\n",
    "SZ_Theta_core = cs.SZ_Theta_core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recalculating the results from the previous stages\n",
    "\n",
    "This creates a map with CMB, point sources, and SZ sources.  It also convolces the map with the beam.  Unlike previous maped based simulations, noise will be added in the time series which are calculated below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Make a CMB map\n",
    "ell, DlTT = np.loadtxt(\"CAMB_fiducial_cosmo_scalCls.dat\", usecols=(0, 1), unpack=True) \n",
    "CMB_T = cmb_modules.make_CMB_T_map(N,pix_size,ell,DlTT)\n",
    "\n",
    "## make a point source map\n",
    "PSMap = cmb_modules.Poisson_source_component(N,pix_size,Number_of_Sources,Amplitude_of_Sources) \n",
    "PSMap += cmb_modules.Exponential_source_component(N,pix_size,Number_of_Sources_EX,Amplitude_of_Sources_EX)\n",
    "\n",
    "## make an SZ map\n",
    "SZMap,SZCat = cmb_modules.SZ_source_component(N,pix_size,Number_of_SZ_Clusters,\\\n",
    "                                  Mean_Amplitude_of_SZ_Clusters,SZ_beta,SZ_Theta_core,False)\n",
    "\n",
    "## add them all together to get the sky map at a single freuqency\n",
    "total_map = CMB_T + PSMap + SZMap\n",
    "\n",
    "\n",
    "## convolve with the beam\n",
    "CMB_T_convolved = cmb_modules.convolve_map_with_gaussian_beam(N,pix_size,beam_size_fwhp,total_map)\n",
    "\n",
    "## plot the result\n",
    "p=cmb_modules.Plot_CMB_Map(CMB_T_convolved,c_min,c_max,X_width,Y_width)\n",
    "\n",
    "\n",
    "\n",
    "## save the result as \"sky\", a variable that represents the simulated sky we will map below\n",
    "sky = total_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating time streams from a simulated sky map\n",
    "\n",
    "This is the first new piece of the notebook.  In this section we simulate the telesacope scanning accross the map.  For simplicity we consider two scan directions: left-right, and up-down. This code generates a list of pixels (x- and y coordinates) in the order in which they are observed.  \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## a function to return the sequence of pixel indices as they are \"observed\" in these simualtions\n",
    "def sim_pointing(map, dir=0):\n",
    "    \"\"\"Simulate a telescope scanning across the given map. The scanning pattern is\n",
    "    as simple as possible: The samples hit the center of each pixel, and we\n",
    "    scan rowwise (dir=0) or columnwise (dir=1).\"\"\"\n",
    "    # The pointing is an [{y,x},nsamp] array of pixel positions\n",
    "    # The einsum stuff is just to swap the second and third axis\n",
    "    # of pixmap, which contains the pixel coordinates of each pixel.\n",
    "    pixmap = np.mgrid[:map.shape[-2],:map.shape[-1]]  ## makes two arrays of the x and y indices in the map\n",
    "    pixmap[1,1::2, :] = pixmap[1,1::2, ::-1]          ## reverse ever other row so the scans go back and forth\n",
    "    if dir == 0: point = pixmap.reshape(2,-1)         ## arranges these for L-R scans\n",
    "    else:        point = np.roll(pixmap,1,axis=0).reshape(2,-1)   ## arranges these for U-D scans\n",
    "    return point\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "## generate left-right scans\n",
    "point_lr = sim_pointing(sky, 0)\n",
    "\n",
    "## generate Up-Down scans\n",
    "point_ud = sim_pointing(sky, 1)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "plt.plot(point_lr[1,:],\"r,\")\n",
    "plt.plot(point_lr[0,:],\"k,\")\n",
    "plt.title(\"Left-Right Scans Poointing (black elevation, red azimuth)\")\n",
    "plt.ylabel(\"pixel number\")\n",
    "plt.xlabel(\"TOD sample\")\n",
    "plt.xlim(0,1e6)\n",
    "plt.show()\n",
    "\n",
    "plt.plot(point_ud[0,:],\"k,\")\n",
    "plt.plot(point_ud[1,:],\"r,\")\n",
    "plt.title(\"Up-Down Scans Poointing (black elevation, red azimuth)\")\n",
    "plt.ylabel(\"pixel number\")\n",
    "plt.xlabel(\"TOD sample\")\n",
    "plt.xlim(0,1e6)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Make a zoomed in plot of the scan strategies above by adjusting the x limit for the plot to be much smaller.     Discribe the scan strategy based on what you see in these plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make your plots here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### discusson goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### noisless TODs\n",
    "\n",
    "This code takes the pointing information genrated above an makes a mock TOD with no noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Observe_map(map, point):\n",
    "    \"\"\"Pointing matrix: Project map to tod\"\"\"\n",
    "    point = np.round(point).astype(int)\n",
    "    return map[point[0],point[1]]  ## return the value of the map at each pointing, this forms the simulated time stream\n",
    "\n",
    "\n",
    "\n",
    "noisless_tod_lr = Observe_map(sky, point_lr)\n",
    "noisless_tod_ud = Observe_map(sky, point_ud)\n",
    "\n",
    "\n",
    "\n",
    "plt.plot(noisless_tod_lr)\n",
    "plt.plot(noisless_tod_ud)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Explain the spikes in the TOD and discuss weather the RMS of the TOD makes sense."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### discussion goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>OPTIONAL EXERCISE: </font> Add in the impact of detector time constants to smear out these TODs.  This excercise is best done aftern completig the rest of the notebook so its impact can be propigatd through this entire code to maps.  This will allow you to model the impact of this sytematic on effect.\n",
    "\n",
    "NOTE: [this wikipedia article](https://en.wikipedia.org/wiki/Exponential_smoothing) is useful in modeling exponetial filters (e.g., detector time cosntants) This apporach can be implemented with arrays using the np.roll funciton."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding noise to the TOD\n",
    "\n",
    "Here we build a noise power spectrum for the atmospheric and instrumental noise.   By default we assume the instrumental noise has an amplitude of 40 $\\mu$-K per sample (1/2 arcmin pixels) which is similar to the full exposure of ACT-MBAC or SPT-SZ which were around 20 $\\mu$ K- arcmin (SPT was slightly lower).   The atmospehric signal is modeled as a simple power law with a 0.1 s 1/f knee and a scan that speends $1/240$ s in each pixel which is equivelent to a $1^\\circ/$s scan speed.   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_noise_spec(nsamp, dt=.00416, fknee=0.1, alpha=3, sigma=40):\n",
    "    \"\"\"Build a simple atmosphere + white noise model, and return it\n",
    "    as a power spectrum.\"\"\"\n",
    "    freq   = np.abs(np.fft.fftfreq(nsamp, dt))\n",
    "    return (1+(np.maximum(freq,freq[1])/fknee)**-alpha)*sigma**2\n",
    "\n",
    "\n",
    "\n",
    "noise_spec_lr = sim_noise_spec(point_lr.shape[-1], dt=.00416, fknee=0.1, alpha=3, sigma=40)\n",
    "noise_spec_ud = sim_noise_spec(point_ud.shape[-1], dt=.00416, fknee=0.1, alpha=3, sigma=40)\n",
    "freq   = np.abs(np.fft.fftfreq(point_lr.shape[-1], .00416))\n",
    "\n",
    "plt.loglog(freq,noise_spec_lr)\n",
    "plt.title(\"noise power spectrum\")\n",
    "plt.xlabel(\"frequency [Hz]\")\n",
    "plt.ylabel(\"Noise amplitude [uK$^{2}$]\")\n",
    "plt.ylim(1e3,1e5)\n",
    "plt.xlim(1e-2,10)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to generate a realziation of this noise and add it to the noisless time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_tod(map, point, noise_spec):\n",
    "    \"\"\"Simulate a noisy TOD using the model d = Pm + n\"\"\"\n",
    "    tod    = Observe_map(map, point)\n",
    "    rand   = np.fft.fft(np.random.standard_normal(tod.shape[-1]))\n",
    "    fnoise = rand * noise_spec**0.5\n",
    "    tod   += np.fft.ifft(fnoise).real\n",
    "    return tod\n",
    "\n",
    "\n",
    "\n",
    "tod_lr = sim_tod(sky, point_lr, noise_spec_lr)\n",
    "tod_ud = sim_tod(sky, point_ud, noise_spec_ud)\n",
    "\n",
    "plt.plot(tod_lr)\n",
    "plt.plot(tod_ud)\n",
    "plt.title(\"simulated time series with noise\")\n",
    "plt.ylabel(\"response [$\\mu$K]\")\n",
    "plt.xlabel(\"TOD sample [integer]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Typical detectors have noises of $\\sim 300 \\mu$K-$\\sqrt{s}$.  Make a plot of what a time stream looks like with this instantaneous sensitivity so you can get a feel for what real data looks like."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### write down the noise and 1/f knee you choose here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## make your plots here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Putting all the pieces together\n",
    "\n",
    "For convenience, the following code bundles the TOD (singal and noise), with pointing information, and the noise spectrum into a class.   This object will be passed as input to the map making codes we develope below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_dataset(map, num_data=2, dt=.00416, fknee=0.1, alpha=3, sigma=40):\n",
    "    \"\"\"Simulate a dataset consisting of num_data scans across the sky.\n",
    "    Returns a list of Data objects, each of which contains the tod,\n",
    "    the pointing and the noise spectrum.\"\"\"\n",
    "    res = []\n",
    "    for i in range(num_data):\n",
    "        point = sim_pointing(map, i % 2)\n",
    "        noise_spec = sim_noise_spec(point.shape[-1], dt=dt, fknee=fknee, alpha=alpha, sigma=sigma)\n",
    "        tod = sim_tod(map, point, noise_spec)\n",
    "        res.append(Data(tod, point, noise_spec))\n",
    "    return res\n",
    "\n",
    "\n",
    "class Data:\n",
    "    def __init__(self, tod, point, noise_spec):\n",
    "        self.tod   = tod\n",
    "        self.point = point\n",
    "        self.noise_spec = noise_spec\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print(\"Generating noisy TOD simulations\")\n",
    "dataset  = sim_dataset(sky,num_data=2, dt=.00416, fknee=0.1, alpha=3, sigma=40)\n",
    "for data in dataset:\n",
    "    plt.plot(data.tod)\n",
    "    \n",
    "plt.title(\"simulated time series with noise\")\n",
    "plt.ylabel(\"response [$\\mu$K]\")\n",
    "plt.xlabel(\"TOD sample [integer]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# naieve mapmaking\n",
    "\n",
    "Making maps from TODs is an inverse problem.  It can be epxressed in matrices as:\n",
    "$$ P m = d.$$\n",
    "Here $d$ is the TOD (whcih is a vector), $m$ is the sky map (whcih is also a vector), and $P$ is a (typically) rectangular matrix that specifies which pixel in the map was observed during each TOD sample.  (NOTE: many TOD sampeles usually correspond to the same pixel in our sky map.   \n",
    "\n",
    "To solve for a map $m$ givein our TODs ($d$) we must invert this equaiton.   The standard technique is to use chi-squared minimizaiton.   Expressed in matries this apporach can be accomplised as derived here: \n",
    "$$ P m = d$$\n",
    "$$ P^tP m  = P^t d      $$\n",
    "$$ m = [P^tP]^{-1} P^t d.  $$\n",
    "\n",
    "In other wordes, we can use the pointing matrix to construct a matrix $[P^tP]^{-1} P^t$ that acts on the TODs to return an estimate for the map.  The difficulty is implementing this when the numnber of smaples in the TODs are large, and in the presence of noise.\n",
    "\n",
    "We have entirely ignored noise at this point.  (Noise correlations will be dealt with below.  \n",
    "However, in the limit where the noise is white the above equation can be solved with a simple binned average.  In this case one simply averages the TOD samples which observae a given pixel and returns a map comprized of these averages.  We implement this below and highlight the value and limitations of such code.  (The value is it is fast)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_binned_average(dataset, shape):\n",
    "    \"\"\"Solve the simplified mapmaking equation Ax=b,\n",
    "    where A = P'P and b = P'd, e.g. ignoring noise\n",
    "    properties such as correlations.\"\"\"\n",
    "    rhs  = np.zeros(shape)\n",
    "    hits = np.zeros(shape)\n",
    "    for data in dataset:\n",
    "        rhs  += PT(data.tod,    data.point, shape)\n",
    "        hits += PT(data.tod*0+1, data.point, shape)\n",
    "    return rhs/hits\n",
    "\n",
    "def PT(tod, point, shape):\n",
    "    \"\"\"Transpose pointing matrix.\"\"\"\n",
    "    point = np.round(point).astype(int)\n",
    "    point_flat = np.ravel_multi_index(point, shape[-2:])\n",
    "    map = np.bincount(point_flat, tod, minlength=shape[-2]*shape[-1])\n",
    "    map = map.reshape(shape[-2:])\n",
    "    return map\n",
    "\n",
    "\n",
    "print(\"Solving for map while ignoring noise correlations\")\n",
    "map_binned = solve_binned_average(dataset, sky.shape)\n",
    "p=cmb_modules.Plot_CMB_Map(map_binned,c_min,c_max,X_width,Y_width)\n",
    "##p=cmb_modules.Plot_CMB_Map(map_binned-sky,c_min,c_max,X_width,Y_width)  #for comparison "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Do you see comon structures in this map and the input sky map?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "comment here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Remake the TOD and this map setting the 1/f knee to be absurly small.  Compare the new result to the originaly input map?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "comment here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  reset the 1/f knee to its original value.  Now filter the time-series using a high-pass filter with a knee frequency set to whiten the TODs.   Remake the map and compare the new result to the input map.  What is different?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "comment here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# maximum liklihood mapmaking including noise correlations\n",
    "\n",
    "Based on what we have seen above we need to develope a map-making apporach that accounts for the noise correlations in our data.   To do this we implement a maximum liklihood mapmaking approiach.  Expressed with matrices this is much like what we did above, except we include weighting by a matrix $N;^{-2}$ whose matdix elements are $1/\\sigma^2_{ij}$ and encode the weight of each sample and correlations between samples in the TOD.   WIth this matrix our mapmaking apporach can be derived as follows: \n",
    "$$ P m = d$$\n",
    "$$ N^{-2} P m = N^{-2} d$$\n",
    "$$ P^tN^{-2}P m  = P^tN^{-2} d      $$\n",
    "$$ m = [P^tN^{-2}P]^{-1} P^tN^{-2} d.  $$\n",
    "\n",
    "In summary, the matrix $[P^tN^{-2}P]^{-1} P^tN^{-2}$ can be applied to the TOD $d$ to produce the map $m$.\n",
    "\n",
    "Given that the matix $N^{-2}$ must have off diagonal terms to encoide noise correlations such as from $1/f$ noise, building this matrix is highly non-trival.  We thus resort to conjugate gradient technique [(see Wikipedia article)](https://en.wikipedia.org/wiki/Conjugate_gradient_method) to invert this equaiton with out ever explicitly building this materix.   These technqiues are iteratative and thus we must run many iterations to arrive at an acceptable estimate for our map. \n",
    "\n",
    "We now run this for both a cross-linked and un-cross-linked maps.  Excercises follow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_full(dataset, shape, niter=100, verbose=True):\n",
    "    \"\"\"Solve the full map-making equation\n",
    "    Ax=b, where A = P'N\"P and b = P'N\"d.\"\"\"\n",
    "    # Set up our A matrix. We don't compute\n",
    "    # explicitly because it's too big. Instead,\n",
    "    # we define it as a function that can be applied\n",
    "    # to a vector x. We will then use Conjugate Gradients\n",
    "    # to invert it.\n",
    "    def A(x):\n",
    "        # x is 1d because the conjugate gradient solver works\n",
    "        # on 1d arrays. So start by expanding it to 2d.\n",
    "        x   = x.reshape(shape)\n",
    "        res = x*0\n",
    "        for data in dataset:\n",
    "            tod  = Observe_map(x, data.point)\n",
    "            tod  = mul_inv_noise(tod, data.noise_spec)\n",
    "            res += PT(tod, data.point, shape)\n",
    "        return res.reshape(-1)\n",
    "    # Build our right-hand side b\n",
    "    b = np.zeros(shape)\n",
    "    for data in dataset:\n",
    "        tod = mul_inv_noise(data.tod, data.noise_spec)\n",
    "        b  += PT(tod, data.point, shape)\n",
    "    # And solve\n",
    "    cg = CG(A, b.reshape(-1))\n",
    "    while cg.i < niter:\n",
    "        cg.step()\n",
    "        if verbose: print(\"iteration: %4d conjugate gradient error: %15.7e\" % (cg.i, cg.err))\n",
    "    return cg.x.reshape(shape)\n",
    "\n",
    "\n",
    "\n",
    "def mul_inv_noise(tod, noise_spec):\n",
    "    \"\"\"Multiply by the inverse noise matrix. We assume that the noise\n",
    "    is stationary, which means that it can be represented by a simple\n",
    "    power spectrum noise_spec. This function is used to apply inverse\n",
    "    variance weighting to the data.\"\"\"\n",
    "    ftod  = np.fft.fft(tod)\n",
    "    ftod /= noise_spec\n",
    "    return np.fft.ifft(ftod).real\n",
    "\n",
    "\n",
    "def default_M(x):     return np.copy(x)\n",
    "def default_dot(a,b): return a.dot(np.conj(b))\n",
    "class CG:\n",
    "    \"\"\"A simple Preconditioned Conjugate gradients solver. Solves\n",
    "    the equation system Ax=b.\"\"\"\n",
    "    def __init__(self, A, b, x0=None, M=default_M, dot=default_dot):\n",
    "        \"\"\"Initialize a solver for the system Ax=b, with a starting guess of x0 (0\n",
    "        if not provided). Vectors b and x0 must provide addition and multiplication,\n",
    "        as well as the .copy() method, such as provided by numpy arrays. The\n",
    "        preconditioner is given by M. A and M must be functors acting on vectors\n",
    "        and returning vectors. The dot product may be manually specified using the\n",
    "        dot argument. This is useful for MPI-parallelization, for example.\"\"\"\n",
    "        # Init parameters\n",
    "        self.A   = A\n",
    "        self.b   = b\n",
    "        self.M   = M\n",
    "        self.dot = dot\n",
    "        if x0 is None:\n",
    "            self.x = b*0\n",
    "            self.r = b\n",
    "        else:\n",
    "            self.x   = x0.copy()\n",
    "            self.r   = b-self.A(self.x)\n",
    "        # Internal work variables\n",
    "        n = b.size\n",
    "        self.z   = self.M(self.r)\n",
    "        self.rz  = self.dot(self.r, self.z)\n",
    "        self.rz0 = float(self.rz)\n",
    "        self.p   = self.z\n",
    "        self.err = np.inf\n",
    "        self.d   = 4\n",
    "        self.arz = []\n",
    "        self.i   = 0\n",
    "    def step(self):\n",
    "        \"\"\"Take a single step in the iteration. Results in .x, .i\n",
    "        and .err being updated. To solve the system, call step() in\n",
    "        a loop until you are satisfied with the accuracy. The result\n",
    "        can then be read off from .x.\"\"\"\n",
    "        Ap = self.A(self.p)\n",
    "        alpha = self.rz/self.dot(self.p, Ap)\n",
    "        self.x += alpha*self.p\n",
    "        self.r -= alpha*Ap\n",
    "        self.z = self.M(self.r)\n",
    "        next_rz = self.dot(self.r, self.z)\n",
    "        self.err = next_rz/self.rz0\n",
    "        beta = next_rz/self.rz\n",
    "        self.rz = next_rz\n",
    "        self.p = self.z + beta*self.p\n",
    "        self.arz.append(self.rz*alpha)\n",
    "        self.i += 1\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print(\"Solving for map while taking noise correlations into account\")\n",
    "map_full  = solve_full(dataset, sky.shape,niter=50)\n",
    "\n",
    "p=cmb_modules.Plot_CMB_Map(map_full - np.mean(map_full),c_min,c_max,X_width,Y_width)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### maps with out cross linking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Solving for map while taking noise correlations into account\")\n",
    "map_single = solve_full(dataset[0:1], sky.shape,niter=50)\n",
    "\n",
    "p=cmb_modules.Plot_CMB_Map(map_single - np.mean(map_single),c_min,c_max,X_width,Y_width)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Compare the cross-linked and non-crosslinked maps.  The horizontal strips in the non-crosslinked maps are refered to (uncreatively) as stripes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Make plots of these maps in 2-d fourier space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Compare these 2d fourier plots here. Be sure to note where modes are missing or de-weighted in each map. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='red'>EXERCISE: </font>  Filter the non-crosslinked map to remove the stripes.  The can be accomplised by masking the relevent modes in your 2d fourier plot before inverse fouerier transforming to create the map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "comment on the difference between the filter map non-crosslinked map, the crosslinked map, and the input map.   Make difference plots if that seems useful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}