LSTMからハイブリッドへ：降雨流出解析プロトタイプ作成 Step 5 微分可能タンクモデルをPyTorchで実装する

はじめに

前回のStep 4では、物理モデル出力をXGBoostやLSTMで後処理する構成を作成しました。

今回はさらに発展させて、タンクモデル自体をPyTorchで微分可能に実装します。

作成するプロトタイプ名は以下です。

RainfallRunoffDifferentiableTank_r001

参考のために、ZIPファイルを置いておきますので、ダウンロードして活用してください。

今回の位置づけは以下です。

第5段階：
  将来的に微分可能水文モデルや物理制約付きLSTMへ発展させる

内容：
  PyTorchで微分可能な降雨流出タンクモデル

微分可能水文モデルとは

微分可能水文モデルとは、タンクモデルやHBVのような物理・概念モデルを、PyTorchなどの自動微分フレームワーク上で実装し、観測流量との誤差を使ってモデルパラメータを学習できるようにしたモデルです。

通常のタンクモデルでは、流出係数や浸透係数を手動で調整します。

一方、微分可能タンクモデルでは、以下のような係数をPyTorchの学習可能パラメータとして扱います。

k_quick
k_perc
k_slow
evap_factor
upper_storage
lower_storage

今回のモデル構成

今回も、簡易2段タンクモデルを使います。

最終流量は以下です。

Q_total = Q_quick + Q_slow

通常のタンクモデルとの違い

項目	通常のタンクモデル	微分可能タンクモデル
パラメータ調整	手動キャリブレーション	勾配降下法で学習
実装	Python / Fortranなど	PyTorch
物理構造	残る	残る
自動微分	なし	あり
深層学習との接続	難しい	接続しやすい

学習対象パラメータ

今回学習するパラメータは以下です。

パラメータ	内容
`k_quick`	上段タンクからの速い流出係数
`k_perc`	上段タンクから下段タンクへの浸透係数
`k_slow`	下段タンクからの遅い流出係数
`evap_factor`	可能蒸発散量をどの程度反映するか
`upper_storage`	上段タンクの初期貯留量
`lower_storage`	下段タンクの初期貯留量

物理制約の入れ方

今回のプロトタイプでは、パラメータを直接自由に学習させず、範囲を設定します。

bounds:
  k_quick: [0.01, 0.90]
  k_perc: [0.001, 0.50]
  k_slow: [0.001, 0.30]
  evap_factor: [0.00, 1.20]
  upper_storage: [0.00, 300.0]
  lower_storage: [0.00, 500.0]

内部では、学習可能なrawパラメータを sigmoid で範囲内に写像します。

これにより、流出係数や貯留量が極端な値になりにくくなります。

また、流量や貯留量は負にならないように制約します。

フォルダー構成

RainfallRunoffDifferentiableTank_r001/
├─ docker/
│  ├─ Dockerfile
│  └─ requirements.txt
├─ docker-compose.yml
├─ .env
├─ configs/
│  ├─ diff_tank_config.yaml
│  └─ diff_tank_config_quick.yaml
├─ data/
│  ├─ raw/
│  └─ processed/
├─ src/
│  ├─ generate_sample_data.py
│  ├─ preprocess.py
│  ├─ model_reference_tank.py
│  ├─ model_differentiable_tank.py
│  ├─ train_differentiable_tank.py
│  ├─ evaluate.py
│  ├─ metrics.py
│  └─ utils.py
├─ outputs/
│  ├─ models/
│  ├─ figures/
│  ├─ metrics/
│  └─ predictions/
├─ scripts/
└─ docs/

実行方法

Docker起動

docker compose up --build

JupyterLabは以下です。

http://localhost:8888

簡易テスト

docker compose run --rm differentiable-tank bash scripts/run_quick_test.sh

このスクリプトでは以下を実行します。

1. サンプルデータ作成
2. 前処理
3. 微分可能タンクモデルの学習
4. 評価

個別実行

docker compose run --rm differentiable-tank python src/generate_sample_data.py
docker compose run --rm differentiable-tank python src/preprocess.py --config configs/diff_tank_config.yaml
docker compose run --rm differentiable-tank python src/train_differentiable_tank.py --config configs/diff_tank_config.yaml
docker compose run --rm differentiable-tank python src/evaluate.py --config configs/diff_tank_config.yaml

出力ファイル

outputs/
├─ models/
│  ├─ differentiable_tank_best.pth
│  ├─ training_history.csv
│  └─ learned_parameters.csv
├─ figures/
│  ├─ loss_curve_differentiable_tank.png
│  ├─ hydrograph_differentiable_tank_test.png
│  ├─ runoff_components_differentiable_tank.png
│  ├─ storage_states_differentiable_tank.png
│  └─ scatter_obs_pred_differentiable_tank.png
├─ metrics/
│  ├─ test_metrics_differentiable_tank.csv
│  └─ learned_parameters_summary.csv
└─ predictions/
   └─ test_predictions_differentiable_tank.csv

今回のポイント

Step 5のポイントは、物理モデルを後処理するだけでなく、物理モデル自体を学習対象にすることです。

Step 4:
  Q_phys をXGBoostやLSTMで補正する

Step 5:
  タンクモデルのパラメータ自体をPyTorchで学習する

つまり、LSTM単独のブラックボックスモデルではなく、物理構造を残したままデータに適応させる方向です。

注意点

今回のモデルは、かなり簡略化した微分可能タンクモデルです。

実務適用には、以下の検討が必要です。

実流域データでの検証
流域ごとのパラメータ範囲の設定
多段タンクモデルへの拡張
水収支制約の強化
積雪・融雪過程の追加
地下水・蒸発散の詳細化
洪水イベント別評価
未経験豪雨への外挿性確認
物理制約付きLSTMとの比較

まとめ

今回は、Step 5として、PyTorchで微分可能な降雨流出タンクモデルを作成しました。

RainfallRunoffDifferentiableTank_r001

このモデルでは、タンクモデルの構造を残しながら、流出係数・浸透係数・初期貯留量を自動微分で学習します。

次の発展としては、以下が考えられます。

1. 多段タンクモデルへの拡張する
2. 水収支制約をより強く入れる
3. 積雪・融雪過程を追加する
4. 流域属性からパラメータを推定する
5. 物理制約付きLSTMへ拡張する
6. 複数流域で学習する

最後にstep 0～ step 5までの記事を紹介しておきます。

LSTMからハイブリッドへ：降雨流出解析プロトタイプ作成 Step 1 LSTM単独モデルを実装する
LSTMからハイブリッドへ：降雨流出解析プロトタイプ作成 Step 2 水文特徴量を追加する
LSTMからハイブリッドへ：降雨流出解析プロトタイプ作成 Step 3 タンクモデル + LSTM残差補正を実装する
LSTMからハイブリッドへ：降雨流出解析プロトタイプ作成 Step 4 物理モデル出力をXGBoost / LSTMで後処理する
LSTMからハイブリッドへ：降雨流出解析プロトタイプ作成 Step 5 微分可能タンクモデルをPyTorchで実装する

参考リンク

本記事に関連する資料として、タンクモデル、PyTorchの自動微分、降雨流出LSTM、微分可能水文モデル、複数流域データセットを中心にまとめました。

Sugawara, M. (1985) - Tank Model: For the Derivation of River Discharge from Rainfall
タンクモデルの基本的な考え方を確認できる資料です。
PyTorch - Automatic differentiation package: torch.autograd
PyTorchで自動微分を扱うための公式ドキュメントです。
PyTorch - Sigmoid
学習パラメータを一定範囲へ写像するときに使用するSigmoid関数の公式ドキュメントです。
Kratzert et al. (2018) - Rainfall-runoff modelling using Long Short-Term Memory (LSTM) networks
LSTMを降雨流出解析へ適用した代表的な研究です。
Feng et al. (2022) - Differentiable, Learnable, Regionalized Process-Based Models With Multiphysical Outputs can Approach State-Of-The-Art Hydrologic Prediction Accuracy
HBVをベースに、物理モデルとニューラルネットワークを微分可能な形で組み合わせた研究です。
Shen et al. (2023) - Differentiable modelling to unify machine learning and physical models for geosciences
微分可能モデルを、物理モデルと機械学習をつなぐ枠組みとして整理したレビュー・展望論文です。
Feng et al. (2023) - The suitability of differentiable, physics-informed machine learning hydrologic models for ungauged regions and climate change impact assessment
未観測流域や気候変化条件への適用可能性を検討した研究です。
Song et al. (2024) - When ancient numerical demons meet physics-informed machine learning: adjoint-based gradients for implicit differentiable modeling
微分可能水文モデルを、陰解法や随伴法まで拡張した研究です。
NeuralHydrology Documentation
LSTMをはじめとした深層学習による水文モデルを扱うオープンソースフレームワークの公式ドキュメントです。
NCAR - CAMELS: Catchment Attributes and Meteorology for Large-sample Studies
米国671流域の気象・流量・流域属性を収録した代表的な大規模水文データセットです。