Publish TropiCycloneNet 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
+*.npz filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,35 @@
+seed: 42
+data:
+  format_version: tcnd_synthetic_v1
+  manifest: data/manifest.json
+  samples: 8
+  history_steps: 8
+  forecast_steps: 4
+  interval_hours: 6
+  data1d_features: 4
+  data3d_channels: 1
+  grid: [81, 81]
+  environment_features: 73
+  basins: [NA, EP, WP, NI, SI, SP]
+model:
+  hidden_dim: 24
+  generators: 6
+  heads: 4
+train:
+  epochs: 1
+  batch_size: 2
+  learning_rate: 0.001
+paths:
+  checkpoint: result/checkpoints/tropicyclonenet.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: [1950, 2016]
+  test_years: [2017, 2021]
+  optimizer: Adam
+  learning_rate: 0.0001
+  batch_size: 96
+  epochs: 102
+  generators: 6

config.json

@@ -0,0 +1 @@
+{"model_name":"TropiCycloneNet","model_type":"tropicyclonenet","architectures":["TropiCycloneNet"],"framework":"PyTorch","domain":"tropical-cyclone","task":"multimodal-track-intensity-forecasting","implementation":{"entry_point":"model/tropicyclonenet.py","scope":"core-method engineering reproduction"},"architecture":{"history_steps":8,"forecast_steps":4,"interval_hours":6,"data1d_variables":["longitude","latitude","pressure","wind"],"data3d_shape":[1,81,81],"environment_features":73,"paper_generators":6,"core":["3D data encoder","1D LSTM encoder","environment-time attention","generator chooser","multiple decoder LSTMs"]},"data":{"name":"TCND","years":[1950,2021],"cyclones":3630,"basins":6,"spatial_resolution_degrees":0.25,"synthetic":true},"configuration_sources":["conf/config.yaml","model/tropicyclonenet.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/tropicyclonenet.py

@@ -0,0 +1,30 @@
+"""Compact multimodal TropiCycloneNet reproduction."""
+import json,math
+from pathlib import Path
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+import yaml
+
+def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text())
+def synthetic_sample(i):
+ torch.manual_seed(1000+i);t=torch.arange(12).float();lon=120+i*.7+1.4*t+.08*t*t;lat=12+i*.3+.65*t;pressure=995-2.4*t-i*.2;wind=18+1.7*t+i*.1;one=torch.stack((lon/180,lat/90,(pressure-960)/50,(wind-40)/25),1);yy,xx=torch.meshgrid(torch.linspace(-1,1,81),torch.linspace(-1,1,81),indexing="ij");grids=[]
+ for k in range(12):grids.append(torch.exp(-((xx-.03*k)**2+(yy+.02*k)**2)/.15)[None])
+ env=torch.zeros(8,73);env[:,0]=1.5;env[:,1+(i%12)]=1;env[:,13+(i%36)]=1;env[:,49+(i%12)]=1;env[:,61+(i%8)]=1;env[:,69+(i%4)]=1
+ return one[:8],torch.stack(grids[:8]),env,one[8:],torch.stack(grids[8:])
+
+class TropiCycloneNet(nn.Module):
+ def __init__(self,hidden_dim=24,generators=6,heads=4):
+  super().__init__();self.generators=generators;self.grid_encoder=nn.Sequential(nn.Conv3d(1,8,3,padding=1),nn.GELU(),nn.AdaptiveAvgPool3d((8,4,4)));self.grid_project=nn.Linear(8*4*4,hidden_dim);self.one_encoder=nn.LSTM(4+hidden_dim,hidden_dim,batch_first=True);self.env_project=nn.Linear(73,hidden_dim);layer=nn.TransformerEncoderLayer(hidden_dim,heads,hidden_dim*2,batch_first=True);self.env_time=nn.TransformerEncoder(layer,1);self.chooser=nn.Linear(hidden_dim*2,generators);self.decoders=nn.ModuleList([nn.LSTMCell(4+hidden_dim*2,hidden_dim) for _ in range(generators)]);self.heads=nn.ModuleList([nn.Linear(hidden_dim,4) for _ in range(generators)]);self.model_config={"hidden_dim":hidden_dim,"generators":generators,"heads":heads}
+ def forward(self,one,grid,env,steps=4):
+  if one.shape[1:]!=(8,4) or grid.shape[1:]!=(8,1,81,81) or env.shape[1:]!=(8,73):raise ValueError("invalid TCND shape")
+  g=self.grid_encoder(grid.transpose(1,2)).permute(0,2,1,3,4).flatten(2);g=self.grid_project(g);_,(h,_)=self.one_encoder(torch.cat((one,g),-1));e=self.env_time(self.env_project(env))[:,-1];context=torch.cat((h[-1],e),-1);prob=self.chooser(context).softmax(-1);all_outputs=[]
+  for decoder,head in zip(self.decoders,self.heads):
+   hx=h[-1];cx=torch.zeros_like(hx);prev=one[:,-1];seq=[]
+   for _ in range(steps):hx,cx=decoder(torch.cat((prev,context),-1),(hx,cx));prev=prev+head(hx);seq.append(prev)
+   all_outputs.append(torch.stack(seq,1))
+  return torch.stack(all_outputs,1),prob
+def haversine_km(a,b):
+ lon1,lat1,lon2,lat2=[torch.deg2rad(x) for x in (a[...,0]*180,a[...,1]*90,b[...,0]*180,b[...,1]*90)];dlat=lat2-lat1;dlon=lon2-lon1;q=torch.sin(dlat/2)**2+torch.cos(lat1)*torch.cos(lat2)*torch.sin(dlon/2)**2;return 6371*2*torch.asin(torch.sqrt(q.clamp(0,1)))
+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,5 @@
+from pathlib import Path
+import sys,json
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.tropicyclonenet import *
+c=load_config(ROOT);o,g,e,y,gy=synthetic_sample(0);assert o.shape==(8,4) and g.shape==(8,1,81,81) and e.shape==(8,73) and y.shape==(4,4);p=ROOT/c["data"]["manifest"];p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps({"format_version":c["data"]["format_version"],"data1d_history_shape":[8,4],"data3d_history_shape":[8,1,81,81],"environment_shape":[8,73],"target_shape":[4,4],"interval_hours":6,"basins":c["data"]["basins"],"synthetic":True},indent=2));print(p)

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.tropicyclonenet import *
+c=load_config(ROOT);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);m=TropiCycloneNet(**ck["model_config"]);m.load_state_dict(ck["model"]);m.eval();pred=[];truth=[];prob=[]
+with torch.no_grad():
+ for i in range(c["data"]["samples"]):o,g,e,y,_=synthetic_sample(100+i);p,q=m(o[None],g[None],e[None]);pred.append(p[0].numpy());prob.append(q[0].numpy());truth.append(y.numpy())
+path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,prediction=np.asarray(pred),target=np.asarray(truth),generator_probability=np.asarray(prob),lead_hours=np.arange(1,5)*6);print(path)

scripts/result.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+import sys,numpy as np,torch
+import matplotlib;matplotlib.use("Agg");import matplotlib.pyplot as plt
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.tropicyclonenet import *
+c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);p=torch.from_numpy(d["prediction"]);t=torch.from_numpy(d["target"]);best=((p-t[:,None])**2).mean((2,3)).argmin(1);chosen=p[torch.arange(len(p)),best];track=haversine_km(chosen[...,:2],t[...,:2]).mean(0);pres=(abs(chosen[...,2]-t[...,2])*50).mean(0);wind=(abs(chosen[...,3]-t[...,3])*25).mean(0);write_json(ROOT/c["paths"]["evaluation"],{"track_mae_km":track.tolist(),"pressure_mae_hpa":pres.tolist(),"wind_mae_ms":wind.tolist(),"generators":int(p.shape[1]),"synthetic":True});fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].plot(d["lead_hours"],track,"o-");ax[0].set(xlabel="Lead (h)",ylabel="Track MAE (km)");ax[1].plot(t[0,:,0]*180,t[0,:,1]*90,"ko-",label="truth");ax[1].plot(chosen[0,:,0]*180,chosen[0,:,1]*90,"r.--",label="prediction");ax[1].legend();ax[1].set(xlabel="Longitude",ylabel="Latitude");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,16 @@
+from pathlib import Path
+import sys,os,math,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.tropicyclonenet 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"]);base=TropiCycloneNet(**c["model"]);m=DDP(base) if distributed else base;opt=torch.optim.Adam(m.parameters(),lr=c["train"]["learning_rate"]);idx=list(range(rank,c["data"]["samples"],world));losses=[]
+for _ in range(c["train"]["epochs"]):
+ for i in idx:o,g,e,y,_=synthetic_sample(i);pred,prob=m(o[None],g[None],e[None]);errors=((pred-y[None,None])**2).mean((2,3));best=errors.min(1).values.mean();diversity=-pred.std(1).mean();loss=best+.01*diversity+.01*(-torch.log(prob+1e-8).mean());opt.zero_grad();loss.backward();opt.step();losses.append(float(loss))
+total=torch.tensor([sum(losses),len(losses)],dtype=torch.float64)
+if distributed:dist.all_reduce(total)
+p=ROOT/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_json(ROOT/c["paths"]["training_metrics"],{"loss":float(total[0]/total[1]),"world_size":world});print(p)
+if distributed:dist.destroy_process_group()