Alignment to dev

.pre-commit-config.yaml

@@ -10,17 +10,17 @@ repos:
       - id: check-toml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.8
+    rev: v0.16.7
     hooks:
       - id: ruff-check
-        types_or: [python, pyi, jupyter]
+        types_or: [python, pyi, jupyter, markdown]
         args: [--fix, --config, pyproject.toml]
       - id: ruff-format
-        types_or: [python, pyi, jupyter]
+        types_or: [python, pyi, jupyter, markdown]
         args: [--config, pyproject.toml]
 
   - repo: https://github.com/kynan/nbstripout
-    rev: 0.8.2
+    rev: 0.9.1
     hooks:
       - id: nbstripout
         types: [jupyter]

docs/setup.md

@@ -80,6 +80,7 @@ git clone https://github.com/mhpi/hydrodl2.git
 
   ```python
   import torch
+
   print(torch.cuda.is_available())
   ```
 

pyproject.toml

@@ -15,6 +15,7 @@ authors = [
 ]
 maintainers = [
     { name = "Leo Lonzarich", email = "lglonzaric@gmail.com" },
+    { name = "Yalan Song", email = "songyalan1@gmail.com" },
 ]
 requires-python = ">=3.9.0"
 dynamic = ["version"]

src/hydrodl2/api/__init__.py

@@ -1,3 +1,4 @@
+# src/hydrodl2/api/__init__.py
 from .methods import available_models, available_modules, load_model, load_module
 
 __all__ = ['available_models', 'available_modules', 'load_model', 'load_module']

src/hydrodl2/api/methods.py

@@ -1,7 +1,4 @@
-"""
-Note: If adding new public methods, please add them to __all__
-at the top of the file and in api/__init__.py.
-"""
+"""API methods for loading and managing models in HydroDL2."""
 
 import importlib.util
 import logging
@@ -117,24 +114,49 @@ def load_model(model: str, ver_name: str = None) -> Module:
     except ImportError as e:
         raise ImportError(f"Model '{model}' not found.") from e
 
+    # Classes defined directly in this module (excludes imported classes like
+    # BasePhysicsModel or, for files like hbv_2_mts.py, other Hbv variants).
+    local_classes = {
+        attr: getattr(module, attr)
+        for attr in dir(module)
+        if isinstance(getattr(module, attr), type)
+        and getattr(module, attr).__module__ == module.__name__
+    }
+
     # Retrieve version name if possible, otherwise get first class in module
     try:
         cls = getattr(module, ver_name)
-    except AttributeError as e:
-        # Find first class in module (NOTE: not guaranteed accurate)
-        classes = [
-            attr
-            for attr in dir(module)
-            if isinstance(getattr(module, attr), type) and attr != 'Any'
+    except AttributeError:
+        # Try case-insensitive match (e.g., hbv_2_hourly -> Hbv_2_hourly)
+        ci_matches = [
+            name for name in local_classes if name.lower() == ver_name.lower()
         ]
-        if not classes:
-            raise ImportError(f"Model version '{model}' not found.") from e
-
-        log.warning(
-            f"Model class '{ver_name}' not found in module '{module.__file__}'. "
-            f"Falling back to the first available: '{classes[0]}'."
-        )
-        cls = getattr(module, classes[0])
+        if ci_matches:
+            cls = local_classes[ci_matches[0]]
+        else:
+            # Try PascalCase conversion (e.g., Prms_Gw0_Triton -> PrmsGw0Triton)
+            pascal_name = ''.join(w.title() for w in ver_name.split('_'))
+            try:
+                cls = getattr(module, pascal_name)
+            except AttributeError as e:
+                # Find first locally-defined class in module (NOTE: not guaranteed
+                # accurate). Prefer local_classes over dir(module) so that classes
+                # merely imported for use by the module's own class(es) - e.g.
+                # BasePhysicsModel, or Hbv_2/Hbv_2_hourly imported into
+                # hbv_2_mts.py - are never picked over the module's own class.
+                classes = sorted(local_classes) or [
+                    attr
+                    for attr in dir(module)
+                    if isinstance(getattr(module, attr), type) and attr != 'Any'
+                ]
+                if not classes:
+                    raise ImportError(f"Model version '{model}' not found.") from e
+
+                log.warning(
+                    f"Model class '{ver_name}' not found in module '{module.__file__}'. "
+                    f"Falling back to the first available: '{classes[0]}'."
+                )
+                cls = getattr(module, classes[0])
 
     return cls
 

src/hydrodl2/core/calc/__init__.py

@@ -1,10 +1,11 @@
 from .fdj import finite_difference_jacobian_p
 from .uh_routing import uh_conv, uh_gamma
-from .utils import change_param_range, param_bounds_2d
+from .utils import change_param_range, param_bounds_2d, trim_warmup
 
 __all__ = [
     'change_param_range',
     'param_bounds_2d',
+    'trim_warmup',
     'uh_gamma',
     'uh_conv',
     'finite_difference_jacobian_p',

src/hydrodl2/core/calc/utils.py

@@ -57,3 +57,47 @@ def param_bounds_2d(
     return (
         out_temp.unsqueeze(0).repeat(ndays, 1, 1).reshape(ndays, params.shape[0], nmul)
     )
+
+
+def trim_warmup(
+    outputs: dict[str, torch.Tensor],
+    pred_cutoff: int,
+    nsteps: int,
+) -> dict[str, torch.Tensor]:
+    """Drop warm-up timesteps from a model's time-major outputs.
+
+    Models spin up their internal states over a warm-up window that should
+    not appear in predictions. Depending on the warm-up strategy, that window
+    is either simulated separately (and never enters the outputs) or simulated
+    inline and removed afterwards. This helper performs the removal so that a
+    model's ``forward`` always returns the same number of timesteps either way.
+
+    Only time-major tensors are trimmed — those whose leading dimension is
+    ``nsteps``. Outputs that have already collapsed the time axis (a baseflow
+    index summed over time, for instance) are passed through untouched, so
+    callers do not have to maintain a list of exceptions.
+
+    Parameters
+    ----------
+    outputs
+        Dictionary of model outputs.
+    pred_cutoff
+        Number of leading timesteps to drop. Values <= 0 are a no-op.
+    nsteps
+        Length of the simulated window, used to identify time-major tensors.
+
+    Returns
+    -------
+    dict[str, torch.Tensor]
+        Outputs with the warm-up period removed from every time-major tensor.
+    """
+    if pred_cutoff <= 0:
+        return outputs
+
+    trimmed = {}
+    for key, value in outputs.items():
+        is_time_major = (
+            torch.is_tensor(value) and value.ndim >= 1 and value.shape[0] == nsteps
+        )
+        trimmed[key] = value[pred_cutoff:] if is_time_major else value
+    return trimmed

src/hydrodl2/core/utils/__init__.py

@@ -1,3 +1,4 @@
+# src/hydrodl2/core/utils/__init__.py
 from .utils import _get_dir, get_model_dirs, get_model_files
 
 __all__ = [

src/hydrodl2/core/utils/utils.py

@@ -1,7 +1,4 @@
-"""
-Note: If adding new public methods, please add them to __all__
-at the top of the file and in utils/__init__.py.
-"""
+"""General utility functions for HydroDL2."""
 
 import os
 from pathlib import Path

src/hydrodl2/models/base.py

@@ -0,0 +1,161 @@
+from typing import Any, Optional
+
+import torch
+
+from hydrodl2.core.calc import change_param_range, trim_warmup
+
+
+class BasePhysicsModel(torch.nn.Module):
+    """Common code for hydrodl2-based models.
+
+    Subclasses keep their own `__init__` and should set their own
+    1. state_names
+    2. nmul
+    3. device
+    4. warmup
+    5. warmup_states
+    6. parameter_bounds
+    7. routing_parameter_bounds
+
+    Forward: returns a timeseries dict (nsteps - warmup); warmup is stripped
+    inside the model -- either simulate it separately (warmup_states=True) or
+    run the full time window and drop the lead.
+    """
+
+    #: Nonzero initial states for safe powers/divisions.
+    initial_state_value: float = 0.001
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.states: Optional[tuple[torch.Tensor, ...]] = None
+        self._state_cache: Optional[tuple[torch.Tensor, ...]] = None
+
+    @staticmethod
+    def trim_warmup(
+        outputs: dict[str, torch.Tensor],
+        pred_cutoff: int,
+        nsteps: int,
+    ) -> dict[str, torch.Tensor]:
+        """Drop `pred_cutoff` leading steps from timeseries outputs."""
+        return trim_warmup(outputs, pred_cutoff, nsteps)
+
+    def _init_states(self, ngrid: int) -> tuple[torch.Tensor, ...]:
+        """One [ngrid, nmul] tensor per state with initial_state_value."""
+
+        def make_state():
+            return torch.full(
+                (ngrid, self.nmul),
+                self.initial_state_value,
+                dtype=torch.float32,
+                device=self.device,
+            )
+
+        return tuple(make_state() for _ in range(len(self.state_names)))
+
+    def get_states(self) -> Optional[tuple[torch.Tensor, ...]]:
+        """States cached by the last forward pass, or None if no yet run."""
+        return self._state_cache
+
+    def load_states(self, states: tuple[torch.Tensor, ...]) -> None:
+        """Load states into the model."""
+        for state in states:
+            if not isinstance(state, torch.Tensor):
+                raise ValueError("Each element in states must be a tensor.")
+        nstates = len(self.state_names)
+        if not (isinstance(states, tuple) and len(states) == nstates):
+            raise ValueError(f"States must be a tuple of {nstates} tensors.")
+
+        self.states = tuple(
+            s.detach().to(self.device, dtype=torch.float32) for s in states
+        )
+
+    def _descale_route_parameters(
+        self,
+        routing_params: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        """Map normalized routing parameters onto their physical ranges.
+
+        Shape: [ngrid, n_route, nmul]
+
+        Parameters
+        ----------
+        routing_params
+            Normalized routing parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary of descaled routing parameters.
+        """
+        parameter_dict = {}
+        for i, name in enumerate(self.routing_parameter_bounds.keys()):
+            parameter_dict[name] = change_param_range(
+                param=routing_params[:, i],
+                bounds=self.routing_parameter_bounds[name],
+            )
+        return parameter_dict
+
+    def _descale_phy_dy_parameters(
+        self,
+        phy_dy_params: torch.Tensor,
+        dy_list: list[str],
+    ) -> dict[str, torch.Tensor]:
+        """Descale the time-dynamic physical parameters.
+
+        Shape: [nsteps, ngrid, n_dynamic, nmul]
+
+        Parameters
+        ----------
+        phy_dy_params
+            Normalized dynamic physical parameters.
+        dy_list
+            List of dynamic parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary of descaled physical parameters.
+        """
+        raise NotImplementedError
+
+    def _descale_phy_stat_parameters(
+        self,
+        phy_stat_params: torch.Tensor,
+        stat_list: list[str],
+    ) -> dict[str, torch.Tensor]:
+        """Descale the time-invariant physical parameters.
+
+        Shape: [ngrid, n_static, nmul]
+
+        Parameters
+        ----------
+        phy_stat_params
+            Normalized static physical parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary of descaled static physical parameters.
+        """
+        raise NotImplementedError
+
+    def forward(
+        self,
+        x_dict: dict[str, torch.Tensor],
+        parameters: Any,
+    ) -> dict[str, torch.Tensor]:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x_dict
+            Dictionary of input forcing data.
+        parameters
+            Unprocessed, learned parameters from a neural network.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            Dictionary of model outputs.
+        """
+        raise NotImplementedError