ENH: Add hybrid models for gr and loieau structures

smash/_constant.py

@@ -45,6 +45,21 @@ def get_neurons_from_hydrological_module(hydrological_module: str, hidden_neuron
         elif hydrological_module == "gr4_ode_mlp":
             # % fixed NN input size = 4 and fixed NN output size 5
             n_in, n_out = (4, 5)
+        elif hydrological_module == "gr5_mlp":
+            # % fixed NN input size = 4 and fixed NN output size 4
+            n_in, n_out = (4, 4)
+        elif hydrological_module == "gr6_mlp":
+            # % fixed NN input size = 5 and fixed NN output size 5
+            n_in, n_out = (5, 5)
+        elif hydrological_module == "grc_mlp":
+            # % fixed NN input size = 5 and fixed NN output size 5
+            n_in, n_out = (5, 5)
+        elif hydrological_module == "grd_mlp":
+            # % fixed NN input size = 4 and fixed NN output size 2
+            n_in, n_out = (4, 2)
+        elif hydrological_module == "loieau_mlp":
+            # % fixed NN input size = 4 and fixed NN output size 3
+            n_in, n_out = (4, 3)
 
         neurons[0] = n_in
         non_zero_hidden_neurons = [val for val in hidden_neuron if val > 0]
@@ -64,16 +79,21 @@ def get_neurons_from_hydrological_module(hydrological_module: str, hidden_neuron
 
 HYDROLOGICAL_MODULE = [
     "gr4",
-    "gr4_ri",
     "gr4_mlp",
+    "gr4_ri",
     "gr4_ode",
     "gr4_ode_mlp",
     "gr5",
+    "gr5_mlp",
     "gr5_ri",
     "gr6",
+    "gr6_mlp",
     "grc",
+    "grc_mlp",
     "grd",
+    "grd_mlp",
     "loieau",
+    "loieau_mlp",
     "vic3l",
 ]
 
@@ -97,20 +117,18 @@ def get_neurons_from_hydrological_module(hydrological_module: str, hidden_neuron
 HYDROLOGICAL_MODULE_RR_PARAMETERS = dict(
     zip(
         HYDROLOGICAL_MODULE,
-        [
-            ["ci", "cp", "ct", "kexc"],  # % gr4
-            ["ci", "cp", "ct", "alpha1", "alpha2", "kexc"],  # % gr4_ri
-            ["ci", "cp", "ct", "kexc"],  # % gr4_mlp
-            ["ci", "cp", "ct", "kexc"],  # % gr4_ode
-            ["ci", "cp", "ct", "kexc"],  # % gr4_ode_mlp
-            ["ci", "cp", "ct", "kexc", "aexc"],  # % gr5
-            ["ci", "cp", "ct", "alpha1", "alpha2", "kexc", "aexc"],  # % gr5_ri
-            ["ci", "cp", "ct", "be", "kexc", "aexc"],  # % gr6
-            ["ci", "cp", "ct", "cl", "kexc"],  # % grc
-            ["cp", "ct"],  # % grd
-            ["ca", "cc", "kb"],  # % loieau
-            ["b", "cusl", "cmsl", "cbsl", "ks", "pbc", "ds", "dsm", "ws"],  # % vic3l
-        ],
+        (
+            [["ci", "cp", "ct", "kexc"]] * 2  # % gr4, gr4_mlp,
+            + [["ci", "cp", "ct", "alpha1", "alpha2", "kexc"]]  # % gr4_ri
+            + [["ci", "cp", "ct", "kexc"]] * 2  # % gr4_ode, gr4_ode_mlp
+            + [["ci", "cp", "ct", "kexc", "aexc"]] * 2  # % gr5, gr5_mlp
+            + [["ci", "cp", "ct", "alpha1", "alpha2", "kexc", "aexc"]]  # % gr5_ri
+            + [["ci", "cp", "ct", "be", "kexc", "aexc"]] * 2  # % gr6, gr6_mlp
+            + [["ci", "cp", "ct", "cl", "kexc"]] * 2  # % grc, grc_mlp
+            + [["cp", "ct"]] * 2  # % grd, grd_mlp
+            + [["ca", "cc", "kb"]] * 2  # % loieau, loieau_mlp
+            + [["b", "cusl", "cmsl", "cbsl", "ks", "pbc", "ds", "dsm", "ws"]]  # % vic3l
+        ),
     )
 )
 
@@ -141,20 +159,14 @@ def get_neurons_from_hydrological_module(hydrological_module: str, hidden_neuron
 HYDROLOGICAL_MODULE_RR_STATES = dict(
     zip(
         HYDROLOGICAL_MODULE,
-        [
-            ["hi", "hp", "ht"],  # % gr4
-            ["hi", "hp", "ht"],  # % gr4_ri
-            ["hi", "hp", "ht"],  # % gr4_mlp
-            ["hi", "hp", "ht"],  # % gr4_ode
-            ["hi", "hp", "ht"],  # % gr4_ode_mlp
-            ["hi", "hp", "ht"],  # % gr5
-            ["hi", "hp", "ht"],  # % gr5_ri
-            ["hi", "hp", "ht", "he"],  # % gr6
-            ["hi", "hp", "ht", "hl"],  # % grc
-            ["hp", "ht"],  # % grd
-            ["ha", "hc"],  # % loieau
-            ["hcl", "husl", "hmsl", "hbsl"],  # % vic3l
-        ],
+        (
+            [["hi", "hp", "ht"]] * 8  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr5_ri
+            + [["hi", "hp", "ht", "he"]] * 2  # % gr6, gr6_mlp
+            + [["hi", "hp", "ht", "hl"]] * 2  # % grc, grc_mlp
+            + [["hp", "ht"]] * 2  # % grd, grd_mlp
+            + [["ha", "hc"]] * 2  # % loieau, loieau_mlp
+            + [["hcl", "husl", "hmsl", "hbsl"]]  # % vic3l
+        ),
     )
 )
 
@@ -187,20 +199,19 @@ def get_neurons_from_hydrological_module(hydrological_module: str, hidden_neuron
 HYDROLOGICAL_MODULE_RR_INTERNAL_FLUXES = dict(
     zip(
         HYDROLOGICAL_MODULE,
-        [
-            ["pn", "en", "pr", "perc", "lexc", "prr", "prd", "qr", "qd", "qt"],  # % gr4
-            ["pn", "en", "pr", "perc", "lexc", "prr", "prd", "qr", "qd", "qt"],  # % gr4-ri
-            ["pn", "en", "pr", "perc", "lexc", "prr", "prd", "qr", "qd", "qt"],  # % gr4_mlp
-            ["pn", "en", "lexc", "qt"],  # % gr4_ode
-            ["pn", "en", "lexc", "qt"],  # % gr4_ode_mlp
-            ["pn", "en", "pr", "perc", "lexc", "prr", "prd", "qr", "qd", "qt"],  # % gr5
-            ["pn", "en", "pr", "perc", "lexc", "prr", "prd", "qr", "qd", "qt"],  # % gr5-ri
-            ["pn", "en", "pr", "perc", "lexc", "prr", "prd", "pre", "qr", "qd", "qe", "qt"],  # % gr6
-            ["pn", "en", "pr", "perc", "lexc", "prr", "prd", "prl", "qr", "qd", "ql", "qt"],  # % grc
-            ["ei", "pn", "en", "pr", "perc", "prr", "qr", "qt"],  # % grd
-            ["ei", "pn", "en", "pr", "perc", "prr", "prd", "qr", "qd", "qt"],  # % loieau
-            ["pn", "en", "qr", "qb", "qt"],  # % vic3l
-        ],
+        (
+            [["pn", "en", "pr", "perc", "lexc", "prr", "prd", "qr", "qd", "qt"]] * 3  # % gr4, gr4_mlp, gr4_ri
+            + [["pn", "en", "lexc", "qt"]] * 2  # % gr4_ode, gr4_ode_mlp
+            + [["pn", "en", "pr", "perc", "lexc", "prr", "prd", "qr", "qd", "qt"]]
+            * 3  # % gr5, gr5_mlp, gr5_ri
+            + [["pn", "en", "pr", "perc", "lexc", "prr", "prd", "pre", "qr", "qd", "qe", "qt"]]
+            * 2  # % gr6, gr6_mlp
+            + [["pn", "en", "pr", "perc", "lexc", "prr", "prd", "prl", "qr", "qd", "ql", "qt"]]
+            * 2  # % grc, grc_mlp
+            + [["ei", "pn", "en", "pr", "perc", "prr", "qr", "qt"]] * 2  # % grd, grd_mlp
+            + [["ei", "pn", "en", "pr", "perc", "prr", "prd", "qr", "qd", "qt"]] * 2  # % loieau, loieau_mlp
+            + [["pn", "en", "qr", "qb", "qt"]]  # % vic3l
+        ),
     )
 )
 
@@ -258,18 +269,20 @@ def get_neurons_from_hydrological_module(hydrological_module: str, hidden_neuron
 
 RR_PARAMETERS = [
     "kmlt",  # % ssn
-    "ci",  # % (gr4, gr4_ri, gr4_mlp, gr4_ode, gr4_ode_mlp, gr5, gr5_ri, grc, gr6)
-    "cp",  # % (gr4, gr4_ri, gr4_mlp, gr4_ode, gr4_ode_mlp, gr5, gr5_ri, grc, gr6, grd)
-    "ct",  # % (gr4, gr4_ri, gr4_mlp, gr4_ode, gr4_ode_mlp, gr5, gr5_ri, grc, gr6, grd)
-    "alpha1",  # % (gr4_ri, gr5_ri)
-    "alpha2",  # % (gr4_ri, gr5_ri)
-    "cl",  # % grc
-    "be",  # % (gr6)
-    "kexc",  # % (gr4, gr4_ri, gr4_mlp, gr4_ode, gr4_ode_mlp, gr5, gr5_ri, grc, gr6)
-    "aexc",  # % (gr5, gr5_ri, gr6)
-    "ca",  # % loieau
-    "cc",  # % loieau
-    "kb",  # % loieau
+    "ci",  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr5_ri, gr6, gr6_mlp, grc, grc_mlp
+    "cp",  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr5_ri, gr6, gr6_mlp, grc, grc_mlp,
+    # % grd, grd_mlp
+    "ct",  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr5_ri, gr6, gr6_mlp, grc, grc_mlp,
+    # %grd, grd_mlp
+    "alpha1",  # % gr4_ri, gr5_ri
+    "alpha2",  # % gr4_ri, gr5_ri
+    "cl",  # % grc, grc_mlp
+    "be",  # % gr6, gr6_mlp
+    "kexc",  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr6, gr6_mlp, grc, grc_mlp
+    "aexc",  # % gr5, gr5_mlp, gr5_ri, gr6, gr6_mlp
+    "ca",  # % loieau, loieau_mlp
+    "cc",  # % loieau, loieau_mlp
+    "kb",  # % loieau, loieau_mlp
     "b",  # % vic3l
     "cusl",  # % vic3l
     "cmsl",  # % vic3l
@@ -286,13 +299,15 @@ def get_neurons_from_hydrological_module(hydrological_module: str, hidden_neuron
 
 RR_STATES = [
     "hs",  # % ssn
-    "hi",  # % (gr4, gr4_ri, gr4_mlp, gr4_ode, gr4_ode_mlp, gr5, gr5_ri, grc, gr6)
-    "hp",  # % (gr4, gr4_ri, gr4_mlp, gr4_ode, gr4_ode_mlp, gr5, gr5_ri, grc, gr6, grd)
-    "ht",  # % (gr4, gr4_ri, gr4_mlp, gr4_ode, gr4_ode_mlp, gr5, gr5_ri, grc, gr6, grd)
-    "hl",  # % grc
-    "he",  # % gr6
-    "ha",  # % loieau
-    "hc",  # % loieau
+    "hi",  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr5_ri, gr6, gr6_mlp, grc, grc_mlp
+    "hp",  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr5_ri, gr6, gr6_mlp, grc, grc_mlp,
+    # % grd, grd_mlp
+    "ht",  # % gr4, gr4_mlp, gr4_ri, gr4_ode, gr4_ode_mlp, gr5, gr5_mlp, gr5_ri, gr6, gr6_mlp, grc, grc_mlp,
+    # % grd, grd_mlp
+    "hl",  # % grc, grc_mlp
+    "he",  # % gr6, gr6_mlp
+    "ha",  # % loieau, loieau_mlp
+    "hc",  # % loieau, loieau_mlp
     "hcl",  # % vic3l
     "husl",  # % vic3l
     "hmsl",  # % vic3l

smash/core/model/model.py

@@ -125,15 +125,12 @@ class Model:
         hydrological_module : `str`, default 'gr4'
             Name of hydrological module. Should be one of:
 
-            - ``'gr4'``
-            - ``'gr4_mlp'``
-            - ``'gr4_ode'``
-            - ``'gr4_ode_mlp'``
-            - ``'gr5'``
-            - ``'gr6'``
-            - ``'grc'``
-            - ``'grd'``
-            - ``'loieau'``
+            - ``'gr4'``, ``'gr4_mlp'``, ``'gr4_ri'``, ``'gr4_ode'``, ``'gr4_ode_mlp'``
+            - ``'gr5'``, ``'gr5_mlp'``, ``'gr5_ri'``
+            - ``'gr6'``, ``'gr6_mlp'``
+            - ``'grc'``, ``'grc_mlp'``
+            - ``'grd'``, ``'grd_mlp'``
+            - ``'loieau'``, ``'loieau_mlp'``
             - ``'vic3l'``
 
             .. hint::