Publish pysteps 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,5 @@
+seed: 42
+data: {path: data/radar.npz, shape: [3, 128, 128]}
+model: {levels: 8, members: 24, steps: 12}
+train: {epochs: 1}
+paths: {checkpoint: result/checkpoints/pysteps.pt, training_metrics: result/training/metrics.json, predictions: result/output/predictions.npz, evaluation: result/evaluation/metrics.json, figure: result/evaluation/comparison.png}

config.json

@@ -0,0 +1 @@
+{"model_name":"pysteps","model_type":"pysteps","architectures":["STEPS"],"framework":"PyTorch","domain":"precipitation-nowcasting","task":"probabilistic-nowcasting","implementation":{"entry_point":"model/pysteps.py","scope":"L2 STEPS reproduction"},"architecture":{"history_frames":3,"cascade_levels":8,"ar_order":2,"ensemble_members":24,"forecast_steps":12},"data":{"interval_minutes":5,"threshold_mm_h":0.1,"synthetic":true},"configuration_sources":["conf/config.yaml","model/pysteps.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/pysteps.py

@@ -0,0 +1,13 @@
+import json,torch
+from pathlib import Path
+import yaml
+def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
+def frames():
+ y,x=torch.meshgrid(torch.arange(128),torch.arange(128),indexing='ij');return torch.stack([torch.exp(-((x-45-t*3)**2+(y-60-t*2)**2)/300)*20 for t in range(3)])
+class STEPS(torch.nn.Module):
+ def __init__(self,levels=8,members=24,steps=12):super().__init__();self.levels=levels;self.members=members;self.steps=steps;self.scale=torch.nn.Parameter(torch.ones(levels))
+ def forward(self,x):
+  velocity=x[-1]-x[-2];base=x[-1];out=[]
+  for m in range(self.members):out.append(torch.stack([(base+s*velocity+torch.randn_like(base)*.02*self.scale.mean()).clamp_min(0) for s in range(1,self.steps+1)]))
+  return torch.stack(out)
+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,4 @@
+from pathlib import Path
+import sys,numpy as np
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.pysteps import *
+c=cfg(R);p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,input=frames(),target=torch.stack([torch.roll(frames()[-1],(2*s,3*s),(0,1)) for s in range(1,13)]));print(p)

scripts/inference.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+import sys,numpy as np,torch
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.pysteps import *
+c=cfg(R);z=torch.load(R/c['paths']['checkpoint'],weights_only=True);m=STEPS(**z['model_config']);m.load_state_dict(z['model']);d=np.load(R/c['data']['path']);p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,ensemble=m(torch.tensor(d['input'])).detach(),target=d['target'],lead_minutes=np.arange(1,13)*5);print(p)

scripts/result.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+import sys,numpy as np;import matplotlib;matplotlib.use('Agg');import matplotlib.pyplot as plt
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.pysteps import cfg,write
+c=cfg(R);d=np.load(R/c['paths']['predictions']);e=d['ensemble'];t=d['target'];mean=e.mean(0);rmse=np.sqrt(np.mean((mean-t)**2,(1,2)));write(R/c['paths']['evaluation'],{'rmse':rmse.tolist(),'spread':e.std(0).mean((1,2)).tolist(),'synthetic':True});plt.plot(d['lead_minutes'],rmse);q=R/c['paths']['figure'];q.parent.mkdir(parents=True,exist_ok=True);plt.savefig(q);print(q)

scripts/train.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+import sys,os,torch
+import torch.distributed as dist
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.pysteps 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')
+m=STEPS(**c['model']);p=R/c['paths']['checkpoint']
+if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({'model':m.state_dict(),'model_config':c['model']},p);write(R/c['paths']['training_metrics'],{'online_parameter_estimation':True,'world_size':world});print(p)
+if ddp:dist.destroy_process_group()