Publish FuXi-Weather 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/fuxi_weather.npz, samples: 8, channels: 20, tile: 16, logical_grid: [721, 1440]}
+model: {channels: 20, hidden: 24}
+train: {epochs: 1, learning_rate: 0.001}
+inference: {steps: 12, paper_steps: 40, switch_step: 6}
+paths: {checkpoint: result/checkpoints/fuxi_weather.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_degrees: 0.25, cycle_hours: 6, observation_window_hours: 8, forecast_days: 10, paper_license: CC-BY-NC-ND-4.0}

config.json

@@ -0,0 +1 @@
+{"model_name":"FuXi-Weather","model_type":"fuxi_weather","architectures":["FuXiWeather"],"framework":"PyTorch","domain":"weather","task":"data-assimilation-and-forecast","implementation":{"entry_point":"model/fuxi_weather.py","scope":"L2 system-flow reproduction"},"architecture":{"logical_grid":[721,1440],"engineering_channels":20,"variables":["Z","T","U","V","R","T2M","MSLP","U10","V10","TP"],"observation_window_hours":8,"forecast_step_hours":6,"cascade_days":[4,10]},"data":{"reference":"ERA5","observations":["FY-3E","Metop-C","NOAA-20","GNSS-RO"],"synthetic_tiles":true},"configuration_sources":["conf/config.yaml","model/fuxi_weather.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/fuxi_weather.py

@@ -0,0 +1,15 @@
+import json
+from pathlib import Path
+import torch
+from torch import nn
+import yaml
+def cfg(r):return yaml.safe_load((Path(r)/"conf/config.yaml").read_text())
+def sample(i,s=0):
+ y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing="ij");c=torch.arange(20)[:,None,None];b=torch.sin((c%5+1)*x+.1*(i+s))*torch.cos((c%4+1)*y);m=(((torch.arange(16)[None,:]+i+s)%4==0).float().expand(16,-1))[None].expand(20,-1,-1);o=b*m;return b.float(),o.float(),m.float()
+class FuXiWeather(nn.Module):
+ def __init__(self,channels=20,hidden=24):super().__init__();self.obs=nn.Conv2d(channels*2,hidden,1);self.bg=nn.Conv2d(channels,hidden,1);self.refine=nn.Sequential(nn.Conv2d(hidden*2,hidden,3,padding=1),nn.GELU(),nn.Conv2d(hidden,channels,1));self.short=nn.Conv2d(channels,channels,3,padding=1);self.medium=nn.Conv2d(channels,channels,3,padding=1);self.model_config={"channels":channels,"hidden":hidden}
+ def analysis(self,b,o,m):return b+self.refine(torch.cat((self.bg(b),self.obs(torch.cat((o,m),1))),1))
+ def forecast(self,x,medium=False):return x+(self.medium(x) if medium else self.short(x))
+ def forward(self,b,o,m,medium=False):
+  a=self.analysis(b,o,m);return a,self.forecast(a,medium)
+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.fuxi_weather import *
+c=cfg(R);b=[];o=[];m=[];t=[]
+for i in range(8):a,z,q=sample(i);b.append(a);o.append(z);m.append(q);t.append(sample(i,1)[0])
+p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,background=b,observation=o,mask=m,target=t,logical_shape=[20,721,1440],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.fuxi_weather import *
+c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);m=FuXiWeather(**z['model_config']);m.load_state_dict(z['model']);x=m.analysis(*map(torch.tensor,(d['background'][:2],d['observation'][:2],d['mask'][:2])));seq=[]
+with torch.no_grad():
+ for s in range(c['inference']['steps']):x=m.forecast(x,s>=c['inference']['switch_step']);seq.append(x.numpy())
+target=np.stack([[sample(i,s+1)[0].numpy() for s in range(12)] for i in range(2)]);p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,prediction=np.stack(seq,1),target=target,lead_hours=np.arange(1,13)*6,logical_shape=np.array([20,721,1440]),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.fuxi_weather import *
+c=cfg(R);d=np.load(R/c['paths']['predictions']);p=d['prediction'];truth=np.stack([[sample(i,s+1)[0].numpy() for s in range(12)] for i in range(2)]);rmse=np.sqrt(np.mean((p-truth)**2,axis=(0,2,3,4)));write(R/c['paths']['evaluation'],{'rmse':rmse.tolist(),'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,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.fuxi_weather 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=FuXiWeather(**c['model']);m=DDP(base,find_unused_parameters=True) if ddp else base;opt=torch.optim.Adam(m.parameters(),lr=c['train']['learning_rate']);ls=[]
+for i in range(rank,8,world):b,o,q,t=map(torch.tensor,(d['background'][i:i+1],d['observation'][i:i+1],d['mask'][i:i+1],d['target'][i:i+1]));a,f=m(b,o,q,i%2==1);loss=((a-t)**2).mean()+((f-t)**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'],{'loss':float(v[0]/v[1]),'world_size':world});print(p)
+if ddp:dist.destroy_process_group()