Integrating poisson and log-normal simulations

delensalot/core/opfilt/MAP_opfilt_aniso_t.py

@@ -28,7 +28,8 @@ def apply_fini(*args, **kwargs):
 
 class alm_filter_ninv_wl(opfilt_base.alm_filter_wl):
     def __init__(self, ninv_geom:utils_geom.Geom, ninv: np.ndarray, ffi:remapping.deflection, transf:np.ndarray,
-                 unlalm_info:tuple, lenalm_info:tuple, sht_threads:int,verbose=False, lmin_dotop=0, tpl:tni.template_tfilt or None =None):
+                 unlalm_info:tuple, lenalm_info:tuple, sht_threads:int,verbose=False, lmin_dotop=0, tpl:tni.template_tfilt or None =None,
+                 extra_ninv: np.ndarray = None):
         r"""CMB inverse-variance and Wiener filtering instance, using unlensed E and lensing deflection
 
             Args:
@@ -40,6 +41,7 @@ def __init__(self, ninv_geom:utils_geom.Geom, ninv: np.ndarray, ffi:remapping.de
                 lenalm_info: tuple of int, lmax and mmax of lensed CMB
                 sht_threads: number of threads for lenspyx SHTs
                 verbose: some printout if set, defaults to False
+                extra_ninv: np.ndarray, extra map for inverse noise filtering to be applied to the observed map
 
         """
         lmax_unl, mmax_unl = unlalm_info
@@ -75,6 +77,9 @@ def __init__(self, ninv_geom:utils_geom.Geom, ninv: np.ndarray, ffi:remapping.de
 
         self.template = tpl
 
+        self.extra_ninv = extra_ninv # I do not use this one for the diagonal pre-conditioned CG
+        ## assume extra_inv will enter in the form: B^t extra_ninv B, then I want to take the inverse of this
+
     def hashdict(self):
         return {'ninv':self._ninv_hash(), 'transf':clhash(self.b_transf_tlm),
                 'deflection':self.ffi.hashdict(),
@@ -101,6 +106,28 @@ def apply_map(self, tmap):
 
         """
         tmap *= self.n_inv
+
+        if self.extra_ninv is not None:
+            """
+            Here I am doing this:
+            N^{-1} - N^{-1}B^t \hat{S}^2 B N^{-1}, applied on X map
+
+            This assumes that we have a small perturbation to the noise.
+
+            Note that for now here the noise is diagonal in pixel space. This is not the case for the extra piece of here.
+            """
+            B = tmap.copy()
+            B = self.geom_.map2alm(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.))
+            B = almxfl(B, self.b_transf_tlm, self.mmax_len, inplace = False)
+            B = self.geom_.alm2map(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.))
+            B *= self.extra_ninv
+            B = self.geom_.map2alm(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.))
+            B = almxfl(B, self.b_transf_tlm, self.mmax_len, inplace = False)
+            B = self.geom_.alm2map(B.copy(), self.lmax_len, self.mmax_len, self.ffi.sht_tr, (-1., 1.))
+            B *= self.n_inv
+            tmap -= B
+
+
         if self.template is not None:
             ts = [self.template] # Hack, this is only meant for one template
             coeffs = np.concatenate(([t.dot(tmap) for t in ts]))

delensalot/sims/foregrounds/pointsources.py

@@ -0,0 +1,110 @@
+"""
+Generates simple point sources, can be correlated with an input lensing convergence map.
+"""
+
+import numpy as np
+import healpy as hp
+
+
+class Foreground(object):
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def randomizing_fg(mappa: np.ndarray):
+        """
+        Randomizes the phase of the input map, preserving the amplitude.
+        """
+        f = lambda z: np.abs(z) * np.exp(1j*np.random.uniform(0., 2.*np.pi, size = z.shape))
+        return f(mappa)
+
+    def randomized_map(self, mappa: np.ndarray, nside: int):
+        """
+        Randomizes the phase of the input map, preserving the amplitude.
+        """
+        alm = hp.map2alm(mappa)
+        alm = self.randomizing_fg(alm)
+        return hp.alm2map(alm, nside)
+
+
+    @staticmethod
+    def matched_filter(input_map_alm: np.ndarray, total_cl: np.ndarray, signal_cl: np.ndarray, nside: int):
+        """
+        Returns the matched filter map.
+        """
+        alm = hp.almxfl(input_map_alm, np.nan_to_num(1/total_cl))
+        alm = hp.almxfl(alm, signal_cl)
+        return hp.alm2map(alm, nside)
+
+    def mask_from_matched_filter(self, input_map_alm: np.ndarray, total_cl: np.ndarray, signal_cl: np.ndarray, nside: int, threshold: float = 0.5):
+        """
+        Returns a mask from the matched filter map.
+        """
+        mappa = self.matched_filter(input_map_alm, total_cl, signal_cl, nside)
+        SN_map = abs(mappa) / np.std(mappa)
+        mask =  np.where(mappa > threshold, 1, 0)
+        return mask
+
+    @staticmethod
+    def smooth_map(mappa: np.ndarray, fwhm: float, nside: int):
+        """
+        Smooths the input map.
+        """
+        return hp.smoothing(mappa, fwhm = np.radians(fwhm))
+
+
+class PointSourcesSimple(Foreground):
+
+    def __init__(self, nside: int = 2048) -> None:
+        self.nside = nside
+
+    @staticmethod
+    def phi_lm_to_kappa_lm(plm: np.ndarray) -> np.ndarray:
+        """
+        Converts the input phi_lm to kappa_lm.
+        """
+        lmax = hp.Alm.getlmax(plm.size)
+        ls = np.arange(0, lmax)
+        factor = (ls * (ls + 1.)) / 2.
+        return hp.almxfl(plm, factor)
+
+
+    @staticmethod
+    def _get_position_from_kappa_default(rng, kappa: np.ndarray, factor: float = 0.5) -> np.ndarray:
+        """
+        Returns the positions of the point sources from the input kappa map.
+        """
+        positions = np.where(rng.poisson(1+kappa*factor) > 0)[0]
+        return positions
+
+    @staticmethod
+    def _get_position_from_kappa_alternative(rng, kappa: np.ndarray, factor: float = 0.5) -> np.ndarray:
+        """
+        Returns the positions of the point sources from the input kappa map.
+        """
+        positions = np.where(rng.poisson(abs(kappa)*factor) > 0)[0]
+        return positions
+
+
+    def generate_ps(self, nsrc: int, amp: float = 100, seed: int = 0, plm: np.ndarray = None, factor: float = 0.5) -> np.ndarray:
+
+        rng = np.random.default_rng(seed)
+
+        mappa = np.zeros(hp.nside2npix(self.nside))
+
+        if plm is not None:
+
+            klm = self.phi_lm_to_kappa_lm(plm)
+            kmap = hp.alm2map(klm, self.nside, verbose = False)
+            positions = self._get_position_from_kappa_default(rng, kmap, factor)
+            nsources = len(positions)
+
+        else:
+
+            nsources = rng.poisson(nsrc)
+            positions = np.random.randint(0, len(mappa), nsources)
+
+        amplitudes = rng.poisson(amp, nsources)
+        mappa[positions] = amplitudes
+
+        return mappa

delensalot/sims/generic.py

@@ -51,9 +51,11 @@ class sims_cmb_len(object):
 
             verbose(defaults to True): lenspyx timing info printout
 
+            extra_tlm: optional extra map to add to the CMB map, e.g. foregrounds
+
     """
     def __init__(self, lib_dir, lmax, cls_unl, lib_pha=None, offsets_plm=None, offsets_cmbunl=None,
-                 dlmax=1024, nside_lens=4096, epsilon=1e-7, nbands=8, cache_plm=True, verbose=True):
+                 dlmax=1024, nside_lens=4096, epsilon=1e-7, nbands=8, cache_plm=True, verbose=True, extra_tlm = None):
 
         fields = _get_fields(cls_unl)
 
@@ -99,6 +101,8 @@ def __init__(self, lib_dir, lmax, cls_unl, lib_pha=None, offsets_plm=None, offse
         self.lens_module = lenspyx
         self.verbose=verbose
 
+        self.extra_tlm = extra_tlm
+
     @staticmethod
     def offset_index(idx, block_size, offset):
         """Offset index by amount 'offset' cyclically within blocks of size block_size
@@ -169,17 +173,32 @@ def _cache_eblm(self, idx):
 
     def get_sim_tlm(self, idx, ret=True):
         fname = os.path.join(self.lib_dir, 'sim_%04d_tlm.fits' % idx)
+        pfname = os.path.join(self.lib_dir, 'sim_%04d_plm.fits' % idx)
         if not os.path.exists(fname):
             tlm = self.unlcmbs.get_sim_tlm(self.offset_index(idx, self.offset_cmb[0], self.offset_cmb[1]))
             dlm = self.get_sim_plm(idx)
+
+            hp.write_alm(pfname, dlm)
+
             assert 'o' not in self.fields, 'not implemented'
 
             lmaxd = hp.Alm.getlmax(dlm.size)
             hp.almxfl(dlm, np.sqrt(np.arange(lmaxd + 1, dtype=float) * np.arange(1, lmaxd + 2)), inplace=True)
             Tlen = self.lens_module.alm2lenmap(np.array(tlm), [dlm, None], epsilon=self.epsilon, verbose=self.verbose)
             hp.write_alm(fname, hp.map2alm(Tlen, lmax=self.lmax, iter=0))
+
+        if (self.extra_tlm is not None):
+            """
+            Adding an extra CMB component to the lensed CMB.
+            """
+            extrafname = os.path.join(self.lib_dir, f'sim_{idx:04}_{self.extra_tlm.get_name()}lm.fits')
+            if (not os.path.exists(extrafname)):
+                extra_tlm = hp.map2alm(self.extra_tlm(idx), lmax=self.lmax, iter=0)
+                hp.write_alm(extrafname, extra_tlm)
+
         if ret:
-            return hp.read_alm(fname)
+            total = hp.read_alm(fname) + (0 if self.extra_tlm is None else hp.read_alm(extrafname))
+            return total
 
     def get_sim_elm(self, idx, ret=True):
         fname = os.path.join(self.lib_dir, 'sim_%04d_elm.fits' % idx)

delensalot/sims/lognormal/lognormal_utils.py

@@ -0,0 +1,170 @@
+"""Module with utilities for the lognormal sims module.
+"""
+
+import numpy as np
+import flt
+import healpy as hp
+
+
+def get_out_quantities_from_a_map(outmap: np.ndarray) -> tuple:
+    """
+    Returns the mean, variance, skewness, lambda, muG and sigmaG from a given map.
+    """
+    skewness = get_skew_from_map(outmap)
+    variance = get_variance_from_map(outmap)
+    mean = get_mean_from_map(outmap)
+    lamb = get_lambda_from_skew(skewness, variance, mean)
+
+    alpha = get_alpha(mean, lamb)
+    sigmaG = get_sigma_gauss(alpha, variance)
+    muG = get_mu_gauss(alpha, variance)
+    return mean, variance, skewness, lamb, muG, sigmaG
+
+
+def create_lognormal_single_map(inputcl: np.ndarray, nside: int, lmax_gen: int, mu: float = 0.0, lamb: float = 0.0):
+    """
+    Creates a lognormal map with a given power spectrum and skewness. These are the parameters of the lognormal distribution that are specified with mu and lamb.
+    """
+
+
+    alpha = get_alpha(mu, lamb)
+    xisinput = cl2xi(inputcl)/alpha/alpha
+
+    xigaussian = np.log(xisinput+1)
+    clgaussian = xi2cl(xigaussian)
+
+    vargauss = np.dot(np.arange(1, 2*len(clgaussian), 2), clgaussian)/(4*np.pi)
+
+    lmax_gen = 2*nside-1 if lmax_gen is None else lmax_gen
+    almgaussian = hp.synalm(clgaussian, lmax = lmax_gen) #GENERATE TO HIGH LMAX
+    maps = hp.alm2map(almgaussian, nside = nside, pol = False)
+
+    #vargauss = np.array([xigaussian[i, i][0] for i in range(Nfields)])
+    #vargauss = np.array([np.var(m) for m in maps])
+
+    expmu = (mu+lamb)*np.exp(-vargauss*0.5)
+    maps = np.array(maps)
+    maps = np.exp(maps) 
+    maps *= expmu
+    maps -= lamb
+    return maps
+
+
+#shifted log-normal distribution
+def shifted_lognormal_zero_mean(x, sigmaG, lamb):
+    #equation (22) of https://www.aanda.org/articles/aa/pdf/2011/12/aa17294-11.pdf
+    return np.exp(-(np.log(x/lamb+1)+sigmaG**2/2)**2./(2.*sigmaG**2.))/(x+lamb)/sigmaG/np.sqrt(2.*np.pi)*(x>-lamb)
+
+
+
+def cl2xi(cl: np.ndarray, closed = False):
+    """
+    This goes from angular power spectrum to a correlation function calculated at theta points.
+    """
+    ls = np.arange(0, len(cl))
+    factorcl = (2*ls+1)/(4*np.pi)
+    coeffs = cl*factorcl
+    return flt.idlt(coeffs, closed = closed)
+
+
+def theta(n, closed = False):
+    """
+    Returns the theta for which the cl2xi are calculated, for a given n
+    """
+    return flt.theta(n, closed = closed)
+
+
+def xi2cl(xi: np.ndarray, closed = False):
+    """
+    This goes from correlation function calculated at theta points to an angular power spectrum.
+    """
+    ls = np.arange(0, len(xi))
+    factorcl = (2*ls+1)/(4*np.pi)
+    return flt.dlt(xi, closed = closed)/factorcl
+
+
+
+def get_mean_from_map(mappa: np.ndarray):
+    """
+    Mean from a map.
+    """
+    return np.mean(mappa)
+
+def get_variance_from_map(mappa: np.ndarray):
+    """
+    Variance from a map.
+    """
+    return np.mean(mappa**2.)-np.mean(mappa)**2.
+
+def get_skew_from_map(mappa: np.ndarray):
+    """
+    Skewness from a map.
+    """
+    return np.mean((mappa-get_mean_from_map(mappa))**3.)/np.mean(mappa**2.)**1.5
+
+def y_skew(skew):
+    """
+    Formula (12) from https://arxiv.org/pdf/1602.08503.pdf
+    It relates the skewnees to some factor that is used to get the lambda parameter for the log-normal generation.
+    """
+    result = 2+skew**2.+skew*np.sqrt(4+skew**2.)
+    result /= 2
+    return np.power(result, 1/3.)
+
+def get_lambda_from_skew(skew, var, mu):
+    lmbda = np.sqrt(var)/skew*(1+y_skew(skew)+1/y_skew(skew))-mu
+    return lmbda 
+
+def get_alpha(mu, lmbda):
+    """
+    Below formula (7) from https://arxiv.org/pdf/1602.08503.pdf
+    """
+    return mu+lmbda
+
+def get_mu_gauss(alpha, var):
+    """
+    Gets the mu parameter for the Gaussian distribution for the log-normal
+    """
+    result = np.log(alpha**2./np.sqrt(var+alpha**2.))
+    return result
+
+def get_sigma_gauss(alpha, var):
+    """
+    Gets the sigma parameter for the Gaussian distribution for the log-normal.
+
+    Here the variance is the variance of the wanted log-normal field.
+    """
+    result = np.log(1+var/alpha**2.)
+    result = np.sqrt(result)
+    return result
+
+
+def suppress(l: np.ndarray, lsup: float = 7000, supindex: float = 10):
+    """
+    Suppression factor at high ell.
+    """
+    return np.exp(-1.0*np.power(l/lsup, supindex))
+
+def suppress_cls(inputcl: np.ndarray, l: np.ndarray, lsup: float = 7000, supindex: float = 10):
+    """
+    Suppresses an input angular power spectrum at high ell.
+    """
+    return inputcl*suppress(l, lsup, supindex)
+
+def process_cl(inputcl: np.ndarray, function: callable = suppress_cls, **kwargs):
+    """
+    Processes an input angular power spectrum with a given function.
+
+    The function must take inputcl and l as arguments.
+
+    Args:
+        inputcl (np.ndarray): Input angular power spectrum.
+        function (callable) default suppress_cls: Function to process the angular power spectrum.
+        **kwargs: Keyword arguments for the function.
+
+    Returns:
+        np.ndarray: Processed angular power spectrum.
+    """
+    ls = np.arange(0, len(inputcl))
+    return function(inputcl = inputcl, l = ls, **kwargs)
+