Skip to content

Commit

Permalink
bringing back tjpcov compatibility
Browse files Browse the repository at this point in the history
  • Loading branch information
Javier Sanchez authored and Javier Sanchez committed Oct 17, 2023
1 parent 8c68c8f commit 08e1874
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 66 deletions.
93 changes: 46 additions & 47 deletions augur/generate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.''')
Expand Down
36 changes: 18 additions & 18 deletions augur/utils/cov_utils.py
View file
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion examples/config_test.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 08e1874

Please sign in to comment.