Exemple #1
0
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')
Exemple #2
0
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()
Exemple #3
0
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')
Exemple #4
0
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')
Exemple #5
0
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')
Exemple #6
0
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()
Exemple #7
0
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()
Exemple #8
0
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')