File size: 2,667 Bytes
186a48a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import json
from pathlib import Path
import numpy as np
import torch
from torch import nn
def load_config(path):
with open(path, encoding="utf-8") as handle:
return json.load(handle)
def write_json(path, value):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as handle:
json.dump(value, handle, indent=2, sort_keys=True)
class StreamflowLSTM(nn.Module):
"""Paper stack with explicit inter-layer activation semantics."""
def __init__(self, input_size=23, hidden_size=50, dropout=0.1):
super().__init__()
self.layers = nn.ModuleList([
nn.LSTM(input_size, hidden_size, batch_first=True),
nn.LSTM(hidden_size, hidden_size, batch_first=True),
nn.LSTM(hidden_size, hidden_size, batch_first=True),
])
self.dropout = nn.Dropout(dropout)
self.dense = nn.Linear(hidden_size, 1)
def forward(self, inputs):
if inputs.ndim != 3 or inputs.shape[-2:] != (28, 23):
raise ValueError(f"expected [B,28,23], got {tuple(inputs.shape)}")
values, _ = self.layers[0](inputs)
values = self.dropout(torch.relu(values))
values, _ = self.layers[1](values)
values = self.dropout(torch.relu(values))
values, _ = self.layers[2](values)
values = torch.tanh(values[:, -1])
return self.dense(values).squeeze(-1)
def nse(prediction, target):
prediction, target = np.asarray(prediction), np.asarray(target)
denominator = np.sum((target - target.mean()) ** 2)
return float(1.0 - np.sum((prediction - target) ** 2) / max(denominator, 1e-12))
def metrics(prediction, target):
prediction = np.asarray(prediction, dtype=np.float64)
target = np.asarray(target, dtype=np.float64)
correlation = float(np.corrcoef(prediction, target)[0, 1]) if prediction.size > 1 else 0.0
if not np.isfinite(correlation):
correlation = 0.0
alpha = float(prediction.std() / max(target.std(), 1e-12))
beta = float(prediction.mean() / max(target.mean(), 1e-12))
kge = float(1.0 - np.sqrt((correlation - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2))
return {
"kge": kge,
"kge_r": correlation,
"kge_alpha": alpha,
"kge_beta": beta,
"nse": nse(prediction, target),
"rmse_m3_s": float(np.sqrt(np.mean((prediction - target) ** 2))),
}
def load_member(payload, device):
model = StreamflowLSTM(hidden_size=payload["hidden_size"], dropout=payload["dropout"]).to(device)
model.load_state_dict(payload["state_dict"])
model.eval()
return model, payload
|