Publish GRACE-SEDA 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/twsa.npz, samples: 20, channels: 9, patch: 32}
+model: {channels: 9, hidden: 16}
+train: {epochs: 2, learning_rate: 0.001, ensemble_members: 5}
+paths: {checkpoint: result/checkpoints/grace_seda.pt, training_metrics: result/training/metrics.json, predictions: result/output/predictions.npz, evaluation: result/evaluation/metrics.json, figure: result/evaluation/comparison.png}
+paper_model: {resolution_degrees: 0.5, patch: [32, 32], channels: 9, encoder_filters: [16, 32, 64], months: 180, optimizer: Adam, batch_size: 512, ensemble_members: 5, paper_license: CC-BY-4.0}

config.json

@@ -0,0 +1 @@
+{"model_name":"GRACE-SEDA","model_type":"grace_seda","architectures":["GRACESEDA"],"framework":"PyTorch","domain":"terrestrial-water-storage","task":"self-supervised-data-assimilation","implementation":{"entry_point":"model/grace_seda.py","scope":"L2 self-supervised patch reproduction"},"architecture":{"input_channels":9,"patch_shape":[32,32],"resolution_degrees":0.5,"encoder_filters":[16,32,64],"deep_ensemble_members":5,"core":"residual encoder-decoder"},"data":{"sources":["JPL GRACE","WGHM","GLDAS"],"period":["2002-04","2019-12"],"months":180,"global_land":true},"configuration_sources":["conf/config.yaml","model/grace_seda.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/grace_seda.py

@@ -0,0 +1,17 @@
+import json,torch
+from pathlib import Path
+from torch import nn
+import torch.nn.functional as F
+import yaml
+FEATURES=('GRACE','WGHM','precipitation','ET','surface_runoff','baseflow','snowmelt','latitude','longitude')
+def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
+def sample(i):
+ y,x=torch.meshgrid(torch.linspace(-1,1,32),torch.linspace(-1,1,32),indexing='ij');w=torch.sin(3*x+.1*i)*torch.cos(2*y);g=F.avg_pool2d(w[None,None],6,1,3)[0,0,:32,:32];g=F.interpolate(g[None,None],size=(32,32),mode='bilinear')[0,0];return torch.stack((g,w,torch.sin(x),torch.cos(y),w*.2,w*.1,w*.05,y,x)).float()
+class Block(nn.Module):
+ def __init__(self,c):super().__init__();self.n=nn.Sequential(nn.Conv2d(c,c,3,padding=1),nn.ReLU(),nn.BatchNorm2d(c),nn.Conv2d(c,c,3,padding=1),nn.BatchNorm2d(c))
+ def forward(self,x):return F.relu(x+self.n(x))
+class GRACESEDA(nn.Module):
+ def __init__(self,channels=9,hidden=16):super().__init__();self.e1=nn.Conv2d(channels,hidden,3,2,1);self.e2=nn.Conv2d(hidden,hidden*2,3,2,1);self.e3=nn.Conv2d(hidden*2,hidden*4,3,2,1);self.b=Block(hidden*4);self.d=nn.Sequential(nn.Conv2d(hidden*4,hidden*2,3,padding=1),nn.ReLU(),nn.Conv2d(hidden*2,hidden,3,padding=1),nn.ReLU(),nn.Conv2d(hidden,1,1));self.model_config={'channels':channels,'hidden':hidden}
+ def forward(self,x):z=F.relu(self.e1(x));z=F.relu(self.e2(z));z=self.b(F.relu(self.e3(z)));return F.interpolate(self.d(z),size=(32,32),mode='bilinear',align_corners=False)[:,0]
+def loss_fn(p,x):g=x[:,0];w=x[:,1];ae=abs(p.mean((1,2))-g.mean((1,2)));pc=((p-p.mean((1,2),keepdim=True))*(w-w.mean((1,2),keepdim=True))).mean((1,2))/(p.std((1,2))*w.std((1,2))+1e-6);mae=abs(p-w).mean((1,2));return (ae+(1-pc)*mae).mean()
+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.grace_seda import *
+c=cfg(R);x=np.stack([sample(i).numpy() for i in range(20)]);p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,input=x,features=FEATURES,shape=[20,9,32,32],resolution=.5);print(p)

scripts/inference.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.grace_seda import *
+c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);x=torch.tensor(d['input']);pred=[]
+with torch.no_grad():
+ for s in z['states']:m=GRACESEDA(**z['model_config']);m.load_state_dict(s);pred.append(m(x).numpy())
+p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,ensemble=pred,grace=d['input'][:,0],wghm=d['input'][:,1],mean=np.mean(pred,0),uncertainty=np.std(pred,0));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.grace_seda import cfg,write
+c=cfg(R);d=np.load(R/c['paths']['predictions']);p=d['mean'];g=d['grace'];w=d['wghm'];corr=float(np.corrcoef(p.ravel(),w.ravel())[0,1]);aggregate=float(abs(p.mean((1,2))-g.mean((1,2))).mean());unc=float(d['uncertainty'].mean());write(R/c['paths']['evaluation'],{'wghm_correlation':corr,'grace_aggregate_error':aggregate,'mean_uncertainty':unc,'synthetic':True});plt.imshow(p[0]-w[0],cmap='coolwarm');plt.colorbar();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,14 @@
+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.grace_seda 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']);states=[];hist=[]
+for member in range(5):
+ torch.manual_seed(c['seed']+member);base=GRACESEDA(**c['model']);m=DDP(base) if ddp else base;opt=torch.optim.Adam(m.parameters(),lr=c['train']['learning_rate']);ls=[]
+ for i in range(rank,20,world):x=torch.tensor(d['input'][i:i+1]);loss=loss_fn(m(x),x);opt.zero_grad();loss.backward();opt.step();ls.append(float(loss))
+ if rank==0:states.append(base.state_dict());hist.append(float(np.mean(ls)))
+p=R/c['paths']['checkpoint']
+if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({'states':states,'model_config':c['model']},p);write(R/c['paths']['training_metrics'],{'ensemble_losses':hist,'world_size':world});print(p)
+if ddp:dist.destroy_process_group()