Publish AI-GAMFS 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/aerosol.npz, samples: 8, channels: 54, tile: 16, logical_grid: [361, 576]}
+model: {channels: 54, hidden: 24, heads: 4}
+train: {epochs: 1, learning_rate: 0.001, leads: [3, 6, 9, 12]}
+inference: {hours: 120}
+paths: {checkpoint: result/checkpoints/ai_gamfs.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: [361, 576], channels: 54, aerosol_variables: 12, upper_levels: 9, leads: [3, 6, 9, 12], epochs: 80, forecast_hours: 120, paper_license: CC-BY-NC-ND-4.0}

config.json

@@ -0,0 +1 @@
+{"model_name":"AI-GAMFS","model_type":"ai_gamfs","architectures":["AIGAMFS"],"framework":"PyTorch","domain":"aerosol-weather","task":"global-aerosol-forecasting","implementation":{"entry_point":"model/ai_gamfs.py","scope":"L2 relay-forecast reproduction"},"architecture":{"logical_grid":[361,576],"aerosol_variables":12,"surface_variables":6,"upper_variables":4,"upper_levels":9,"total_channels":54,"lead_models":[3,6,9,12]},"data":{"source":"MERRA-2","years":[1980,2023],"test_year":2022,"operational_initialization":"GEOS-FP","synthetic_tiles":true},"configuration_sources":["conf/config.yaml","model/ai_gamfs.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/ai_gamfs.py

@@ -0,0 +1,15 @@
+import json,torch
+from pathlib import Path
+from torch import nn
+import yaml
+AEROSOL=('AOD','TSAOD','SUAOD','DUAOD','BCAOD','OCAOD','SSAOD','SUSMC','DUSMC','BCSMC','OCSMC','SSSMC')
+def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text())
+def state(i,h=0):
+ y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing='ij');c=torch.arange(54)[:,None,None];return (torch.sin((c%8+1)*x+.02*(i+h))*torch.cos((c%6+1)*y)).float()
+class Relay(nn.Module):
+ def __init__(self,c,h,heads):super().__init__();self.e=nn.Conv2d(c,h,2,2);self.a=nn.MultiheadAttention(h,heads,batch_first=True);self.d=nn.ConvTranspose2d(h,c,2,2)
+ def forward(self,x):z=self.e(x);b,c,y,w=z.shape;t=z.permute(0,2,3,1).reshape(b,y*w,c);t=self.a(t,t,t,need_weights=False)[0];return x+self.d(t.reshape(b,y,w,c).permute(0,3,1,2))
+class AIGAMFS(nn.Module):
+ def __init__(self,channels=54,hidden=24,heads=4):super().__init__();self.models=nn.ModuleDict({str(h):Relay(channels,hidden,heads) for h in (3,6,9,12)});self.model_config={'channels':channels,'hidden':hidden,'heads':heads}
+ def forward(self,x,lead):return self.models[str(lead)](x)
+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.ai_gamfs import *
+c=cfg(R);x=np.stack([state(i).numpy() for i in range(8)]);t=np.stack([[state(i,h).numpy() for h in (3,6,9,12)] for i in range(8)]);p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,input=x,targets=t,leads=[3,6,9,12],aerosol_variables=AEROSOL,logical_shape=[54,361,576],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.ai_gamfs import *
+c=cfg(R);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);m=AIGAMFS(**z['model_config']);m.load_state_dict(z['model']);x=state(20)[None];seq=[]
+with torch.no_grad():
+ for h in range(3,121,3):lead=max(v for v in (12,9,6,3) if v<=120-(h-3));x=m(x,lead);seq.append(x[0].numpy())
+p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,prediction=seq,target=[state(20,h).numpy() for h in range(3,121,3)],lead_hours=np.arange(3,121,3),logical_shape=[54,361,576],is_complete_global=False);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.ai_gamfs import cfg,write
+c=cfg(R);d=np.load(R/c['paths']['predictions']);p=d['prediction'];t=d['target'];rmse=np.sqrt(np.mean((p-t)**2,axis=(1,2,3)));corr=[float(np.corrcoef(p[i,0].ravel(),t[i,0].ravel())[0,1]) for i in range(40)];write(R/c['paths']['evaluation'],{'aod_rmse':rmse.tolist(),'aod_correlation':corr,'is_complete_global':False,'synthetic':True});plt.plot(d['lead_hours'],rmse);plt.xlabel('Lead (h)');plt.ylabel('RMSE');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.ai_gamfs 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=AIGAMFS(**c['model']);m=DDP(base,find_unused_parameters=True) if ddp else base;opt=torch.optim.AdamW(m.parameters(),lr=c['train']['learning_rate']);ls=[]
+for i in range(rank,8,world):
+ for j,h in enumerate(c['train']['leads']):loss=((m(torch.tensor(d['input'][i:i+1]),h)-torch.tensor(d['targets'][i:i+1,j]))**2).mean();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'],{'mse':float(v[0]/v[1]),'world_size':world});print(p)
+if ddp:dist.destroy_process_group()