sceua calibration could be called

hydromodel/models/gr4j.py

@@ -86,7 +86,8 @@ def production(
  # The precip. for routing is the sum of the rainfall which
  # did not make it to storage and the percolation from the store
  current_runoff = perc + (precip_net - precip_store)
- return current_runoff, s_update
+ # TODO: check if evap_store is the real ET
+ return current_runoff, evap_store, s_update
 
 
 # @jit
@@ -189,8 +190,8 @@ def gr4j(p_and_e, parameters, warmup_length: int, return_state=False, **kwargs):
 
  Parameters
  ----------
- p_and_e
-
+ p_and_e: ndarray
+ 3-dim input -- [time, basin, variable]: precipitation and potential evaporation
  parameters
  2-dim variable -- [basin, parameter]:
  the parameters are x1, x2, x3 and x4
@@ -216,7 +217,7 @@ def gr4j(p_and_e, parameters, warmup_length: int, return_state=False, **kwargs):
  if warmup_length > 0:
  # set no_grad for warmup periods
  p_and_e_warmup = p_and_e[0:warmup_length, :, :]
- _, s0, r0 = gr4j(
+ _, _, s0, r0 = gr4j(
  p_and_e_warmup, parameters, warmup_length=0, return_state=True, **kwargs
  )
  else:
@@ -225,12 +226,14 @@ def gr4j(p_and_e, parameters, warmup_length: int, return_state=False, **kwargs):
  inputs = p_and_e[warmup_length:, :, :]
  streamflow_ = np.full(inputs.shape[:2], 0.0)
  prs = np.full(inputs.shape[:2], 0.0)
+ ets = np.full(inputs.shape[:2], 0.0)
  for i in range(inputs.shape[0]):
  if i == 0:
- pr, s = production(inputs[i, :, :], x1, s0)
+ pr, et, s = production(inputs[i, :, :], x1, s0)
  else:
- pr, s = production(inputs[i, :, :], x1, s)
+ pr, et, s = production(inputs[i, :, :], x1, s)
  prs[i, :] = pr
+ ets[i, :] = et
  prs_x = np.expand_dims(prs, axis=2)
  conv_q9, conv_q1 = uh_gr4j(x4)
  q9 = np.full([inputs.shape[0], inputs.shape[1], 1], 0.0)
@@ -250,5 +253,5 @@ def gr4j(p_and_e, parameters, warmup_length: int, return_state=False, **kwargs):
  streamflow_[i, :] = q
  streamflow = np.expand_dims(streamflow_, axis=2)
  if return_state:
- return streamflow, s, r
- return streamflow
+ return streamflow, ets, s, r
+ return streamflow, ets

hydromodel/models/hymod.py

@@ -44,7 +44,7 @@ def hymod(p_and_e, parameters, warmup_length=30, return_state=False, **kwargs):
  if warmup_length > 0:
  # set no_grad for warmup periods
  p_and_e_warmup = p_and_e[0:warmup_length, :, :]
- _, x_slow, x_quick, x_loss = hymod(
+ _, _, x_slow, x_quick, x_loss = hymod(
  p_and_e_warmup,
  parameters,
  warmup_length=0,
@@ -88,11 +88,11 @@ def hymod(p_and_e, parameters, warmup_length=30, return_state=False, **kwargs):
 
  # Compute total flow for timestep
  output[t, :] = qs + outflow
- t = t + 1
+ t += 1
  streamflow = np.expand_dims(output, axis=2)
  if return_state:
- return streamflow, x_slow, x_quick, x_loss
- return streamflow
+ return streamflow, et, x_slow, x_quick, x_loss
+ return streamflow, et
 
 
 # @jit

hydromodel/models/model_dict.py

@@ -0,0 +1,15 @@
+from spotpy.objectivefunctions import rmse
+from hydromodel.models.xaj import xaj
+from hydromodel.models.gr4j import gr4j
+from hydromodel.models.hymod import hymod
+
+CRITERION_DICT = {
+ "RMSE": rmse,
+}
+
+MODEL_DICT = {
+ "xaj_mz": xaj,
+ "xaj": xaj,
+ "gr4j": gr4j,
+ "hymod": hymod,
+}

hydromodel/models/xaj.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2021-12-10 23:01:02
-LastEditTime: 2024-02-10 11:51:35
+LastEditTime: 2024-03-22 17:05:22
 LastEditors: Wenyu Ouyang
 Description: Core code for XinAnJiang model
 FilePath: /hydro-model-xaj/hydromodel/models/xaj.py
@@ -767,7 +767,7 @@ def xaj(
  (value[1] - value[0]) * params[:, i] + value[0]
  for i, (key, value) in enumerate(param_ranges.items())
  ]
- # xaj_params = [param_ranges[key] for key in param_ranges] 
+ # xaj_params = [param_ranges[key] for key in param_ranges]
 
  k = xaj_params[0]
  b = xaj_params[1]
@@ -783,7 +783,7 @@ def xaj(
  # ki+kg should be smaller than 1; if not, we scale them
  ki = np.where(ki_ + kg_ < 1.0, ki_, (1.0 - PRECISION) / (ki_ + kg_) * ki_)
  kg = np.where(ki_ + kg_ < 1.0, kg_, (1.0 - PRECISION) / (ki_ + kg_) * kg_)
- 
+
  if route_method == "CSL":
  cs = xaj_params[11]
  l = xaj_params[12]
@@ -817,7 +817,6 @@ def xaj(
  qi0 = np.full(ci.shape, 0.1)
  qg0 = np.full(cg.shape, 0.1)
 
-
  # state_variables
  inputs = p_and_e[warmup_length:, :, :]
  runoff_ims_ = np.full(inputs.shape[:2], 0.0)

hydromodel/models/xaj_bmi.py

@@ -12,7 +12,9 @@
 
 
 class xajBmi(Bmi):
- """Empty model wrapped in a BMI interface."""
+ """Empty model wrapped in a BMI interface.
+ TODO: not implemented yet
+ """
 
  name = "hydro-model-xaj"
  input_var_names = ("precipitation", "ETp")

hydromodel/trainers/calibrate_ga.py

@@ -25,7 +25,7 @@
 import definitions
 from hydromodel.models.model_config import MODEL_PARAM_DICT
 from hydromodel.utils import units
-from hydromodel.trainers.plots import plot_sim_and_obs, plot_train_iteration
+from hydromodel.trainers.train_utils import plot_sim_and_obs, plot_train_iteration
 from hydromodel.models.gr4j import gr4j
 from hydromodel.models.hymod import hymod
 from hydromodel.models.xaj import xaj