add param_range file so that you can set your own param range

hydromodel/datasets/data_preprocess.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-10-25 21:16:22
-LastEditTime: 2024-03-26 21:20:18
+LastEditTime: 2024-03-27 14:30:15
 LastEditors: Wenyu Ouyang
 Description: preprocess data for models in hydro-model-xaj
 FilePath: \hydro-model-xaj\hydromodel\datasets\data_preprocess.py
@@ -390,7 +390,23 @@ def cross_valid_data(ts_data, period, warmup, cv_fold, freq="1D"):
  return train_test_data
 
 
-def get_basin_area(data_type, data_dir, basin_ids):
+def get_basin_area(data_type, data_dir, basin_ids) -> xr.Dataset:
+ """_summary_
+
+ Parameters
+ ----------
+ data_type : _type_
+ _description_
+ data_dir : _type_
+ _description_
+ basin_ids : _type_
+ _description_
+
+ Returns
+ -------
+ xr.Dataset
+ _description_
+ """
  area_name = remove_unit_from_name(AREA_NAME)
  if data_type == "camels":
  camels_data_dir = os.path.join(
@@ -399,10 +415,9 @@ def get_basin_area(data_type, data_dir, basin_ids):
  camels = Camels(camels_data_dir)
  basin_area = camels.read_area(basin_ids)
  elif data_type == "owndata":
- attr_data = xr.open_dataset(
- os.path.join(os.path.dirname(data_dir), "attributes.nc")
- )
- basin_area = attr_data[area_name].values
+ attr_data = xr.open_dataset(os.path.join(data_dir, "attributes.nc"))
+ # to guarantee the column name is same as the column name in the time series data
+ basin_area = attr_data[[area_name]].rename({"id": "basin"})
  return basin_area
 
 
@@ -451,9 +466,7 @@ def get_ts_from_diffsource(data_type, data_dir, periods, basin_ids):
  ts_data = ts_data.rename({"PET": pet_name})
  # ts_data = ts_data.drop_vars('streamflow')
  elif data_type == "owndata":
- ts_data = xr.open_dataset(
- os.path.join(os.path.dirname(data_dir), "timeseries.nc")
- )
+ ts_data = xr.open_dataset(os.path.join(data_dir, "timeseries.nc"))
  target_unit = ts_data[prcp_name].attrs.get("units", "unknown")
  qobs_ = ts_data[[flow_name]]
  r_mmd = streamflow_unit_conv(qobs_, basin_area, target_unit=target_unit)

hydromodel/models/gr4j.py

@@ -3,7 +3,7 @@
 import numpy as np
 from numba import jit
 
-from hydromodel.models.model_config import MODEL_PARAM_DICT
+from hydromodel.models.model_config import read_model_param_dict
 from hydromodel.models.xaj import uh_conv
 
 
@@ -204,10 +204,12 @@ def gr4j(p_and_e, parameters, warmup_length: int, return_state=False, **kwargs):
  Union[np.array, tuple]
  streamflow or (streamflow, states)
  """
- x1_scale = MODEL_PARAM_DICT["gr4j"]["param_range"]["x1"]
- x2_sacle = MODEL_PARAM_DICT["gr4j"]["param_range"]["x2"]
- x3_scale = MODEL_PARAM_DICT["gr4j"]["param_range"]["x3"]
- x4_scale = MODEL_PARAM_DICT["gr4j"]["param_range"]["x4"]
+ pr_file = kwargs.get("param_range_file", None)
+ model_param_dict = read_model_param_dict(pr_file)
+ x1_scale = model_param_dict["gr4j"]["param_range"]["x1"]
+ x2_sacle = model_param_dict["gr4j"]["param_range"]["x2"]
+ x3_scale = model_param_dict["gr4j"]["param_range"]["x3"]
+ x4_scale = model_param_dict["gr4j"]["param_range"]["x4"]
  x1 = x1_scale[0] + parameters[:, 0] * (x1_scale[1] - x1_scale[0])
  x2 = x2_sacle[0] + parameters[:, 1] * (x2_sacle[1] - x2_sacle[0])
  x3 = x3_scale[0] + parameters[:, 2] * (x3_scale[1] - x3_scale[0])

hydromodel/models/hymod.py

@@ -2,7 +2,7 @@
 import numpy as np
 from numba import jit
 
-from hydromodel.models.model_config import MODEL_PARAM_DICT
+from hydromodel.models.model_config import read_model_param_dict
 
 
 def hymod(p_and_e, parameters, warmup_length=30, return_state=False, **kwargs):
@@ -30,12 +30,14 @@ def hymod(p_and_e, parameters, warmup_length=30, return_state=False, **kwargs):
  Union[list, np.array]
  streamflow, x_slow, x_quick, x_loss or streamflow
  """
+ pr_file = kwargs.get("param_range_file", None)
+ model_param_dict = read_model_param_dict(pr_file)
  # parameter, 2-dim variable: [parameter=1, basin]
- cmax_scale = MODEL_PARAM_DICT["hymod"]["param_range"]["cmax"]
- bexp_sacle = MODEL_PARAM_DICT["hymod"]["param_range"]["bexp"]
- alpha_scale = MODEL_PARAM_DICT["hymod"]["param_range"]["alpha"]
- ks_scale = MODEL_PARAM_DICT["hymod"]["param_range"]["ks"]
- kq_scale = MODEL_PARAM_DICT["hymod"]["param_range"]["kq"]
+ cmax_scale = model_param_dict["hymod"]["param_range"]["cmax"]
+ bexp_sacle = model_param_dict["hymod"]["param_range"]["bexp"]
+ alpha_scale = model_param_dict["hymod"]["param_range"]["alpha"]
+ ks_scale = model_param_dict["hymod"]["param_range"]["ks"]
+ kq_scale = model_param_dict["hymod"]["param_range"]["kq"]
  cmax = cmax_scale[0] + parameters[:, 0] * (cmax_scale[1] - cmax_scale[0])
  bexp = bexp_sacle[0] + parameters[:, 1] * (bexp_sacle[1] - bexp_sacle[0])
  alpha = alpha_scale[0] + parameters[:, 2] * (alpha_scale[1] - alpha_scale[0])

hydromodel/models/model_config.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-10-25 21:16:22
-LastEditTime: 2024-03-26 11:13:08
+LastEditTime: 2024-03-27 15:03:41
 LastEditors: Wenyu Ouyang
 Description: some basic config for hydro-model-xaj models
 FilePath: \hydro-model-xaj\hydromodel\models\model_config.py
@@ -10,7 +10,27 @@
 
 from collections import OrderedDict
 
-# NOTE: Don't change the parameter settings
+import yaml
+
+
+def read_model_param_dict(file_path="param.yaml"):
+ try:
+ with open(file_path, "r") as file:
+ data = yaml.safe_load(file)
+
+ return {
+ model: {
+ "param_name": contents["param_name"],
+ "param_range": OrderedDict(contents["param_range"]),
+ }
+ for model, contents in data.items()
+ }
+ except FileNotFoundError:
+ print(
+ f"File not found: {file_path}, we directly use the default MODEL_PARAM_DICT."
+ )
+ return MODEL_PARAM_DICT
+
 
 MODEL_PARAM_DICT = {
  "xaj": {
@@ -49,7 +69,7 @@
  "KI": [0.0, 0.7],
  "KG": [0.0, 0.7],
  "CS": [0.0, 1.0],
- "L": [1.0, 10.0], # unit is day
+ "L": [1.0, 10.0], # unit is same as your time step
  "CI": [0.0, 0.9],
  "CG": [0.98, 0.998],
  }

hydromodel/models/param.yaml

@@ -0,0 +1,172 @@
+xaj:
+ param_name:
+ - K
+ - B
+ - IM
+ - UM
+ - LM
+ - DM
+ - C
+ - SM
+ - EX
+ - KI
+ - KG
+ - CS
+ - L
+ - CI
+ - CG
+ param_range:
+ K:
+ - 0.1
+ - 1.0
+ B:
+ - 0.1
+ - 0.4
+ IM:
+ - 0.01
+ - 0.1
+ UM:
+ - 0.0
+ - 20.0
+ LM:
+ - 60.0
+ - 90.0
+ DM:
+ - 60.0
+ - 120.0
+ C:
+ - 0.0
+ - 0.2
+ SM:
+ - 1
+ - 100.0
+ EX:
+ - 1.0
+ - 1.5
+ KI:
+ - 0.0
+ - 0.7
+ KG:
+ - 0.0
+ - 0.7
+ CS:
+ - 0.0
+ - 1.0
+ L:
+ - 1.0
+ - 10.0
+ CI:
+ - 0.0
+ - 0.9
+ CG:
+ - 0.98
+ - 0.998
+xaj_mz:
+ param_name:
+ - K
+ - B
+ - IM
+ - UM
+ - LM
+ - DM
+ - C
+ - SM
+ - EX
+ - KI
+ - KG
+ - A
+ - THETA
+ - CI
+ - CG
+ - KERNEL
+ param_range:
+ K:
+ - 0.1
+ - 1.0
+ B:
+ - 0.1
+ - 0.4
+ IM:
+ - 0.01
+ - 0.1
+ UM:
+ - 0.0
+ - 20.0
+ LM:
+ - 60.0
+ - 90.0
+ DM:
+ - 60.0
+ - 120.0
+ C:
+ - 0.0
+ - 0.2
+ SM:
+ - 1.0
+ - 100.0
+ EX:
+ - 1.0
+ - 1.5
+ KI:
+ - 0.0
+ - 0.7
+ KG:
+ - 0.0
+ - 0.7
+ A:
+ - 0.0
+ - 2.9
+ THETA:
+ - 0.0
+ - 6.5
+ CI:
+ - 0.0
+ - 0.9
+ CG:
+ - 0.98
+ - 0.998
+ KERNEL:
+ - 1
+ - 15
+gr4j:
+ param_name:
+ - x1
+ - x2
+ - x3
+ - x4
+ param_range:
+ x1:
+ - 100.0
+ - 1200.0
+ x2:
+ - -5.0
+ - 3.0
+ x3:
+ - 20.0
+ - 300.0
+ x4:
+ - 1.1
+ - 2.9
+hymod:
+ param_name:
+ - cmax
+ - bexp
+ - alpha
+ - ks
+ - kq
+ param_range:
+ cmax:
+ - 1.0
+ - 500.0
+ bexp:
+ - 0.1
+ - 2.0
+ alpha:
+ - 0.1
+ - 0.99
+ ks:
+ - 0.001
+ - 0.1
+ kq:
+ - 0.1
+ - 0.99

hydromodel/models/xaj.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2021-12-10 23:01:02
-LastEditTime: 2024-03-26 11:06:12
+LastEditTime: 2024-03-27 15:09:11
 LastEditors: Wenyu Ouyang
 Description: Core code for XinAnJiang model
 FilePath: /hydro-model-xaj/hydromodel/models/xaj.py
@@ -14,7 +14,7 @@
 from numba import jit
 from scipy.special import gamma
 
-from hydromodel.models.model_config import MODEL_PARAM_DICT
+from hydromodel.models.model_config import read_model_param_dict
 
 PRECISION = 1e-5
 
@@ -741,11 +741,13 @@ def xaj(
  streamflow or (streamflow, states)
  """
  # default values for some function parameters
- model_name = kwargs["name"] if "name" in kwargs else "xaj"
- source_type = kwargs["source_type"] if "source_type" in kwargs else "sources"
- source_book = kwargs["source_book"] if "source_book" in kwargs else "HF"
+ model_name = kwargs.get("name", "xaj")
+ source_type = kwargs.get("source_type", "sources")
+ source_book = kwargs.get("source_book", "HF")
+ pr_file = kwargs.get("param_range_file", None)
+ model_param_dict = read_model_param_dict(pr_file)
  # params
- param_ranges = MODEL_PARAM_DICT[model_name]["param_range"]
+ param_ranges = model_param_dict[model_name]["param_range"]
  if model_name == "xaj":
  route_method = "CSL"
  elif model_name == "xaj_mz":