def write_analysis_grid():
"""
Write a P(k,mu) grid for analysis purposes
"""
from pyRSD.rsd import PkmuGrid
parser = argparse.ArgumentParser()
parser.formatter_class = argparse.RawTextHelpFormatter
# required arguments
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 (optional) data keys to slice the data on; specified as `a = '0.6452'`"
parser.add_argument('--key', type=str, nargs='+', action=parse_tools.StoreDataKeys, help=h)
h = 'the output file name'
parser.add_argument('-o', '--output', required=True, type=str, help=h)
args = parser.parse_args()
# get the data from the parent data and function
data = getattr(args.data, args.callable['name'])(**args.callable['kwargs'])
if args.key is not None:
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
if args.key is not None:
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)))
# now make the grid and save to plaintext file
grid = PkmuGrid.from_pkmuresult(data)
grid.to_plaintext(args.output)
