fix import bug; start to refactor

hydromodel/calibrate/calibrate_ga_xaj_bmi.py

@@ -1,4 +1,5 @@
 """Calibrate XAJ model using DEAP"""
+
 import os
 import pickle
 import random
@@ -11,14 +12,15 @@
 from deap import tools
 from tqdm import tqdm
 
+import HydroErr as he
+from hydroutils import hydro_stat, hydro_file
+
 sys.path.append(os.path.dirname(Path(os.path.abspath(__file__)).parent.parent))
 import definitions
 from hydromodel.models.model_config import MODEL_PARAM_DICT
-from hydromodel.utils import hydro_constant, hydro_utils
-from hydromodel.utils import stat
-from hydromodel.utils.stat import statRmse
-from hydromodel.visual.hydro_plot import plot_sim_and_obs, plot_train_iteration
+from hydromodel.utils.plots import plot_sim_and_obs, plot_train_iteration
 from hydromodel.models.xaj_bmi import xajBmi
+from hydromodel.utils import units
 
 
 def evaluate(individual, x_input, y_true, warmup_length, model):
@@ -50,16 +52,16 @@ def evaluate(individual, x_input, y_true, warmup_length, model):
         # xaj model's output include streamflow and evaporation now,
         # but now we only calibrate the model with streamflow
         model = xajBmi()
-        model.initialize(os.path.relpath('runxaj.yaml'), params, x_input)
-        while model.get_current_time() <= model.get_end_time('train'):
+        model.initialize(os.path.relpath("runxaj.yaml"), params, x_input)
+        while model.get_current_time() <= model.get_end_time("train"):
             model.update()
         sim = model.get_value("discharge")
         sim = np.expand_dims(sim, 0)
         sim = np.expand_dims(sim, 1)
         sim = np.transpose(sim, [2, 1, 0])
     else:
         raise NotImplementedError("We don't provide this model now")
-    rmses = statRmse(y_true[warmup_length:, :, :], sim)
+    rmses = he.rmse(y_true[warmup_length:, :, :], sim)
     rmse = rmses.mean(axis=0)
     # print(f"-----------------RMSE: {str(rmse)}------------------------")
     return rmse
@@ -103,7 +105,7 @@ def wrapper(*args, **kargs):
 
 
 def calibrate_by_ga(
-        input_data, observed_output, deap_dir, warmup_length=30, model=None, ga_param=None
+    input_data, observed_output, deap_dir, warmup_length=30, model=None, ga_param=None
 ):
     """
     Use GA algorithm to find optimal parameters for hydrologic models
@@ -230,8 +232,8 @@ def calibrate_by_ga(
         invalid_ind = [ind for ind in offspring if not ind.fitness.valid]
         fitnesses = map(toolbox.evaluate, invalid_ind)
         for ind, fit in tqdm(
-                zip(invalid_ind, fitnesses),
-                desc=f"{str(gen + 1)} generation fitness calculating",
+            zip(invalid_ind, fitnesses),
+            desc=f"{str(gen + 1)} generation fitness calculating",
         ):
             ind.fitness.values = fit
 
@@ -256,7 +258,7 @@ def calibrate_by_ga(
             )
 
             with open(
-                    os.path.join(deap_dir, f"epoch{str(gen + 1)}.pkl"), "wb"
+                os.path.join(deap_dir, f"epoch{str(gen + 1)}.pkl"), "wb"
             ) as cp_file:
                 pickle.dump(cp, cp_file)
             print(f"Files of generation {gen} saved.")
@@ -265,15 +267,15 @@ def calibrate_by_ga(
 
 
 def show_ga_result(
-        deap_dir,
-        warmup_length,
-        basin_id,
-        the_data,
-        the_period,
-        basin_area,
-        model_info,
-        result_unit="mm/day",
-        train_mode=True,
+    deap_dir,
+    warmup_length,
+    basin_id,
+    the_data,
+    the_period,
+    basin_area,
+    model_info,
+    result_unit="mm/day",
+    train_mode=True,
 ):
     """
     show the result of GA
@@ -290,22 +292,24 @@ def show_ga_result(
     train_test_flag = "train" if train_mode else "test"
 
     model = xajBmi()
-    model.initialize("runxaj.yaml", np.array(list(halloffame[0])).reshape(1, -1), the_data[:, :, 0:2])
-    while model.get_current_time() <= model.get_end_time('train'):
+    model.initialize(
+        "runxaj.yaml", np.array(list(halloffame[0])).reshape(1, -1), the_data[:, :, 0:2]
+    )
+    while model.get_current_time() <= model.get_end_time("train"):
         model.update()
     best_simulation = model.get_value("discharge")
 
-    convert_unit_sim = hydro_constant.convert_unit(
+    convert_unit_sim = units.convert_unit(
         np.array(best_simulation).reshape(1, -1),
         # best_simulation,
         result_unit,
-        hydro_constant.unit["streamflow"],
+        units.unit["streamflow"],
         basin_area=basin_area,
     )
-    convert_unit_obs = hydro_constant.convert_unit(
+    convert_unit_obs = units.convert_unit(
         np.array(the_data[warmup_length:, :, -1:]).reshape(1, -1),
         result_unit,
-        hydro_constant.unit["streamflow"],
+        units.unit["streamflow"],
         basin_area=basin_area,
     )
     # save calibrated results of calibration period
@@ -320,12 +324,12 @@ def show_ga_result(
         header=False,
     )
     # calculation rmse、nashsutcliffe and bias for training period
-    stat_error = stat.statError(
+    stat_error = hydro_stat.stat_error(
         convert_unit_obs,
         convert_unit_sim,
     )
     print(f"{train_test_flag}ing metrics:", basin_id, stat_error)
-    hydro_utils.serialize_json_np(
+    hydro_file.serialize_json_np(
         stat_error, os.path.join(deap_dir, f"{train_test_flag}_metrics.json")
     )
     t_range = pd.to_datetime(the_period[warmup_length:]).values.astype("datetime64[D]")
@@ -359,8 +363,8 @@ def show_ga_result(
     )
     train_data_info_file = os.path.join(data_dir, "data_info_fold0_train.json")
     train_data_file = os.path.join(data_dir, "basins_lump_p_pe_q_fold0_train.npy")
-    data_train = hydro_utils.unserialize_numpy(train_data_file)
-    data_info_train = hydro_utils.unserialize_json_ordered(train_data_info_file)
+    data_train = hydro_file.unserialize_numpy(train_data_file)
+    data_info_train = hydro_file.unserialize_json_ordered(train_data_info_file)
     model_info = {
         "name": "xaj_mz",
         "source_type": "sources",

hydromodel/models/configuration.py

@@ -1,10 +1,10 @@
 import yaml
 import numpy as np
-from hydromodel.models.xaj import xaj_state as xaj_state
+from hydromodel.models.xaj import xaj
 
 
 def read_config(config_file):
-    with open(config_file, encoding='utf8') as f:
+    with open(config_file, encoding="utf8") as f:
         config = yaml.safe_load(f)
     return config
 
@@ -17,7 +17,7 @@ def extract_forcing(forcing_data):
 
 
 def warmup(p_and_e_warmup, params_state, warmup_length, model_info):
-    q_sim, es, *w0, w1, w2, s0, fr0, qi0, qg0 = xaj_state(
+    q_sim, es, *w0, w1, w2, s0, fr0, qi0, qg0 = xaj(
         p_and_e_warmup,
         params_state,
         warmup_length=warmup_length,

hydromodel/models/xaj.py

@@ -1,6 +1,13 @@
 """
-Core code for XinAnJiang model
+Author: Wenyu Ouyang
+Date: 2021-12-10 23:01:02
+LastEditTime: 2024-02-10 11:51:35
+LastEditors: Wenyu Ouyang
+Description: Core code for XinAnJiang model
+FilePath: /hydro-model-xaj/hydromodel/models/xaj.py
+Copyright (c) 2023-2024 Wenyu Ouyang. All rights reserved.
 """
+
 import logging
 from typing import Union
 import numpy as np
@@ -693,6 +700,14 @@ def uh_gamma(a, theta, len_uh=15):
     return w
 
 
+def xaj_runoff():
+    pass
+
+
+def xaj_route():
+    pass
+
+
 def xaj(
     p_and_e,
     params: Union[np.array, list],

test/test_rr_event_iden.py

@@ -1,12 +1,13 @@
 """
 Author: Wenyu Ouyang
 Date: 2023-10-28 09:23:22
-LastEditTime: 2023-10-28 16:19:22
+LastEditTime: 2024-02-10 11:17:37
 LastEditors: Wenyu Ouyang
 Description: Test for rainfall-runoff event identification
 FilePath: \hydro-model-xaj\test\test_rr_event_iden.py
 Copyright (c) 2023-2024 Wenyu Ouyang. All rights reserved.
 """
+
 import os
 import pandas as pd
 import definitions

test/test_xaj_bmi.py

@@ -11,7 +11,9 @@
 import socket
 from datetime import datetime
 
-from hydromodel.utils import hydro_utils
+from hydroutils import hydro_file
+
+from hydromodel.utils import units
 from hydromodel.data.data_preprocess import (
     cross_valid_data,
     split_train_test,
@@ -21,15 +23,14 @@
     calibrate_by_ga,
     show_ga_result,
 )
-from hydromodel.visual.pyspot_plots import show_calibrate_result, show_test_result
+from hydromodel.utils.plots import show_calibrate_result, show_test_result
 from hydromodel.data.data_postprocess import (
     renormalize_params,
     read_save_sceua_calibrated_params,
     save_streamflow,
     summarize_metrics,
     summarize_parameters,
 )
-from hydromodel.utils import hydro_constant
 
 logging.basicConfig(level=logging.INFO)
 
@@ -121,10 +122,10 @@ def test_bmi():
             raise FileNotFoundError(
                 "The data files are not found, please run datapreprocess4calibrate.py first."
             )
-        data_train = hydro_utils.unserialize_numpy(train_data_file)
-        data_test = hydro_utils.unserialize_numpy(test_data_file)
-        data_info_train = hydro_utils.unserialize_json_ordered(train_data_info_file)
-        data_info_test = hydro_utils.unserialize_json_ordered(test_data_info_file)
+        data_train = hydro_file.unserialize_numpy(train_data_file)
+        data_test = hydro_file.unserialize_numpy(test_data_file)
+        data_info_train = hydro_file.unserialize_json_ordered(train_data_info_file)
+        data_info_test = hydro_file.unserialize_json_ordered(test_data_info_file)
         current_time = datetime.now().strftime("%b%d_%H-%M-%S")
         # one directory for one model + one hyperparam setting and one basin
         save_dir = os.path.join(
@@ -194,17 +195,17 @@ def test_bmi():
                     j += 1
                 q_sim = model.get_value("discharge")
 
-                qsim = hydro_constant.convert_unit(
+                qsim = units.convert_unit(
                     q_sim,
                     unit_now="mm/day",
-                    unit_final=hydro_constant.unit["streamflow"],
+                    unit_final=units.unit["streamflow"],
                     basin_area=basin_area,
                 )
 
-                qobs = hydro_constant.convert_unit(
+                qobs = units.convert_unit(
                     data_test[warmup_length:, i : i + 1, -1:],
                     unit_now="mm/day",
-                    unit_final=hydro_constant.unit["streamflow"],
+                    unit_final=units.unit["streamflow"],
                     basin_area=basin_area,
                 )
                 test_result_file = os.path.join(