From 08e1874cf704168b4350b8eb3ac56b3eed04684f Mon Sep 17 00:00:00 2001 From: Javier Sanchez Date: Tue, 17 Oct 2023 14:55:39 -0400 Subject: [PATCH] bringing back tjpcov compatibility --- augur/generate.py | 93 ++++++++++++++++++++-------------------- augur/utils/cov_utils.py | 36 ++++++++-------- examples/config_test.yml | 2 +- 3 files changed, 65 insertions(+), 66 deletions(-) diff --git a/augur/generate.py b/augur/generate.py index 8e6b1af..92e8d02 100644 --- a/augur/generate.py +++ b/augur/generate.py @@ -11,7 +11,7 @@ import sacc from augur.tracers.two_point import ZDist, LensSRD2018, SourceSRD2018 from augur.utils.cov_utils import get_gaus_cov, get_SRD_cov, get_noise_power -# from augur.utils.cov_utils import TJPCovGaus +from augur.utils.cov_utils import TJPCovGaus import firecrown.likelihood.gauss_family.statistic.source.weak_lensing as wl import firecrown.likelihood.gauss_family.statistic.source.number_counts as nc from firecrown.likelihood.gauss_family.statistic.two_point import TwoPoint @@ -315,52 +315,51 @@ def generate(config, return_all_outputs=False, write_sacc=True, force_read=True) S.add_covariance(cov) # The option using TJPCov takes a while. TODO: Use some sort of parallelization. elif config['cov_options']['cov_type'] == 'tjpcov': - raise NotImplementedError('TJPCov not supported currently.') - # tjpcov_config = dict() # Create a config dictionary to instantiate TJPCov - # tjpcov_config['tjpcov'] = dict() - # tjpcov_config['tjpcov']['cosmo'] = tools.ccl_cosmo - # ccl_tracers = dict() - # bias_all = dict() - # for i, myst1 in enumerate(lk.statistics): - # trname1 = myst1.source0.sacc_tracer - # trname2 = myst1.source1.sacc_tracer - # tr1 = myst1.source0.tracers[0].ccl_tracer # Pulling out the tracers - # tr2 = myst1.source1.tracers[0].ccl_tracer - # ccl_tracers[trname1] = tr1 - # ccl_tracers[trname2] = tr2 - # if 'lens' in trname1: - # bias_all[trname1] = myst1.source0.bias - # if 'lens' in trname2: - # bias_all[trname2] = myst1.source1.bias - # for key in bias_all.keys(): - # tjpcov_config['tjpcov'][f'bias_{key}'] = bias_all[key] - # tjpcov_config['tjpcov']['sacc_file'] = S - # tjpcov_config['tjpcov']['IA'] = config['cov_options'].get('IA', None) - # tjpcov_config['GaussianFsky'] = {} - # tjpcov_config['GaussianFsky']['fsky'] = config['cov_options']['fsky'] - # tjpcov_config['tjpcov']['binning_info'] = dict() - # tjpcov_config['tjpcov']['binning_info']['ell_edges'] = \ - # eval(config['cov_options']['binning_info']['ell_edges']) - # for tr in S.tracers: - # _, ndens = get_noise_power(config, S, tr, return_ndens=True) - # tjpcov_config['tjpcov'][f'Ngal_{tr}'] = ndens - # if 'src' in tr: - # tjpcov_config['tjpcov'][f'sigma_e_{tr}'] = config['sources']['ellipticity_error'] - # cov_calc = TJPCovGaus(tjpcov_config) - # if config['general']['ignore_scale_cuts']: - # cov_all = cov_calc.get_covariance() - # else: - # ndata = len(S.mean) - # cov_all = np.zeros((ndata, ndata)) - # for i, trcombs1 in enumerate(S.get_tracer_combinations()): - # ii = S.indices(tracers=trcombs1) - # for trcombs2 in S.get_tracer_combinations()[i:]: - # jj = S.indices(tracers=trcombs2) - # ii_all, jj_all = np.meshgrid(ii, jj, indexing='ij') - # cov_here = cov_calc.get_covariance_block(trcombs1, trcombs2) - # cov_all[ii_all, jj_all] = cov_here[:len(ii), :len(jj)] - # cov_all[jj_all.T, ii_all.T] = cov_here[:len(ii), :len(jj)].T - # S.add_covariance(cov_all) + tjpcov_config = dict() # Create a config dictionary to instantiate TJPCov + tjpcov_config['tjpcov'] = dict() + tjpcov_config['tjpcov']['cosmo'] = tools.ccl_cosmo + ccl_tracers = dict() + bias_all = dict() + for i, myst1 in enumerate(lk.statistics): + trname1 = myst1.source0.sacc_tracer + trname2 = myst1.source1.sacc_tracer + tr1 = myst1.source0.tracers[0].ccl_tracer # Pulling out the tracers + tr2 = myst1.source1.tracers[0].ccl_tracer + ccl_tracers[trname1] = tr1 + ccl_tracers[trname2] = tr2 + if 'lens' in trname1: + bias_all[trname1] = myst1.source0.bias + if 'lens' in trname2: + bias_all[trname2] = myst1.source1.bias + for key in bias_all.keys(): + tjpcov_config['tjpcov'][f'bias_{key}'] = bias_all[key] + tjpcov_config['tjpcov']['sacc_file'] = S + tjpcov_config['tjpcov']['IA'] = config['cov_options'].get('IA', None) + tjpcov_config['GaussianFsky'] = {} + tjpcov_config['GaussianFsky']['fsky'] = config['cov_options']['fsky'] + tjpcov_config['tjpcov']['binning_info'] = dict() + tjpcov_config['tjpcov']['binning_info']['ell_edges'] = \ + eval(config['cov_options']['binning_info']['ell_edges']) + for tr in S.tracers: + _, ndens = get_noise_power(config, S, tr, return_ndens=True) + tjpcov_config['tjpcov'][f'Ngal_{tr}'] = ndens + if 'src' in tr: + tjpcov_config['tjpcov'][f'sigma_e_{tr}'] = config['sources']['ellipticity_error'] + cov_calc = TJPCovGaus(tjpcov_config) + if config['general']['ignore_scale_cuts']: + cov_all = cov_calc.get_covariance() + else: + ndata = len(S.mean) + cov_all = np.zeros((ndata, ndata)) + for i, trcombs1 in enumerate(S.get_tracer_combinations()): + ii = S.indices(tracers=trcombs1) + for trcombs2 in S.get_tracer_combinations()[i:]: + jj = S.indices(tracers=trcombs2) + ii_all, jj_all = np.meshgrid(ii, jj, indexing='ij') + cov_here = cov_calc.get_covariance_block(trcombs1, trcombs2) + cov_all[ii_all, jj_all] = cov_here[:len(ii), :len(jj)] + cov_all[jj_all.T, ii_all.T] = cov_here[:len(ii), :len(jj)].T + S.add_covariance(cov_all) else: raise Warning('''Currently only internal Gaussian covariance and SRD has been implemented, cov_type is not understood. Using identity matrix as covariance.''') diff --git a/augur/utils/cov_utils.py b/augur/utils/cov_utils.py index a123d54..dad48c0 100644 --- a/augur/utils/cov_utils.py +++ b/augur/utils/cov_utils.py @@ -1,6 +1,6 @@ import numpy as np import pyccl as ccl -# from tjpcov.covariance_gaussian_fsky import FourierGaussianFsky +from tjpcov.covariance_gaussian_fsky import FourierGaussianFsky def get_noise_power(config, S, tracer_name, return_ndens=False): @@ -193,20 +193,20 @@ def get_SRD_cov(config, S): return cov_sacc_all -# class TJPCovGaus(FourierGaussianFsky): -# """ -# Class to patch FourierGaussianFsky to work with Augur -# """ -# def __init__(self, config): -# super().__init__(config) -# self.tracer_Noise = self.tracer_Noise_coupled - -# def get_binning_info(self): -# ell_eff = self.get_ell_eff() -# if 'ell_edges' in self.config['tjpcov']['binning_info'].keys(): -# ell_edges = self.config['tjpcov']['binning_info']['ell_edges'] -# ell_min = np.min(ell_edges) -# ell_max = np.max(ell_edges) -# nbpw = ell_max - ell_min -# ell = np.linspace(ell_min, ell_max, nbpw+1).astype(np.int32) -# return ell, ell_eff, ell_edges +class TJPCovGaus(FourierGaussianFsky): + """ + Class to patch FourierGaussianFsky to work with Augur + """ + def __init__(self, config): + super().__init__(config) + self.tracer_Noise = self.tracer_Noise_coupled + + def get_binning_info(self): + ell_eff = self.get_ell_eff() + if 'ell_edges' in self.config['tjpcov']['binning_info'].keys(): + ell_edges = self.config['tjpcov']['binning_info']['ell_edges'] + ell_min = np.min(ell_edges) + ell_max = np.max(ell_edges) + nbpw = ell_max - ell_min + ell = np.linspace(ell_min, ell_max, nbpw+1).astype(np.int32) + return ell, ell_eff, ell_edges diff --git a/examples/config_test.yml b/examples/config_test.yml index 1a901fc..7d05db9 100644 --- a/examples/config_test.yml +++ b/examples/config_test.yml @@ -75,7 +75,7 @@ cov_options: # ------ # Or you can also get it from a file cov_type : 'SRD' - SRD_cov_path : '/Users/jsanchez/Software/augur/data/Y1_3x2_SRD_cov.npy' + SRD_cov_path : '../data/Y1_3x2_SRD_cov.npy' # Or using TJPCov # cov_type : 'tjpcov' # IA : None