Merge 6159aed into edda069

environment-dev.yml

@@ -16,6 +16,7 @@ dependencies:
   - pip
   - pip:
     - xdatasets
+    - tensorflow
   # Dev
   - black ==24.1.1
   - blackdoc ==0.3.9

environment.yml

@@ -16,3 +16,4 @@ dependencies:
   - pip
   - pip:
     - xdatasets
+    - tensorflow

pyproject.toml

@@ -45,7 +45,8 @@ dependencies = [
   "xarray",
   "xclim>=0.47.0",
   "xdatasets>=0.3.1",
-  "xscen>=0.7.1"
+  "xscen>=0.7.1",
+  "tensorflow>=2.10"
 ]
 
 [project.optional-dependencies]

tests/test_lstm.py

@@ -0,0 +1,96 @@
+"""Test suite for LSTM model implementations"""
+
+import os
+
+from xhydro.lstm_tools.lstm_controller import control_regional_lstm_training
+
+
+def test_lstm_controller():
+    """Test the regional LSTM model implementation."""
+    batch_size = 64  # batch size used in the training - multiple of 32
+    epochs = 2  # Number of epoch to train the LSTM model
+    window_size = 5  # Number of time step (days) to use in the LSTM model
+    train_pct = 10  # Percentage of watersheds used for the training
+    valid_pct = 5  # Percentage of watersheds used for the validation
+    use_cpu = (
+        True  # Use CPU as GPU is not guaranteed to be installed with CUDA/CuDNN etc.
+    )
+    use_parallel = False
+    do_train = True
+    do_simulation = True
+    input_data_filename = "LSTM_test_data.nc"
+    filename_base = "LSTM_results"
+    simulation_phases = ["test"]
+
+    dynamic_var_tags = ["tasmax_MELCC", "rf", "Qsim"]
+
+    # Scale variable according to area. Used for simulated flow inputs.
+    qsim_pos = [
+        False,
+        False,
+        False,
+        False,
+        True,
+    ]
+
+    # static variables used to condition flows on catchment properties
+    static_var_tags = [
+        "drainage_area",
+        "elevation",
+        "slope",
+        "deciduous_forest",
+        "coniferous_forest",
+        "impervious",
+        "loamy_sand",
+        "silt",
+        "silty_clay_loam",
+        "meanPrecip",
+        "centroid_lat",
+        "centroid_lon",
+    ]
+
+    kge_results, flow_results, name_of_saved_model = control_regional_lstm_training(
+        input_data_filename,
+        dynamic_var_tags,
+        qsim_pos,
+        static_var_tags,
+        batch_size=batch_size,
+        epochs=epochs,
+        window_size=window_size,
+        train_pct=train_pct,
+        valid_pct=valid_pct,
+        use_cpu=use_cpu,
+        use_parallel=use_parallel,
+        do_train=do_train,
+        do_simulation=do_simulation,
+        filename_base=filename_base,
+        simulation_phases=simulation_phases,
+    )
+
+    assert len(kge_results[0]) == 4
+    assert len(flow_results[0]) == 4
+    assert os.path.isfile(name_of_saved_model)
+
+    # Do a sim with no training
+    kge_results, flow_results, name_of_saved_model = control_regional_lstm_training(
+        input_data_filename,
+        dynamic_var_tags,
+        qsim_pos,
+        static_var_tags,
+        batch_size=batch_size,
+        epochs=epochs,
+        window_size=window_size,
+        train_pct=train_pct,
+        valid_pct=valid_pct,
+        use_cpu=use_cpu,
+        use_parallel=use_parallel,
+        do_train=False,
+        do_simulation=do_simulation,
+        filename_base=filename_base,
+        simulation_phases=simulation_phases,
+        name_of_saved_model=name_of_saved_model,
+    )
+
+    assert len(kge_results[0]) == 4
+    assert len(flow_results[0]) == 4
+    assert os.path.isfile(name_of_saved_model)