def write_covariance(): """ Write out a P(k,mu) or multipoles covariance matrix """ from pyRSD.rsdfit.data import PkmuCovarianceMatrix, PoleCovarianceMatrix parser = argparse.ArgumentParser() parser.formatter_class = argparse.RawTextHelpFormatter # required arguments h = 'the mode, either `pkmu` or `poles`' parser.add_argument('mode', choices=['pkmu', 'poles'], help=h) h = parse_tools.PowerSpectraParser.format_help() parser.add_argument('data', type=parse_tools.PowerSpectraParser.data, help=h) h = parse_tools.PowerSpectraCallable.format_help() parser.add_argument('callable', type=parse_tools.PowerSpectraCallable.data, help=h) h = 'the data format to use, either `pickle` or `plaintext`' parser.add_argument('--format', choices=['pickle', 'plaintext'], default='plaintext', help=h) h = 'the output file name' parser.add_argument('-o', '--output', required=True, type=str, help=h) h = 'the multipole numbers to compute' parser.add_argument('--ells', nargs='*', type=int, help=h) # options h = "the minimum wavenumber to use" parser.add_argument('--kmin', nargs='+', type=float, default=[-np.inf], help=h) h = "the maximum wavenumber to use" parser.add_argument('--kmax', nargs='+', type=float, default=[np.inf], help=h) h = "set off-diagonal elements to zero" parser.add_argument('--force_diagonal', action='store_true', help=h) # parse args = parser.parse_args() if len(args.kmin) == 1: args.kmin = args.kmin[0] if len(args.kmax) == 1: args.kmax = args.kmax[0] # get the data from the parent data and function data = getattr(args.data, args.callable['name'])(**args.callable['kwargs']) # compute the covariance matrix kwargs = {} kwargs['kmin'] = args.kmin kwargs['kmax'] = args.kmax kwargs['force_diagonal'] = args.force_diagonal if args.mode == 'pkmu': C = PkmuCovarianceMatrix.from_spectra_set(data, **kwargs) elif args.mode == 'poles': if args.ells is not None: kwargs['ells'] = args.ells C = PoleCovarianceMatrix.from_spectra_set(data, **kwargs) # now output if args.format == 'pickle': C.to_pickle(args.output) else: C.to_plaintext(args.output)
def write_data_gaussian_covariance(): """ Write out the gaussian covariance matrix from data measurements, for either P(k,mu) or multipoles """ from pyRSD.rsdfit.data import PkmuCovarianceMatrix, PoleCovarianceMatrix parser = argparse.ArgumentParser() parser.formatter_class = argparse.RawTextHelpFormatter subparsers = parser.add_subparsers(dest='subparser_name') # pkmu parser pkmu_parser = subparsers.add_parser('pkmu') h = parse_tools.PowerSpectraParser.format_help() pkmu_parser.add_argument('data', type=parse_tools.PowerSpectraParser.data, help=h) h = parse_tools.PowerSpectraCallable.format_help() pkmu_parser.add_argument('callable', type=parse_tools.PowerSpectraCallable.data, help=h) h = 'a list of (lower, upper) specifying the mu bin bounds' pkmu_parser.add_argument('mu_bounds', type=str, help=h) h = "the data keys to slice the data on; specified as `a = '0.6452'`" pkmu_parser.add_argument('key', type=str, nargs='+', action=parse_tools.StoreDataKeys, help=h) # poles parser pole_parser = subparsers.add_parser('poles') h = parse_tools.PowerSpectraParser.format_help() pole_parser.add_argument('data', type=parse_tools.PowerSpectraParser.data, help=h) h = parse_tools.PowerSpectraCallable.format_help() pole_parser.add_argument('callable', type=parse_tools.PowerSpectraCallable.data, help=h) h = 'the multipole numbers' pole_parser.add_argument('ells', type=str, help=h) h = "the data keys to slice the data on; specified as `a = '0.6452'`" pole_parser.add_argument('key', type=str, nargs='+', action=parse_tools.StoreDataKeys, help=h) # options for p in [pkmu_parser, pole_parser]: h = 'the data format to use, either `pickle` or `plaintext`' p.add_argument('--format', choices=['pickle', 'plaintext'], default='plaintext', help=h) h = 'the output file name' p.add_argument('-o', '--output', required=True, type=str, help=h) h = "the minimum wavenumber to use" p.add_argument('--kmin', nargs='+', type=float, default=[-np.inf], help=h) h = "the maximum wavenumber to use" p.add_argument('--kmax', nargs='+', type=float, default=[np.inf], help=h) # parse args = parser.parse_args() if len(args.kmin) == 1: args.kmin = args.kmin[0] if len(args.kmax) == 1: args.kmax = args.kmax[0] # get the data from the parent data and function data = getattr(args.data, args.callable['name'])(**args.callable['kwargs']) coords = data.coords for c in coords: if c in args.key: dt = data.coords[c].dtype.type args.key[c] = [dt(x) for x in args.key[c]] # now slice for k in args.key: if len(args.key[k]) != 1: raise ValueError("must specify exactly one key for each dimension") args.key[k] = args.key[k][0] try: data = data.sel(**args.key) if data.size == 1: data = data.get() except Exception as e: raise RuntimeError("error slicing data with key %s: %s" %(str(args.key), str(e))) # compute the covariance matrix if args.subparser_name == 'pkmu': mu_bounds = eval(args.mu_bounds) C, coords = covariance.data_pkmu_gausscov(data, mu_bounds, kmin=args.kmin, kmax=args.kmax) C = PkmuCovarianceMatrix(C, coords[0], coords[1]) else: ells = np.array(eval(args.ells), dtype=float) C, coords = covariance.data_pole_gausscov(data, ells, kmin=args.kmin, kmax=args.kmax) C = PoleCovarianceMatrix(C, coords[0], coords[1]) # now output if args.format == 'pickle': C.to_pickle(args.output) else: C.to_plaintext(args.output)