ngen mts (#33)

* update hbv_2 for mts

* update dhbv2 models for ngen mts

* add mts hbv2

* update mts to make daily parameters optional

* fix caching

* updates to latest code release

* cleanup

* update gh actions

* remove print statement

* update license check

* new license check

* patch license agreement

Author: leoglonz

.github/workflows/lint.yaml

@@ -16,10 +16,10 @@ jobs:
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@v5.0.0
+        uses: actions/checkout@v6.0.1
 
       - name: Install uv + Python
-        uses: astral-sh/[email protected].2
+        uses: astral-sh/[email protected].6
         with:
           python-version: ${{ matrix.python-version }}
           enable-cache: true

.github/workflows/pytest.yaml

@@ -16,10 +16,10 @@ jobs:
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@v5.0.0
+        uses: actions/checkout@v6.0.1
 
       - name: Install uv + Python
-        uses: astral-sh/[email protected].2
+        uses: astral-sh/[email protected].6
         with:
           python-version: ${{ matrix.python-version }}
           enable-cache: true

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
       - id: check-toml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.2
+    rev: v0.14.8
     hooks:
       - id: ruff-check
         types_or: [python, pyi, jupyter]
@@ -20,7 +20,7 @@ repos:
         args: [--config, pyproject.toml]
 
   - repo: https://github.com/kynan/nbstripout
-    rev: 0.8.1
+    rev: 0.8.2
     hooks:
       - id: nbstripout
         types: [jupyter]

README.md

@@ -1,10 +1,10 @@
 # HydroDL2: Differentiable Hydrological Models
 
-[![Python](https://img.shields.io/badge/python-3.9%20%7C%203.12%20%7C%203.13-blue)](https://www.python.org/downloads/)
-[![PyTorch](https://img.shields.io/badge/PyTorch-2.9.0-EE4C2C?logo=pytorch)](https://pytorch.org/)
+[![Python](https://img.shields.io/badge/python-3.9%20%7C%203.12%20%7C%203.13-blue?labelColor=333333)](https://www.python.org/downloads/)
+[![PyTorch Version](https://img.shields.io/badge/dynamic/json?label=PyTorch&query=info.version&url=https%3A%2F%2Fpypi.org%2Fpypi%2Ftorch%2Fjson&logo=pytorch&color=EE4C2C&logoColor=F900FF&labelColor=333333)](https://pypi.org/project/torch/)
 
-[![Build](https://github.com/mhpi/hydrodl2/actions/workflows/pytest.yaml/badge.svg?branch=master)](https://github.com/mhpi/hydrodl2/actions/workflows/pytest.yaml/)
-[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![Build](https://img.shields.io/github/actions/workflow/status/mhpi/generic_deltamodel/pytest.yaml?branch=master&logo=github&label=tests&labelColor=333333)](https://github.com/mhpi/generic_deltamodel/actions/workflows/pytest.yaml)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json&labelColor=333333)](https://github.com/astral-sh/ruff)
 
 ---
 

pyproject.toml

@@ -12,6 +12,7 @@ license = { file = "LICENSE" }
 authors = [
     { name = "Leo Lonzarich", email = "[email protected]" },
     { name = "Yalan Song", email = "[email protected]" },
+    { name = "Wencong Yang", email = "[email protected]" },
     { name = "Tadd Bindas", email = "[email protected]" },
 ]
 maintainers = [

src/hydrodl2/__init__.py

@@ -1,5 +1,7 @@
+# src/hydrodl2/__init__.py
 import logging
 import os
+import sys
 from datetime import datetime
 from importlib.resources import files
 from pathlib import Path
@@ -69,7 +71,13 @@ def _check_license_agreement():
 
         print("-" * 40)
 
-        response = input("Do you agree to these terms? Type 'Yes' to continue: ")
+        try:
+            response = input("Do you agree to these terms? Type 'Yes' to continue: ")
+        except EOFError:
+            # If we get here, it means we're in an environment that can't take
+            # input -- default to no agreement.
+            print("\n[!] No terminal detected. Skipping license prompt.")
+            return
 
         if response.strip().lower() in ['yes', 'y']:
             try:
@@ -92,7 +100,23 @@ def _check_license_agreement():
             raise SystemExit(1)
 
 
+def _should_skip_license():
+    """Returns True if the license check should be bypassed."""
+    # 1. Check if we are in a Non-Interactive shell (No user to type 'Yes')
+    if not sys.stdin.isatty():
+        return True
+
+    # 2. Check for common 'Silent' or 'CI' flags
+    if os.environ.get('CI') or os.environ.get('NGEN_SILENT'):
+        return True
+
+    # 3. Check for the Docker flag (as a fallback)
+    if os.path.exists('/.dockerenv'):
+        return True
+
+    return False
+
+
 # This only runs once when package is first imported.
-if not os.environ.get('CI'):
-    # Skip license check in CI envs (e.g., GitHub Actions)
+if not _should_skip_license():
     _check_license_agreement()

src/hydrodl2/models/hbv/hbv_2.py

@@ -6,14 +6,14 @@
 
 
 class Hbv_2(torch.nn.Module):
-    """HBV 2.0 ~.
+    """HBV 2.0.
 
-    Multi-component, multiscale, differentiable PyTorch HBV model with rainfall
+    Multi-component, multi-scale, differentiable PyTorch HBV model with rainfall
     runoff simulation on unit basins.
 
     Authors
     -------
-    -   Yalan Song, Leo Lonzarich
+    -   Yalan Song, Leo Lonzarich, Wencong Yang
     -   (Original NumPy HBV ver.) Beck et al., 2020 (http://www.gloh2o.org/hbv/).
     -   (HBV-light Version 2) Seibert, 2005
         (https://www.geo.uzh.ch/dam/jcr:c8afa73c-ac90-478e-a8c7-929eed7b1b62/HBV_manual_2005.pdf).
@@ -48,14 +48,16 @@ def __init__(
         self.dynamic_params = []
         self.dy_drop = 0.0
         self.variables = ['prcp', 'tmean', 'pet']
-        self.routing = True
+        self.routing = False
+        self.lenF = 15
         self.comprout = False
+        self.muwts = None
         self.nearzero = 1e-5
         self.nmul = 1
         self.cache_states = False
         self.device = device
 
-        self.states, self._states_cache = None, None
+        self.states, self._state_cache = None, None
 
         self.state_names = [
             'SNOWPACK',  # Snowpack storage
@@ -124,7 +126,7 @@ def __init__(
             self.comprout = config.get('comprout', self.comprout)
             self.nearzero = config.get('nearzero', self.nearzero)
             self.nmul = config.get('nmul', self.nmul)
-            self.cache_states = config.get('cache_states', False)
+            self.cache_states = config.get('cache_states', self.cache_states)
         self._set_parameters()
 
     def _init_states(self, ngrid: int) -> tuple[torch.Tensor]:
@@ -145,7 +147,7 @@ def get_states(self) -> Optional[tuple[torch.Tensor, ...]]:
         tuple[torch.Tensor, ...]
             A tuple containing the states (SNOWPACK, MELTWATER, SM, SUZ, SLZ).
         """
-        return self._states_cache
+        return self._state_cache
 
     def load_states(
         self,
@@ -380,10 +382,10 @@ def forward(
         )
 
         # State caching
-        self._states_cache = [s.detach() for s in states]
+        self._state_cache = states
 
         if self.cache_states:
-            self.states = self._states_cache
+            self.states = tuple(s[-1].detach() for s in self._state_cache)
 
         return fluxes
 
@@ -398,6 +400,8 @@ def _PBM(
     ) -> Union[tuple, dict[str, torch.Tensor]]:
         """Run through process-based model (PBM).
 
+        Flux outputs are in mm/day.
+
         Parameters
         ----------
         forcing
@@ -449,6 +453,13 @@ def _PBM(
         SWE_sim = torch.zeros(Pm.size(), dtype=torch.float32, device=self.device)
         capillary_sim = torch.zeros(Pm.size(), dtype=torch.float32, device=self.device)
 
+        # NOTE: new for MTS -- Save model states for all time steps.
+        SNOWPACK_sim = torch.zeros(Pm.size(), dtype=torch.float32, device=self.device)
+        MELTWATER_sim = torch.zeros(Pm.size(), dtype=torch.float32, device=self.device)
+        SM_sim = torch.zeros(Pm.size(), dtype=torch.float32, device=self.device)
+        SUZ_sim = torch.zeros(Pm.size(), dtype=torch.float32, device=self.device)
+        SLZ_sim = torch.zeros(Pm.size(), dtype=torch.float32, device=self.device)
+
         param_dict = {}
         for t in range(nsteps):
             # Get dynamic parameter values per timestep.
@@ -541,6 +552,7 @@ def _PBM(
             Q2 = param_dict['parK2'] * SLZ
             SLZ = SLZ - Q2
 
+            # --- Outputs ---
             Qsimmu[t, :, :] = Q0 + Q1 + Q2
             Q0_sim[t, :, :] = Q0
             Q1_sim[t, :, :] = Q1
@@ -555,6 +567,13 @@ def _PBM(
             tosoil_sim[t, :, :] = tosoil
             PERC_sim[t, :, :] = PERC
 
+            # NOTE: new for MTS -- Save model states for all time steps.
+            SNOWPACK_sim[t, :, :] = SNOWPACK
+            MELTWATER_sim[t, :, :] = MELTWATER
+            SM_sim[t, :, :] = SM
+            SUZ_sim[t, :, :] = SUZ
+            SLZ_sim[t, :, :] = SLZ
+
         # Get the average or weighted average using learned weights.
         if self.muwts is None:
             Qsimavg = Qsimmu.mean(-1)
@@ -574,7 +593,7 @@ def _PBM(
             UH = uh_gamma(
                 self.routing_param_dict['route_a'].repeat(nsteps, 1).unsqueeze(-1),
                 self.routing_param_dict['route_b'].repeat(nsteps, 1).unsqueeze(-1),
-                lenF=15,
+                lenF=self.lenF,
             )
             rf = torch.unsqueeze(Qsim, -1).permute([1, 2, 0])  # [gages,vars,time]
             UH = UH.permute([1, 2, 0])  # [gages,vars,time]
@@ -603,11 +622,11 @@ def _PBM(
             Qs = torch.unsqueeze(Qsimavg, -1)
             Q0_rout = Q1_rout = Q2_rout = None
 
-        states = (SNOWPACK, MELTWATER, SM, SUZ, SLZ)
+        states = (SNOWPACK_sim, MELTWATER_sim, SM_sim, SUZ_sim, SLZ_sim)
 
         if self.initialize:
             # If initialize is True, only return warmed-up storages.
-            return states
+            return {}, states
         else:
             # Baseflow index (BFI) calculation
             BFI_sim = (