Added ellwise likelihood

Stages.md

@@ -55,7 +55,7 @@ This stage takes a set of coadded multi-frequency polarization power spectra and
 - `cells_fiducial`: a fits file containing multi-frequency and multi-polarization power spectra for a model that should roughly resemble the data (e.g. the model used to generate the simulations in the previous stages). This is only needed if you're using the Hamimeche & Lewis likelihood.
 
 ### Output
-- `param_chains`: a numpy `npz` file containing the parameter MCMC chain.
+- `output_dir`: a directory containing the parameter MCMC chain.
 
 ### Parameters
 See [test/test_config_sampling.yml](test/test_config_sampling.yml)
@@ -78,3 +78,16 @@ This stage generates a webpage containing a set of plots illustrating the main p
 
 ### Parameters
 See [test/test_config_sampling.yml](test/test_config_sampling.yml)
+
+## 5. BBEllwise
+### Stage summary
+This stage saves the CMB bandpowers inferred by the ellwise sampler, implemented in BBCompSep. It optionally validates the results by re-inserting the bandpowers in a CMB-only likelihood, inferring r and A_lens, and comparing those results with r and A_lens inferred by the default pipeline.
+
+### Inputs
+- `default_chain`: either a `npz`, or a `stats` file containing the output from the default (non-ellwise) sampler.
+
+### Outputs
+- `output_dir`: a directory containing the input (!) and output chains, as well as the resulting bandpowers from the ellwise CMB sampler.
+
+### Parameters
+See [test/test_config_ellwise.yml](test/test_config_ellwise.yml)

bbpower/__init__.py

@@ -3,3 +3,4 @@
 from .power_summarizer import BBPowerSummarizer  # noqa
 from .compsep import BBCompSep  # noqa
 from .plotter import BBPlotter  # noqa
+from .ellwise import BBEllwise

bbpower/compsep.py

@@ -333,10 +333,19 @@ def model(self, params):
         Defines the total model and integrates over
         the bandpasses and windows.
         """
-        # [npol,npol,nell]
-        cmb_cell = (params['r_tensor'] * self.cmb_tens +
-                    params['A_lens'] * self.cmb_lens +
-                    self.cmb_scal) * self.dl2cl
+        if self.config['cmb_model'].get('use_ellwise'):
+            cmb_amps = np.asarray([params[f'A_cmb_{str(i).zfill(2)}'] \
+                                   for i in range(1, self.n_bpws+1)])
+            delta_ell = np.mean(np.diff(self.ell_b))
+            cmb_amps_ell = np.dot(cmb_amps, self.windows[0,:,:]) * delta_ell * self.dl2cl
+
+            # [npol,npol,nell]
+            cmb_cell = (cmb_amps_ell * self.cmb_lens + self.cmb_scal) * self.dl2cl
+        else:
+            # [npol,npol,nell]
+            cmb_cell = (params['r_tensor'] * self.cmb_tens +
+                        params['A_lens'] * self.cmb_lens +
+                        self.cmb_scal) * self.dl2cl
         # [nell,npol,npol]
         cmb_cell = np.transpose(cmb_cell, axes=[2, 0, 1])
         if self.config['cmb_model'].get('use_birefringence'):
@@ -610,8 +619,10 @@ def emcee_sampler(self):
         """
         import emcee
         from multiprocessing import Pool
-
-        fname_temp = self.get_output('output_dir')+'/emcee.npz.h5'
+        if self.config['cmb_model'].get('use_ellwise'):
+            fname_temp = self.get_output('output_dir')+'/emcee_ellwise.npz.h5'
+        else:
+            fname_temp = self.get_output('output_dir')+'/emcee.npz.h5'
         backend = emcee.backends.HDFBackend(fname_temp)
 
         nwalkers = self.config['nwalkers']
@@ -642,7 +653,7 @@ def emcee_sampler(self):
                                             backend=backend)
             if nsteps_use > 0:
                 sampler.run_mcmc(pos, nsteps_use, store=True, progress=False)
-                end = time.time()
+            end = time.time()
 
         return sampler, end-start
 
@@ -673,7 +684,10 @@ def dumper(live, dead, logweights, logZ, logZerr):
             print("Last dead point:", dead[-1])
 
         settings = PolyChordSettings(ndim, nder)
-        settings.base_dir = self.get_output('output_dir')+'/polychord'
+        if self.config['cmb_model'].get('use_ellwise'):
+            settings.base_dir = self.get_output('output_dir')+'/polychord_ellwise'
+        else:
+            settings.base_dir = self.get_output('output_dir')+'/polychord'
         settings.file_root = 'pch'
         settings.nlive = self.config['nlive']
         settings.num_repeats = self.config['nrepeat']
@@ -684,10 +698,13 @@ def dumper(live, dead, logweights, logZ, logZerr):
         settings.read_resume = False   # Read from resume file of earlier run
         settings.feedback = 2          # Verbosity {0,1,2,3}
 
+        import time
+        start = time.time()
         output = pypolychord.run_polychord(likelihood, ndim, nder, settings, 
                                            prior, dumper)
+        end = time.time()
 
-        return output
+        return output, end-start
 
     def minimizer(self):
         """
@@ -769,16 +786,34 @@ def run(self):
         from shutil import copyfile
         copyfile(self.get_input('config'), self.get_output('config_copy'))
         self.setup_compsep()
+        if self.config['cmb_model'].get('use_ellwise'):
+            if self.config.get('sampler') == 'emcee':
+                print('Running compsep: ellwise CMB (emcee)')
+                sampler, timing = self.emcee_sampler()
+                np.savez(self.get_output('output_dir')+'/emcee_ellwise.npz',
+                         chain=sampler.chain,
+                         names=self.params.p_free_names,
+                         time=timing)
+                print("Finished sampling", timing)
+            elif self.config.get('sampler') == 'polychord':
+                print('Running compsep: ellwise CMB (polychord)')
+                sampler, timing = self.polychord_sampler()
+                print("Finished sampling", timing)
+            else:
+                raise ValueError("Ellwise CMB: Unknown sampler.")
+            return
         if self.config.get('sampler') == 'emcee':
+            print('Running compsep: default (emcee)')
             sampler, timing = self.emcee_sampler()
             np.savez(self.get_output('output_dir')+'/emcee.npz',
                      chain=sampler.chain,
                      names=self.params.p_free_names,
                      time=timing)
             print("Finished sampling", timing)
         elif self.config.get('sampler')=='polychord':
-            sampler = self.polychord_sampler()
-            print("Finished sampling")
+            print('Running compsep: default (polychord)')
+            sampler, timing = self.polychord_sampler()
+            print("Finished sampling", timing)
         elif self.config.get('sampler') == 'fisher':
             p0, fisher = self.fisher()
             cov = np.linalg.inv(fisher)

bbpower/ellwise.py

@@ -0,0 +1,226 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+from bbpipe import PipelineStage
+from .types import NpzFile, FitsFile, YamlFile, DirFile
+from .param_manager import ParameterManager
+
+
+class BBEllwise(PipelineStage):
+    """
+    CMB bandpower stage
+    This stage saves CMB bandpowers after ellwise component separation.
+    """
+    name = "BBEllwise"
+    inputs = [('default_chain', DirFile)]
+    outputs = [('output_dir', DirFile)]
+    config_options = {'sampler': 'emcee', 'validation': True}
+
+    def bin_cells(self, bin_edges, ells, cells):
+        """
+        Bins power spectra into bandpower windows, given bin edges.
+        bin_edges contains [lo_1, lo_2, lo_3, lo_4, ...,lo_n, hi_n+1],
+        where n-th bin has range [lo_n, hi_n=lo_n+1 - 1].
+        """
+        cells_b = np.zeros(len(bin_edges)-1)
+        for i_b in range(len(bin_edges)-1):
+            target = bin_edges[i_b] + np.diff(bin_edges)[i_b]/2
+            value = cells[np.where(ells==target)]
+            if value:
+                cells_b[i_b] = value[-1]
+        assert(len(cells_b) == len(bin_edges)-1)
+        return np.asarray(cells_b)
+
+    def read_ell_mc(self, mcfile):
+        """
+        Reads CMB bandpowers from emcee output.
+        Assumes bandpowers = first nell parameters.
+        """
+        mcfile = np.load(mcfile)
+        pnames = mcfile['names']
+        chain = mcfile['chain']
+        A = np.zeros(self.nell)
+        A_std = np.zeros(self.nell)
+        for il in range(self.nell):
+            pname = f'A_cmb_{str(il+1).zfill(2)}'
+            assert(pname == pnames[il])
+            A[il] = np.mean(chain[:,:,il], axis=(0,1))
+            A_std[il] = np.std(chain[:,:,il], axis=(0,1))
+
+        return A, A_std
+
+    def read_ell_pc(self, pcfile):
+        """
+        Reads CMB bandpowers from PolyChord output.
+        Assumes bandpowers = first nell parameters.
+        """
+        A = np.zeros(self.nell)
+        A_std = np.zeros(self.nell)
+        with open(pcfile) as f:
+            for line in f:
+                if line.startswith("Dim No."):
+                    break
+            for il in range(self.nell):
+                l = f.readline()
+                A[il] += float(l.split()[1])
+                A_std[il] += float(l.split()[3])
+
+        return A, A_std
+
+    def read_r_AL_mc(self, mc_npz_dir):
+        """
+        Reads r, A_lens from emcee output.
+        Assumes A_lens, r are first two parameters.
+        """
+        mc_chain = np.load(mc_npz_dir)['chain']
+        AL_fid = np.mean(mc_chain[:,:,0], axis=(0,1))
+        AL_std_fid = np.std(mc_chain[:,:,0], axis=(0,1))
+        r_fid = np.mean(mc_chain[:,:,1], axis=(0,1))
+        r_std_fid = np.std(mc_chain[:,:,1], axis=(0,1))
+
+        return r_fid, r_std_fid, AL_fid, AL_std_fid
+
+    def read_r_AL_pc(self, pch_stats_dir):
+        """
+        Reads r, A_lens from PolyChord output.
+        Assumes A_lens, r are first two parameters.
+        """
+        with open(pch_stats_dir) as fid:
+            for line in fid:
+                if line.startswith("Dim No."):
+                    break
+            l = fid.readline() 
+            AL_fid = float(l.split()[1])
+            AL_std_fid = float(l.split()[3])
+            l = fid.readline() 
+            r_fid = float(l.split()[1])
+            r_std_fid = float(l.split()[3])
+
+        return r_fid, r_std_fid, AL_fid, AL_std_fid
+
+    def validate(self, A, A_std):
+        """
+        Inserts CMB bandpowers and standard errors (in units of
+        lensing B-modes, assuming no correlations) in likelihood
+        parameterized by (r, A_lens), and outputs marginal posteriors.
+        """
+        import emcee, os
+        from multiprocessing import Pool
+
+        def lnprob_cmb(theta):
+            res = (A - theta[1])*self.cmb_lens_b - theta[0]*self.cmb_tens_b
+            return -0.5*np.sum(np.square(res)/self.cmb_lens_b**2 / A_std**2)
+
+        def emcee_sampler_valid():
+            fname_temp = self.get_output('output_dir')+'/valid_ellwise.h5'
+            backend = emcee.backends.HDFBackend(fname_temp)
+
+            nwalkers = 5
+            n_iters = 1000
+            ndim = 2
+            found_file = os.path.isfile(fname_temp)
+
+            try:
+                nchain = len(backend.get_chain())
+            except AttributeError:
+                found_file = False
+
+            if not found_file:
+                backend.reset(nwalkers, ndim)
+                pos = [np.asarray([0.01,1.]) + 1.e-3*np.random.randn(ndim)
+                       for i in range(nwalkers)]
+                nsteps_use = n_iters
+            else:
+                print("Restarting from previous run")
+                pos = None
+                nsteps_use = max(n_iters - nchain, 0)
+
+            with Pool() as pool:
+                sampler = emcee.EnsembleSampler(nwalkers, ndim,
+                                                lnprob_cmb,
+                                                backend=backend)
+                if nsteps_use > 0:
+                    sampler.run_mcmc(pos, nsteps_use, store=True, progress=False)
+
+            return sampler
+
+        print('Validating CMB bandpowers')
+        import time
+        start = time.time()
+        sampler_valid = emcee_sampler_valid()
+        end = time.time()
+        print(f" Finished validation in {(end-start):.0f} s")
+
+        chain_valid = sampler_valid.chain
+        r_valid = np.mean(chain_valid[:,:,0], axis=(0,1))
+        r_std_valid = np.std(chain_valid[:,:,0], axis=(0,1))
+        AL_valid = np.mean(chain_valid[:,:,1], axis=(0,1))
+        AL_std_valid = np.std(chain_valid[:,:,1], axis=(0,1))
+
+        return r_valid, r_std_valid, AL_valid, AL_std_valid
+
+    def run(self):
+        # Load bandpower metadata
+        bin_edges = np.loadtxt(self.config.get('bin_edges_path'))
+        lmin = self.config.get('lmin')
+        lmax = self.config.get('lmax')
+        ell_mask = (bin_edges>lmin) * (bin_edges<lmax+3)
+        bin_edges = bin_edges[np.where(ell_mask)]
+        dell = np.mean(np.diff(bin_edges))
+
+        # Load fiducial power spectra
+        cmb_ells = np.loadtxt(self.config.get('fid_spec_r1'))[:,0]
+        mask = (cmb_ells <= lmax+3) & (cmb_ells > max(1,lmin))
+        cmb_ells = cmb_ells[mask]
+        cmb_dl2cl = 2*np.pi/cmb_ells/(cmb_ells+1)
+        cmb_r1 = np.loadtxt(self.config.get('fid_spec_r1'))[:,3][mask]
+        cmb_tens = (np.loadtxt(self.config.get('fid_spec_r1'))[:,3] - \
+                    np.loadtxt(self.config.get('fid_spec_r0'))[:,3])[mask]
+        cmb_lens = np.loadtxt(self.config.get('fid_spec_r0'))[:,3][mask]
+        self.cmb_tens_b = self.bin_cells(bin_edges, cmb_ells, cmb_tens)
+        self.cmb_lens_b = self.bin_cells(bin_edges, cmb_ells, cmb_lens)
+        ells_b = self.bin_cells(bin_edges, cmb_ells, cmb_ells)
+        self.nell = len(ells_b)
+
+        # Load CMB bandpowers
+        if self.config.get('sampler') == 'emcee':
+            A, A_std = self.read_ell_mc(self.get_output('output_dir')+'/emcee_ellwise.npz')
+        elif self.config.get('sampler') == 'polychord':
+            A, A_std = self.read_ell_pc(self.get_output('output_dir')+'/polychord_ellwise/pch.stats')
+        else:
+            raise ValueError('Unknown sampler')
+        print('CMB bandpowers, best fit in units of input CMB:\n', A)
+        np.savez(self.get_output('output_dir')+'/cl_cmb.npz', ells=ells_b, 
+                 dells=self.cmb_lens_b*A, errs=self.cmb_lens_b*A_std)
+        print('Saved under "'+self.get_output('output_dir')+'/cl_cmb.npz"')
+
+
+        if self.config.get('validation') == True:
+            # Load marginal validated posteriors on r and A_lens
+            r, r_std, AL, AL_std = self.validate(A, A_std)
+
+            # Load marginal fiducial posteriors on r and A_lens
+            if '.npz' in self.get_input('default_chain'):
+                r_fid, r_std_fid, AL_fid, AL_std_fid = self.read_r_AL_mc(self.get_input('default_chain'))
+            elif '.stats' in self.get_input('default_chain'):
+                r_fid, r_std_fid, AL_fid, AL_std_fid = self.read_r_AL_pc(self.get_input('default_chain'))
+            else:
+                print('Default chain file has the wrong format - skipping.')
+
+            print('*** Validation: ***')
+            print(f'r      = {r:.5f} +/- {r_std:.5f} from CMB bandpowers')
+            try:
+                print(f'r      = {r_fid:.5f} +/- {r_std_fid:.5f} from default likelihood')
+            except:
+                pass
+            print(f'A_lens = {AL:.3f} +/- {AL_std:.3f} from CMB bandpowers')
+            try:
+                print(f'A_lens = {AL_fid:.3f} +/- {AL_std_fid:.3f}  from default likelihood')
+            except:
+                pass
+
+        return
+
+
+if __name__ == '__main__':
+    cls = PipelineStage.main()

test/run_compsep_ellwise_test.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+mkdir -p test/test_out
+
+# Generate fake data
+python ./examples/generate_SO_spectra.py test/test_out
+
+# Run fiducial compsep (Cori login node: 6 min)
+python -m bbpower BBCompSep   --cells_coadded="./test/test_out/cls_coadd.fits"   --cells_noise="./test/test_out/cls_noise.fits"   --cells_fiducial="./test/test_out/cls_fid.fits"   --output_dir="./test/test_out"   --config_copy="./test/test_out/config_copy.yml"   --config="./test/test_config_emcee.yml"
+
+# Run ellwise compsep (Cori login node: 119 min)
+python -m bbpower BBCompSep   --cells_coadded="./test/test_out/cls_coadd.fits"   --cells_noise="./test/test_out/cls_noise.fits"   --cells_fiducial="./test/test_out/cls_fid.fits"   --output_dir="./test/test_out"   --config_copy="./test/test_out/config_copy_ellwise.yml"   --config="./test/test_config_ellwise.yml"
+
+# Save and validate CMB bandpowers
+python -m bbpower BBEllwise   --default_chain="./test/test_out/emcee.npz"   --output_dir="./test/test_out"   --config="./test/test_config_ellwise.yml"
+
+# Cleanup
+rm -r test/test_out

test/run_compsep_test.sh

@@ -16,4 +16,4 @@ else
 fi
 
 # Cleanup
-rm -r test/test_out
+rm -r test/test_out