| """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") |
|
|
