diff --git a/Stages.md b/Stages.md index b3e461e..2dfcba4 100644 --- a/Stages.md +++ b/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) diff --git a/bbpower/__init__.py b/bbpower/__init__.py index 0f86ea4..a6a4ba7 100644 --- a/bbpower/__init__.py +++ b/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 diff --git a/bbpower/compsep.py b/bbpower/compsep.py index 7f6760e..e163c0e 100644 --- a/bbpower/compsep.py +++ b/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,7 +786,24 @@ 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, @@ -777,8 +811,9 @@ def run(self): 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) diff --git a/bbpower/ellwise.py b/bbpower/ellwise.py new file mode 100644 index 0000000..a0667e5 --- /dev/null +++ b/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 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() \ No newline at end of file diff --git a/test/run_compsep_ellwise_test.sh b/test/run_compsep_ellwise_test.sh new file mode 100644 index 0000000..c4267bf --- /dev/null +++ b/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 \ No newline at end of file diff --git a/test/run_compsep_test.sh b/test/run_compsep_test.sh index 2ada19c..70edad4 100755 --- a/test/run_compsep_test.sh +++ b/test/run_compsep_test.sh @@ -16,4 +16,4 @@ else fi # Cleanup -rm -r test/test_out +rm -r test/test_out \ No newline at end of file diff --git a/test/test_config_ellwise.yml b/test/test_config_ellwise.yml new file mode 100644 index 0000000..004a5d4 --- /dev/null +++ b/test/test_config_ellwise.yml @@ -0,0 +1,136 @@ +# These parameters are accessible to all stages +global: + nside: 64 + compute_dell: True + +BBPowerSpecter: + # Bandpower definition + bpw_edges: "./examples/test_data/bpw_edges.txt" + purify_B: True + n_iter : 3 + +BBPowerSummarizer: + # Covariance types + nulls_covar_type: "diagonal" + nulls_covar_diag_order: 0 + data_covar_type: "block_diagonal" + data_covar_diag_order: 0 + +BBCompSep: + # Sampler type (choose 'emcee' or 'polychord' for ellwise compsep) + sampler: 'emcee' + # If you chose polychord: + nlive: 50 + nrepeat: 50 + # If you chose emcee (note higher nwalkers for ellwise compsep): + nwalkers: 96 + n_iters: 5000 + # Likelihood type (choose 'chi2' or 'h&l') + likelihood_type: 'chi2' + # What is the starting point? + r_init: 1.e-3 + # Which polarization channels do you want to include? + pol_channels: ['B'] + # Scale cuts (will apply to all frequencies) + l_min: 30 + l_max: 120 + + # CMB model + cmb_model: + # Inference targets CMB bandpowers + use_ellwise: True + # Template power spectrum. Should contained the lensed power spectra + # with r=0 and r=1 respectively. + cmb_templates: ["./examples/data/camb_lens_nobb.dat", + "./examples/data/camb_lens_r1.dat"] + # Free parameters + params: + # CMB bandpower amplitudes. This assumes 9 bandpowers with + # delta_ell = 10, ell in [30, 120]. + # Append more parameters if needed. + A_cmb_01: ['A_cmb_01', 'tophat', [0.,1.,2.]] + A_cmb_02: ['A_cmb_02', 'tophat', [0.,1.,2.]] + A_cmb_03: ['A_cmb_03', 'tophat', [0.,1.,2.]] + A_cmb_04: ['A_cmb_04', 'tophat', [0.,1.,2.]] + A_cmb_05: ['A_cmb_05', 'tophat', [0.,1.,2.]] + A_cmb_06: ['A_cmb_06', 'tophat', [0.,1.,2.]] + A_cmb_07: ['A_cmb_07', 'tophat', [0.,1.,2.]] + A_cmb_08: ['A_cmb_08', 'tophat', [0.,1.,2.]] + A_cmb_09: ['A_cmb_09', 'tophat', [0.,1.,2.]] + + # Foreground model + fg_model: + # Add one section per component. They should be called `component_X`, + # starting with X=1 + component_1: + # Name for this component + name: Dust + # Type of SED. Should be one of the classes stored in fgbuster.components + # https://github.com/fgbuster/fgbuster/blob/master/fgbuster/component_model.py + sed: Dust + # Type of power spectra for all possible polarization channel combinations. + # Any combinations not added here will be assumed to be zero. + # The names should be one of the classes in bbpower/fgcls.py + cl: + EE: ClPowerLaw + BB: ClPowerLaw + # Parameters of the SED + sed_parameters: + # The key can be anything you want, but each parameter in the model + # must have a different name. + # The first item in the list is the name of the parameter used by fgbuster + # The second item is the type of prior. The last item are the numbers + # necessary to define the prior. They should be: + # - Gaussian: [mean,sigma] + # - tophat: [lower edge, start value, upper edge] + # - fixed: [parameter value] + # nu0-type parameters can only be fixed. + beta_d: ['beta_d', 'Gaussian', [1.59, 0.11]] + temp_d: ['temp', 'fixed', [19.6]] + nu0_d: ['nu0', 'fixed', [353.]] + cl_parameters: + # Same for power spectrum parameters + # (broken down by polarization channel combinations) + EE: + amp_d_ee: ['amp', 'tophat', [0., 10., 20.]] + alpha_d_ee: ['alpha', 'tophat', [-1., -0.42, 0.]] + l0_d_ee: ['ell0', 'fixed', [80.]] + BB: + amp_d_bb: ['amp', 'tophat', [0., 5., 10.]] + alpha_d_bb: ['alpha', 'tophat', [-1., -0.2, 0.]] + l0_d_bb: ['ell0', 'fixed', [80.]] + # If this component should be correlated with any other, list them here + cross: + # In this case the list should contain: + # [component name, prior type, prior parameters] + epsilon_ds: ['component_2', 'tophat', [-1., 0., 1.]] + + component_2: + name: Synchrotron + sed: Synchrotron + cl: + EE: ClPowerLaw + BB: ClPowerLaw + sed_parameters: + beta_s: ['beta_pl', 'Gaussian', [-3.0, 0.3]] + nu0_s: ['nu0', 'fixed', [23.]] + cl_parameters: + EE: + amp_s_ee: ['amp', 'tophat', [0., 4., 8.]] + alpha_s_ee: ['alpha', 'tophat', [-1., -0.6, 0.]] + l0_s_ee: ['ell0', 'fixed', [80.]] + BB: + amp_s_bb: ['amp', 'tophat', [0., 2., 4.]] + alpha_s_bb: ['alpha', 'tophat', [-1., -0.4, 0.]] + l0_s_bb: ['ell0', 'fixed', [80.]] + +BBEllwise: + # Sampler used by BBCompSep to produce ellwise chains ('emcee' or 'polychord') + sampler: 'emcee' + # Re-insert CMB amplitudes in sampler to check (r, A_lens)? + validation: True + fid_spec_r0: './examples/data/camb_lens_nobb.dat' + fid_spec_r1: './examples/data/camb_lens_r1.dat' + bin_edges_path: './examples/test_data/bpw_edges.txt' + lmin: 30 + lmax: 120 diff --git a/test/test_config_emcee.yml b/test/test_config_emcee.yml index fddf3f3..84c1195 100644 --- a/test/test_config_emcee.yml +++ b/test/test_config_emcee.yml @@ -34,7 +34,7 @@ BBCompSep: pol_channels: ['B'] # Scale cuts (will apply to all frequencies) l_min: 30 - l_max: 300 + l_max: 120 # CMB model cmb_model: @@ -117,7 +117,7 @@ BBCompSep: l0_s_bb: ['ell0', 'fixed', [80.]] BBPlotter: - lmax_plot: 300 + lmax_plot: 128 plot_coadded_total: False plot_noise: False plot_nulls: False diff --git a/test/test_config_polychord.yml b/test/test_config_polychord.yml index c5d3fda..7723e05 100644 --- a/test/test_config_polychord.yml +++ b/test/test_config_polychord.yml @@ -34,7 +34,7 @@ BBCompSep: pol_channels: ['B'] # Scale cuts (will apply to all frequencies) l_min: 30 - l_max: 300 + l_max: 120 # CMB model cmb_model: