fixed bug in unit tests

Hydrology-IFH · Sep 17, 2024 · 156123e · 156123e
1 parent d1df2a0
commit 156123e
Show file tree

Hide file tree

Showing 6 changed files with 160 additions and 112 deletions.
diff --git a/examples/plot_scale/boadkh/svat_crop/svat_crop.py b/examples/plot_scale/boadkh/svat_crop/svat_crop.py
@@ -303,7 +303,7 @@ def set_boundary_conditions(self, state):
 
         @roger_routine(
             dist_safe=False,
-            local_variables=["PREC", "TA", "TA_MIN", "TA_MAX", "PET", "RS"],
+            local_variables=["PREC", "TA", "TA_MIN", "TA_MAX", "PET"],
         )
         def set_forcing_setup(self, state):
             vs = state.variables

diff --git a/examples/plot_scale/boadkh/svat_crop_nitrate/svat_crop_nitrate.py b/examples/plot_scale/boadkh/svat_crop_nitrate/svat_crop_nitrate.py
@@ -842,14 +842,13 @@ def set_diagnostics(self, state, base_path=tmp_dir):
                 "rt50_s",
                 "rt90_s",
                 "rtavg_s",
-                "C_q_ss"
             ]
             diagnostics["average"].output_frequency = 24 * 60 * 60
             diagnostics["average"].sampling_frequency = 1
             if base_path:
                 diagnostics["average"].base_output_path = base_path
 
-            diagnostics["collect"].output_variables = ["M_s", "Nmin_s", "C_s"]
+            diagnostics["collect"].output_variables = ["M_s", "Nmin_s", "C_s", "C_q_ss"]
             diagnostics["collect"].output_frequency = 24 * 60 * 60
             diagnostics["collect"].sampling_frequency = 1
             if base_path:

diff --git a/roger/models/oneD/oneD.py b/roger/models/oneD/oneD.py
@@ -64,6 +64,8 @@ def set_settings(self, state):
         settings.nx, settings.ny = 1, 1
         # length of simulation (in seconds)
         settings.runlen = self._get_runlen(self._input_dir, 'forcing.nc')
+        # number of time steps in meteorological data
+        settings.nitt_forc = len(self._read_var_from_nc("Time", self._input_dir, "forcing.nc"))
 
         # spatial discretization (in meters)
         settings.dx = 1
@@ -196,36 +198,42 @@ def set_boundary_conditions_setup(self, state):
     def set_boundary_conditions(self, state):
         pass
 
-    @roger_routine
-    def set_forcing_setup(self, state):
-        pass
-
     @roger_routine(
         dist_safe=False,
-        local_variables=[
-            "time",
-            "itt_day",
-            "itt_forc",
-            "prec_day",
-            "ta_day",
-            "pet_day",
-            "year",
-            "month",
-            "doy",
-        ],
+        local_variables=["PREC", "TA", "PET"],
     )
+    def set_forcing_setup(self, state):
+        vs = state.variables
+
+        vs.PREC = update(vs.PREC, at[:], self._read_var_from_nc("PREC", self._input_dir, "forcing.nc")[0, 0, :])
+        vs.TA = update(vs.TA, at[:], self._read_var_from_nc("TA", self._input_dir, "forcing.nc")[0, 0, :])
+        vs.PET = update(vs.PET, at[:], self._read_var_from_nc("PET", self._input_dir, "forcing.nc")[0, 0, :])
+
+    @roger_routine
     def set_forcing(self, state):
         vs = state.variables
 
-        condt = (vs.time % (24 * 60 * 60) == 0)
+        condt = vs.time % (24 * 60 * 60) == 0
         if condt:
             vs.itt_day = 0
-            vs.year = update(vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.month = update(vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.doy = update(vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.prec_day = update(vs.prec_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PREC", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
-            vs.ta_day = update(vs.ta_day, at[2:-2, 2:-2, :], self._read_var_from_nc("TA", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
-            vs.pet_day = update(vs.pet_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PET", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
+            vs.year = update(
+                vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.month = update(
+                vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.doy = update(
+                vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.prec_day = update(
+                vs.prec_day, at[:, :, :], vs.PREC[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
+            vs.ta_day = update(
+                vs.ta_day, at[:, :, :], vs.TA[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
+            vs.pet_day = update(
+                vs.pet_day, at[:, :, :], vs.PET[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
             vs.itt_forc = vs.itt_forc + 6 * 24
 
     @roger_routine

diff --git a/roger/models/oneD_crop/oneD_crop.py b/roger/models/oneD_crop/oneD_crop.py
@@ -61,6 +61,8 @@ def set_settings(self, state):
         settings.nx, settings.ny = 1, 1
         # length of simulation (in seconds)
         settings.runlen = self._get_runlen(self._input_dir, 'forcing.nc')
+        # number of time steps in meteorological data
+        settings.nitt_forc = len(self._read_var_from_nc("Time", self._input_dir, "forcing.nc"))
 
         # spatial discretization (in meters)
         settings.dx = 1
@@ -231,46 +233,74 @@ def set_boundary_conditions_setup(self, state):
     def set_boundary_conditions(self, state):
         pass
 
-    @roger_routine
-    def set_forcing_setup(self, state):
-        pass
-
     @roger_routine(
         dist_safe=False,
-        local_variables=[
-            "time",
-            "itt_day",
-            "itt_forc",
-            "prec_day",
-            "ta_day",
-            "pet_day",
-            "year",
-            "month",
-            "doy",
-            "itt",
-            "itt_cr",
-            "crop_type",
-        ],
+        local_variables=["PREC", "TA", "TA_MIN", "TA_MAX", "PET"],
     )
+    def set_forcing_setup(self, state):
+        vs = state.variables
+
+        vs.PREC = update(vs.PREC, at[:], self._read_var_from_nc("PREC", self._input_dir, "forcing.nc")[0, 0, :])
+        vs.TA = update(vs.TA, at[:], self._read_var_from_nc("TA", self._input_dir, "forcing.nc")[0, 0, :])
+        vs.TA_MIN = update(
+            vs.TA_MIN, at[:], self._read_var_from_nc("TA_min", self._input_dir, "forcing.nc")[0, 0, :]
+        )
+        vs.TA_MAX = update(
+            vs.TA_MAX, at[:], self._read_var_from_nc("TA_max", self._input_dir, "forcing.nc")[0, 0, :]
+        )
+        vs.PET = update(vs.PET, at[:], self._read_var_from_nc("PET", self._input_dir, "forcing.nc")[0, 0, :])
+
+    @roger_routine
     def set_forcing(self, state):
         vs = state.variables
 
-        condt = (vs.time % (24 * 60 * 60) == 0)
+        condt = vs.time % (24 * 60 * 60) == 0
         if condt:
             vs.itt_day = 0
-            vs.year = update(vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.month = update(vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.doy = update(vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.prec_day = update(vs.prec_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PREC", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
-            vs.ta_day = update(vs.ta_day, at[2:-2, 2:-2, :], self._read_var_from_nc("TA", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
-            vs.pet_day = update(vs.pet_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PET", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
+            vs.year = update(
+                vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.month = update(
+                vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.doy = update(
+                vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.prec_day = update(
+                vs.prec_day, at[:, :, :], vs.PREC[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
+            vs.ta_day = update(
+                vs.ta_day, at[:, :, :], vs.TA[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
+            vs.ta_min = update(
+                vs.ta_min,
+                at[:, :],
+                npx.min(vs.TA_MIN[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24], axis=-1),
+            )
+            vs.ta_max = update(
+                vs.ta_max,
+                at[:, :],
+                npx.max(vs.TA_MAX[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24], axis=-1),
+            )
+            vs.pet_day = update(
+                vs.pet_day, at[:, :, :], vs.PET[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
             vs.itt_forc = vs.itt_forc + 6 * 24
 
         if (vs.year[1] != vs.year[0]) & (vs.itt > 1):
             vs.itt_cr = vs.itt_cr + 2
-            vs.crop_type = update(vs.crop_type, at[2:-2, 2:-2, 0], 564)
-            vs.crop_type = update(vs.crop_type, at[2:-2, 2:-2, 1], 539)
-            vs.crop_type = update(vs.crop_type, at[2:-2, 2:-2, 2], 564)
+            vs.crop_type = update(vs.crop_type, at[2:-2, 2:-2, 0], 599)
+            vs.crop_type = update(
+                vs.crop_type,
+                at[2:-2, 2:-2, 1],
+                539,
+            )
+            vs.crop_type = update(
+                vs.crop_type,
+                at[2:-2, 2:-2, 2],
+                599,
+            )
+
 
     @roger_routine
     def set_diagnostics(self, state):

diff --git a/roger/models/svat/svat.py b/roger/models/svat/svat.py
@@ -46,6 +46,8 @@ def set_settings(self, state):
         settings.nx, settings.ny = 1, 1
         # length of simulation (in seconds)
         settings.runlen = self._get_runlen(self._input_dir, 'forcing.nc')
+        # number of time steps in meteorological data
+        settings.nitt_forc = len(self._read_var_from_nc("Time", self._input_dir, "forcing.nc"))
 
         # spatial discretization (in meters)
         settings.dx = 1
@@ -132,36 +134,42 @@ def set_boundary_conditions_setup(self, state):
     def set_boundary_conditions(self, state):
         pass
 
-    @roger_routine
-    def set_forcing_setup(self, state):
-        pass
-
     @roger_routine(
         dist_safe=False,
-        local_variables=[
-            "time",
-            "itt_day",
-            "itt_forc",
-            "prec_day",
-            "ta_day",
-            "pet_day",
-            "year",
-            "month",
-            "doy",
-        ],
+        local_variables=["PREC", "TA", "PET"],
     )
+    def set_forcing_setup(self, state):
+        vs = state.variables
+
+        vs.PREC = update(vs.PREC, at[:], self._read_var_from_nc("PREC", self._input_dir, "forcing.nc")[0, 0, :])
+        vs.TA = update(vs.TA, at[:], self._read_var_from_nc("TA", self._input_dir, "forcing.nc")[0, 0, :])
+        vs.PET = update(vs.PET, at[:], self._read_var_from_nc("PET", self._input_dir, "forcing.nc")[0, 0, :])
+
+    @roger_routine
     def set_forcing(self, state):
         vs = state.variables
 
-        condt = (vs.time % (24 * 60 * 60) == 0)
+        condt = vs.time % (24 * 60 * 60) == 0
         if condt:
             vs.itt_day = 0
-            vs.year = update(vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.month = update(vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.doy = update(vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, 'forcing.nc')[vs.itt_forc])
-            vs.prec_day = update(vs.prec_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PREC", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
-            vs.ta_day = update(vs.ta_day, at[2:-2, 2:-2, :], self._read_var_from_nc("TA", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
-            vs.pet_day = update(vs.pet_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PET", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
+            vs.year = update(
+                vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.month = update(
+                vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.doy = update(
+                vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, "forcing.nc")[vs.itt_forc]
+            )
+            vs.prec_day = update(
+                vs.prec_day, at[:, :, :], vs.PREC[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
+            vs.ta_day = update(
+                vs.ta_day, at[:, :, :], vs.TA[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
+            vs.pet_day = update(
+                vs.pet_day, at[:, :, :], vs.PET[npx.newaxis, npx.newaxis, vs.itt_forc : vs.itt_forc + 6 * 24]
+            )
             vs.itt_forc = vs.itt_forc + 6 * 24
 
     @roger_routine