#!/usr/bin/env python
# coding: utf-8

# # Bayesian signal reconstruction with Gaussian Random Fields (Wiener Filtering) - CMB example

# Florent Leclercq,<br/>
# Imperial Centre for Inference and Cosmology, Imperial College London,<br/>
# florent.leclercq@polytechnique.org

# This notebook requires the CLASS code (https://class-code.net) and its python wrapper classy, as well as the Planck colormap (https://raw.githubusercontent.com/zonca/paperplots/master/data/Planck_Parchment_RGB.txt).
# Optionally, pre-computed data are available:<br/>
# * http://www.florent-leclercq.eu/data/Bayes_InfoTheory/sqrtsignalcovarPix.npy (2.4 GB)
# * http://www.florent-leclercq.eu/data/Bayes_InfoTheory/sqrtCovWF.npy (2.4 GB)
# * http://www.florent-leclercq.eu/data/Bayes_InfoTheory/CovWFinvnoisecovarmat.npy (2.4 GB)
# * http://www.florent-leclercq.eu/data/Bayes_InfoTheory/invnoisecovar.npy (98.4 kB)

# In[1]:


import numpy as np
import scipy.linalg
import os.path
import matplotlib.pyplot as plt
from classy import Class
import healpy as hp
np.random.seed(123457)
get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


# Download and define the Planck color map
from matplotlib.colors import ListedColormap
if not os.path.isfile("data/Planck_Parchment_RGB.txt"):
    get_ipython().system('wget https://raw.githubusercontent.com/zonca/paperplots/master/data/Planck_Parchment_RGB.txt --directory-prefix=data/')
planck = ListedColormap(np.loadtxt("data/Planck_Parchment_RGB.txt")/255.)
planck.set_bad("gray") # color of missing pixels
planck.set_under("white") # color of background


# ## Setup Cls using CLASS

# In[3]:


# Define cosmology (what is not specified will be set to CLASS default parameters)
paramsLCDM = {
    'h':0.702,
    'n_s':0.9619,
    'Omega_b':0.045,
    'Omega_cdm':0.272-0.045,
    'output':'tCl lCl',
    'l_max_scalars': 1000,
    'ic':'ad',
}
# Create an instance of the CLASS wrapper
cosmoLCDM = Class()
# Set the parameters to the cosmological code
cosmoLCDM.set(paramsLCDM)
# Run the whole code.
cosmoLCDM.compute()


# In[4]:


# Compute the power spectrum for the current set of cosmological parameters
res=cosmoLCDM.raw_cl(1000)
ell=res['ell']
Cl=res['tt']


# In[5]:


# Clean CLASS
cosmoLCDM.struct_cleanup()


# In[6]:


# Define a function of l from the arrays
from scipy.interpolate import InterpolatedUnivariateSpline
Clfunc = InterpolatedUnivariateSpline(ell, Cl, k=2)


# In[7]:


# Plot l*(l+1)*Cl/(2*pi)
plt.xlabel("$\ell$",size=18)
plt.ylabel("$\ell(\ell+1) C_\ell$",size=18)
plt.loglog(ell,Cl*(ell*(ell+1))/(2*np.pi),label='CLASS output')
plt.loglog(np.arange(0,999,2),Clfunc(np.arange(0,999,2))*(np.arange(0,999,2)*(np.arange(0,999,2)+1))/(2*np.pi)
           ,label='Spline function')
plt.legend(loc='best')
plt.show()


# In[8]:


# Setup notations
nside=32
npix=hp.pixelfunc.nside2npix(nside)
lmax=min(len(Cl)-1, 3*nside-1)
nalm=hp.Alm.getsize(lmax)


# Above $\mathrm{nside}=64$ the covariance matrix can't be stored...

# ## Setup signal covariance matrix

# In[9]:


# Setup covariance matrix in harmonic space
covarFourier=np.zeros(nalm,dtype=complex)
for l in range(lmax+1):
    for m in range(lmax+1):
        idx=hp.Alm.getidx(lmax,l,m)
        covarFourier[idx]=Clfunc(l)


# In[10]:


# Plot diagonal of the harmonic-space covariance matrix
plt.xlabel("$a_{\ell m}$",size=18)
plt.ylabel("$C_{\ell}$",size=18)
plt.loglog(np.arange(nalm),covarFourier.real)
plt.show()


# Each row of the covariance matrix in pixel space is given by:
# $C_i = \mathcal{F}^{-1} \left[ C_{a_{\ell m}a_{\ell m}}a_{\ell m}(i) \right]$
# where
# $a_{\ell m}(i) = \mathcal{F} \left[ e(i) \right]$.<br/>
# $e(i) = [\delta^{i,j}]_{0 \leq j < N_\mathrm{pix}}^\mathrm{T}$ is the unit vector of the pixel space basis.<br/>
# $\mathcal{F}$ and $\mathcal{F}^{-1}$ are transformations from harmonic space to pixel space:
# \begin{eqnarray}
#     m(n) & = & \sum_{\ell m} a_{\ell m} Y_{\ell m}(n)\\
#     a_{\ell m} & = & \sum_n m(n) Y_{\ell m}^*(n).
# \end{eqnarray}

# In[11]:


get_ipython().run_cell_magic('capture', '', 'if os.path.isfile("data/sqrtsignalcovarPix.npy"):\n    # Load precomputed sqrt covariance matrix in pixel space\n    sqrtsignalcovarPix=np.load("data/sqrtsignalcovarPix.npy")\nelse:\n    # Setup covariance matrix in pixel space\n    covarFouriermat=np.diagflat(covarFourier)\n    signalcovarPix=np.zeros((npix,npix))\n    for i in xrange(npix):\n        ei=np.zeros(npix)\n        ei[i]=1.\n        alm=hp.sphtfunc.map2alm(ei)\n        Calm=covarFouriermat.dot(alm)\n        signalcovarPix[i]=hp.sphtfunc.alm2map(Calm,nside)\n    signalcovarPix=signalcovarPix.T\n    sqrtsignalcovarPix=scipy.linalg.sqrtm(signalcovarPix)\n    del covarFouriermat, signalcovarPix\n    np.save("data/sqrtsignalcovarPix",signalcovarPix)\n')


# The signal covariance matrix is diagonal in harmonic space, but dense in pixel space...

# In[12]:


# Show sqrt covariance matrix in pixel space
i=1000
ei=np.zeros(npix)
ei[i]=1.
hp.mollview(ei, coord='G', title='one pixel', unit='mK',cmap="Blues")
hp.mollview(sqrtsignalcovarPix.real[i], coord='G', title='one row of the sqrt covariance matrix', unit='mK',cmap="Blues")
plt.show()


# ## Setup mask and noise covariance

# In[13]:


get_ipython().run_cell_magic('capture', '', '# Load the mask\nfn = "data/TMASK_32_R2.01_full.fits"\ntmask_map = hp.read_map(fn)\n# Simplify to binary\ntmask_map[np.where(tmask_map>0.)]=1.\n')


# In[14]:


hp.mollview(tmask_map, coord='G', title='mask', unit='mK', cmap="magma")
plt.show()


# In[15]:


get_ipython().run_cell_magic('capture', '', 'if os.path.isfile("data/invnoisecovar.npy"):\n    # Load precomputed inverse noise covariance matrix in pixel space\n    invnoisecovar=np.load("data/invnoisecovar.npy")\nelse:\n    # Setup the noise covariance (infinite in masked regions)\n    noisepower=1e-12\n    noisecovar=noisepower*tmask_map\n    noisecovar[np.where(tmask_map==0.)]=10000.\n    invnoisecovar=1./noisecovar\n    np.save("data/invnoisecovar",invnoisecovar)\n')


# The noise covariance matrix is assumed diagonal in pixel space.

# ## Generate mock data

# In[16]:


get_ipython().run_cell_magic('capture', '', '# Generate the true CMB signal\nsignal_map, signal_alms = hp.synfast(Cl, nside, alm=True, lmax=lmax, new=True)\n')


# In[17]:


hp.mollview(signal_map, coord='G', title='true signal', unit='mK', cmap=planck)
plt.show()


# Data model: $d=s+n$

# In[18]:


# Generate simulated data
normalsim = np.random.normal(0., 1., len(signal_map))
noise_map = normalsim/np.sqrt(invnoisecovar+1e-12)
data_map = signal_map + noise_map
data_map_vis = hp.ma(data_map)
data_map_vis.mask = 1-tmask_map


# In[19]:


hp.mollview(data_map_vis, coord='G', title='simulated data', unit='mK', cmap=planck)
plt.show()


# ## Setup Wiener filter

# Covariance of the Wiener Filter:
# \begin{equation}
# \mathrm{Cov}_\mathrm{WF} = (N^{-1}+S^{-1})^{-1} = S^{1/2}(I+S^{1/2}N^{-1}S^{1/2})^{-1}S^{1/2}
# \end{equation}
# For applications we only need to keep in memory
# $\mathrm{Cov}_\mathrm{WF}N^{-1}$ and $\sqrt{\mathrm{Cov}_\mathrm{WF}}$.

# In[20]:


get_ipython().run_cell_magic('capture', '', 'if os.path.isfile("data/CovWFinvnoisecovarmat.npy") and os.path.isfile("data/sqrtCovWF.npy"):\n    # Load precomputed covariance of the Wiener filter\n    del sqrtsignalcovarPix\n    CovWFinvnoisecovarmat=np.load("data/CovWFinvnoisecovarmat.npy")\n    sqrtCovWF=np.load("data/sqrtCovWF.npy")\nelse:\n    # Setup covariance of the Wiener filter\n    invnoisecovarmat=np.diagflat(tmask_map*invnoisecovar)\n    M=np.identity(npix)+sqrtsignalcovarPix.dot(invnoisecovarmat).dot(sqrtsignalcovarPix)\n    CovWF=sqrtsignalcovarPix.dot(np.linalg.inv(M)).dot(sqrtsignalcovarPix)\n    del sqrtsignalcovarPix, M\n    # Setup CovWF*N^-1\n    CovWFinvnoisecovarmat=CovWF.dot(invnoisecovarmat)\n    np.save("data/CovWFinvnoisecovarmat",CovWFinvnoisecovarmat)\n    # Setup sqrtCovWF for simulating signals\n    CovWFsym=(CovWF+CovWF.T)/2\n    del CovWF\n    sqrtCovWF=scipy.linalg.sqrtm(CovWFsym)\n    del CovWFsim\n    np.save("data/sqrtCovWF",sqrtCovWF)\n')


# ## Perform signal reconstruction (apply Wiener Filtering)

# Mean of the Wiener posterior:
# \begin{equation}
# s_\mathrm{WF} = \mathrm{Cov}_\mathrm{WF} N^{-1} d
# \end{equation}

# In[21]:


# Wiener filtered map (posterior mean)
sWF=CovWFinvnoisecovarmat.dot(data_map).real


# In[22]:


hp.mollview(sWF, coord='G', title='Wiener-filtered data', unit='mK', cmap=planck)
plt.show()


# ## Generate constrained realizations (draw samples from the Wiener posterior)

# Samples of the Wiener posterior:
# \begin{equation}
# s=s_\mathrm{WF}+\sqrt{C_\mathrm{WF}} \, G(0,1)
# \end{equation}
# so that $\left\langle s \right\rangle = s_\mathrm{WF}$ and $\mathrm{Cov}(s) = C_\mathrm{WF}$

# In[23]:


# Constrained realizations (samples of the Wiener posterior)
cr1=sqrtCovWF.dot(np.random.normal(0.,1.,npix)).real+sWF
cr2=sqrtCovWF.dot(np.random.normal(0.,1.,npix)).real+sWF
cr3=sqrtCovWF.dot(np.random.normal(0.,1.,npix)).real+sWF


# In[24]:


hp.mollview(cr1.real, coord='G', title='One simulated signal', unit='mK', cmap=planck)
hp.mollview(cr2.real, coord='G', title='One simulated signal', unit='mK', cmap=planck)
hp.mollview(cr3.real, coord='G', title='One simulated signal', unit='mK', cmap=planck)
plt.show()


# In[25]:


hp.mollview(cr1.real-sWF, coord='G', title='Simulated noise contribution', unit='mK', cmap=planck)


# In[26]:


hp.mollview(cr1.real-signal_map, coord='G', title='Reconstructed - true signal', unit='mK', cmap=planck)


# Of course, this is a toy example. At higher resolution, it is impossible to deal with (invert, or even store) dense covariance matrices in either pixel or harmonic space. Techniques have been devised to perform Wiener filtering in such situations: (preconditioned) conjugate gradients and messenger-field approaches. They are beyond the scope of these lectures.