Beispiel #1
0
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)