Merge pull request #49 from /issues/48

Issues/48 - Redshift distributions from file

augur/analyze.py

@@ -191,7 +191,7 @@ def get_fisher_matrix(self, method='5pt_stencil', save_txt=True):
  if save_txt:
  np.savetxt(self.config['output'], self.Fij)
  tab_out = Table(self.x.T, names=self.var_pars)
- tab_out.write(self.config['fid_output'], format='ascii')
+ tab_out.write(self.config['fid_output'], format='ascii', overwrite=True)
  return self.Fij
  else:
  return self.Fij

augur/generate.py

@@ -10,6 +10,7 @@
 import pyccl as ccl
 import sacc
 from augur.tracers.two_point import ZDist, LensSRD2018, SourceSRD2018
+from augur.tracers.two_point import ZDistFromFile
 from augur.utils.cov_utils import get_gaus_cov, get_SRD_cov, get_noise_power
 from augur.utils.cov_utils import TJPCovGaus
 from packaging.version import Version
@@ -29,7 +30,7 @@
 from augur.utils.config_io import parse_config
 
 
-implemented_nzs = [ZDist, LensSRD2018, SourceSRD2018]
+implemented_nzs = [ZDist, LensSRD2018, SourceSRD2018, ZDistFromFile]
 
 
 def _get_tracers(statistic, comb):
@@ -132,14 +133,22 @@ def generate_sacc_and_stats(config):
  sacc_tracer = f'{src_root}{i}'
  if isinstance(src_cfg['Nz_type'], list):
  if eval(src_cfg['Nz_type'][i]) in implemented_nzs:
- dndz[sacc_tracer] = eval(src_cfg['Nz_type'][i])(z, Nz_nbins=nbins, Nz_ibin=i,
- **src_cfg['Nz_kwargs'])
+ if 'ZDistFromFile' not in src_cfg['Nz_type'][i]:
+ dndz[sacc_tracer] = eval(src_cfg['Nz_type'][i])(z, Nz_nbins=nbins,
+ Nz_ibin=i,
+ **src_cfg['Nz_kwargs'])
+ else:
+ dndz[sacc_tracer] = ZDistFromFile(**src_cfg['Nz_kwargs'], ibin=i)
  else:
  raise NotImplementedError('The selected N(z) is yet not implemented')
  else:
  if eval(src_cfg['Nz_type']) in implemented_nzs:
- dndz[sacc_tracer] = eval(src_cfg['Nz_type'])(z, Nz_nbins=nbins, Nz_ibin=i,
- **src_cfg['Nz_kwargs'])
+ if 'ZDistFromFile' not in src_cfg['Nz_type']:
+ dndz[sacc_tracer] = eval(src_cfg['Nz_type'])(z, Nz_nbins=nbins,
+ Nz_ibin=i,
+ **src_cfg['Nz_kwargs'])
+ else:
+ dndz[sacc_tracer] = ZDistFromFile(**src_cfg['Nz_kwargs'], ibin=i)
  else:
  raise NotImplementedError('The selected N(z) is yet not implemented')
  S.add_tracer('NZ', sacc_tracer, dndz[sacc_tracer].z, dndz[sacc_tracer].Nz)
@@ -164,16 +173,17 @@ def generate_sacc_and_stats(config):
  lns_cfg = config['lenses']
  nbins = lns_cfg['nbins']
  lns_root = 'lens'
- Nz_centers = eval(lns_cfg['Nz_kwargs']['Nz_center'])
- lns_cfg['Nz_kwargs'].pop('Nz_center')
-
- if np.isscalar(Nz_centers):
- Nz_centers = [Nz_centers]
- if nbins != 1:
- raise ValueError('Nz_centers should have the same length as the number of bins')
- else:
- if len(Nz_centers) != nbins:
- raise ValueError('Nz_centers should have the same length as the number of bins')
+ if 'Nz_center' in lns_cfg['Nz_kwargs'].keys():
+ Nz_centers = eval(lns_cfg['Nz_kwargs']['Nz_center'])
+ lns_cfg['Nz_kwargs'].pop('Nz_center')
+
+ if np.isscalar(Nz_centers):
+ Nz_centers = [Nz_centers]
+ if nbins != 1:
+ raise ValueError('Nz_centers should have the same length as the number of bins')
+ else:
+ if len(Nz_centers) != nbins:
+ raise ValueError('Nz_centers should have the same length as the number of bins')
  # Checking if there's photo-z shift systematics in the config
  if 'delta_z' in lns_cfg.keys():
  if np.isscalar(lns_cfg['delta_z']):
@@ -185,16 +195,24 @@ def generate_sacc_and_stats(config):
  sacc_tracer = f'{lns_root}{i}'
  if isinstance(lns_cfg['Nz_type'], list):
  if eval(lns_cfg['Nz_type'][i]) in implemented_nzs:
- dndz[sacc_tracer] = eval(lns_cfg['Nz_type'][i])(z, Nz_center=Nz_centers[i],
- Nz_nbins=nbins,
- **lns_cfg['Nz_kwargs'])
+ if 'ZDistFromFile' not in lns_cfg['Nz_type'][i]:
+ dndz[sacc_tracer] = eval(lns_cfg['Nz_type'][i])(z,
+ Nz_center=Nz_centers[i],
+ Nz_nbins=nbins,
+ **lns_cfg['Nz_kwargs'])
+ else:
+ dndz[sacc_tracer] = ZDistFromFile(**lns_cfg['Nz_kwargs'], ibin=i)
  else:
  raise NotImplementedError('The selected N(z) is yet not implemented')
  else:
  if eval(lns_cfg['Nz_type']) in implemented_nzs:
- dndz[sacc_tracer] = eval(lns_cfg['Nz_type'])(z, Nz_center=Nz_centers[i],
- Nz_nbins=nbins,
- **lns_cfg['Nz_kwargs'])
+ if 'ZDistFromFile' not in lns_cfg['Nz_type']:
+ dndz[sacc_tracer] = eval(lns_cfg['Nz_type'])(z,
+ Nz_center=Nz_centers[i],
+ Nz_nbins=nbins,
+ **lns_cfg['Nz_kwargs'])
+ else:
+ dndz[sacc_tracer] = ZDistFromFile(**lns_cfg['Nz_kwargs'], ibin=i)
  else:
  raise NotImplementedError('The selected N(z) is yet not implemented')
  S.add_tracer('NZ', sacc_tracer, dndz[sacc_tracer].z, dndz[sacc_tracer].Nz)

augur/tracers/two_point.py

@@ -221,3 +221,34 @@ def to_dict(self):
  dict_out = self.__dict__
  dict_out['kind'] = 'WLSource'
  return dict_out
+
+
+class ZDistFromFile(ZDist):
+ """
+ Class to pass an arbitrary redshift distribution to a tracer
+ """
+ def __init__(self, input_file, ibin=None, format='npy', **kwargs):
+ """
+ Parameters:
+ -----------
+ input_file : str
+ Path to input file with z and dndz
+ """
+ if 'npy' in format:
+ data = np.load(input_file, allow_pickle=True)
+ self.z = data[()]['redshift_range']
+ if ibin is None:
+ raise ValueError('Expected bin number for this format')
+ else:
+ self.Nz = data[()]['bins'][ibin]
+ elif 'ascii' in format:
+ data = np.loadtxt(input_file)
+ self.z = data[:, 0]
+ if ibin is None:
+ self.Nz = data[:, 1]
+ else:
+ self.Nz = data[:, ibin+1]
+ else:
+ raise NotImplementedError('Only ascii files or npy files are currently supported')
+
+ self.zav = np.average(self.z, weights=self.Nz/np.sum(self.Nz))