Publish NNCAM-Stable 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: 128, input: 122, output: 68}
+model: {hidden: 64}
+train: {lr: 0.001}
+paths: {checkpoint: result/checkpoints/nncam_stable.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":"NNCAM-Stable","model_type":"nncam_stable","architectures":["NNCAM"],"framework":"PyTorch","domain":"climate-parameterization","task":"column-emulation","implementation":{"entry_point":"model/nncam_stable.py"},"architecture":{"input_features":122,"output_features":68,"vertical_levels":30,"networks":3},"data":{"source":"SPCAM","synthetic":true},"configuration_sources":["conf/config.yaml","model/nncam_stable.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/nncam_stable.py

@@ -0,0 +1,11 @@
+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 ResidualDNN(nn.Module):
+ def __init__(self,i,o,h):super().__init__();self.a=nn.Linear(i,h);self.b=nn.Linear(h,h);self.o=nn.Linear(h,o)
+ def forward(self,x):z=torch.relu(self.a(x));return self.o(z+torch.relu(self.b(z)))
+class NNCAM(nn.Module):
+ def __init__(self,hidden=64):super().__init__();self.q=ResidualDNN(122,30,hidden);self.s=ResidualDNN(122,30,hidden);self.rad=ResidualDNN(122,8,hidden);self.model_config={'hidden':hidden}
+ def forward(self,x):return torch.cat((self.q(x),self.s(x),self.rad(x)),1)
+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(2);x=g.normal(size=(128,122)).astype('f');y=x[:,:68]*.2;p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez(p,input=x,target=y);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.nncam_stable import *
+c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],weights_only=True);m=NNCAM(**z['model_config']);m.load_state_dict(z['model']);p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez(p,prediction=m(torch.tensor(d['input'])).detach(),target=d['target']);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.nncam_stable import *
+c=cfg(R);d=np.load(R/c['paths']['predictions']);rmse=float(np.sqrt(np.mean((d['prediction']-d['target'])**2)));write(R/c['paths']['evaluation'],{'rmse':rmse});plt.plot(d['prediction'][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.nncam_stable 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=NNCAM(**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()