Commit 5a1b6ee8 authored by Klaus Zimmermann's avatar Klaus Zimmermann
Browse files

Add better CLI (closes #24)

parent fad01c0d
# -*- coding: utf-8 -*-
from .precip_dry_aggregator import main
from .main import main
from ._version import get_versions
__version__ = get_versions()['version']
......
......@@ -6,7 +6,8 @@ drop:
- history
indices:
- var_name: max_dry_spell2_amjj
cdd:
var_name: max_dry_spell2_amjj
season: AMJJ
threshold:
data: 2.31481481e-05
......
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import os
import time
from dask.distributed import Client, LocalCluster
# from dask_jobqueue import SLURMCluster
import iris
import yaml
import climix
from .precip_dry_aggregator import loop_body
MISSVAL = 1.0e20
def parse_args():
parser = argparse.ArgumentParser(
description=('A climate index thing, version {}.'
''.format(climix.__version__)))
parser.add_argument('datafiles', nargs='+', metavar="DATAFILE",
help='the input data files')
parser.add_argument('-x', '--index', action='append',
required=True, metavar='INDEX', dest='indices',
help='the index to calculcate (currently only cdd)')
parser.add_argument('-o', '--output', dest='output_template',
help='output filename')
return parser.parse_args()
def setup_cluster():
cluster = LocalCluster()
# cluster = SLURMCluster()
# cluster.adapt(minimum=10, maximum=100)
# cluster.scale(10)
return cluster
def ignore_cb(cube, field, filename):
cube.attributes.pop('creation_date')
cube.attributes.pop('tracking_id')
def load_metadata():
metadata_filename = os.path.join(os.path.dirname(__file__),
'etc', 'metadata.yml')
with open(metadata_filename) as md_file:
metadata = yaml.safe_load(md_file)
return metadata
def prepare_indices(index_definitions, requested_indices):
return [index_definitions[index] for index in requested_indices]
def prepare_input_data(datafiles):
datacube = iris.load(datafiles, callback=ignore_cb)
iris.util.unify_time_units(datacube)
cube = datacube.concatenate_cube()
return cube
def guess_output_template(datafiles):
output_template = '{var_name}.nc'
def filename_stripper(path):
# remove directory part...
basename = os.path.basename(path)
# ...and extension
root, ext = os.path.splitext(basename)
# split at _...
parts = root.split('_')
# and remove
# first part (usually variable)
# and last part (usually time)
base = '_'.join(parts[1:-1])
try:
time = [int(t) for t in parts[-1].split('-')]
if len(time) == 1:
time *= 2
except ValueError:
time = [None, None]
return (base, time[0], time[1])
files = [filename_stripper(p) for p in datafiles]
bases, starts, ends = zip(*files)
unique_bases = set(bases)
if len(unique_bases) == 1:
start = min(starts)
end = max(ends)
output_template = '{{var_name}}_{base}_{start}-{end}.nc'.format(
base=unique_bases.pop(), start=start, end=end)
return output_template
def build_output_filename(index, datafiles, output_template):
if output_template is None:
output_template = guess_output_template(datafiles)
return output_template.format(**index)
def do_main(indices, datafiles, output_template):
metadata = load_metadata()
indices = prepare_indices(metadata['indices'], indices)
for index in indices:
output_filename = build_output_filename(index, datafiles, output_template)
input_data = prepare_input_data(datafiles)
result = loop_body(input_data, index)
result.data
iris.save(result, output_filename, fill_value=MISSVAL)
def main():
args = parse_args()
cluster = setup_cluster()
client = Client(cluster)
try:
start = time.time()
do_main(args.indices, args.datafiles, args.output_template)
end = time.time()
input("Calculation took {:.4f} seconds.\n"
"Press enter to close the cluster ".format(end-start))
finally:
client.close()
cluster.close()
if __name__ == "__main__":
main()
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import glob
import numpy as np
import os
import time
from cf_units import Unit
import dask.array as da
from dask.distributed import Client, LocalCluster
from dask_jobqueue import SLURMCluster
import iris
from iris.analysis import Aggregator
import iris.coords
import iris.coord_categorisation
import yaml
MISSVAL = 1.0e20
def non_lazy(data, axis, threshold, condition):
......@@ -132,57 +123,3 @@ def loop_body(datacube, ix):
season_cube = datacube.extract(season.constraint)
cube = season_cube.aggregated_by('year', aggregator)
return cube
def ignore_cb(cube, field, filename):
cube.attributes.pop('creation_date')
cube.attributes.pop('tracking_id')
def do_main():
var = 'pr'
indir = '/home/rossby/imports/obs/UERRA/SURFEX-MESCAN/input/day/'
inroot = '_EUR-055_UERRA-HARMONIE_RegRean_v1d1-v1d2_SURFEX-MESCAN_v1_day_*.nc' # noqa
filepattern = '{}{}{}'.format(indir, var, inroot)
filelist = sorted(glob.glob(filepattern))[:2]
metadata_filename = os.path.join(os.path.dirname(__file__),
'etc', 'metadata.yml')
with open(metadata_filename) as md_file:
metadata = yaml.safe_load(md_file)
ix = metadata['indices'][0]
datacube = iris.load(filelist, callback=ignore_cb)
iris.util.unify_time_units(datacube)
cube = datacube.concatenate_cube()
max_dry_spell2_amjj = loop_body(cube, ix)
max_dry_spell2_amjj.data
out_template = '{}_EUR-055_UERRA-HARMONIE_RegRean_v1d1-v1d2_SURFEX-MESCAN_v1_ann_1961-2015.nc' # noqa
iris.save(max_dry_spell2_amjj,
out_template.format('maxDrySpell2_AMJJ'),
fill_value=MISSVAL)
def main():
cluster = LocalCluster()
# cluster = SLURMCluster()
# cluster.adapt(minimum=10, maximum=100)
# cluster.scale(10)
client = Client(cluster)
try:
start = time.time()
do_main()
end = time.time()
input("Calculation took {:.4f} seconds.\n"
"Press enter to close the cluster ".format(end-start))
finally:
client.close()
cluster.close()
if __name__ == "__main__":
main()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment