Publish WeatherBench reproduction

.gitattributes

@@ -1,35 +1,6 @@
-*.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
+*.py text eol=lf
+*.md text eol=lf
+*.json text eol=lf
+*.yaml text eol=lf
+*.pt binary
+*.npz binary

conf/config.yaml

@@ -0,0 +1,30 @@
+seed: 42
+data:
+  format_version: weatherbench_v1
+  path: data/weatherbench.npz
+  grid: [32, 64]
+  samples: 24
+  channels: [Z500, T850]
+model:
+  hidden_channels: 16
+  layers: 5
+  kernel_size: 5
+train:
+  epochs: 2
+  learning_rate: 0.001
+paths:
+  checkpoint: result/checkpoints/weatherbench_cnn.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png
+paper_model:
+  grid: [32, 64]
+  input_channels: [Z500, T850]
+  conv_layers: 5
+  hidden_channels: 64
+  kernel_size: 5
+  activation: ELU
+  optimizer: Adam
+  loss: MSE
+  test_years: [2017, 2018]

config.json

@@ -0,0 +1 @@
+{"model_name":"WeatherBench-CNN","model_type":"weatherbench_cnn","architectures":["WeatherBenchCNN"],"framework":"PyTorch","domain":"weather","task":"medium-range-forecasting-benchmark","implementation":{"entry_point":"model/weatherbench.py","scope":"L2 benchmark and paper CNN reproduction"},"architecture":{"input_channels":["Z500","T850"],"output_channels":["Z500","T850"],"grid":[32,64],"conv_layers":5,"paper_hidden_channels":64,"kernel_size":5,"activation":"ELU","longitude_boundary":"periodic"},"data":{"source":"ERA5","years":[1979,2018],"training_years":[1979,2015],"validation_year":2016,"test_years":[2017,2018],"pressure_levels":13,"synthetic":true},"configuration_sources":["conf/config.yaml","model/weatherbench.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/weatherbench.py

@@ -0,0 +1,24 @@
+"""WeatherBench five-layer fully convolutional baseline."""
+import json
+from pathlib import Path
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+import yaml
+
+PRESSURE_LEVELS=(50,100,150,200,250,300,400,500,600,700,850,925,1000)
+def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text())
+def periodic_conv(x,conv,pad=2):x=F.pad(x,(pad,pad,0,0),mode="circular");x=F.pad(x,(0,0,pad,pad),mode="replicate");return conv(x)
+class WeatherBenchCNN(nn.Module):
+ def __init__(self,hidden_channels=16,layers=5,kernel_size=5):
+  super().__init__();chs=[2]+[hidden_channels]*(layers-1)+[2];self.convs=nn.ModuleList([nn.Conv2d(chs[i],chs[i+1],kernel_size) for i in range(layers)]);self.model_config={"hidden_channels":hidden_channels,"layers":layers,"kernel_size":kernel_size}
+ def forward(self,x):
+  if x.shape[1:]!=(2,32,64):raise ValueError("expected [B,2,32,64]")
+  for c in self.convs[:-1]:x=F.elu(periodic_conv(x,c))
+  return periodic_conv(x,self.convs[-1])
+def synthetic_state(i,lead=0):
+ lat,lon=torch.meshgrid(torch.linspace(-90,90,32),torch.arange(64).float()*360/64,indexing="ij");phase=.12*(i+lead);z=50000+3000*torch.cos(torch.deg2rad(lat))*torch.sin(torch.deg2rad(lon)+phase);t=270-35*abs(lat)/90+4*torch.cos(torch.deg2rad(lon*2)-phase);return torch.stack((z/50000,t/270)).float()
+def weighted_rmse(p,t,lat):w=np.cos(np.deg2rad(lat));w=w/w.mean();return np.mean(np.sqrt(np.mean((p-t)**2*w[None,None,:,None],axis=(1,2,3))))
+def weighted_acc(p,t,clim,lat):w=np.cos(np.deg2rad(lat))[None,None,:,None];a=p-clim;b=t-clim;return float((w*a*b).sum()/np.sqrt((w*a*a).sum()*(w*b*b).sum()))
+def write_json(path,obj):path=Path(path);path.parent.mkdir(parents=True,exist_ok=True);path.write_text(json.dumps(obj,indent=2)+"\n")

scripts/fake_data.py

@@ -0,0 +1,5 @@
+from pathlib import Path
+import sys,numpy as np
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.weatherbench import *
+c=load_config(ROOT);x=np.stack([synthetic_state(i).numpy() for i in range(24)]);y6=np.stack([synthetic_state(i,1).numpy() for i in range(24)]);y3=np.stack([synthetic_state(i,12).numpy() for i in range(24)]);y5=np.stack([synthetic_state(i,20).numpy() for i in range(24)]);split=np.array([0]*16+[1]*4+[2]*4);path=ROOT/c["data"]["path"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,input=x,target_6h=y6,target_3d=y3,target_5d=y5,split=split,latitude=np.linspace(-90,90,32),longitude=np.linspace(0,360,64,endpoint=False),pressure_levels=np.array(PRESSURE_LEVELS));print(path)

scripts/inference.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+import sys,numpy as np,torch
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.weatherbench import *
+c=load_config(ROOT);d=np.load(ROOT/c["data"]["path"]);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);m=WeatherBenchCNN(**ck["model_config"]);m.load_state_dict(ck["model"]);x=torch.tensor(d["input"][d["split"]==2]);p=[]
+with torch.no_grad():
+ for step in range(20):
+  x=m(x)
+  if step in (11,19):p.append(x.numpy())
+path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,prediction=np.stack(p,1),target=np.stack((d["target_3d"][d["split"]==2],d["target_5d"][d["split"]==2]),1),initial=d["input"][d["split"]==2],lead_days=np.array([3,5]),latitude=d["latitude"],longitude=d["longitude"]);print(path)

scripts/result.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+import sys,numpy as np
+import matplotlib;matplotlib.use("Agg");import matplotlib.pyplot as plt
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.weatherbench import load_config,weighted_rmse,weighted_acc,write_json
+c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);p=d["prediction"];t=d["target"];clim=t.mean((0,1),keepdims=True);metrics={}
+for k,name in enumerate(("Z500","T850")):metrics[name]={"rmse":[weighted_rmse(p[:,j,k:k+1],t[:,j,k:k+1],d["latitude"]) for j in range(2)],"acc":[weighted_acc(p[:,j,k:k+1],t[:,j,k:k+1],clim[:,:,k:k+1],d["latitude"]) for j in range(2)],"persistence_rmse":[weighted_rmse(d["initial"][:,k:k+1],t[:,j,k:k+1],d["latitude"]) for j in range(2)]}
+metrics["synthetic"]=True;write_json(ROOT/c["paths"]["evaluation"],metrics);fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].plot(d["lead_days"],metrics["Z500"]["rmse"],"o-",label="CNN");ax[0].plot(d["lead_days"],metrics["Z500"]["persistence_rmse"],"s-",label="persistence");ax[0].legend();im=ax[1].imshow(p[0,1,0]-t[0,1,0],cmap="coolwarm");ax[1].set_title("5-day Z500 error");fig.colorbar(im,ax=ax[1]);fig.tight_layout();path=ROOT/c["paths"]["figure"];path.parent.mkdir(parents=True,exist_ok=True);fig.savefig(path,dpi=150);print(path)

scripts/train.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+import sys,os,numpy as np,torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.weatherbench import *
+c=load_config(ROOT);rank=int(os.getenv("RANK",0));world=int(os.getenv("WORLD_SIZE",1));distributed=world>1
+if distributed:dist.init_process_group("gloo")
+torch.manual_seed(c["seed"]);d=np.load(ROOT/c["data"]["path"]);ids=np.where(d["split"]==0)[0];base=WeatherBenchCNN(**c["model"]);m=DDP(base) if distributed else base;opt=torch.optim.Adam(m.parameters(),lr=c["train"]["learning_rate"]);losses=[]
+for _ in range(c["train"]["epochs"]):
+ for i in ids[rank::world]:x=torch.tensor(d["input"][i:i+1]);y=torch.tensor(d["target_6h"][i:i+1]);loss=((m(x)-y)**2).mean();opt.zero_grad();loss.backward();opt.step();losses.append(float(loss))
+v=torch.tensor([sum(losses),len(losses)],dtype=torch.float64)
+if distributed:dist.all_reduce(v)
+p=ROOT/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_json(ROOT/c["paths"]["training_metrics"],{"mse":float(v[0]/v[1]),"world_size":world});print(p)
+if distributed:dist.destroy_process_group()