Publish CLM5-Emulator reproduction

.gitattributes

@@ -1,35 +1,7 @@
-*.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
+*.png binary

conf/config.yaml

@@ -0,0 +1,31 @@
+seed: 42
+data:
+  format_version: clm5_emulator_v1
+  path: data/clm5_ppe.npz
+  samples: 40
+  grid: [46, 72]
+  input_parameters: 6
+  targets: [GPP, LHF]
+  eof_modes: 3
+model:
+  hidden_dim: 12
+train:
+  epochs: 4
+  batch_size: 10
+  learning_rate: 0.005
+  weight_decay: 0.0001
+paths:
+  checkpoint: result/checkpoints/clm5_emulator.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png
+paper_model:
+  ppe_members: 100
+  input_parameters: [medlynslope, dleaf, kmax, fff, dint, baseflow_scalar]
+  hidden_layers: 2
+  hidden_candidates: [5, 15]
+  optimizer: RMSprop
+  learning_rate: 0.005
+  batch_size: 30
+  split: [0.6, 0.2, 0.2]

config.json

@@ -0,0 +1 @@
+{"model_name":"CLM5-Emulator","model_type":"clm5_emulator","architectures":["CLM5Emulator"],"framework":"PyTorch","domain":"land-surface","task":"emulation-and-parameter-estimation","implementation":{"entry_point":"model/clm5_emulator.py","scope":"L2 core-method engineering reproduction"},"architecture":{"input_parameters":6,"independent_targets":["GPP","LHF"],"eof_modes_per_target":3,"hidden_layers":2,"activations":["relu","tanh","linear"]},"data":{"driver":"GSWP3","forcing_interval_hours":3,"simulation_years":20,"spinup_years":15,"analysis_years":5,"paper_grid_degrees":[4,5],"engineering_grid_ledger":[46,72],"synthetic":true},"configuration_sources":["conf/config.yaml","model/clm5_emulator.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/clm5_emulator.py

@@ -0,0 +1,24 @@
+"""CLM5 GPP/LHF EOF emulators and bounded parameter search."""
+import json
+from pathlib import Path
+import numpy as np
+import torch
+from torch import nn
+import yaml
+
+PARAMETERS=("medlynslope","dleaf","kmax","fff","dint","baseflow_scalar")
+def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text())
+class TargetEmulator(nn.Module):
+ def __init__(self,hidden_dim=12):super().__init__();self.net=nn.Sequential(nn.Linear(6,hidden_dim),nn.ReLU(),nn.Linear(hidden_dim,hidden_dim),nn.Tanh(),nn.Linear(hidden_dim,3))
+ def forward(self,x):return self.net(x)
+class CLM5Emulator(nn.Module):
+ def __init__(self,hidden_dim=12):super().__init__();self.gpp=TargetEmulator(hidden_dim);self.lhf=TargetEmulator(hidden_dim);self.model_config={"hidden_dim":hidden_dim}
+ def forward(self,x):
+  if x.ndim!=2 or x.shape[1]!=6:raise ValueError("expected [B,6] normalized parameters")
+  return torch.stack((self.gpp(x),self.lhf(x)),1)
+def reconstruct(pc,eof,mean):return np.einsum("btm,tmg->btg",pc,eof)+mean[None]
+def bounded_search(model,observed_pc,scales,weights,candidates=4096,seed=0):
+ rng=np.random.default_rng(seed);p=torch.from_numpy(rng.random((candidates,6)).astype("float32"));
+ with torch.no_grad():pred=model(p).numpy()
+ cost=(((pred-observed_pc[None])/scales[None])**2*weights[None]).sum((1,2));i=int(cost.argmin());return p[i].numpy(),float(cost[i])
+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,numpy as np
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.clm5_emulator import *
+c=load_config(ROOT);rng=np.random.default_rng(c["seed"]);n=c["data"]["samples"];p=rng.random((n,6)).astype("float32");lat=np.linspace(-90,90,46);lon=np.arange(72)*5;la,lo=np.meshgrid(lat,lon,indexing="ij");eof=np.stack([[np.cos(np.deg2rad(la)),np.sin(np.deg2rad(lo)),np.sin(np.deg2rad(la*2+lo))],[np.cos(np.deg2rad(la))*np.cos(np.deg2rad(lo)),np.sin(np.deg2rad(la)),np.cos(np.deg2rad(lo*2))]]).astype("float32");eof=eof.reshape(2,3,-1);pc=np.stack((np.stack((2*p[:,0]-p[:,3],p[:,1]*p[:,2],np.sin(np.pi*p[:,4])),1),np.stack((p[:,2]+p[:,5],p[:,0]*p[:,4],np.cos(np.pi*p[:,3])),1)),1).astype("float32");mean=np.stack((2+np.cos(np.deg2rad(la)),50+5*np.sin(np.deg2rad(la)))).reshape(2,-1).astype("float32");fields=reconstruct(pc,eof,mean).reshape(n,2,46,72);split=np.array([0]*24+[1]*8+[2]*8);path=ROOT/c["data"]["path"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,parameters=p,pc=pc,eof=eof,mean=mean,fields=fields,split=split,latitude=lat,longitude=lon,format_version=c["data"]["format_version"]);print(path)

scripts/inference.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+import sys,numpy as np,torch
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.clm5_emulator 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=CLM5Emulator(**ck["model_config"]);m.load_state_dict(ck["model"]);mask=d["split"]==2
+with torch.no_grad():pc=m(torch.tensor(d["parameters"][mask])).numpy()
+fields=reconstruct(pc,d["eof"],d["mean"]).reshape(-1,2,46,72);obs=d["pc"].mean(0);best,cost=bounded_search(m,obs,np.maximum(d["pc"].std(0),.1),np.ones((2,3)),seed=c["seed"]);path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,predicted_pc=pc,target_pc=d["pc"][mask],predicted_fields=fields,target_fields=d["fields"][mask],best_parameters=best,cost=cost,latitude=d["latitude"],longitude=d["longitude"]);print(path)

scripts/result.py

@@ -0,0 +1,8 @@
+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.clm5_emulator import load_config,write_json
+c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);p=d["predicted_pc"];t=d["target_pc"];r2=[]
+for v in range(2):r2.append([float(np.corrcoef(p[:,v,k],t[:,v,k])[0,1]**2) for k in range(3)])
+rmse=np.sqrt(np.mean((d["predicted_fields"]-d["target_fields"])**2,axis=(0,2,3)));write_json(ROOT/c["paths"]["evaluation"],{"pc_r2":r2,"spatial_rmse":{"GPP":float(rmse[0]),"LHF":float(rmse[1])},"estimated_parameters":d["best_parameters"].tolist(),"calibration_cost":float(d["cost"]),"synthetic":True});fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].scatter(t[:,0,0],p[:,0,0]);ax[0].set(xlabel="CLM PC1",ylabel="Emulated PC1",title="GPP emulator");im=ax[1].imshow((d["predicted_fields"][0,1]-d["target_fields"][0,1]),cmap="coolwarm");ax[1].set_title("LHF spatial error");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,16 @@
+from pathlib import Path
+import sys,os,torch,numpy as np
+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.clm5_emulator 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"]);d=np.load(ROOT/c["data"]["path"]);x=torch.tensor(d["parameters"][d["split"]==0]);y=torch.tensor(d["pc"][d["split"]==0]);base=CLM5Emulator(**c["model"]);m=DDP(base) if distributed else base;opt=torch.optim.RMSprop(m.parameters(),lr=c["train"]["learning_rate"],weight_decay=c["train"]["weight_decay"]);losses=[]
+for _ in range(c["train"]["epochs"]):
+ for i in range(rank,len(x),world):loss=((m(x[i:i+1])-y[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"],{"mse":float(v[0]/v[1]),"world_size":world});print(path)
+if distributed:dist.destroy_process_group()