mHM: add calculated outputs and outputs per horizon

src/finam_mhm/component.py

@@ -6,27 +6,32 @@ from datetime import datetime, timedelta
 
 import finam as fm
 import mhm
+import numpy as np
 
 OUTPUT_META = {
     "L0_GRIDDED_LAI": dict(units="1", long_name="leaf area index"),
     "L1_FSEALED": dict(units="1", long_name="Fraction of sealed area"),
     "L1_FNOTSEALED": dict(units="1", long_name="Fraction of unsealed area"),
-    "L1_INTER": dict(units="mm / h", long_name="Interception"),
-    "L1_SNOWPACK": dict(units="mm / h", long_name="Snowpack"),
+    "L1_INTER": dict(units="mm / h", long_name="Interception"),  # interception (1)
+    "L1_SNOWPACK": dict(units="mm / h", long_name="Snowpack"),  # snowpack (2)
     "L1_SEALSTW": dict(
         units="mm / h", long_name="Retention storage of impervious areas"
-    ),
-    "L1_UNSATSTW": dict(units="mm / h", long_name="upper soil storage"),
-    "L1_SATSTW": dict(units="mm / h", long_name="groundwater storage"),
-    "L1_NEUTRONS": dict(units="mm / h", long_name="Ground Albedo Neutrons"),
-    "L1_PET_CALC": dict(units="mm / h", long_name="potential evapotranspiration"),
+    ),  # sealedSTW (6)
+    "L1_UNSATSTW": dict(units="mm / h", long_name="upper soil storage"),  # unsatSTW (7)
+    "L1_SATSTW": dict(units="mm / h", long_name="groundwater storage"),  # satSTW (8)
+    "L1_NEUTRONS": dict(
+        units="mm / h", long_name="Ground Albedo Neutrons"
+    ),  # neutrons (18)
+    "L1_PET_CALC": dict(
+        units="mm / h", long_name="potential evapotranspiration"
+    ),  # PET (9)
     "L1_AETCANOPY": dict(
         units="mm / h", long_name="Real evaporation intensity from canopy"
     ),
     "L1_AETSEALED": dict(
         units="mm / h", long_name="Real evap. from free water surfaces"
     ),
-    "L1_TOTAL_RUNOFF": dict(units="m^3 / h", long_name="Generated runoff"),
+    "L1_TOTAL_RUNOFF": dict(units="m^3 / h", long_name="Generated runoff"),  # Q (11)
     "L1_RUNOFFSEAL": dict(
         units="mm / h", long_name="Direct runoff from impervious areas"
     ),
@@ -34,30 +39,77 @@ OUTPUT_META = {
     "L1_SLOWRUNOFF": dict(units="mm / h", long_name="Slow runoff component"),
     "L1_BASEFLOW": dict(units="mm / h", long_name="Baseflow"),
     "L1_PERCOL": dict(units="mm / h", long_name="Percolation"),
-    "L1_PREEFFECT": dict(units="mm / h", long_name="Effective precip. depth"),
+    "L1_PREEFFECT": dict(
+        units="mm / h", long_name="Effective precip. depth"
+    ),  # preEffect (20)
     "L1_SOILMOIST_VOL_ALL": dict(
         units="1", long_name="average soil moisture over all layers"
-    ),
+    ),  # SM_Lall (5)
     "L11_QMOD": dict(units="m^3 / s", long_name="Simulated discharge"),
     "L11_QOUT": dict(units="m^3 / s", long_name="Total outflow from cells"),
 }
 """dict: meta information about available outputs in mHM."""
 
 OUTPUT_HORIZONS_META = {
-    "L1_SOILMOIST": dict(units="mm", long_name="soil water content of soil layer"),
+    "L1_SOILMOIST": dict(
+        units="mm", long_name="soil water content of soil layer {n}"
+    ),  # SWC_Lxx (3)
     "L1_SOILMOIST_VOL": dict(
-        units="mm / mm", long_name="volumetric soil moisture of soil layer"
-    ),
+        units="mm / mm", long_name="volumetric soil moisture of soil layer {n}"
+    ),  # SM_Lxx (4)
     "L1_SOILMOISTSAT": dict(
-        units="mm", long_name="Saturation soil moisture for each horizon"
+        units="mm", long_name="saturation soil moisture of soil layer {n}"
     ),
-    "L1_AETSOIL": dict(units="mm / h", long_name="Actual ET from soil layers"),
+    "L1_AETSOIL": dict(units="mm / h", long_name="actual ET  of soil layer {n}"),
     "L1_INFILSOIL": dict(
-        units="mm / h", long_name="Infiltration intensity each soil horizon"
+        units="mm / h", long_name="infiltration intensity of soil layer {n}"
     ),
+    # missing
+    # L1_MELT
+    # L1_RAIN
+    # L1_SNOW
+    # L1_THROUGHFALL
 }
 """dict: meta information about available outputs per horizon in mHM."""
 
+OUTPUT_CALC_META = {
+    # sum(aETSoil(horizons)) * fNotSealed + aETCanopy + aETSealed * fSealed
+    "L1_AET": dict(units="mm / h", long_name="actual Evapotranspiration"),  # aET (10)
+    # runoffSeal * fSealed
+    "L1_QD": dict(
+        units="mm / h", long_name="direct runoff generated by every cell (runoffSeal)"
+    ),  # QD (12)
+    # fastRunoff * fNotSealed
+    "L1_QIF": dict(
+        units="mm / h", long_name="fast interflow generated by every cell (fastRunoff)"
+    ),  # QIf (13)
+    # slowRunoff * fNotSealed
+    "L1_QIS": dict(
+        units="mm / h", long_name="slow interflow generated by every cell (slowRunoff)"
+    ),  # QIs (14)
+    # baseflow * fNotSealed
+    "L1_QB": dict(
+        units="mm / h", long_name="baseflow generated by every cell"
+    ),  # QB (15)
+    # percol * fNotSealed
+    "L1_RECHARGE": dict(
+        units="mm / h", long_name="groundwater recharge"
+    ),  # recharge (16)
+}
+"""dict: meta information about calculated outputs in mHM."""
+
+OUTPUT_CALC_HORIZONS_META = {
+    # infilSoil(horizon) * fNotSealed
+    "L1_SOIL_INFIL": dict(
+        units="mm / h", long_name="infiltration flux from soil layer {n}"
+    ),  # soil_infil_Lxx (17)
+    # aETSoil(horizon) * fNotSealed
+    "L1_AET": dict(
+        units="mm / h", long_name="actual Evapotranspiration from soil layer {n}"
+    ),  # aET_Lxx (19)
+}
+"""dict: meta information about calculated outputs per horizon in mHM."""
+
 INPUT_UNITS = {
     "L0_GRIDDED_LAI": "1",
     "METEO_PRE": "mm / {ts}",
@@ -78,6 +130,101 @@ HOURS_TO_TIMESTEP = {1: "h", 24: "d"}
 """dict: timestep string from hours."""
 
 
+def _fill_var(var, grid="l1"):
+    grid_info = getattr(mhm.get, grid + "_domain_info")()
+    # mask in mHM is the opposite in numpy (selection)
+    sel = mhm.get_mask(grid, indexing="xy", selection=True)
+    sel = sel.ravel(order="F")
+    output = np.ma.empty_like(sel, dtype=float)
+    output.fill_value = grid_info[-1]
+    output.mask = ~sel
+    output[sel] = var
+    return output.reshape((grid_info[0], grid_info[1]), order="C")
+
+
+def _L1_AET():
+    # sum(aETSoil(horizons)) * fNotSealed + aETCanopy + aETSealed * fSealed
+    fsealed = mhm.get_variable("L1_FSEALED", compressed=True)
+    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
+    aetcanopy = mhm.get_variable("L1_AETCANOPY", compressed=True)
+    aetsealed = mhm.get_variable("L1_AETSEALED", compressed=True)
+    horizons = mhm.get.number_of_horizons()
+
+    sum_aetsoil = np.zeros_like(fsealed, dtype=float)
+    for n in range(1, horizons + 1):
+        sum_aetsoil += mhm.get_variable("L1_AETSOIL", index=n, compressed=True)
+
+    return _fill_var(sum_aetsoil * fnotsealed + aetcanopy + aetsealed * fsealed)
+
+
+def _L1_QD():
+    # runoffSeal * fSealed
+    fsealed = mhm.get_variable("L1_FSEALED", compressed=True)
+    runoffseal = mhm.get_variable("L1_RUNOFFSEAL", compressed=True)
+    return _fill_var(runoffseal * fsealed)
+
+
+def _L1_QIF():
+    # fastRunoff * fNotSealed
+    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
+    fastrunoff = mhm.get_variable("L1_FASTRUNOFF", compressed=True)
+    return _fill_var(fastrunoff * fnotsealed)
+
+
+def _L1_QIS():
+    # slowRunoff * fNotSealed
+    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
+    slowrunoff = mhm.get_variable("L1_SLOWRUNOFF", compressed=True)
+    return _fill_var(slowrunoff * fnotsealed)
+
+
+def _L1_QB():
+    # baseflow * fNotSealed
+    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
+    baseflow = mhm.get_variable("L1_BASEFLOW", compressed=True)
+    return _fill_var(baseflow * fnotsealed)
+
+
+def _L1_RECHARGE():
+    # percol * fNotSealed
+    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
+    percol = mhm.get_variable("L1_PERCOL", compressed=True)
+    return _fill_var(percol * fnotsealed)
+
+
+def _L1_SOIL_INFIL_N(n):
+    # infilSoil(horizon) * fNotSealed
+    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
+    infilsoil = mhm.get_variable("L1_INFILSOIL", index=n, compressed=True)
+    return _fill_var(infilsoil * fnotsealed)
+
+
+def _L1_AET_N(n):
+    # aETSoil(horizon) * fNotSealed
+    fnotsealed = mhm.get_variable("L1_FNOTSEALED", compressed=True)
+    aetsoil = mhm.get_variable("L1_AETSOIL", index=n, compressed=True)
+    return _fill_var(aetsoil * fnotsealed)
+
+
+OUTPUT_CALC = {
+    "L1_AET": _L1_AET,
+    "L1_QD": _L1_QD,
+    "L1_QIF": _L1_QIF,
+    "L1_QIS": _L1_QIS,
+    "L1_QB": _L1_QB,
+    "L1_RECHARGE": _L1_RECHARGE,
+}
+
+OUTPUT_CALC_HORIZON = {
+    "L1_SOIL_INFIL": _L1_SOIL_INFIL_N,
+    "L1_AET": _L1_AET_N,
+}
+
+
+def _horizon_name(name, horizon):
+    return name + "_L" + str(horizon).zfill(2)
+
+
 def _get_grid_name(var):
     grid_name = var.split("_")[0]
     return "L1" if grid_name == "METEO" else grid_name
@@ -93,10 +240,6 @@ def _get_meteo_inputs(inputs):
     }
 
 
-def _horizon_name(name, horizon):
-    return name + "_L" + str(horizon).zfill(2)
-
-
 class MHM(fm.TimeComponent):
     """
     mHM FINAM compoment.
@@ -142,9 +285,9 @@ class MHM(fm.TimeComponent):
         super().__init__()
         self.gridspec = {}
         self.no_data = None
-        self.horizons = None
+        self.number_of_horizons = None
 
-        self.OUTPUT_NAMES = list(OUTPUT_META)
+        self.OUTPUT_NAMES = None
         self.INPUT_NAMES = (
             [] if input_names is None else [n.upper() for n in input_names]
         )
@@ -175,6 +318,11 @@ class MHM(fm.TimeComponent):
         """Next pull time."""
         return self.time + self.step
 
+    @property
+    def horizons(self):
+        """Iterator for all horizons starting at 1."""
+        return range(1, self.number_of_horizons + 1)
+
     @fm.tools.execute_in_cwd
     def _initialize(self):
         # only show errors
@@ -198,13 +346,21 @@ class MHM(fm.TimeComponent):
         mhm.run.prepare()
         # only one domain possible
         mhm.run.prepare_domain()
-        self.horizons = mhm.get.number_of_horizons()
+        self.number_of_horizons = mhm.get.number_of_horizons()
+        # prepare outputs
+        self.OUTPUT_NAMES = list(OUTPUT_META)
         self.OUTPUT_NAMES += [
             _horizon_name(var, horizon)
             for var in OUTPUT_HORIZONS_META
-            for horizon in range(1, self.horizons)
+            for horizon in self.horizons
+        ]
+        self.OUTPUT_NAMES += list(OUTPUT_CALC_META)
+        self.OUTPUT_NAMES += [
+            _horizon_name(var, horizon)
+            for var in OUTPUT_CALC_HORIZONS_META
+            for horizon in self.horizons
         ]
-        # set time
+        # get start time
         year, month, day, hour = mhm.run.current_time()
         self.time = datetime(year=year, month=month, day=max(day, 0), hour=max(hour, 0))
         # first time step compensate by negative values in mHM
@@ -243,16 +399,47 @@ class MHM(fm.TimeComponent):
                 _FillValue=self.no_data,
                 **meta,
             )
+        for var, meta in OUTPUT_CALC_META.items():
+            grid_name = _get_grid_name(var)
+            self.outputs.add(
+                name=var,
+                time=self.time,
+                grid=self.gridspec[grid_name],
+                missing_value=self.no_data,
+                _FillValue=self.no_data,
+                **meta,
+            )
         for var, meta in OUTPUT_HORIZONS_META.items():
             grid_name = _get_grid_name(var)
-            for horizon in range(1, self.horizons + 1):
+            for horizon in self.horizons:
+                # add horizon number to long name
+                n_meta = {
+                    att: val.format(n=horizon) if att == "long_name" else val
+                    for att, val in meta.items()
+                }
+                self.outputs.add(
+                    name=_horizon_name(var, horizon),
+                    time=self.time,
+                    grid=self.gridspec[grid_name],
+                    missing_value=self.no_data,
+                    _FillValue=self.no_data,
+                    **n_meta,
+                )
+        for var, meta in OUTPUT_CALC_HORIZONS_META.items():
+            grid_name = _get_grid_name(var)
+            for horizon in self.horizons:
+                # add horizon number to long name
+                n_meta = {
+                    att: val.format(n=horizon) if att == "long_name" else val
+                    for att, val in meta.items()
+                }
                 self.outputs.add(
                     name=_horizon_name(var, horizon),
                     time=self.time,
                     grid=self.gridspec[grid_name],
                     missing_value=self.no_data,
                     _FillValue=self.no_data,
-                    **meta,
+                    **n_meta,
                 )
         for var in self.INPUT_NAMES:
             grid_name = _get_grid_name(var)
@@ -270,11 +457,19 @@ class MHM(fm.TimeComponent):
 
     def _connect(self, start_time):
         push_data = {var: mhm.get_variable(var) for var in OUTPUT_META}
+        push_data.update({var: func() for var, func in OUTPUT_CALC.items()})
         push_data.update(
             {
                 _horizon_name(var, horizon): mhm.get_variable(var, index=horizon)
                 for var in OUTPUT_HORIZONS_META
-                for horizon in range(1, self.horizons + 1)
+                for horizon in self.horizons
+            }
+        )
+        push_data.update(
+            {
+                _horizon_name(var, horizon): func(horizon)
+                for var, func in OUTPUT_CALC_HORIZON.items()
+                for horizon in self.horizons
             }
         )
         self.try_connect(start_time=start_time, push_data=push_data)
@@ -289,7 +484,7 @@ class MHM(fm.TimeComponent):
             # every hour or every 24 hours
             if self.time.hour % self.meteo_timestep == 0:
                 kwargs = {
-                    var: self.inputs[name].pull_data(self.time)[0].magnitude
+                    var: self.inputs[name].pull_data(self.next_time)[0].magnitude
                     for var, name in self.meteo_inputs.items()
                 }
                 kwargs["time"] = self.time
@@ -307,8 +502,15 @@ class MHM(fm.TimeComponent):
                 data=mhm.get_variable(var),
                 time=self.time,
             )
+        for var, func in OUTPUT_CALC.items():
+            if not self.outputs[var].has_targets:
+                continue
+            self.outputs[var].push_data(
+                data=func(),
+                time=self.time,
+            )
         for var in OUTPUT_HORIZONS_META:
-            for horizon in range(1, self.horizons + 1):
+            for horizon in self.horizons:
                 name = _horizon_name(var, horizon)
                 if not self.outputs[name].has_targets:
                     continue
@@ -316,6 +518,15 @@ class MHM(fm.TimeComponent):
                     data=mhm.get_variable(var, index=horizon),
                     time=self.time,
                 )
+        for var, func in OUTPUT_CALC_HORIZON.items():
+            for horizon in self.horizons:
+                name = _horizon_name(var, horizon)
+                if not self.outputs[name].has_targets:
+                    continue
+                self.outputs[name].push_data(
+                    data=func(horizon),
+                    time=self.time,
+                )
 
     @fm.tools.execute_in_cwd
     def _finalize(self):