Publish DLESyM 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,7 @@
+seed: 42
+data: {path: data/dlesym.npz, samples: 6, tile: 16, atmosphere_channels: 9, logical_grid: [180, 360]}
+model: {atmosphere_channels: 9, hidden: 16}
+train: {epochs: 1, learning_rate: 0.001}
+inference: {cycles: 4}
+paths: {checkpoint: result/checkpoints/dlesym.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_km: 110, atmosphere_channels: 9, ocean_channels: 1, atmosphere_step_hours: 6, ocean_step_hours: 48, coupling_cycle_hours: 96, atmosphere_epochs: 250, ocean_epochs: 300}

config.json

@@ -0,0 +1 @@
+{"model_name":"DLESyM","model_type":"dlesym","architectures":["DLESyM"],"framework":"PyTorch","domain":"coupled-climate","task":"earth-system-simulation","implementation":{"entry_point":"model/dlesym.py","scope":"L2 coupled core reproduction"},"architecture":{"atmosphere_fields":9,"ocean_fields":["SST"],"atmosphere_step_hours":6,"ocean_step_hours":48,"coupling_cycle_hours":96,"logical_grid":[180,360],"modules":["DLWP","DLOM","precipitation"]},"data":{"sources":["ERA5","ISCCP"],"train_years":[1983,2016],"synthetic_tiles":true},"configuration_sources":["conf/config.yaml","model/dlesym.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/dlesym.py

@@ -0,0 +1,20 @@
+import json,torch
+from pathlib import Path
+from torch import nn
+import yaml
+FIELDS=('Z1000','Z500','Z250','THICK300_700','T2M','T850','TCWV','WS10','OLR')
+def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
+def state(i,s=0):
+ y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing='ij');c=torch.arange(9)[:,None,None];a=torch.sin((c+1)*x+.04*(i+s))*torch.cos((c%4+1)*y);sst=torch.cos(x+.01*(i+s))*torch.cos(y);return a.float(),sst[None].float()
+class Net(nn.Module):
+ def __init__(self,ci,co,h):super().__init__();self.n=nn.Sequential(nn.Conv2d(ci,h,3,padding=1),nn.GELU(),nn.Conv2d(h,h,3,padding=1),nn.GELU(),nn.Conv2d(h,co,1))
+ def forward(self,x):return self.n(x)
+class DLESyM(nn.Module):
+ def __init__(self,atmosphere_channels=9,hidden=16):super().__init__();self.atm=Net(10,9,hidden);self.ocean=Net(4,1,hidden//2);self.precip=Net(9,1,hidden//2);self.model_config={'atmosphere_channels':atmosphere_channels,'hidden':hidden}
+ def atmosphere(self,a,sst):return a+self.atm(torch.cat((a,sst),1))
+ def cycle(self,a,sst):
+  seq=[]
+  for _ in range(16):a=self.atmosphere(a,sst);seq.append(a)
+  f=torch.stack(seq);forcing=torch.stack((f[:,:,7].mean(0),f[:,:,0].mean(0),f[:,:,8].mean(0)),1);sst=sst+self.ocean(torch.cat((sst,forcing),1));return a,sst,self.precip(a)
+ def forward(self,a,sst):return self.cycle(a,sst)
+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,6 @@
+from pathlib import Path
+import sys,numpy as np
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.dlesym import *
+c=cfg(R);a=[];s=[];ta=[];ts=[]
+for i in range(6):x,y=state(i);u,v=state(i,1);a.append(x);s.append(y);ta.append(u);ts.append(v)
+p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,atmosphere=a,sst=s,target_atmosphere=ta,target_sst=ts,fields=FIELDS,logical_shape=[10,180,360],is_complete_global=False);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.dlesym import *
+c=cfg(R);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);m=DLESyM(**z['model_config']);m.load_state_dict(z['model']);a,s=state(20);a=a[None];s=s[None];aa=[];ss=[];pp=[]
+with torch.no_grad():
+ for _ in range(4):a,s,p=m(a,s);aa.append(a[0].numpy());ss.append(s[0].numpy());pp.append(p[0].numpy())
+q=R/c['paths']['predictions'];q.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(q,atmosphere=aa,sst=ss,precipitation=pp,lead_days=np.arange(1,5)*4,logical_shape=[10,180,360],is_complete_global=False);print(q)

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.dlesym import cfg,write
+c=cfg(R);d=np.load(R/c['paths']['predictions']);a=d['atmosphere'];s=d['sst'];drift=np.sqrt(np.mean((a-a[:1])**2,axis=(1,2,3)));sst=s.mean((1,2,3));prec=d['precipitation'].mean((1,2,3));write(R/c['paths']['evaluation'],{'atmosphere_drift_rmse':drift.tolist(),'sst_mean':sst.tolist(),'precipitation_mean':prec.tolist(),'is_complete_global':False,'synthetic':True});plt.plot(d['lead_days'],drift);plt.xlabel('Lead (days)');plt.ylabel('Atmosphere drift RMSE');p=R/c['paths']['figure'];p.parent.mkdir(parents=True,exist_ok=True);plt.savefig(p,dpi=150);print(p)

scripts/train.py

@@ -0,0 +1,13 @@
+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.dlesym 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=DLESyM(**c['model']);m=DDP(base) if ddp else base;opt=torch.optim.AdamW(m.parameters(),lr=c['train']['learning_rate']);ls=[]
+for i in range(rank,6,world):a,s=map(torch.tensor,(d['atmosphere'][i:i+1],d['sst'][i:i+1]));pa,ps,pr=m(a,s);loss=((pa-torch.tensor(d['target_atmosphere'][i:i+1]))**2).mean()+((ps-torch.tensor(d['target_sst'][i:i+1]))**2).mean()+pr.pow(2).mean()*.01;opt.zero_grad();loss.backward();opt.step();ls.append(float(loss))
+v=torch.tensor([sum(ls),len(ls)],dtype=torch.float64)
+if ddp:dist.all_reduce(v)
+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'],{'loss':float(v[0]/v[1]),'world_size':world});print(p)
+if ddp:dist.destroy_process_group()