Publish Delta-HBV reproduction

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pt binary
+*.npz binary

conf/config.yaml

@@ -0,0 +1,6 @@
+seed: 42
+data: {path: data/camels.npz, basins: 12, days: 730, forcing_features: 3, static_features: 35}
+model: {static_features: 35, hidden: 16, components: 16}
+train: {epochs: 1, learning_rate: 0.001, warmup_days: 365}
+paths: {checkpoint: result/checkpoints/delta_hbv.pt, training_metrics: result/training/metrics.json, predictions: result/output/predictions.npz, evaluation: result/evaluation/metrics.json, figure: result/evaluation/comparison.png}
+paper_model: {basins: 671, reported_basins: 531, attributes: 35, components: 16, minibatch: 100, warmup_years: 1, PUB_folds: 10, PUR_regions: 7}

config.json

@@ -0,0 +1 @@
+{"model_name":"Delta-HBV","model_type":"delta_hbv","architectures":["DeltaHBV"],"framework":"PyTorch","domain":"hydrology","task":"ungauged-basin-and-trend-modeling","implementation":{"entry_point":"model/delta_hbv.py","scope":"L2 differentiable HBV reproduction"},"architecture":{"dynamic_inputs":["precipitation","temperature","PET"],"static_attributes":35,"parallel_components":16,"dynamic_parameters":["beta","gamma"],"states":["snow","soil","upper","lower"]},"data":{"source":"CAMELS","basins":671,"reported_basins":531,"period":["1989-10-01","1999-09-30"],"synthetic":true},"configuration_sources":["conf/config.yaml","model/delta_hbv.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/delta_hbv.py

@@ -0,0 +1,16 @@
+import json,torch
+from pathlib import Path
+from torch import nn
+import yaml
+def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
+def sample(i,days=730):
+ t=torch.arange(days).float();p=(.5+.5*torch.sin(t/20+i)).clamp_min(0);temp=10+12*torch.sin(t/58+i);pet=(2+.1*temp).clamp_min(0);f=torch.stack((p,temp,pet),1);a=torch.linspace(0,1,35)+i*.01;return f,a
+class DeltaHBV(nn.Module):
+ def __init__(self,static_features=35,hidden=16,components=16):super().__init__();self.components=components;self.rnn=nn.LSTM(3+static_features,hidden,batch_first=True);self.head=nn.Linear(hidden,components*4);self.model_config={'static_features':static_features,'hidden':hidden,'components':components}
+ def forward(self,forcing,attrs):
+  b,t,_=forcing.shape;x=torch.cat((forcing,attrs[:,None].expand(-1,t,-1)),-1);h,_=self.rnn(x);par=torch.sigmoid(self.head(h)).reshape(b,t,self.components,4);snow=torch.zeros(b,self.components);soil=torch.ones_like(snow);upper=torch.zeros_like(snow);lower=torch.zeros_like(snow);flows=[];ets=[]
+  for k in range(t):
+   p,temp,pet=forcing[:,k].unbind(-1);rain=p[:,None]*(temp[:,None]>0);snow=snow+p[:,None]-rain;melt=torch.minimum(snow,torch.relu(temp[:,None])*.02);snow=snow-melt;water=rain+melt;beta=1+4*par[:,k,:,0];gamma=.5+2*par[:,k,:,1];fc=1+4*par[:,k,:,2];recharge=torch.minimum((soil/fc).clamp(0,1).pow(beta)*water,water);soil=soil+water-recharge;et=torch.minimum((soil/fc).clamp(0,1).pow(gamma)*pet[:,None],soil);soil=soil-et;upper=upper+recharge;q0=upper*.2;upper=upper-q0;lower=lower+upper*.05;upper=upper*.95;q2=lower*.02;lower=lower-q2;flows.append((q0+q2).mean(1));ets.append(et.mean(1))
+  return torch.stack(flows,1),torch.stack(ets,1)
+def nse(p,t):return 1-((p-t)**2).sum()/((t-t.mean())**2).sum().clamp_min(1e-8)
+def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n')

scripts/fake_data.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+import sys,numpy as np,torch
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.delta_hbv import *
+c=cfg(R);fs=[];aa=[];qq=[];ee=[];teacher=DeltaHBV()
+with torch.no_grad():
+ for i in range(12):f,a=sample(i);q,e=teacher(f[None],a[None]);fs.append(f);aa.append(a);qq.append(q[0]);ee.append(e[0])
+p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,forcing=np.stack(fs),attributes=np.stack(aa),streamflow=np.stack(qq),et=np.stack(ee),basin_ids=np.arange(12));print(p)

scripts/inference.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+import sys,numpy as np,torch
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.delta_hbv import *
+c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);m=DeltaHBV(**z['model_config']);m.load_state_dict(z['model'])
+with torch.no_grad():q,e=m(torch.tensor(d['forcing']),torch.tensor(d['attributes']))
+p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,streamflow=q.numpy(),target=d['streamflow'],et=e.numpy(),et_target=d['et'],basin_ids=d['basin_ids']);print(p)

scripts/result.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+import sys,numpy as np,torch;import matplotlib;matplotlib.use('Agg');import matplotlib.pyplot as plt
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.delta_hbv import cfg,nse,write
+c=cfg(R);d=np.load(R/c['paths']['predictions']);p=d['streamflow'][:,365:];t=d['target'][:,365:];scores=[float(nse(torch.tensor(p[i]),torch.tensor(t[i]))) for i in range(12)];et_rmse=float(np.sqrt(np.mean((d['et']-d['et_target'])**2)));trend_p=np.polyfit(np.arange(p.shape[1]),p.mean(0),1)[0];trend_t=np.polyfit(np.arange(t.shape[1]),t.mean(0),1)[0];write(R/c['paths']['evaluation'],{'median_nse':float(np.median(scores)),'et_rmse':et_rmse,'trend_error':float(abs(trend_p-trend_t)),'synthetic':True});plt.plot(t[0],label='target');plt.plot(p[0],label='prediction');plt.legend();q=R/c['paths']['figure'];q.parent.mkdir(parents=True,exist_ok=True);plt.savefig(q,dpi=150);print(q)

scripts/train.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+import sys,os,numpy as np,torch;import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.delta_hbv import *
+c=cfg(R);rank=int(os.getenv('RANK',0));world=int(os.getenv('WORLD_SIZE',1));ddp=world>1
+if ddp:dist.init_process_group('gloo')
+d=np.load(R/c['data']['path']);torch.set_num_threads(2);base=DeltaHBV(**c['model']);m=DDP(base) if ddp else base;opt=torch.optim.Adam(m.parameters(),lr=c['train']['learning_rate']);ids=np.arange(rank,12,world);q,_=m(torch.tensor(d['forcing'][ids]),torch.tensor(d['attributes'][ids]));t=torch.tensor(d['streamflow'][ids]);loss=torch.sqrt(((q[:,365:]-t[:,365:])**2).mean()+1e-8);opt.zero_grad();loss.backward();opt.step();ls=[float(loss)]
+v=torch.tensor([sum(ls),len(ls)],dtype=torch.float64)
+if ddp:dist.all_reduce(v)
+p=R/c['paths']['checkpoint']
+if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({'model':base.state_dict(),'model_config':c['model']},p);write(R/c['paths']['training_metrics'],{'rmse':float(v[0]/v[1]),'world_size':world});print(p)
+if ddp:dist.destroy_process_group()