Publish WeatherBench2 reproduction

.gitattributes

@@ -1,35 +1 @@
-*.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

conf/config.yaml

@@ -0,0 +1,5 @@
+seed: 42
+data: {path: data/data.npz, samples: 8, variables: 8, grid: [24, 48]}
+model: {variables: 8, hidden: 16}
+train: {lr: 0.001}
+paths: {checkpoint: result/checkpoints/weatherbench2.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":"WeatherBench2","model_type":"weatherbench2","architectures":["WB2Baseline"],"framework":"PyTorch","domain":"weather-benchmark","task":"forecast-evaluation","implementation":{"entry_point":"model/weatherbench2.py"},"architecture":{"headline_variables":8,"grid":[120,240],"ensemble_members":8},"data":{"source":"ERA5","evaluation_year":2020},"configuration_sources":["conf/config.yaml","model/weatherbench2.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/weatherbench2.py

@@ -0,0 +1,8 @@
+import json,torch,yaml
+from pathlib import Path
+from torch import nn
+def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
+class WB2Baseline(nn.Module):
+ def __init__(self,variables=8,hidden=16):super().__init__();self.n=nn.Sequential(nn.Conv2d(variables,hidden,3,padding=1),nn.ReLU(),nn.Conv2d(hidden,variables,3,padding=1));self.model_config={'variables':variables,'hidden':hidden}
+ def forward(self,x):return x+self.n(x)
+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,2 @@
+from pathlib import Path
+import numpy as np,yaml;R=Path(__file__).resolve().parents[1];c=yaml.safe_load((R/'conf/config.yaml').read_text());g=np.random.default_rng(3);x=g.normal(size=(8,8,24,48)).astype('f');y=x+.1;p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez(p,input=x,target=y,latitude=np.linspace(-90,90,24));print(p)

scripts/inference.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+import sys,numpy as np,torch;R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.weatherbench2 import *
+c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],weights_only=True);m=WB2Baseline(**z['model_config']);m.load_state_dict(z['model']);base=m(torch.tensor(d['input'])).detach().numpy();ens=np.stack([base+np.random.default_rng(i).normal(0,.02,base.shape) for i in range(8)]);p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez(p,ensemble=ens,target=d['target'],latitude=d['latitude']);print(p)

scripts/result.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+import sys,numpy as np,matplotlib;matplotlib.use('Agg');import matplotlib.pyplot as plt;R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.weatherbench2 import *
+c=cfg(R);d=np.load(R/c['paths']['predictions']);e=d['ensemble'];t=d['target'];mean=e.mean(0);rmse=float(np.sqrt(np.mean((mean-t)**2)));crps=float(np.mean(abs(e-t[None]))-.5*np.mean(abs(e[:,None]-e[None,:])));spread=float(e.std(0).mean());write(R/c['paths']['evaluation'],{'rmse':rmse,'crps':crps,'spread_skill':spread/(rmse+1e-8)});plt.imshow(mean[0,0]-t[0,0]);p=R/c['paths']['figure'];p.parent.mkdir(parents=True,exist_ok=True);plt.savefig(p);print(p)

scripts/train.py

@@ -0,0 +1,11 @@
+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.weatherbench2 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']);base=WB2Baseline(**c['model']);m=DDP(base) if ddp else base;opt=torch.optim.Adam(m.parameters(),lr=c['train']['lr']);ids=np.arange(rank,len(d['input']),world);loss=((m(torch.tensor(d['input'][ids]))-torch.tensor(d['target'][ids]))**2).mean();opt.zero_grad();loss.backward();opt.step();v=loss.detach().double()
+if ddp:dist.all_reduce(v);v/=world
+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'],{'mse':float(v),'world_size':world});print(p)
+if ddp:dist.destroy_process_group()