Add better CLI (closes #24)

climix/__init__.py

+# -*- coding: utf-8 -*-
 
-from .precip_dry_aggregator import main
+from .main import main
 
 from ._version import get_versions
 __version__ = get_versions()['version']

climix/etc/metadata.yml

@@ -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

climix/main.py

+#!/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()

climix/precip_dry_aggregator.py

-#!/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()