Publish RISE-UNet 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
+*.npz filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,31 @@
+seed: 42
+data:
+  format_version: rise_unet_synthetic_v1
+  manifest: data/manifest.json
+  initializations: 4
+  ensemble_members: 11
+  grid: [48, 96]
+  input_channels: 8
+  forecast_weeks: 5
+model:
+  input_channels: 8
+  width: 12
+  dropout: 0.15
+train:
+  epochs: 1
+  learning_rate: 0.001
+  crps_spread_factor: 0.08
+paths:
+  checkpoint: result/checkpoints/rise_unet.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: [48, 96]
+  ensemble_members: 11
+  training_years: [2000, 2015]
+  validation_years: [2016, 2017]
+  test_years: [2018, 2019]
+  forecast_weeks: 5
+  encoder_filters: [32, 64, 128, 256]

config.json

@@ -0,0 +1 @@
+{"model_name":"RISE-UNet","model_type":"rise_unet","architectures":["RISEUNet"],"framework":"PyTorch","domain":"soil-moisture-drought","task":"subseasonal-ensemble-forecasting","implementation":{"entry_point":"model/rise_unet.py","scope":"core-method engineering reproduction"},"architecture":{"grid":[48,96],"ensemble_members":11,"forecast_weeks":5,"engineering_input_channels":8,"core":["residual inception","squeeze excitation","UNet encoder-decoder","deep supervision","recursive forecast"]},"data":{"target":"GLEAM 0-100 cm root-zone soil moisture anomaly","sources":["GLEAM","ERA5","GEFSv12","ECMWF S2S"],"training_years":[2000,2015],"validation_years":[2016,2017],"test_years":[2018,2019],"resolution_degrees":0.5,"synthetic":true},"configuration_sources":["conf/config.yaml","model/rise_unet.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/rise_unet.py

@@ -0,0 +1,29 @@
+"""Compact residual-inception squeeze-excitation UNet."""
+import json,math
+from pathlib import Path
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+import yaml
+
+def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text())
+def synthetic_initialization(i,members=11):
+ y,x=torch.meshgrid(torch.linspace(-1,1,48),torch.linspace(-2,2,96),indexing="ij");base=torch.sin(2*x)*torch.cos(2*y)+.2*i;xs=[]
+ for m in range(members):xs.append(torch.stack([base+.03*m+.05*j+torch.sin(x*(j+1))*.1 for j in range(8)]))
+ targets=torch.stack([base-.12*w+.08*torch.sin(x*w+y) for w in range(1,6)]);return torch.stack(xs),targets
+
+class RISEBlock(nn.Module):
+ def __init__(self,cin,cout,dropout):
+  super().__init__();q=max(1,cout//4);self.paths=nn.ModuleList([nn.Conv2d(cin,q,k,padding=k//2) for k in (3,5,7)]);self.pool=nn.Sequential(nn.MaxPool2d(5,1,2),nn.Conv2d(cin,cout-3*q,1));self.norm=nn.BatchNorm2d(cout);self.skip=nn.Conv2d(cin,cout,1);self.se=nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(cout,max(1,cout//4),1),nn.ReLU(),nn.Conv2d(max(1,cout//4),cout,1),nn.Sigmoid());self.drop=nn.Dropout2d(dropout)
+ def forward(self,x):
+  z=torch.cat([p(x) for p in self.paths]+[self.pool(x)],1);z=F.relu(self.norm(z));z=self.drop(z);return F.relu(z+self.skip(x))*self.se(z)
+class RISEUNet(nn.Module):
+ def __init__(self,input_channels=8,width=12,dropout=.15):
+  super().__init__();self.e1=RISEBlock(input_channels,width,dropout);self.e2=RISEBlock(width,width*2,dropout);self.e3=RISEBlock(width*2,width*4,dropout);self.d2=RISEBlock(width*6,width*2,dropout);self.d1=RISEBlock(width*3,width,dropout);self.deep=nn.Conv2d(width*2,1,1);self.out=nn.Conv2d(width,1,1);self.model_config={"input_channels":input_channels,"width":width,"dropout":dropout}
+ def forward(self,x):
+  if x.shape[1:]!=(8,48,96):raise ValueError("expected [B,8,48,96]")
+  a=self.e1(x);b=self.e2(F.max_pool2d(a,2));c=self.e3(F.max_pool2d(b,2));u=self.d2(torch.cat((F.interpolate(c,size=b.shape[-2:],mode="bilinear"),b),1));v=self.d1(torch.cat((F.interpolate(u,size=a.shape[-2:],mode="bilinear"),a),1));d=F.interpolate(self.deep(u),size=(48,96),mode="bilinear");o=self.out(v);return (d,o,(d+o)/2)
+def crps_exp(outputs,target,members=11,f=.08):
+ pred=outputs[-1].reshape(-1,members,48,96);truth=target.reshape(-1,members,48,96);return (abs(pred-truth).mean(1)-f*pred.std(1)).mean()+sum(abs(o-target).mean() for o in outputs[:-1])*.1
+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,json
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.rise_unet import *
+c=load_config(ROOT);x,y=synthetic_initialization(0);assert x.shape==(11,8,48,96) and y.shape==(5,48,96);p=ROOT/c["data"]["manifest"];p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps({"format_version":c["data"]["format_version"],"input_shape":[11,8,48,96],"target_shape":[5,48,96],"forecast_weeks":[1,2,3,4,5],"resolution_degrees":0.5,"sources":["GLEAM","ERA5","GEFSv12","ECMWF S2S"],"synthetic":True},indent=2));print(p)

scripts/inference.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+import sys,numpy as np,torch
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.rise_unet import *
+c=load_config(ROOT);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);m=RISEUNet(**ck["model_config"]);m.load_state_dict(ck["model"]);m.train();x,t=synthetic_initialization(20);pred=[]
+with torch.no_grad():
+ for w in range(5):o=m(x)[-1];pred.append(o[:,0].numpy());x=torch.cat((x[:,1:],o),1)
+p=ROOT/c["paths"]["predictions"];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,ensemble=np.stack(pred,1),target=t.numpy(),lead_weeks=np.arange(1,6),dropout_ensemble=np.bool_(True));print(p)

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.rise_unet import load_config,write_json
+c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);e=d["ensemble"];t=d["target"];mean=e.mean(0);acc=[]
+for w in range(5):acc.append(float(np.corrcoef(mean[w].ravel(),t[w].ravel())[0,1]))
+crps=(abs(e-t[None]).mean(0)-.5*abs(e[:,None]-e[None,:]).mean((0,1))).mean((1,2));q=np.quantile(t,.2);obs=t<q;hit=mean<q;tp=(obs&hit).sum((1,2));fp=(~obs&hit).sum((1,2));fn=(obs&~hit).sum((1,2));gss=(tp-(tp+fp)*(tp+fn)/obs[0].size)/(tp+fp+fn-(tp+fp)*(tp+fn)/obs[0].size+1e-8);write_json(ROOT/c["paths"]["evaluation"],{"acc":acc,"crps":crps.tolist(),"drought_gss":gss.tolist(),"threshold":"synthetic 20th percentile","synthetic":True});fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].plot(d["lead_weeks"],acc,"o-",label="ACC");ax[0].plot(d["lead_weeks"],crps,"s-",label="CRPS");ax[0].legend();im=ax[1].imshow(mean[2]-t[2],cmap="coolwarm");ax[1].set_title("Week-3 error");fig.colorbar(im,ax=ax[1]);fig.tight_layout();p=ROOT/c["paths"]["figure"];p.parent.mkdir(parents=True,exist_ok=True);fig.savefig(p,dpi=150);print(p)

scripts/train.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+import sys,os,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.rise_unet 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"]);torch.set_num_threads(2);base=RISEUNet(**c["model"]);m=DDP(base) if distributed else base;opt=torch.optim.Adam(m.parameters(),lr=c["train"]["learning_rate"]);losses=[]
+for i in range(rank,c["data"]["initializations"],world):
+ x,t=synthetic_initialization(i);current=x
+ for w in range(5):out=m(current);target=t[w][None,None].expand(11,1,-1,-1);loss=crps_exp(out,target,f=c["train"]["crps_spread_factor"]);opt.zero_grad();loss.backward();opt.step();losses.append(float(loss));current=torch.cat((current[:,1:],out[-1].detach()),1)
+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"],{"crps_exp":float(v[0]/v[1]),"world_size":world});print(p)
+if distributed:dist.destroy_process_group()