Publish SamudrACE 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,31 @@
+seed: 42
+data:
+  format_version: samudrace_tiles_v1
+  path: data/coupled_tiles.npz
+  logical_grid: [180, 360]
+  tile_size: 8
+  tile_origins: [[20, 40], [80, 160], [140, 280], [50, 320]]
+  atmosphere_channels: 46
+  ocean_channels: 80
+model:
+  hidden_dim: 16
+train:
+  epochs: 1
+  learning_rate: 0.001
+  atmosphere_steps_per_ocean: 20
+paths:
+  checkpoint: result/checkpoints/samudrace.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png
+paper_model:
+  training_years: 155
+  validation_years: 5
+  test_years: 40
+  atmosphere_step_hours: 6
+  ocean_step_days: 5
+  atmosphere_levels: 8
+  ocean_levels: 19
+  optimizer: AdamW
+  batch_size: 16

config.json

@@ -0,0 +1 @@
+{"model_name":"SamudrACE","model_type":"samudrace","architectures":["SamudrACE"],"framework":"PyTorch","domain":"coupled-climate","task":"atmosphere-ocean-emulation","implementation":{"entry_point":"model/samudrace.py","scope":"L2 core coupling reproduction"},"architecture":{"logical_grid":[180,360],"atmosphere_channels":46,"atmosphere_levels":8,"ocean_channels":80,"ocean_levels":19,"coupling_ratio":20,"atmosphere_step_hours":6,"ocean_step_days":5},"data":{"source":"GFDL CM4 piControl","years":200,"split_years":[155,5,40],"synthetic_tiles":true,"enumerated_fields":126,"paper_abstract_fields":145,"unresolved_fields":19},"configuration_sources":["conf/config.yaml","model/samudrace.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/samudrace.py

@@ -0,0 +1,26 @@
+"""Compact physical-state coupled atmosphere-ocean emulator."""
+import json
+from pathlib import Path
+import numpy as np
+import torch
+from torch import nn
+import yaml
+
+ATM3D=("T","qT","U","V");ATM2D=("Ts","ps");DIAG=("RSW","OLR","USWsfc","ULWsfc","DSWsfc","DLWsfc","LHF","SHF","P","dTWP_adv","tau_u","tau_v")
+OCEAN3D=("thetao","so","uo","vo");OCEAN2D=("SIC","HI","SST","ZOS")
+def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text())
+def synthetic_state(origin,size,i,step=0):
+ y,x=torch.meshgrid(torch.arange(origin[0],origin[0]+size),torch.arange(origin[1],origin[1]+size),indexing="ij");lat=torch.deg2rad(90-(y+.5));lon=torch.deg2rad(-180+(x+.5));a=torch.arange(46)[:,None,None];o=torch.arange(80)[:,None,None];atm=torch.sin((a%6+1)*lat+.03*a+.05*(i+step))*torch.cos(lon);oce=torch.cos((o%5+1)*lat-.01*o+.02*(i+step))*torch.sin(lon);return atm.float(),oce.float()
+class ResidualEmulator(nn.Module):
+ def __init__(self,cin,cout,h):super().__init__();self.net=nn.Sequential(nn.Conv2d(cin,h,1),nn.GELU(),nn.Conv2d(h,h,3,padding=1),nn.GELU(),nn.Conv2d(h,cout,1))
+ def forward(self,x):return self.net(x)
+class SamudrACE(nn.Module):
+ def __init__(self,hidden_dim=16):super().__init__();self.atmosphere=ResidualEmulator(48,46,hidden_dim);self.ocean=ResidualEmulator(92,80,hidden_dim);self.model_config={"hidden_dim":hidden_dim}
+ def atmosphere_step(self,atm,ocean):return atm+self.atmosphere(torch.cat((atm,ocean[:,[76,0]]),1))
+ def ocean_step(self,ocean,flux):return ocean+self.ocean(torch.cat((ocean,flux),1))
+ def forward(self,atm,ocean,n=20):return self.coupled_step(atm,ocean,n)
+ def coupled_step(self,atm,ocean,n=20):
+  flux=[]
+  for _ in range(n):atm=self.atmosphere_step(atm,ocean);flux.append(atm[:,-12:])
+  ocean=self.ocean_step(ocean,torch.stack(flux).mean(0));return atm,ocean
+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,7 @@
+from pathlib import Path
+import sys,json,numpy as np
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.samudrace import *
+c=load_config(ROOT);aa=[];oo=[];at=[];ot=[]
+for i,p in enumerate(c["data"]["tile_origins"]):a,o=synthetic_state(p,8,i);ta,to=synthetic_state(p,8,i,1);aa.append(a);oo.append(o);at.append(ta);ot.append(to)
+path=ROOT/c["data"]["path"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,atmosphere=np.stack(aa),ocean=np.stack(oo),atmosphere_target=np.stack(at),ocean_target=np.stack(ot),origins=np.array(c["data"]["tile_origins"]),logical_shape=np.array([126,180,360]),is_complete_global=False);print(path)

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.samudrace 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=SamudrACE(**ck["model_config"]);m.load_state_dict(ck["model"]);a=torch.tensor(d["atmosphere"]);o=torch.tensor(d["ocean"]);ap=[];op=[]
+with torch.no_grad():
+ for _ in range(4):a,o=m.coupled_step(a,o,20);ap.append(a.numpy());op.append(o.numpy())
+path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,atmosphere=np.stack(ap,1),ocean=np.stack(op,1),origins=d["origins"],lead_days=np.arange(1,5)*5,logical_shape=d["logical_shape"],is_complete_global=False);print(path)

scripts/result.py

@@ -0,0 +1,6 @@
+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.samudrace import load_config,write_json
+c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);a=d["atmosphere"];o=d["ocean"];armse=np.sqrt(np.mean((a-a[:,0:1])**2,axis=(0,2,3,4)));ormse=np.sqrt(np.mean((o-o[:,0:1])**2,axis=(0,2,3,4)));heat=o[:,:,0].mean((0,2,3));salt=o[:,:,19].mean((0,2,3));write_json(ROOT/c["paths"]["evaluation"],{"atmosphere_drift_rmse":armse.tolist(),"ocean_drift_rmse":ormse.tolist(),"ocean_heat_proxy":heat.tolist(),"ocean_salt_proxy":salt.tolist(),"is_complete_global":False,"synthetic":True});fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].plot(d["lead_days"],armse,label="atmosphere");ax[0].plot(d["lead_days"],ormse,label="ocean");ax[0].legend();ax[0].set(xlabel="Lead (days)",ylabel="Drift RMSE");im=ax[1].imshow(o[0,-1,76]-o[0,0,76],cmap="coolwarm");ax[1].set_title("SST drift proxy");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,15 @@
+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.samudrace 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);d=np.load(ROOT/c["data"]["path"]);base=SamudrACE(**c["model"]);m=DDP(base) if distributed else base;opt=torch.optim.AdamW(m.parameters(),lr=c["train"]["learning_rate"]);losses=[]
+for i in range(rank,len(d["atmosphere"]),world):a=torch.tensor(d["atmosphere"][i:i+1]);o=torch.tensor(d["ocean"][i:i+1]);pa,po=m(a,o,20);loss=((pa-torch.tensor(d["atmosphere_target"][i:i+1]))**2).mean()+((po-torch.tensor(d["ocean_target"][i:i+1]))**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)
+path=ROOT/c["paths"]["checkpoint"]
+if rank==0:path.parent.mkdir(parents=True,exist_ok=True);torch.save({"model":base.state_dict(),"model_config":c["model"]},path);write_json(ROOT/c["paths"]["training_metrics"],{"coupled_mse":float(v[0]/v[1]),"world_size":world});print(path)
+if distributed:dist.destroy_process_group()