def main(cmdargs): """Compute the cross-correlation between a catalog of objects and a delta field.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=('Compute the cross-correlation between a catalog of ' 'objects and a delta field.')) parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--from-image', type=str, default=None, required=False, help='Read delta from image format', nargs='*') parser.add_argument('--drq', type=str, default=None, required=True, help='Catalog of objects in format selected by mode') parser.add_argument('--mode', type=str, required=False, choices=['desi', 'sdss'], default='sdss', help='Mode for reading the catalog (default sdss)') parser.add_argument('--rp-min', type=float, default=-200., required=False, help='Min r-parallel [h^-1 Mpc]') parser.add_argument('--rp-max', type=float, default=200., required=False, help='Max r-parallel [h^-1 Mpc]') parser.add_argument('--rt-max', type=float, default=200., required=False, help='Max r-transverse [h^-1 Mpc]') parser.add_argument('--np', type=int, default=100, required=False, help='Number of r-parallel bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of r-transverse bins') parser.add_argument('--z-min-obj', type=float, default=0, required=False, help='Min redshift for object field') parser.add_argument('--z-max-obj', type=float, default=10, required=False, help='Max redshift for object field') parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift larger than ' 'z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift smaller than ' 'z-cut-max')) parser.add_argument( '--lambda-abs', type=str, default='LYA', required=False, help=( 'Name of the absorption in picca.constants defining the redshift ' 'of the delta')) parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol-del', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the delta field') parser.add_argument( '--z-evol-obj', type=float, default=1., required=False, help='Exponent of the redshift evolution of the object field') parser.add_argument( '--fid-Om', type=float, default=0.315, required=False, help='Omega_matter(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-Or', type=float, default=0., required=False, help='Omega_radiation(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-wl', type=float, default=-1., required=False, help='Equation of state of dark energy of fiducial LambdaCDM cosmology' ) parser.add_argument('--no-project', action='store_true', required=False, help='Do not project out continuum fitting modes') parser.add_argument('--no-remove-mean-lambda-obs', action='store_true', required=False, help='Do not remove mean delta versus lambda_obs') parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') parser.add_argument( '--shuffle-distrib-obj-seed', type=int, default=None, required=False, help=('Shuffle the distribution of objects on the sky following the ' 'given seed. Do not shuffle if None')) parser.add_argument( '--shuffle-distrib-forest-seed', type=int, default=None, required=False, help=('Shuffle the distribution of forests on the sky following the ' 'given seed. Do not shuffle if None')) args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 # setup variables in module xcf xcf.r_par_max = args.rp_max xcf.r_par_min = args.rp_min xcf.z_cut_max = args.z_cut_max xcf.z_cut_min = args.z_cut_min xcf.r_trans_max = args.rt_max xcf.num_bins_r_par = args.np xcf.num_bins_r_trans = args.nt xcf.nside = args.nside xcf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] # read blinding keyword blinding = io.read_blinding(args.in_dir) # load fiducial cosmology cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) t0 = time.time() # Find the redshift range if args.z_min_obj is None: r_comov_min = cosmo.get_r_comov(z_min) r_comov_min = max(0., r_comov_min + xcf.r_par_min) args.z_min_obj = cosmo.distance_to_redshift(r_comov_min) userprint("z_min_obj = {}".format(args.z_min_obj), end="") if args.z_max_obj is None: r_comov_max = cosmo.get_r_comov(z_max) r_comov_max = max(0., r_comov_max + xcf.r_par_max) args.z_max_obj = cosmo.distance_to_redshift(r_comov_max) userprint("z_max_obj = {}".format(args.z_max_obj), end="") ### Read objects objs, z_min2 = io.read_objects(args.drq, args.nside, args.z_min_obj, args.z_max_obj, args.z_evol_obj, args.z_ref, cosmo, mode=args.mode) xcf.objs = objs ### Read deltas data, num_data, z_min, z_max = io.read_deltas(args.in_dir, args.nside, xcf.lambda_abs, args.z_evol_del, args.z_ref, cosmo=cosmo, max_num_spec=args.nspec, no_project=args.no_project, from_image=args.from_image) xcf.data = data xcf.num_data = num_data userprint("") userprint("done, npix = {}\n".format(len(data))) ### Remove <delta> vs. lambda_obs if not args.no_remove_mean_lambda_obs: Forest.delta_log_lambda = None for healpix in xcf.data: for delta in xcf.data[healpix]: delta_log_lambda = np.asarray([ delta.log_lambda[index] - delta.log_lambda[index - 1] for index in range(1, delta.log_lambda.size) ]).min() if Forest.delta_log_lambda is None: Forest.delta_log_lambda = delta_log_lambda else: Forest.delta_log_lambda = min(delta_log_lambda, Forest.delta_log_lambda) Forest.log_lambda_min = (np.log10( (z_min + 1.) * xcf.lambda_abs) - Forest.delta_log_lambda / 2.) Forest.log_lambda_max = (np.log10( (z_max + 1.) * xcf.lambda_abs) + Forest.delta_log_lambda / 2.) log_lambda, mean_delta, stack_weight = prep_del.stack( xcf.data, stack_from_deltas=True) del log_lambda, stack_weight for healpix in xcf.data: for delta in xcf.data[healpix]: bins = ((delta.log_lambda - Forest.log_lambda_min) / Forest.delta_log_lambda + 0.5).astype(int) delta.delta -= mean_delta[bins] # shuffle forests and objects if not args.shuffle_distrib_obj_seed is None: xcf.objs = utils.shuffle_distrib_forests(objs, args.shuffle_distrib_obj_seed) if not args.shuffle_distrib_forest_seed is None: xcf.data = utils.shuffle_distrib_forests( xcf.data, args.shuffle_distrib_forest_seed) userprint("") # compute maximum angular separation xcf.ang_max = utils.compute_ang_max(cosmo, xcf.r_trans_max, z_min, z_min2) t1 = time.time() userprint(f'picca_xcf.py - Time reading data: {(t1-t0)/60:.3f} minutes') # compute correlation function, use pool to parallelize xcf.counter = Value('i', 0) xcf.lock = Lock() cpu_data = {healpix: [healpix] for healpix in data} context = multiprocessing.get_context('fork') pool = context.Pool(processes=args.nproc) correlation_function_data = pool.map(corr_func, sorted(cpu_data.values())) pool.close() t2 = time.time() userprint( f'picca_xcf.py - Time computing cross-correlation function: {(t2-t1)/60:.3f} minutes' ) # group data from parallelisation correlation_function_data = np.array(correlation_function_data) weights_list = correlation_function_data[:, 0, :] xi_list = correlation_function_data[:, 1, :] r_par_list = correlation_function_data[:, 2, :] r_trans_list = correlation_function_data[:, 3, :] z_list = correlation_function_data[:, 4, :] num_pairs_list = correlation_function_data[:, 5, :].astype(np.int64) healpix_list = np.array(sorted(list(cpu_data.keys()))) w = (weights_list.sum(axis=0) > 0.) r_par = (r_par_list * weights_list).sum(axis=0) r_par[w] /= weights_list.sum(axis=0)[w] r_trans = (r_trans_list * weights_list).sum(axis=0) r_trans[w] /= weights_list.sum(axis=0)[w] z = (z_list * weights_list).sum(axis=0) z[w] /= weights_list.sum(axis=0)[w] num_pairs = num_pairs_list.sum(axis=0) results = fitsio.FITS(args.out, 'rw', clobber=True) header = [{ 'name': 'RPMIN', 'value': xcf.r_par_min, 'comment': 'Minimum r-parallel [h^-1 Mpc]' }, { 'name': 'RPMAX', 'value': xcf.r_par_max, 'comment': 'Maximum r-parallel [h^-1 Mpc]' }, { 'name': 'RTMAX', 'value': xcf.r_trans_max, 'comment': 'Maximum r-transverse [h^-1 Mpc]' }, { 'name': 'NP', 'value': xcf.num_bins_r_par, 'comment': 'Number of bins in r-parallel' }, { 'name': 'NT', 'value': xcf.num_bins_r_trans, 'comment': 'Number of bins in r-transverse' }, { 'name': 'ZCUTMIN', 'value': xcf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': xcf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'NSIDE', 'value': xcf.nside, 'comment': 'Healpix nside' }, { 'name': 'OMEGAM', 'value': args.fid_Om, 'comment': 'Omega_matter(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAR', 'value': args.fid_Or, 'comment': 'Omega_radiation(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAK', 'value': args.fid_Ok, 'comment': 'Omega_k(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'WL', 'value': args.fid_wl, 'comment': 'Equation of state of dark energy of fiducial LambdaCDM cosmology' }, { 'name': "BLINDING", 'value': blinding, 'comment': 'String specifying the blinding strategy' }] results.write( [r_par, r_trans, z, num_pairs], names=['RP', 'RT', 'Z', 'NB'], comment=['R-parallel', 'R-transverse', 'Redshift', 'Number of pairs'], units=['h^-1 Mpc', 'h^-1 Mpc', '', ''], header=header, extname='ATTRI') header2 = [{ 'name': 'HLPXSCHM', 'value': 'RING', 'comment': 'Healpix scheme' }] da_name = "DA" if blinding != "none": da_name += "_BLIND" results.write([healpix_list, weights_list, xi_list], names=['HEALPID', 'WE', da_name], comment=['Healpix index', 'Sum of weight', 'Correlation'], header=header2, extname='COR') results.close() t3 = time.time() userprint(f'picca_xcf.py - Time total: {(t3-t0)/60:.3f} minutes')
def main(cmdargs): # pylint: disable-msg=too-many-locals,too-many-branches,too-many-statements """Computes the wick covariance for the cross-correlation of object x forests.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=( 'Compute the wick covariance for the cross-correlation of ' 'object x forests.')) parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--from-image', type=str, default=None, required=False, help='Read delta from image format', nargs='*') parser.add_argument('--drq', type=str, default=None, required=True, help='Catalog of objects in DRQ format') parser.add_argument('--mode', type=str, default='sdss', choices=['sdss', 'desi'], required=False, help='type of catalog supplied, default sdss') parser.add_argument('--rp-min', type=float, default=-200., required=False, help='Min r-parallel [h^-1 Mpc]') parser.add_argument('--rp-max', type=float, default=200., required=False, help='Max r-parallel [h^-1 Mpc]') parser.add_argument('--rt-max', type=float, default=200., required=False, help='Max r-transverse [h^-1 Mpc]') parser.add_argument('--np', type=int, default=100, required=False, help='Number of r-parallel bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of r-transverse bins') parser.add_argument('--z-min-obj', type=float, default=0, required=False, help='Min redshift for object field') parser.add_argument('--z-max-obj', type=float, default=10, required=False, help='Max redshift for object field') parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift larger than ' 'z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift smaller than ' 'z-cut-max')) parser.add_argument( '--lambda-abs', type=str, default='LYA', required=False, help=( 'Name of the absorption in picca.constants defining the redshift ' 'of the delta')) parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol-del', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the delta field') parser.add_argument( '--z-evol-obj', type=float, default=1., required=False, help='Exponent of the redshift evolution of the object field') parser.add_argument( '--fid-Om', type=float, default=0.315, required=False, help='Omega_matter(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-Or', type=float, default=0., required=False, help='Omega_radiation(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-wl', type=float, default=-1., required=False, help='Equation of state of dark energy of fiducial LambdaCDM cosmology' ) parser.add_argument('--max-diagram', type=int, default=4, required=False, help='Maximum diagram to compute') parser.add_argument( '--cf1d', type=str, required=True, help=('1D auto-correlation of pixels from the same forest file: ' 'picca_cf1d.py')) parser.add_argument( '--cf', type=str, default=None, required=False, help=('3D auto-correlation of pixels from different forests: ' 'picca_cf.py')) parser.add_argument( '--rej', type=float, default=1., required=False, help=('Fraction of rejected object-forests pairs: -1=no rejection, ' '1=all rejection')) parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 # setup variables in module xcf xcf.r_par_min = args.rp_min xcf.r_par_max = args.rp_max xcf.r_trans_max = args.rt_max xcf.z_cut_min = args.z_cut_min xcf.z_cut_max = args.z_cut_max xcf.num_bins_r_par = args.np xcf.num_bins_r_trans = args.nt xcf.nside = args.nside xcf.reject = args.rej xcf.z_ref = args.z_ref xcf.alpha = args.z_evol_del xcf.alpha_obj = args.z_evol_obj xcf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] xcf.max_diagram = args.max_diagram # read blinding keyword blinding = io.read_blinding(args.in_dir) # load fiducial cosmology if (args.fid_Or != 0.) or (args.fid_Ok != 0.) or (args.fid_wl != -1.): userprint(("ERROR: Cosmology with other than Omega_m set are not yet " "implemented")) sys.exit() cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) ### Read deltas data, num_data, z_min, z_max = io.read_deltas(args.in_dir, args.nside, xcf.lambda_abs, args.z_evol_del, args.z_ref, cosmo=cosmo, max_num_spec=args.nspec) for deltas in data.values(): for delta in deltas: delta.fname = 'D1' for item in [ 'cont', 'delta', 'order', 'ivar', 'exposures_diff', 'mean_snr', 'mean_reso', 'mean_z', 'delta_log_lambda' ]: setattr(delta, item, None) xcf.data = data xcf.num_data = num_data sys.stderr.write("\n") userprint("done, npix = {}, ndels = {}".format(len(data), xcf.num_data)) sys.stderr.write("\n") ### Find the redshift range if args.z_min_obj is None: r_comov_min = cosmo.get_r_comov(z_min) r_comov_min = max(0., r_comov_min + xcf.r_par_min) args.z_min_obj = cosmo.distance_to_redshift(r_comov_min) sys.stderr.write("\r z_min_obj = {}\r".format(args.z_min_obj)) if args.z_max_obj is None: r_comov_max = cosmo.get_r_comov(z_max) r_comov_max = max(0., r_comov_max + xcf.r_par_max) args.z_max_obj = cosmo.distance_to_redshift(r_comov_max) sys.stderr.write("\r z_max_obj = {}\r".format(args.z_max_obj)) ### Read objects objs, z_min2 = io.read_objects(args.drq, args.nside, args.z_min_obj, args.z_max_obj, args.z_evol_obj, args.z_ref, cosmo, mode=args.mode) xcf.objs = objs sys.stderr.write("\n") userprint("done, npix = {}".format(len(objs))) sys.stderr.write("\n") # compute maximum angular separation xcf.ang_max = utils.compute_ang_max(cosmo, xcf.r_trans_max, z_min, z_min2) # Load 1d correlation functions hdul = fitsio.FITS(args.cf1d) header = hdul[1].read_header() log_lambda_min = header['LLMIN'] delta_log_lambda = header['DLL'] num_pairs_variance_1d = hdul[1]['nv1d'][:] variance_1d = hdul[1]['v1d'][:] log_lambda = log_lambda_min + delta_log_lambda * np.arange( variance_1d.size) xcf.get_variance_1d['D1'] = interp1d( log_lambda[num_pairs_variance_1d > 0], variance_1d[num_pairs_variance_1d > 0], kind='nearest', fill_value='extrapolate') num_pairs1d = hdul[1]['nb1d'][:] xi_1d = hdul[1]['c1d'][:] xcf.xi_1d['D1'] = interp1d((log_lambda - log_lambda_min)[num_pairs1d > 0], xi_1d[num_pairs1d > 0], kind='nearest', fill_value='extrapolate') hdul.close() # Load correlation functions if not args.cf is None: hdul = fitsio.FITS(args.cf) header = hdul[1].read_header() assert cf.num_bins_r_par == header['NP'] assert cf.num_bins_r_trans == header['NT'] assert cf.r_par_min == header['RPMIN'] assert cf.r_par_max == header['RPMAX'] assert cf.r_trans_max == header['RTMAX'] xi = hdul[2]['DA'][:] weights = hdul[2]['WE'][:] xi = (xi * weights).sum(axis=0) weights = weights.sum(axis=0) w = weights > 0. xi[w] /= weights[w] xcf.xi_wick = xi.copy() hdul.close() cf.data = xcf.data cf.ang_max = xcf.ang_max cf.nside = xcf.nside cf.z_cut_max = xcf.z_cut_max cf.z_cut_min = xcf.z_cut_min ### Send xcf.counter = Value('i', 0) xcf.lock = Lock() cpu_data = {} for index, healpix in enumerate(sorted(data)): num_processor = index % args.nproc if not num_processor in cpu_data: cpu_data[num_processor] = [] cpu_data[num_processor].append(healpix) # Find neighbours for healpixs in cpu_data.values(): xcf.fill_neighs(healpixs) if not xcf.xi_wick is None: cf.fill_neighs(healpixs) # compute the covariance matrix context = multiprocessing.get_context('fork') pool = context.Pool(processes=min(args.nproc, len(cpu_data.values()))) userprint(" \nStarting\n") wick_data = pool.map(calc_wick_terms, sorted(cpu_data.values())) userprint(" \nFinished\n") pool.close() wick_data = np.array(wick_data) weights_wick = wick_data[:, 0].sum(axis=0) num_pairs_wick = wick_data[:, 1].sum(axis=0) npairs = wick_data[:, 2].sum(axis=0) npairs_used = wick_data[:, 3].sum(axis=0) t1 = wick_data[:, 4].sum(axis=0) t2 = wick_data[:, 5].sum(axis=0) t3 = wick_data[:, 6].sum(axis=0) t4 = wick_data[:, 7].sum(axis=0) t5 = wick_data[:, 8].sum(axis=0) t6 = wick_data[:, 9].sum(axis=0) weights = weights_wick * weights_wick[:, None] w = weights > 0. t1[w] /= weights[w] t2[w] /= weights[w] t3[w] /= weights[w] t4[w] /= weights[w] t5[w] /= weights[w] t6[w] /= weights[w] t1 *= 1. * npairs_used / npairs t2 *= 1. * npairs_used / npairs t3 *= 1. * npairs_used / npairs t4 *= 1. * npairs_used / npairs t5 *= 1. * npairs_used / npairs t6 *= 1. * npairs_used / npairs t_tot = t1 + t2 + t3 + t4 + t5 + t6 results = fitsio.FITS(args.out, 'rw', clobber=True) header = [{ 'name': 'RPMIN', 'value': xcf.r_par_min, 'comment': 'Minimum r-parallel [h^-1 Mpc]' }, { 'name': 'RPMAX', 'value': xcf.r_par_max, 'comment': 'Maximum r-parallel [h^-1 Mpc]' }, { 'name': 'RTMAX', 'value': xcf.r_trans_max, 'comment': 'Maximum r-transverse [h^-1 Mpc]' }, { 'name': 'NP', 'value': xcf.num_bins_r_par, 'comment': 'Number of bins in r-parallel' }, { 'name': 'NT', 'value': xcf.num_bins_r_trans, 'comment': 'Number of bins in r-transverse' }, { 'name': 'ZCUTMIN', 'value': xcf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': xcf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'REJ', 'value': xcf.reject, 'comment': 'Rejection factor' }, { 'name': 'NPALL', 'value': npairs, 'comment': 'Number of pairs' }, { 'name': 'NPUSED', 'value': npairs_used, 'comment': 'Number of used pairs' }, { 'name': 'OMEGAM', 'value': args.fid_Om, 'comment': 'Omega_matter(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAR', 'value': args.fid_Or, 'comment': 'Omega_radiation(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAK', 'value': args.fid_Ok, 'comment': 'Omega_k(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'WL', 'value': args.fid_wl, 'comment': 'Equation of state of dark energy of fiducial LambdaCDM cosmology' }, { 'name': "BLINDING", 'value': blinding, 'comment': 'String specifying the blinding strategy' }] comment = [ 'Sum of weight', 'Covariance', 'Nomber of pairs', 'T1', 'T2', 'T3', 'T4', 'T5', 'T6' ] results.write( [t_tot, weights_wick, num_pairs_wick, t1, t2, t3, t4, t5, t6], names=['CO', 'WALL', 'NB', 'T1', 'T2', 'T3', 'T4', 'T5', 'T6'], comment=comment, header=header, extname='COV') results.close()
def main(cmdargs): """Compute the distortion matrix of the cross-correlation delta x object for a list of IGM absorption.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=('Compute the distortion matrix of the cross-correlation ' 'delta x object for a list of IGM absorption.')) parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--drq', type=str, default=None, required=True, help='Catalog of objects in DRQ format') parser.add_argument('--mode', type=str, default='sdss', choices=['sdss', 'desi', 'desi_mocks', 'desi_healpix'], required=False, help='type of catalog supplied, default sdss') parser.add_argument('--rp-min', type=float, default=-200., required=False, help='Min r-parallel [h^-1 Mpc]') parser.add_argument('--rp-max', type=float, default=200., required=False, help='Max r-parallel [h^-1 Mpc]') parser.add_argument('--rt-max', type=float, default=200., required=False, help='Max r-transverse [h^-1 Mpc]') parser.add_argument('--np', type=int, default=100, required=False, help='Number of r-parallel bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of r-transverse bins') parser.add_argument( '--coef-binning-model', type=int, default=1, required=False, help=( 'Coefficient multiplying np and nt to get finner binning for the ' 'model')) parser.add_argument('--z-min-obj', type=float, default=0, required=False, help='Min redshift for object field') parser.add_argument('--z-max-obj', type=float, default=10, required=False, help='Max redshift for object field') parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift larger than ' 'z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift smaller than ' 'z-cut-max')) parser.add_argument( '--lambda-abs', type=str, default='LYA', required=False, help=( 'Name of the absorption in picca.constants defining the redshift ' 'of the delta')) parser.add_argument('--obj-name', type=str, default='QSO', required=False, help='Name of the object tracer') parser.add_argument( '--abs-igm', type=str, default=None, required=False, nargs='*', help='List of names of metal absorption in picca.constants') parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol-del', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the delta field') parser.add_argument( '--z-evol-obj', type=float, default=1., required=False, help='Exponent of the redshift evolution of the object field') parser.add_argument( '--metal-alpha', type=float, default=1., required=False, help='Exponent of the redshift evolution of the metal delta field') parser.add_argument( '--fid-Om', type=float, default=0.315, required=False, help='Omega_matter(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-Or', type=float, default=0., required=False, help='Omega_radiation(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-wl', type=float, default=-1., required=False, help='Equation of state of dark energy of fiducial LambdaCDM cosmology' ) parser.add_argument( '--rej', type=float, default=1., required=False, help=('Fraction of rejected object-forests pairs: -1=no rejection, ' '1=all rejection')) parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 # setup variables in module xcf xcf.r_par_max = args.rp_max xcf.r_par_min = args.rp_min xcf.r_trans_max = args.rt_max xcf.z_cut_max = args.z_cut_max xcf.z_cut_min = args.z_cut_min xcf.num_bins_r_par = args.np * args.coef_binning_model xcf.num_bins_r_trans = args.nt * args.coef_binning_model xcf.num_model_bins_r_par = args.np * args.coef_binning_model xcf.num_model_bins_r_trans = args.nt * args.coef_binning_model xcf.nside = args.nside xcf.z_ref = args.z_ref xcf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] xcf.reject = args.rej xcf.alpha_abs = {} xcf.alpha_abs[args.lambda_abs] = args.z_evol_del for metal in args.abs_igm: xcf.alpha_abs[metal] = args.metal_alpha # read blinding keyword blinding = io.read_blinding(args.in_dir) # load fiducial cosmology cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) xcf.cosmo = cosmo t0 = time.time() # read deltas data, num_data, z_min, z_max = io.read_deltas(args.in_dir, args.nside, xcf.lambda_abs, args.z_evol_del, args.z_ref, cosmo, max_num_spec=args.nspec) xcf.data = data xcf.num_data = num_data ### Find the redshift range if args.z_min_obj is None: r_comov_min = cosmo.get_r_comov(z_min) r_comov_min = max(0., r_comov_min + xcf.r_par_min) args.z_min_obj = cosmo.distance_to_redshift(r_comov_min) userprint("z_min_obj = {}".format(args.z_min_obj), end="") if args.z_max_obj is None: r_comov_max = cosmo.get_r_comov(z_max) r_comov_max = max(0., r_comov_max + xcf.r_par_max) args.z_max_obj = cosmo.distance_to_redshift(r_comov_max) userprint("z_max_obj = {}".format(args.z_max_obj), end="") # read objets objs, z_min2 = io.read_objects(args.drq, args.nside, args.z_min_obj, args.z_max_obj, args.z_evol_obj, args.z_ref, cosmo, mode=args.mode) xcf.objs = objs # compute maximum angular separation xcf.ang_max = utils.compute_ang_max(cosmo, xcf.r_trans_max, z_min, z_min2) t1 = time.time() userprint( f'picca_metal_xdmat.py - Time reading data: {(t1-t0)/60:.3f} minutes') xcf.counter = Value('i', 0) xcf.lock = Lock() cpu_data = {} for index, healpix in enumerate(sorted(list(data.keys()))): num_processor = index % args.nproc if not num_processor in cpu_data: cpu_data[num_processor] = [] cpu_data[num_processor].append(healpix) # intiialize arrays to store the results for the different metal absorption dmat_all = [] weights_dmat_all = [] r_par_all = [] r_trans_all = [] z_all = [] names = [] num_pairs_all = [] num_pairs_used_all = [] # loop over metals for index, abs_igm in enumerate(args.abs_igm): xcf.counter.value = 0 calc_metal_dmat_wrapper = partial(calc_metal_dmat, abs_igm) userprint("") if args.nproc > 1: context = multiprocessing.get_context('fork') pool = context.Pool(processes=args.nproc) dmat_data = pool.map(calc_metal_dmat_wrapper, sorted(cpu_data.values())) pool.close() elif args.nproc == 1: dmat_data = map(calc_metal_dmat_wrapper, sorted(cpu_data.values())) dmat_data = list(dmat_data) # merge the results from different CPUs dmat_data = np.array(dmat_data) weights_dmat = dmat_data[:, 0].sum(axis=0) dmat = dmat_data[:, 1].sum(axis=0) r_par = dmat_data[:, 2].sum(axis=0) r_trans = dmat_data[:, 3].sum(axis=0) z = dmat_data[:, 4].sum(axis=0) weights = dmat_data[:, 5].sum(axis=0) # normalize_values w = weights > 0 r_par[w] /= weights[w] r_trans[w] /= weights[w] z[w] /= weights[w] num_pairs = dmat_data[:, 6].sum(axis=0) num_pairs_used = dmat_data[:, 7].sum(axis=0) w = weights_dmat > 0 dmat[w, :] /= weights_dmat[w, None] # add these results to the list ofor the different metal absorption dmat_all.append(dmat) weights_dmat_all.append(weights_dmat) r_par_all.append(r_par) r_trans_all.append(r_trans) z_all.append(z) names.append(abs_igm) num_pairs_all.append(num_pairs) num_pairs_used_all.append(num_pairs_used) t2 = time.time() userprint( f'picca_metal_xdmat.py - Time computing all metal matrix: {(t2-t1)/60:.3f} minutes' ) # save the results results = fitsio.FITS(args.out, 'rw', clobber=True) header = [{ 'name': 'RPMIN', 'value': xcf.r_par_min, 'comment': 'Minimum r-parallel [h^-1 Mpc]' }, { 'name': 'RPMAX', 'value': xcf.r_par_max, 'comment': 'Maximum r-parallel [h^-1 Mpc]' }, { 'name': 'RTMAX', 'value': xcf.r_trans_max, 'comment': 'Maximum r-transverse [h^-1 Mpc]' }, { 'name': 'NP', 'value': xcf.num_bins_r_par, 'comment': 'Number of bins in r-parallel' }, { 'name': 'NT', 'value': xcf.num_bins_r_trans, 'comment': 'Number of bins in r-transverse' }, { 'name': 'COEFMOD', 'value': args.coef_binning_model, 'comment': 'Coefficient for model binning' }, { 'name': 'ZCUTMIN', 'value': xcf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': xcf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'REJ', 'value': xcf.reject, 'comment': 'Rejection factor' }, { 'name': 'OMEGAM', 'value': args.fid_Om, 'comment': 'Omega_matter(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAR', 'value': args.fid_Or, 'comment': 'Omega_radiation(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAK', 'value': args.fid_Ok, 'comment': 'Omega_k(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'WL', 'value': args.fid_wl, 'comment': 'Equation of state of dark energy of fiducial LambdaCDM cosmology' }, { 'name': "BLINDING", 'value': blinding, 'comment': 'String specifying the blinding strategy' }] len_names = np.array([len(name) for name in names]).max() names = np.array(names, dtype='S' + str(len_names)) results.write( [ np.array(num_pairs_all), np.array(num_pairs_used_all), np.array(names) ], names=['NPALL', 'NPUSED', 'ABS_IGM'], header=header, comment=['Number of pairs', 'Number of used pairs', 'Absorption name'], extname='ATTRI') dmat_name = "DM_" if blinding != "none": dmat_name += "BLIND_" names = names.astype(str) out_list = [] out_names = [] out_comment = [] out_units = [] for index, name in enumerate(names): out_names += ['RP_' + args.obj_name + '_' + name] out_list += [r_par_all[index]] out_comment += ['R-parallel'] out_units += ['h^-1 Mpc'] out_names += ['RT_' + args.obj_name + '_' + name] out_list += [r_trans_all[index]] out_comment += ['R-transverse'] out_units += ['h^-1 Mpc'] out_names += ['Z_' + args.obj_name + '_' + name] out_list += [z_all[index]] out_comment += ['Redshift'] out_units += [''] out_names += [dmat_name + args.obj_name + '_' + name] out_list += [dmat_all[index]] out_comment += ['Distortion matrix'] out_units += [''] out_names += ['WDM_' + args.obj_name + '_' + name] out_list += [weights_dmat_all[index]] out_comment += ['Sum of weight'] out_units += [''] results.write(out_list, names=out_names, comment=out_comment, units=out_units, extname='MDMAT') results.close() t3 = time.time() userprint(f'picca_metal_xdmat.py - Time total: {(t3-t0)/60:.3f} minutes')
def main(cmdargs): """Compute the auto and cross-correlation of delta fields""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='Compute the auto and cross-correlation of delta fields') parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--from-image', type=str, default=None, required=False, help='Read delta from image format', nargs='*') parser.add_argument('--in-dir2', type=str, default=None, required=False, help='Directory to 2nd delta files') parser.add_argument('--rp-min', type=float, default=0., required=False, help='Min r-parallel [h^-1 Mpc]') parser.add_argument('--rp-max', type=float, default=200., required=False, help='Max r-parallel [h^-1 Mpc]') parser.add_argument('--rt-max', type=float, default=200., required=False, help='Max r-transverse [h^-1 Mpc]') parser.add_argument('--np', type=int, default=50, required=False, help='Number of r-parallel bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of r-transverse bins') parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean ' 'of the last absorber redshift and the object ' 'redshift larger than z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean ' 'of the last absorber redshift and the object ' 'redshift smaller than z-cut-max')) parser.add_argument('--lambda-abs', type=str, default='LYA', required=False, help=('Name of the absorption in picca.constants ' 'defining the redshift of the delta')) parser.add_argument('--lambda-abs2', type=str, default=None, required=False, help=('Name of the absorption in picca.constants ' 'defining the redshift of the 2nd delta')) parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol', type=float, default=2.9, required=False, help=('Exponent of the redshift evolution of the delta ' 'field')) parser.add_argument('--z-evol2', type=float, default=2.9, required=False, help=('Exponent of the redshift evolution of the 2nd ' 'delta field')) parser.add_argument('--fid-Om', type=float, default=0.315, required=False, help=('Omega_matter(z=0) of fiducial LambdaCDM ' 'cosmology')) parser.add_argument('--fid-Or', type=float, default=0., required=False, help=('Omega_radiation(z=0) of fiducial LambdaCDM ' 'cosmology')) parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-wl', type=float, default=-1., required=False, help=('Equation of state of dark energy of fiducial ' 'LambdaCDM cosmology')) parser.add_argument('--no-project', action='store_true', required=False, help='Do not project out continuum fitting modes') parser.add_argument('--remove-same-half-plate-close-pairs', action='store_true', required=False, help=('Reject pairs in the first bin in r-parallel ' 'from same half plate')) parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') parser.add_argument( '--unfold-cf', action='store_true', required=False, help=('rp can be positive or negative depending on the ' 'relative position between absorber1 and ' 'absorber2')) parser.add_argument('--shuffle-distrib-forest-seed', type=int, default=None, required=False, help=('Shuffle the distribution of forests on the sky ' 'following the given seed. Do not shuffle if ' 'None')) args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 # setup variables in module cf cf.r_par_max = args.rp_max cf.r_trans_max = args.rt_max cf.r_par_min = args.rp_min cf.z_cut_max = args.z_cut_max cf.z_cut_min = args.z_cut_min cf.num_bins_r_par = args.np cf.num_bins_r_trans = args.nt cf.nside = args.nside cf.z_ref = args.z_ref cf.alpha = args.z_evol cf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] cf.remove_same_half_plate_close_pairs = args.remove_same_half_plate_close_pairs # read blinding keyword blinding = io.read_blinding(args.in_dir) # load fiducial cosmology cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) t0 = time.time() ### Read data 1 data, num_data, z_min, z_max = io.read_deltas(args.in_dir, cf.nside, cf.lambda_abs, cf.alpha, cf.z_ref, cosmo, max_num_spec=args.nspec, no_project=args.no_project, from_image=args.from_image) del z_max cf.data = data cf.num_data = num_data cf.ang_max = utils.compute_ang_max(cosmo, cf.r_trans_max, z_min) userprint("") userprint("done, npix = {}".format(len(data))) ### Read data 2 if args.in_dir2 or args.lambda_abs2: if args.lambda_abs2 or args.unfold_cf: cf.x_correlation = True cf.alpha2 = args.z_evol2 if args.in_dir2 is None: args.in_dir2 = args.in_dir if args.lambda_abs2: cf.lambda_abs2 = constants.ABSORBER_IGM[args.lambda_abs2] else: cf.lambda_abs2 = cf.lambda_abs data2, num_data2, z_min2, z_max2 = io.read_deltas( args.in_dir2, cf.nside, cf.lambda_abs2, cf.alpha2, cf.z_ref, cosmo, max_num_spec=args.nspec, no_project=args.no_project, from_image=args.from_image) del z_max2 cf.data2 = data2 cf.num_data2 = num_data2 cf.ang_max = utils.compute_ang_max(cosmo, cf.r_trans_max, z_min, z_min2) userprint("") userprint("done, npix = {}".format(len(data2))) # shuffle forests if args.shuffle_distrib_forest_seed is not None: cf.data = utils.shuffle_distrib_forests( cf.data, args.shuffle_distrib_forest_seed) t1 = time.time() userprint(f'picca_cf.py - Time reading data: {(t1-t0)/60:.3f} minutes') # compute correlation function, use pool to parallelize cf.counter = Value('i', 0) cf.lock = Lock() cpu_data = {healpix: [healpix] for healpix in data} context = multiprocessing.get_context('fork') pool = context.Pool(processes=args.nproc) correlation_function_data = pool.map(corr_func, sorted(cpu_data.values())) pool.close() t2 = time.time() userprint( f'picca_cf.py - Time computing correlation function: {(t2-t1)/60:.3f} minutes' ) # group data from parallelisation correlation_function_data = np.array(correlation_function_data) weights_list = correlation_function_data[:, 0, :] xi_list = correlation_function_data[:, 1, :] r_par_list = correlation_function_data[:, 2, :] r_trans_list = correlation_function_data[:, 3, :] z_list = correlation_function_data[:, 4, :] num_pairs_list = correlation_function_data[:, 5, :].astype(np.int64) healpix_list = np.array(sorted(list(cpu_data.keys()))) # normalize values w = (weights_list.sum(axis=0) > 0.) r_par = (r_par_list * weights_list).sum(axis=0) r_par[w] /= weights_list.sum(axis=0)[w] r_trans = (r_trans_list * weights_list).sum(axis=0) r_trans[w] /= weights_list.sum(axis=0)[w] z = (z_list * weights_list).sum(axis=0) z[w] /= weights_list.sum(axis=0)[w] num_pairs = num_pairs_list.sum(axis=0) # save data results = fitsio.FITS(args.out, 'rw', clobber=True) header = [{ 'name': 'RPMIN', 'value': cf.r_par_min, 'comment': 'Minimum r-parallel [h^-1 Mpc]' }, { 'name': 'RPMAX', 'value': cf.r_par_max, 'comment': 'Maximum r-parallel [h^-1 Mpc]' }, { 'name': 'RTMAX', 'value': cf.r_trans_max, 'comment': 'Maximum r-transverse [h^-1 Mpc]' }, { 'name': 'NP', 'value': cf.num_bins_r_par, 'comment': 'Number of bins in r-parallel' }, { 'name': 'NT', 'value': cf.num_bins_r_trans, 'comment': 'Number of bins in r-transverse' }, { 'name': 'ZCUTMIN', 'value': cf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': cf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'NSIDE', 'value': cf.nside, 'comment': 'Healpix nside' }, { 'name': 'OMEGAM', 'value': args.fid_Om, 'comment': 'Omega_matter(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAR', 'value': args.fid_Or, 'comment': 'Omega_radiation(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAK', 'value': args.fid_Ok, 'comment': 'Omega_k(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'WL', 'value': args.fid_wl, 'comment': 'Equation of state of dark energy of fiducial LambdaCDM cosmology' }, { 'name': "BLINDING", 'value': blinding, 'comment': 'String specifying the blinding strategy' }] results.write( [r_par, r_trans, z, num_pairs], names=['RP', 'RT', 'Z', 'NB'], comment=['R-parallel', 'R-transverse', 'Redshift', 'Number of pairs'], units=['h^-1 Mpc', 'h^-1 Mpc', '', ''], header=header, extname='ATTRI') header2 = [{ 'name': 'HLPXSCHM', 'value': 'RING', 'comment': 'Healpix scheme' }] xi_list_name = "DA" if blinding != "none": xi_list_name += "_BLIND" results.write([healpix_list, weights_list, xi_list], names=['HEALPID', 'WE', xi_list_name], comment=['Healpix index', 'Sum of weight', 'Correlation'], header=header2, extname='COR') results.close() t3 = time.time() userprint(f'picca_cf.py - Time total : {(t3-t0)/60:.3f} minutes')
def main(cmdargs): # pylint: disable-msg=too-many-locals,too-many-branches,too-many-statements """Compute the auto and cross-correlation of delta fields for a list of IGM absorption.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=('Compute the auto and cross-correlation of delta fields ' 'for a list of IGM absorption.')) parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--in-dir2', type=str, default=None, required=False, help='Directory to 2nd delta files') parser.add_argument('--rp-min', type=float, default=0., required=False, help='Min r-parallel [h^-1 Mpc]') parser.add_argument('--rp-max', type=float, default=200., required=False, help='Max r-parallel [h^-1 Mpc]') parser.add_argument('--rt-max', type=float, default=200., required=False, help='Max r-transverse [h^-1 Mpc]') parser.add_argument('--np', type=int, default=50, required=False, help='Number of r-parallel bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of r-transverse bins') parser.add_argument( '--coef-binning-model', type=int, default=1, required=False, help=('Coefficient multiplying np and nt to get finner binning for the ' 'model')) parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift larger than ' 'z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift smaller than ' 'z-cut-max')) parser.add_argument( '--lambda-abs', type=str, default='LYA', required=False, help=('Name of the absorption in picca.constants defining the redshift ' 'of the delta')) parser.add_argument( '--lambda-abs2', type=str, default=None, required=False, help=('Name of the absorption in picca.constants defining the redshift ' 'of the 2nd delta')) parser.add_argument( '--abs-igm', type=str, default=[], required=False, nargs='*', help=('List of names of metal absorption in picca.constants present in ' 'forest')) parser.add_argument( '--abs-igm2', type=str, default=[], required=False, nargs='*', help=('List of names of metal absorption in picca.constants present in ' '2nd forest')) parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the delta field') parser.add_argument( '--z-evol2', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the 2nd delta field') parser.add_argument( '--metal-alpha', type=float, default=1., required=False, help='Exponent of the redshift evolution of the metal delta field') parser.add_argument( '--fid-Om', type=float, default=0.315, required=False, help='Omega_matter(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-Or', type=float, default=0., required=False, help='Omega_radiation(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-wl', type=float, default=-1., required=False, help='Equation of state of dark energy of fiducial LambdaCDM cosmology') parser.add_argument( '--remove-same-half-plate-close-pairs', action='store_true', required=False, help='Reject pairs in the first bin in r-parallel from same half plate') parser.add_argument( '--rej', type=float, default=1., required=False, help=('Fraction of rejected forest-forest pairs: -1=no rejection, ' '1=all rejection')) parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') parser.add_argument( '--unfold-cf', action='store_true', required=False, help=('rp can be positive or negative depending on the relative ' 'position between absorber1 and absorber2')) args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 userprint("nproc", args.nproc) # setup variables in module cf cf.r_par_max = args.rp_max cf.r_trans_max = args.rt_max cf.r_par_min = args.rp_min cf.z_cut_max = args.z_cut_max cf.z_cut_min = args.z_cut_min cf.num_bins_r_par = args.np * args.coef_binning_model cf.num_bins_r_trans = args.nt * args.coef_binning_model cf.num_model_bins_r_par = args.np * args.coef_binning_model cf.num_model_bins_r_trans = args.nt * args.coef_binning_model cf.nside = args.nside cf.z_ref = args.z_ref cf.alpha = args.z_evol cf.reject = args.rej cf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] cf.remove_same_half_plate_close_pairs = args.remove_same_half_plate_close_pairs cf.alpha_abs = {} cf.alpha_abs[args.lambda_abs] = cf.alpha for metal in args.abs_igm: cf.alpha_abs[metal] = args.metal_alpha # read blinding keyword blinding = io.read_blinding(args.in_dir) # load fiducial cosmology cf.cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) t0 = time.time() ### Read data 1 data, num_data, z_min, z_max = io.read_deltas(args.in_dir, cf.nside, cf.lambda_abs, cf.alpha, cf.z_ref, cf.cosmo, max_num_spec=args.nspec) del z_max cf.data = data cf.num_data = num_data cf.ang_max = utils.compute_ang_max(cf.cosmo, cf.r_trans_max, z_min) userprint("") userprint("done, npix = {}".format(len(data))) ### Read data 2 if args.in_dir2 or args.lambda_abs2: if args.lambda_abs2 or args.unfold_cf: cf.x_correlation = True cf.alpha2 = args.z_evol2 if args.in_dir2 is None: args.in_dir2 = args.in_dir if args.lambda_abs2: cf.lambda_abs2 = constants.ABSORBER_IGM[args.lambda_abs2] else: cf.lambda_abs2 = cf.lambda_abs cf.alpha_abs[args.lambda_abs2] = cf.alpha2 for m in args.abs_igm2: cf.alpha_abs[m] = args.metal_alpha data2, num_data2, z_min2, z_max2 = io.read_deltas( args.in_dir2, cf.nside, cf.lambda_abs2, cf.alpha2, cf.z_ref, cf.cosmo, max_num_spec=args.nspec) del z_max2 cf.data2 = data2 cf.num_data2 = num_data2 cf.ang_max = utils.compute_ang_max(cf.cosmo, cf.r_trans_max, z_min, z_min2) userprint("") userprint("done, npix = {}".format(len(data2))) t1 = time.time() userprint(f'picca_metal_dmat.py - Time reading data: {(t1-t0)/60:.3f} minutes') cf.counter = Value('i', 0) cf.lock = Lock() cpu_data = {} for index, healpix in enumerate(sorted(list(data.keys()))): num_processor = index % args.nproc if not num_processor in cpu_data: cpu_data[num_processor] = [] cpu_data[num_processor].append(healpix) # intiialize arrays to store the results for the different metal absorption dmat_all = [] weights_dmat_all = [] r_par_all = [] r_trans_all = [] z_all = [] names = [] num_pairs_all = [] num_pairs_used_all = [] abs_igm = [args.lambda_abs] + args.abs_igm userprint("abs_igm = {}".format(abs_igm)) if args.lambda_abs2 is None: args.lambda_abs2 = args.lambda_abs args.abs_igm2 = args.abs_igm abs_igm_2 = [args.lambda_abs2] + args.abs_igm2 if cf.x_correlation: userprint("abs_igm2 = {}".format(abs_igm_2)) # loop over metals for index1, abs_igm1 in enumerate(abs_igm): index0 = index1 if args.lambda_abs != args.lambda_abs2: index0 = 0 for index2, abs_igm2 in enumerate(abs_igm_2[index0:]): if index1 == 0 and index2 == 0: continue cf.counter.value = 0 calc_metal_dmat_wrapper = partial(calc_metal_dmat, abs_igm1, abs_igm2) userprint("") # compute the distortion matrix if args.nproc > 1: context = multiprocessing.get_context('fork') pool = context.Pool(processes=args.nproc) dmat_data = pool.map(calc_metal_dmat_wrapper, sorted(cpu_data.values())) pool.close() elif args.nproc == 1: dmat_data = map(calc_metal_dmat_wrapper, sorted(cpu_data.values())) dmat_data = list(dmat_data) # merge the results from different CPUs dmat_data = np.array(dmat_data) weights_dmat = dmat_data[:, 0].sum(axis=0) dmat = dmat_data[:, 1].sum(axis=0) r_par = dmat_data[:, 2].sum(axis=0) r_trans = dmat_data[:, 3].sum(axis=0) z = dmat_data[:, 4].sum(axis=0) weights = dmat_data[:, 5].sum(axis=0) num_pairs = dmat_data[:, 6].sum(axis=0) num_pairs_used = dmat_data[:, 7].sum(axis=0) # normalize_values w = weights > 0 r_par[w] /= weights[w] r_trans[w] /= weights[w] z[w] /= weights[w] w = weights_dmat > 0 dmat[w, :] /= weights_dmat[w, None] # add these results to the list ofor the different metal absorption dmat_all.append(dmat) weights_dmat_all.append(weights_dmat) r_par_all.append(r_par) r_trans_all.append(r_trans) z_all.append(z) names.append(abs_igm1 + "_" + abs_igm2) num_pairs_all.append(num_pairs) num_pairs_used_all.append(num_pairs_used) t2 = time.time() userprint(f'picca_metal_dmat.py - Time computing all metal matrices : {(t2-t1)/60:.3f} minutes') # save the results results = fitsio.FITS(args.out, 'rw', clobber=True) header = [ { 'name': 'RPMIN', 'value': cf.r_par_min, 'comment': 'Minimum r-parallel [h^-1 Mpc]' }, { 'name': 'RPMAX', 'value': cf.r_par_max, 'comment': 'Maximum r-parallel [h^-1 Mpc]' }, { 'name': 'RTMAX', 'value': cf.r_trans_max, 'comment': 'Maximum r-transverse [h^-1 Mpc]' }, { 'name': 'NP', 'value': cf.num_bins_r_par, 'comment': 'Number of bins in r-parallel' }, { 'name': 'NT', 'value': cf.num_bins_r_trans, 'comment': ' Number of bins in r-transverse' }, { 'name': 'COEFMOD', 'value': args.coef_binning_model, 'comment': 'Coefficient for model binning' }, { 'name': 'ZCUTMIN', 'value': cf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': cf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'REJ', 'value': cf.reject, 'comment': 'Rejection factor' }, { 'name': 'ALPHAMET', 'value': args.metal_alpha, 'comment': 'Evolution of metal bias' }, { 'name': 'OMEGAM', 'value': args.fid_Om, 'comment': 'Omega_matter(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAR', 'value': args.fid_Or, 'comment': 'Omega_radiation(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAK', 'value': args.fid_Ok, 'comment': 'Omega_k(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'WL', 'value': args.fid_wl, 'comment': 'Equation of state of dark energy of fiducial LambdaCDM cosmology' }, { 'name': "BLINDING", 'value': blinding, 'comment': 'String specifying the blinding strategy' } ] len_names = np.array([len(name) for name in names]).max() names = np.array(names, dtype='S' + str(len_names)) results.write( [ np.array(num_pairs_all), np.array(num_pairs_used_all), np.array(names) ], names=['NPALL', 'NPUSED', 'ABS_IGM'], header=header, comment=['Number of pairs', 'Number of used pairs', 'Absorption name'], extname='ATTRI') dmat_name = "DM_" if blinding != "none": dmat_name += "BLIND_" names = names.astype(str) out_list = [] out_names = [] out_comment = [] out_units = [] for index, name in enumerate(names): out_names += ['RP_' + name] out_list += [r_par_all[index]] out_comment += ['R-parallel'] out_units += ['h^-1 Mpc'] out_names += ['RT_' + name] out_list += [r_trans_all[index]] out_comment += ['R-transverse'] out_units += ['h^-1 Mpc'] out_names += ['Z_' + name] out_list += [z_all[index]] out_comment += ['Redshift'] out_units += [''] out_names += [dmat_name + name] out_list += [dmat_all[index]] out_comment += ['Distortion matrix'] out_units += [''] out_names += ['WDM_' + name] out_list += [weights_dmat_all[index]] out_comment += ['Sum of weight'] out_units += [''] results.write(out_list, names=out_names, comment=out_comment, units=out_units, extname='MDMAT') results.close() t3 = time.time() userprint(f'picca_metal_dmat.py - Time total : {(t3-t0)/60:.3f} minutes')
def main(cmdargs): """Computes the cross-correlation between a catalog of objects and a delta field as a function of angle and wavelength ratio""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=('Compute the cross-correlation between a catalog of ' 'objects and a delta field as a function of angle and ' 'wavelength ratio')) parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--drq', type=str, default=None, required=True, help='Catalog of objects in DRQ format') parser.add_argument('--mode', type=str, default='sdss', choices=['sdss', 'desi'], required=False, help='type of catalog supplied, default sdss') parser.add_argument('--wr-min', type=float, default=0.9, required=False, help='Min of wavelength ratio') parser.add_argument('--wr-max', type=float, default=1.1, required=False, help='Max of wavelength ratio') parser.add_argument('--ang-max', type=float, default=0.02, required=False, help='Max angle (rad)') parser.add_argument('--np', type=int, default=100, required=False, help='Number of wavelength ratio bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of angular bins') parser.add_argument('--z-min-obj', type=float, default=0, required=False, help='Min redshift for object field') parser.add_argument('--z-max-obj', type=float, default=10, required=False, help='Max redshift for object field') parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift larger than ' 'z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift smaller than ' 'z-cut-max')) parser.add_argument( '--lambda-abs', type=str, default='LYA', required=False, help=( 'Name of the absorption in picca.constants defining the redshift ' 'of the delta')) parser.add_argument( '--lambda-abs-obj', type=str, default='LYA', required=False, help=('Name of the absorption in picca.constants the object is ' 'considered as')) parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol-del', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the delta field') parser.add_argument( '--z-evol-obj', type=float, default=1., required=False, help='Exponent of the redshift evolution of the object field') parser.add_argument( '--fid-Om', type=float, default=0.315, required=False, help='Omega_matter(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-Or', type=float, default=0., required=False, help='Omega_radiation(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-wl', type=float, default=-1., required=False, help='Equation of state of dark energy of fiducial LambdaCDM cosmology' ) parser.add_argument('--no-project', action='store_true', required=False, help='Do not project out continuum fitting modes') parser.add_argument('--no-remove-mean-lambda-obs', action='store_true', required=False, help='Do not remove mean delta versus lambda_obs') parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 # setup variables in module xcf xcf.r_par_min = args.wr_min xcf.r_par_max = args.wr_max xcf.r_trans_max = args.ang_max xcf.z_cut_min = args.z_cut_min xcf.z_cut_max = args.z_cut_max xcf.num_bins_r_par = args.np xcf.num_bins_r_trans = args.nt xcf.nside = args.nside xcf.ang_correlation = True xcf.ang_max = args.ang_max xcf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] # read blinding keyword blinding = io.read_blinding(args.in_dir) # load fiducial cosmology cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) ### Read deltas data, num_data, z_min, z_max = io.read_deltas( args.in_dir, args.nside, constants.ABSORBER_IGM[args.lambda_abs], args.z_evol_del, args.z_ref, cosmo=cosmo, max_num_spec=args.nspec, no_project=args.no_project) xcf.data = data xcf.num_data = num_data userprint("") userprint("done, npix = {}".format(len(data))) ### Remove <delta> vs. lambda_obs if not args.no_remove_mean_lambda_obs: Forest.delta_log_lambda = None for healpix in xcf.data: for delta in xcf.data[healpix]: delta_log_lambda = np.asarray([ delta.log_lambda[index] - delta.log_lambda[index - 1] for index in range(1, delta.log_lambda.size) ]).min() if Forest.delta_log_lambda is None: Forest.delta_log_lambda = delta_log_lambda else: Forest.delta_log_lambda = min(delta_log_lambda, Forest.delta_log_lambda) Forest.log_lambda_min = (np.log10( (z_min + 1.) * xcf.lambda_abs) - Forest.delta_log_lambda / 2.) Forest.log_lambda_max = (np.log10( (z_max + 1.) * xcf.lambda_abs) + Forest.delta_log_lambda / 2.) log_lambda, mean_delta, stack_weight = prep_del.stack( xcf.data, stack_from_deltas=True) del log_lambda, stack_weight for healpix in xcf.data: for delta in xcf.data[healpix]: bins = ((delta.log_lambda - Forest.log_lambda_min) / Forest.delta_log_lambda + 0.5).astype(int) delta.delta -= mean_delta[bins] ### Read objects objs, z_min2 = io.read_objects(args.drq, args.nside, args.z_min_obj, args.z_max_obj, args.z_evol_obj, args.z_ref, cosmo, mode=args.mode) del z_min2 for index, healpix in enumerate(sorted(objs)): for obj in objs[healpix]: obj.log_lambda = np.log10( (1. + obj.z_qso) * constants.ABSORBER_IGM[args.lambda_abs_obj]) userprint("") xcf.objs = objs # compute correlation function, use pool to parallelize xcf.counter = Value('i', 0) xcf.lock = Lock() cpu_data = {healpix: [healpix] for healpix in data} context = multiprocessing.get_context('fork') pool = context.Pool(processes=args.nproc) correlation_function_data = pool.map(corr_func, sorted(list(cpu_data.values()))) pool.close() # group data from parallelisation correlation_function_data = np.array(correlation_function_data) weights_list = correlation_function_data[:, 0, :] xi_list = correlation_function_data[:, 1, :] r_par_list = correlation_function_data[:, 2, :] r_trans_list = correlation_function_data[:, 3, :] z_list = correlation_function_data[:, 4, :] num_pairs_list = correlation_function_data[:, 5, :].astype(np.int64) healpix_list = np.array(sorted(list(cpu_data.keys()))) w = (weights_list.sum(axis=0) > 0.) r_par = (r_par_list * weights_list).sum(axis=0) r_par[w] /= weights_list.sum(axis=0)[w] r_trans = (r_trans_list * weights_list).sum(axis=0) r_trans[w] /= weights_list.sum(axis=0)[w] z = (z_list * weights_list).sum(axis=0) z[w] /= weights_list.sum(axis=0)[w] num_pairs = num_pairs_list.sum(axis=0) # save results results = fitsio.FITS(args.out, 'rw', clobber=True) header = [{ 'name': 'RPMIN', 'value': xcf.r_par_min, 'comment': 'Minimum wavelength ratio' }, { 'name': 'RPMAX', 'value': xcf.r_par_max, 'comment': 'Maximum wavelength ratio' }, { 'name': 'RTMAX', 'value': xcf.r_trans_max, 'comment': 'Maximum angle [rad]' }, { 'name': 'NP', 'value': xcf.num_bins_r_par, 'comment': 'Number of bins in wavelength ratio' }, { 'name': 'NT', 'value': xcf.num_bins_r_trans, 'comment': 'Number of bins in angle' }, { 'name': 'ZCUTMIN', 'value': xcf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': xcf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'NSIDE', 'value': xcf.nside, 'comment': 'Healpix nside' }] results.write( [r_par, r_trans, z, num_pairs], names=['RP', 'RT', 'Z', 'NB'], units=['', 'rad', '', ''], comment=['Wavelength ratio', 'Angle', 'Redshift', 'Number of pairs'], header=header, extname='ATTRI') header2 = [{ 'name': 'HLPXSCHM', 'value': 'RING', 'comment': ' Healpix scheme' }] xi_list_name = "DA" if blinding != "none": xi_list_name += "_BLIND" results.write([healpix_list, weights_list, xi_list], names=['HEALPID', 'WE', xi_list_name], comment=['Healpix index', 'Sum of weight', 'Correlation'], header=header2, extname='COR') results.close()
def main(cmdargs): # pylint: disable-msg=too-many-locals,too-many-branches,too-many-statements """Computes the wick covariance for the auto-correlation of forests""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=('Compute the wick covariance for the auto-correlation of ' 'forests')) parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--in-dir2', type=str, default=None, required=False, help='Directory to 2nd delta files') parser.add_argument('--rp-min', type=float, default=0., required=False, help='Min r-parallel [h^-1 Mpc]') parser.add_argument('--rp-max', type=float, default=200., required=False, help='Max r-parallel [h^-1 Mpc]') parser.add_argument('--rt-max', type=float, default=200., required=False, help='Max r-transverse [h^-1 Mpc]') parser.add_argument('--np', type=int, default=50, required=False, help='Number of r-parallel bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of r-transverse bins') parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift larger than ' 'z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift smaller than ' 'z-cut-max')) parser.add_argument( '--lambda-abs', type=str, default='LYA', required=False, help=('Name of the absorption in picca.constants defining the redshift ' 'of the delta')) parser.add_argument( '--lambda-abs2', type=str, default=None, required=False, help=('Name of the absorption in picca.constants defining the redshift ' 'of the 2nd delta')) parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the delta field') parser.add_argument( '--z-evol2', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the 2nd delta field') parser.add_argument( '--fid-Om', type=float, default=0.315, required=False, help='Omega_matter(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-Or', type=float, default=0., required=False, help='Omega_radiation(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-wl', type=float, default=-1., required=False, help='Equation of state of dark energy of fiducial LambdaCDM cosmology') parser.add_argument('--no-project', action='store_true', required=False, help='Do not project out continuum fitting modes') parser.add_argument('--max-diagram', type=int, default=3, required=False, help='Maximum diagram to compute') parser.add_argument( '--cf1d', type=str, required=True, help=('1D auto-correlation of pixels from the same forest file: ' 'picca_cf1d.py')) parser.add_argument( '--cf1d2', type=str, default=None, required=False, help=('1D auto-correlation of pixels from the same forest file of the ' '2nd delta field: picca_cf1d.py')) parser.add_argument( '--cf', type=str, default=None, required=False, help=('3D auto-correlation of pixels from different forests: ' 'picca_cf.py')) parser.add_argument( '--cf2', type=str, default=None, required=False, help=('3D auto-correlation of pixels from different forests for 2nd ' 'catalog: picca_cf.py')) parser.add_argument( '--cf12', type=str, default=None, required=False, help=('3D auto-correlation of pixels from different forests for cross ' '1st and 2nd catalog: picca_cf.py')) parser.add_argument( '--unfold-cf', action='store_true', required=False, help=('rp can be positive or negative depending on the relative ' 'position between absorber1 and absorber2')) parser.add_argument( '--rej', type=float, default=1., required=False, help='Fraction of rejected pairs: -1=no rejection, 1=all rejection') parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 userprint("nproc", args.nproc) # setup variables in module cf cf.r_par_max = args.rp_max cf.r_trans_max = args.rt_max cf.r_par_min = args.rp_min cf.z_cut_max = args.z_cut_max cf.z_cut_min = args.z_cut_min cf.num_bins_r_par = args.np cf.num_bins_r_trans = args.nt cf.nside = args.nside cf.z_ref = args.z_ref cf.alpha = args.z_evol cf.alpha2 = args.z_evol cf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] cf.reject = args.rej cf.max_diagram = args.max_diagram # read blinding keyword blinding = io.read_blinding(args.in_dir) # load cosmology if (args.fid_Or != 0.) or (args.fid_Ok != 0.) or (args.fid_wl != -1.): userprint(("ERROR: Cosmology with other than Omega_m set are not yet " "implemented")) sys.exit() cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) # read data 1 data, num_data, z_min, z_max = io.read_deltas(args.in_dir, cf.nside, cf.lambda_abs, cf.alpha, cf.z_ref, cosmo, max_num_spec=args.nspec) for deltas in data.values(): for delta in deltas: delta.fname = 'D1' for item in [ 'cont', 'delta', 'order', 'ivar', 'exposures_diff', 'mean_snr', 'mean_reso', 'mean_z', 'delta_log_lambda' ]: setattr(delta, item, None) del z_max cf.data = data cf.num_data = num_data cf.ang_max = utils.compute_ang_max(cosmo, cf.r_trans_max, z_min) sys.stderr.write("\n") userprint("done, npix = {}".format(len(data))) # Load 1d correlation functions dic_xi1d = {"D1": args.cf1d, "D2": args.cf1d2} for fname, filename in dic_xi1d.items(): if filename is None: continue hdul = fitsio.FITS(filename) header = hdul[1].read_header() log_lambda_min = header['LLMIN'] delta_log_lambda = header['DLL'] num_pairs_variance_1d = hdul[1]['nv1d'][:] variance_1d = hdul[1]['v1d'][:] log_lambda = (log_lambda_min + delta_log_lambda * np.arange(len(variance_1d))) cf.get_variance_1d[fname] = interp1d( log_lambda[num_pairs_variance_1d > 0], variance_1d[num_pairs_variance_1d > 0], kind='nearest', fill_value='extrapolate') num_pairs1d = hdul[1]['nb1d'][:] xi_1d = hdul[1]['c1d'][:] cf.xi_1d[fname] = interp1d( (log_lambda - log_lambda_min)[num_pairs1d > 0], xi_1d[num_pairs1d > 0], kind='nearest', fill_value='extrapolate') hdul.close() # Load correlation functions dic_xi = { "D1_D1": args.cf, "D2_D2": args.cf2, "D1_D2": args.cf12, "D2_D1": args.cf12 } for fname, filename in dic_xi.items(): if filename is None: continue hdul = fitsio.FITS(filename) header = hdul[1].read_header() assert cf.num_bins_r_par == header['NP'] assert cf.num_bins_r_trans == header['NT'] assert cf.r_par_min == header['RPMIN'] assert cf.r_par_max == header['RPMAX'] assert cf.r_trans_max == header['RTMAX'] xi = hdul[2]['DA'][:] weights = hdul[2]['WE'][:] xi = (xi * weights).sum(axis=0) weights = weights.sum(axis=0) w = weights > 0. xi[w] /= weights[w] cf.xi_wick[fname] = xi.copy() hdul.close() ### Read data 2 if args.in_dir2 or args.lambda_abs2: if args.lambda_abs2 or args.unfold_cf: cf.x_correlation = True cf.alpha2 = args.z_evol2 if args.in_dir2 is None: args.in_dir2 = args.in_dir if args.lambda_abs2: cf.lambda_abs2 = constants.ABSORBER_IGM[args.lambda_abs2] else: cf.lambda_abs2 = cf.lambda_abs data2, num_data2, z_min2, z_max2 = io.read_deltas( args.in_dir2, cf.nside, cf.lambda_abs2, cf.alpha2, cf.z_ref, cosmo, max_num_spec=args.nspec) for deltas in data.values(): for delta in deltas: delta.fname = 'D2' for item in [ 'cont', 'delta', 'order', 'ivar', 'exposures_diff', 'mean_snr', 'mean_reso', 'mean_z', 'delta_log_lambda' ]: setattr(delta, item, None) del z_max2 cf.data2 = data2 cf.num_data2 = num_data2 cf.ang_max = utils.compute_ang_max(cosmo, cf.r_trans_max, z_min, z_min2) userprint("") userprint("done, npix = {}".format(len(data2))) cf.counter = Value('i', 0) cf.lock = Lock() cpu_data = {} for index, healpix in enumerate(sorted(data)): num_processor = index % args.nproc if not num_processor in cpu_data: cpu_data[num_processor] = [] cpu_data[num_processor].append(healpix) # compute the covariance matrix context = multiprocessing.get_context('fork') pool = context.Pool(processes=min(args.nproc, len(cpu_data.values()))) userprint(" \nStarting\n") if args.nproc>1: wick_data = pool.map(calc_wick_terms, sorted(cpu_data.values())) else: wick_data = [calc_wick_terms(arg) for arg in sorted(cpu_data.values())] userprint(" \nFinished\n") pool.close() # merge the results from the different CPUs wick_data = np.array(wick_data) weights_wick = wick_data[:, 0].sum(axis=0) num_pairs_wick = wick_data[:, 1].sum(axis=0) num_pairs = wick_data[:, 2].sum(axis=0) num_pairs_used = wick_data[:, 3].sum(axis=0) t1 = wick_data[:, 4].sum(axis=0) t2 = wick_data[:, 5].sum(axis=0) t3 = wick_data[:, 6].sum(axis=0) t4 = wick_data[:, 7].sum(axis=0) t5 = wick_data[:, 8].sum(axis=0) t6 = wick_data[:, 9].sum(axis=0) weights = weights_wick * weights_wick[:, None] w = weights > 0. t1[w] /= weights[w] t2[w] /= weights[w] t3[w] /= weights[w] t4[w] /= weights[w] t5[w] /= weights[w] t6[w] /= weights[w] t1 *= 1. * num_pairs_used / num_pairs t2 *= 1. * num_pairs_used / num_pairs t3 *= 1. * num_pairs_used / num_pairs t4 *= 1. * num_pairs_used / num_pairs t5 *= 1. * num_pairs_used / num_pairs t6 *= 1. * num_pairs_used / num_pairs t_tot = t1 + t2 + t3 + t4 + t5 + t6 # save results results = fitsio.FITS(args.out, 'rw', clobber=True) header = [ { 'name': 'RPMIN', 'value': cf.r_par_min, 'comment': 'Minimum r-parallel [h^-1 Mpc]' }, { 'name': 'RPMAX', 'value': cf.r_par_max, 'comment': 'Maximum r-parallel [h^-1 Mpc]' }, { 'name': 'RTMAX', 'value': cf.r_trans_max, 'comment': 'Maximum r-transverse [h^-1 Mpc]' }, { 'name': 'NP', 'value': cf.num_bins_r_par, 'comment': 'Number of bins in r-parallel' }, { 'name': 'NT', 'value': cf.num_bins_r_trans, 'comment': 'Number of bins in r-transverse' }, { 'name': 'ZCUTMIN', 'value': cf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': cf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'REJ', 'value': cf.reject, 'comment': 'Rejection factor' }, { 'name': 'NPALL', 'value': num_pairs, 'comment': 'Number of pairs' }, { 'name': 'NPUSED', 'value': num_pairs_used, 'comment': 'Number of used pairs' }, { 'name': 'OMEGAM', 'value': args.fid_Om, 'comment': 'Omega_matter(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAR', 'value': args.fid_Or, 'comment': 'Omega_radiation(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAK', 'value': args.fid_Ok, 'comment': 'Omega_k(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'WL', 'value': args.fid_wl, 'comment': 'Equation of state of dark energy of fiducial LambdaCDM cosmology' }, { 'name': "BLINDING", 'value': blinding, 'comment': 'String specifying the blinding strategy' } ] comment = [ 'Sum of weight', 'Covariance', 'Nomber of pairs', 'T1', 'T2', 'T3', 'T4', 'T5', 'T6' ] results.write( [t_tot, weights_wick, num_pairs_wick, t1, t2, t3, t4, t5, t6], names=['CO', 'WALL', 'NB', 'T1', 'T2', 'T3', 'T4', 'T5', 'T6'], comment=comment, header=header, extname='COV') results.close()
def main(cmdargs): # pylint: disable-msg=too-many-locals,too-many-branches,too-many-statements """Computes the distortion matrix""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=('Compute the distortion matrix of the auto and ' 'cross-correlation of delta fields')) parser.add_argument('--out', type=str, default=None, required=True, help='Output file name') parser.add_argument('--in-dir', type=str, default=None, required=True, help='Directory to delta files') parser.add_argument('--in-dir2', type=str, default=None, required=False, help='Directory to 2nd delta files') parser.add_argument('--rp-min', type=float, default=0., required=False, help='Min r-parallel [h^-1 Mpc]') parser.add_argument('--rp-max', type=float, default=200., required=False, help='Max r-parallel [h^-1 Mpc]') parser.add_argument('--rt-max', type=float, default=200., required=False, help='Max r-transverse [h^-1 Mpc]') parser.add_argument('--np', type=int, default=50, required=False, help='Number of r-parallel bins') parser.add_argument('--nt', type=int, default=50, required=False, help='Number of r-transverse bins') parser.add_argument( '--coef-binning-model', type=int, default=1, required=False, help=( 'Coefficient multiplying np and nt to get finner binning for the ' 'model')) parser.add_argument( '--z-cut-min', type=float, default=0., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift larger than ' 'z-cut-min')) parser.add_argument( '--z-cut-max', type=float, default=10., required=False, help=('Use only pairs of forest x object with the mean of the last ' 'absorber redshift and the object redshift smaller than ' 'z-cut-max')) parser.add_argument( '--lambda-abs', type=str, default='LYA', required=False, help=( 'Name of the absorption in picca.constants defining the redshift ' 'of the delta')) parser.add_argument( '--lambda-abs2', type=str, default=None, required=False, help=( 'Name of the absorption in picca.constants defining the redshift ' 'of the 2nd delta')) parser.add_argument('--z-ref', type=float, default=2.25, required=False, help='Reference redshift') parser.add_argument( '--z-evol', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the delta field') parser.add_argument( '--z-evol2', type=float, default=2.9, required=False, help='Exponent of the redshift evolution of the 2nd delta field') parser.add_argument( '--fid-Om', type=float, default=0.315, required=False, help='Omega_matter(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-Or', type=float, default=0., required=False, help='Omega_radiation(z=0) of fiducial LambdaCDM cosmology') parser.add_argument('--fid-Ok', type=float, default=0., required=False, help='Omega_k(z=0) of fiducial LambdaCDM cosmology') parser.add_argument( '--fid-wl', type=float, default=-1., required=False, help='Equation of state of dark energy of fiducial LambdaCDM cosmology' ) parser.add_argument('--no-project', action='store_true', required=False, help='Do not project out continuum fitting modes') parser.add_argument( '--remove-same-half-plate-close-pairs', action='store_true', required=False, help='Reject pairs in the first bin in r-parallel from same half plate' ) parser.add_argument( '--rej', type=float, default=1., required=False, help=('Fraction of rejected forest-forest pairs: -1=no rejection, ' '1=all rejection')) parser.add_argument('--nside', type=int, default=16, required=False, help='Healpix nside') parser.add_argument('--nproc', type=int, default=None, required=False, help='Number of processors') parser.add_argument('--nspec', type=int, default=None, required=False, help='Maximum number of spectra to read') parser.add_argument( '--unfold-cf', action='store_true', required=False, help=('rp can be positive or negative depending on the relative ' 'position between absorber1 and absorber2')) args = parser.parse_args(cmdargs) if args.nproc is None: args.nproc = cpu_count() // 2 userprint("nproc", args.nproc) # setup variables in module cf cf.r_par_max = args.rp_max cf.r_par_min = args.rp_min cf.r_trans_max = args.rt_max cf.z_cut_max = args.z_cut_max cf.z_cut_min = args.z_cut_min cf.num_bins_r_par = args.np cf.num_bins_r_trans = args.nt cf.num_model_bins_r_par = args.np * args.coef_binning_model cf.num_model_bins_r_trans = args.nt * args.coef_binning_model cf.nside = args.nside cf.z_ref = args.z_ref cf.alpha = args.z_evol cf.reject = args.rej cf.lambda_abs = constants.ABSORBER_IGM[args.lambda_abs] cf.remove_same_half_plate_close_pairs = args.remove_same_half_plate_close_pairs # read blinding keyword blinding = io.read_blinding(args.in_dir) # load fiducial cosmology cosmo = constants.Cosmo(Om=args.fid_Om, Or=args.fid_Or, Ok=args.fid_Ok, wl=args.fid_wl, blinding=blinding) t0 = time.time() ### Read data 1 data, num_data, z_min, z_max = io.read_deltas(args.in_dir, cf.nside, cf.lambda_abs, cf.alpha, cf.z_ref, cosmo, max_num_spec=args.nspec, no_project=args.no_project) del z_max cf.data = data cf.num_data = num_data cf.ang_max = utils.compute_ang_max(cosmo, cf.r_trans_max, z_min) userprint("") userprint("done, npix = {}".format(len(data))) ### Read data 2 if args.in_dir2 or args.lambda_abs2: if args.lambda_abs2 or args.unfold_cf: cf.x_correlation = True cf.alpha2 = args.z_evol2 if args.in_dir2 is None: args.in_dir2 = args.in_dir if args.lambda_abs2: cf.lambda_abs2 = constants.ABSORBER_IGM[args.lambda_abs2] else: cf.lambda_abs2 = cf.lambda_abs data2, num_data2, z_min2, z_max2 = io.read_deltas( args.in_dir2, cf.nside, cf.lambda_abs2, cf.alpha2, cf.z_ref, cosmo, max_num_spec=args.nspec, no_project=args.no_project) del z_max2 cf.data2 = data2 cf.num_data2 = num_data2 cf.ang_max = utils.compute_ang_max(cosmo, cf.r_trans_max, z_min, z_min2) userprint("") userprint("done, npix = {}".format(len(data2))) t1 = time.time() userprint(f'picca_dmat.py - Time reading data: {(t1-t0)/60:.3f} minutes') cf.counter = Value('i', 0) cf.lock = Lock() cpu_data = {} for index, healpix in enumerate(sorted(data)): num_processor = index % args.nproc if num_processor not in cpu_data: cpu_data[num_processor] = [] cpu_data[num_processor].append(healpix) # compute the distortion matrix if args.nproc > 1: context = multiprocessing.get_context('fork') pool = context.Pool(processes=args.nproc) dmat_data = pool.map(calc_dmat, sorted(cpu_data.values())) pool.close() elif args.nproc == 1: dmat_data = map(calc_dmat, sorted(cpu_data.values())) dmat_data = list(dmat_data) t2 = time.time() userprint( f'picca_dmat.py - Time computing distortion matrix: {(t2-t1)/60:.3f} minutes' ) # merge the results from different CPUs dmat_data = np.array(dmat_data) weights_dmat = dmat_data[:, 0].sum(axis=0) dmat = dmat_data[:, 1].sum(axis=0) r_par = dmat_data[:, 2].sum(axis=0) r_trans = dmat_data[:, 3].sum(axis=0) z = dmat_data[:, 4].sum(axis=0) weights = dmat_data[:, 5].sum(axis=0) num_pairs = dmat_data[:, 6].sum(axis=0) num_pairs_used = dmat_data[:, 7].sum(axis=0) # normalize values w = weights > 0. r_par[w] /= weights[w] r_trans[w] /= weights[w] z[w] /= weights[w] w = weights_dmat > 0 dmat[w] /= weights_dmat[w, None] # save results results = fitsio.FITS(args.out, 'rw', clobber=True) header = [{ 'name': 'RPMIN', 'value': cf.r_par_min, 'comment': 'Minimum r-parallel [h^-1 Mpc]' }, { 'name': 'RPMAX', 'value': cf.r_par_max, 'comment': 'Maximum r-parallel [h^-1 Mpc]' }, { 'name': 'RTMAX', 'value': cf.r_trans_max, 'comment': 'Maximum r-transverse [h^-1 Mpc]' }, { 'name': 'NP', 'value': cf.num_bins_r_par, 'comment': 'Number of bins in r-parallel' }, { 'name': 'NT', 'value': cf.num_bins_r_trans, 'comment': 'Number of bins in r-transverse' }, { 'name': 'COEFMOD', 'value': args.coef_binning_model, 'comment': 'Coefficient for model binning' }, { 'name': 'ZCUTMIN', 'value': cf.z_cut_min, 'comment': 'Minimum redshift of pairs' }, { 'name': 'ZCUTMAX', 'value': cf.z_cut_max, 'comment': 'Maximum redshift of pairs' }, { 'name': 'REJ', 'value': cf.reject, 'comment': 'Rejection factor' }, { 'name': 'NPALL', 'value': num_pairs, 'comment': 'Number of pairs' }, { 'name': 'NPUSED', 'value': num_pairs_used, 'comment': 'Number of used pairs' }, { 'name': 'OMEGAM', 'value': args.fid_Om, 'comment': 'Omega_matter(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAR', 'value': args.fid_Or, 'comment': 'Omega_radiation(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'OMEGAK', 'value': args.fid_Ok, 'comment': 'Omega_k(z=0) of fiducial LambdaCDM cosmology' }, { 'name': 'WL', 'value': args.fid_wl, 'comment': 'Equation of state of dark energy of fiducial LambdaCDM cosmology' }, { 'name': "BLINDING", 'value': blinding, 'comment': 'String specifying the blinding strategy' }] dmat_name = "DM" if blinding != "none": dmat_name += "_BLIND" results.write([weights_dmat, dmat], names=['WDM', dmat_name], comment=['Sum of weight', 'Distortion matrix'], units=['', ''], header=header, extname='DMAT') results.write([r_par, r_trans, z], names=['RP', 'RT', 'Z'], comment=['R-parallel', 'R-transverse', 'Redshift'], units=['h^-1 Mpc', 'h^-1 Mpc', ''], extname='ATTRI') results.close() t3 = time.time() userprint(f'picca_dmat.py - Time total : {(t3-t0)/60:.3f} minutes')