Add persistent state caching (Fixes #31)

Co-authored-by: ZhennanShi1 <[email protected]>

Author: leoglonz

.github/PULL_REQUEST_TEMPLATE.md

@@ -5,8 +5,6 @@
 
 ## Description
 
-Fixes # (issue number)
-
 <!-- Describe how you addressed the bug/feature request, what choices you made and why. Changes can be listed as bullet point;
 
 - ...

.github/workflows/lint.yaml

@@ -12,14 +12,14 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: ['3.12', '3.13']
+        python-version: ['3.9', '3.12', '3.13']
 
     steps:
       - name: Checkout code
         uses: actions/[email protected]
 
       - name: Install uv + Python
-        uses: astral-sh/setup-uv@v6.7.0
+        uses: astral-sh/setup-uv@v7.1.2
         with:
           python-version: ${{ matrix.python-version }}
           enable-cache: true

.github/workflows/pytest.yaml

@@ -12,14 +12,14 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: ['3.12', '3.13']
+        python-version: ['3.9', '3.12', '3.13']
 
     steps:
       - name: Checkout code
         uses: actions/[email protected]
-  
+
       - name: Install uv + Python
-        uses: astral-sh/setup-uv@v6.7.0
+        uses: astral-sh/setup-uv@v7.1.2
         with:
           python-version: ${{ matrix.python-version }}
           enable-cache: true

.pre-commit-config.yaml

@@ -1,17 +1,27 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v6.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-yaml
-        exclude: "mkdocs.*\\.yml$"  # Exclude mkdocs YAML files
+        exclude: "mkdocs.*\\.yml$"
       - id: debug-statements
+      - id: check-toml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.5  # Use whatever version you're currently using
+    rev: v0.14.2
     hooks:
-      - id: ruff
-        args: [--fix]
-        exclude: "hooks/.*\\.py$"  # Exclude the hooks directory from ruff
+      - id: ruff-check
+        types_or: [python, pyi, jupyter]
+        args: [--fix, --config, pyproject.toml]
       - id: ruff-format
+        types_or: [python, pyi, jupyter]
+        args: [--config, pyproject.toml]
+
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.8.1
+    hooks:
+      - id: nbstripout
+        types: [jupyter]
+        language_version: python3

README.md

@@ -1,7 +1,7 @@
 # HydroDL2: Differentiable Hydrological Models
 
-[![Python](https://img.shields.io/badge/python-3.12%20%7C%203.13-blue)](https://www.python.org/downloads/)
-[![PyTorch](https://img.shields.io/badge/PyTorch-2.8.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)](https://www.python.org/downloads/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.9.0-EE4C2C?logo=pytorch)](https://pytorch.org/)
 
 [![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)
@@ -39,15 +39,15 @@ Optionally, add flag `-e` to install in editable mode.
     ```text
     .
     ├── src/
-    |   └── hydrodl2/ 
+    |   └── hydrodl2/
     │       ├── api/                   # Main API
-    │       |   ├── __init__.py        
+    │       |   ├── __init__.py
     │       |   └── methods.py         # Methods exposed to end-users
     |       ├── core/                  # Methods used internally
     │       ├── models/                # Shared models directory
     │       |   └── hbv/               # HBV models
     |       └── modules/               # Augmentations for δMG models
-    └── docs/                          
+    └── docs/
     ```
 
 </br>

docs/CONTRIBUTING.md

@@ -1,34 +1,40 @@
-# Contributing to *hydroDL2*
+# Contributing to *HydroDL2*
 
 Thank you for considering contributing to this project! Whether it's fixing a bug, improving documentation, or adding a new feature, we welcome contributions.
 
-We have a minimal set of standards we would ask you to consider to speed up the review process.
+There is a minimal set of standards we would ask you to consider to speed up the review process.
 
 ## 🧭 How to Contribute
 
 1. **Fork the repository**
-   - If you have not already done so, create a fork of the hydroDL2 repo (master branch) and make changes to this copy.
+   - If you have not already done so, create a fork of the `hydrodl2` repo (master branch) and make changes to this copy.
 
 2. **Lint & test your code**
-   - If you have not already, install development packages for hydroDL2.
+   - Make sure development packages for HydroDL2 are installed. This can be done by flagging dev packages during pip install like `uv pip install "./hydrodl2[dev]`.
 
-   - When your changes are ready, run
+   - Once your changes are complete, run the following in your Python environement:
 
       ```bash
       cd ./hydrodl2
-      
-      isort .
-      
-      ruff check .
 
       pytest tests
+
+      pre-commit install
+
+      git add .
+
+      git commit -m 'your commit message'
       ```
 
-   - If there ruff or pytest report any errors, please try to correct these if possible. (We can also help in the next step).
+     Upon committing, pre-commit will run a series of checks according to `.pre-commit-config.yaml` lint and format your code. This will block your commit if changes are made or requested. If manual changes are required, you will be notified. If only automatic changes are made, simply perform the git add and commit once more to push your code.
+
+     Note: if pytest does not work, try `python -m pytest`.
+
+   - If ruff or pytest report any errors, please try to correct these if possible. Otherwise, do `pre-commit uninstall` to proceed with committing your code and we can help in the next step.
 
 3. **Make a pull request (PR)**
-    - When you are ready make a PR to the hydroDL2 repository master branch.
+    - When you are ready, make a PR of your fork to the HydroDL2 repository master branch.
 
-    - In the PR description, make sure to include enough detail so that we can understand the changes made and rationale if necessary.
+    - In the PR description, include enough detail so that we can understand the changes made and rationale if necessary.
 
-    - If hydroDL2 master branch has new commits not included in your forked version, we would ask you to merge these new changes into your fork before we accept the PR. We can assist with this if necessary.
+    - If the HydroDL2 master branch has new commits not included in your forked version, we would ask you to merge these new changes into your fork before we accept the PR. We can assist with this if necessary.

docs/adding_models_and_modules.md

@@ -1,12 +1,12 @@
-# Adding Models and Modules to *hydroDL2*
+# Adding Models and Modules to *HydroDL2*
 
 We illustrate this with a hydrology model HBV 1.2.
 
 - In the `models/` directory, create a folder for the model type if it does not already exist, using only lowercase.
   - e.g., we should create `models/hbv/`, or confirm it exists.
 
 - Within `models/<your_model>/`, create a `.py` file for the model taking the name that will be exposed to users, converting to lowercase.
-  - e.g., HBV 1.2 goes to `hbv1_2.py`.
+  - e.g., HBV 1.2 goes to `hbv_1_2.py`.
 
 - The model file should only contain one model class. If you want to place multiple variants of your model in the same file, make sure to add a variant specification to your config and indicate
 with the `ver_name` flag when loading your model with `load_model()`. Otherwise, the first model listed in your file will be loaded by default.

docs/index.md

@@ -17,14 +17,13 @@ authors = [
 maintainers = [
     { name = "Leo Lonzarich", email = "[email protected]" },
 ]
-requires-python = ">=3.12.0"
+requires-python = ">=3.9.0"
 dynamic = ["version"]
 dependencies = [
-    "numpy>=2.3.3",
-    "pre-commit>=4.3.0",
-    "sourcedefender>=15.0.14",
-    "torch>=2.8.0",
-    "uv>=0.8.22",
+    "numpy",
+    "sourcedefender",
+    "torch",
+    "uv",
 ]
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -34,9 +33,10 @@ classifiers = [
 
 [project.optional-dependencies]
 dev = [
-    "ipywidgets>=8.1.7",
-    "pytest>=8.4.2",
-    "ruff>=0.13.2",
+    "ipywidgets",
+    "pre-commit",
+    "pytest",
+    "ruff",
 ]
 
 [project.urls]

pyproject.toml

@@ -18,11 +18,10 @@ def change_param_range(param: torch.Tensor, bounds: list[float]) -> torch.Tensor
 
     Returns
     -------
-    out
+    torch.Tensor
         The parameter with the specified bounds.
     """
-    out = param * (bounds[1] - bounds[0]) + bounds[0]
-    return out
+    return param * (bounds[1] - bounds[0]) + bounds[0]
 
 
 def param_bounds_2d(
@@ -49,19 +48,12 @@ def param_bounds_2d(
 
     Returns
     -------
-    out : torch.Tensor
+    torch.Tensor
         The 3D parameter array.
     """
     out_temp = (
         params[:, num * nmul : (num + 1) * nmul] * (bounds[1] - bounds[0]) + bounds[0]
     )
-    out = (
-        out_temp.unsqueeze(0)
-        .repeat(ndays, 1, 1)
-        .reshape(
-            ndays,
-            params.shape[0],
-            nmul,
-        )
+    return (
+        out_temp.unsqueeze(0).repeat(ndays, 1, 1).reshape(ndays, params.shape[0], nmul)
     )
-    return out