add modifier

3 files changed, 34 insertions, 1 deletions

diff --git a/aggregator/dwd_icon.py b/aggregator/dwd_icon.py
index 0827469..6df905c 100755
--- a/aggregator/dwd_icon.py
+++ b/aggregator/dwd_icon.py
@@ -113,7 +113,9 @@ def load_data(name, output, description = None, clean = False, force_filename =
         if clean:
             clean_output_dir(output, filename)
 
-    ds = xr.load_dataset(target, engine='cfgrib')
+    # we drop heightAboveGround to allow 2m and 10m values to be merged down to one dataset
+    ds = xr.load_dataset(target, engine='cfgrib', drop_variables='heightAboveGround')
+
     if description is not None:
         ds.attrs['_description'] = description
     return ds
diff --git a/modifier/merge.py b/modifier/merge.py
new file mode 100644
index 0000000..692287d
--- /dev/null
+++ b/modifier/merge.py
@@ -0,0 +1,4 @@
+import xarray as xr 
+
+def run(data):
+    return xr.merge(data)
diff --git a/run.py b/run.py
index f239f45..28362d9 100755
--- a/run.py
+++ b/run.py
@@ -22,6 +22,30 @@ def create_aggregators(cfg):
 
     return ret
 
+def create_modifiers(aggregators, cfg):
+    # naming is scuffed
+    ret = {}
+    for modifier in cfg:
+        mod = cfg[modifier]
+        modname = mod['module']
+        del mod['module']
+
+        if 'aggregator' in mod:
+            if type(mod['aggregator']) == list:
+                mod['data'] = []
+                for ag in mod['aggregator']:
+                    mod['data'].append(aggregators[ag])
+
+                del mod['aggregator']
+            else:
+                mod['data'] = aggregators[mod['aggregator']]
+                del mod['aggregator']
+
+        pymod = __import__(modname, fromlist=[None])
+        ret[modifier] = pymod.run(**mod)
+
+    return ret
+
 mpl.use('agg')
 
 # Define custom gpm and gpdm units. The default gpm in metpy is aliased to meter.
@@ -60,6 +84,9 @@ with open(FILE, 'r') as f:
 
 aggregators = create_aggregators(conf['aggregator'])
 
+if 'modifier' in conf:
+    aggregators.update(create_modifiers(aggregators, conf['modifier']))
+
 index = []
 
 for plotter in conf['plotter']: