| """WRF column layout adapter and compact Model-D radiation emulator.""" |
| import json,math |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch import nn |
| import yaml |
|
|
| INPUT_NAMES=("pressure","temperature","water_vapor","ozone","cloud_water","cloud_ice","cloud_fraction","albedo","solar_zenith","emissivity") |
| OUTPUT_NAMES=("sw_flux","lw_flux","sw_heating_rate","lw_heating_rate","surface_sw_flux","toa_sw_flux") |
| def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text()) |
| def synthetic_columns(i,n=1): |
| level=torch.linspace(1,0,57);x=torch.stack([level,250+45*level,0.012*level**2,.001*(1-level),.002*torch.sin(math.pi*level)**2,.001*torch.cos(math.pi*level)**2,.5*torch.sin(math.pi*level)**2,torch.full_like(level,.2),torch.full_like(level,.6),torch.full_like(level,.95)],1);x=x[None].repeat(n,1,1);x+=torch.randn_like(x)*.001;sw=500*torch.exp(-1.2*(1-level));lw=220+80*level;shr=torch.gradient(sw,spacing=(level,))[0]/100;lhr=torch.gradient(lw,spacing=(level,))[0]/100;profiles=torch.stack((sw,lw,shr,lhr),1)[None].repeat(n,1,1);boundary=torch.tensor([sw[-1],sw[0]])[None].repeat(n,1);return x.float(),profiles.float(),boundary.float() |
| def preprocess_layout(wrf): |
| if wrf.ndim!=4 or wrf.shape[1]!=57 or wrf.shape[-1]!=10:raise ValueError("expected WRF [i,57,j,10]") |
| return wrf.permute(0,2,1,3).reshape(-1,57,10) |
| class RadiationBiLSTM(nn.Module): |
| def __init__(self,input_features=10,hidden_size=16,layers=3,profile_outputs=4,boundary_outputs=2): |
| super().__init__();self.rnn=nn.LSTM(input_features,hidden_size,layers,batch_first=True,bidirectional=True);self.profile=nn.Linear(hidden_size*2,profile_outputs);self.boundary=nn.Linear(hidden_size*4,boundary_outputs);self.model_config={"input_features":input_features,"hidden_size":hidden_size,"layers":layers,"profile_outputs":profile_outputs,"boundary_outputs":boundary_outputs} |
| def forward(self,x): |
| if x.shape[1:]!=(57,10):raise ValueError("expected [B,57,10]") |
| h,_=self.rnn(x);return self.profile(h),self.boundary(torch.cat((h[:,0],h[:,-1]),-1)) |
| class SynchronousCoupler: |
| def __init__(self,model):self.model=model |
| def infer_run(self,wrf_layout): |
| columns=preprocess_layout(wrf_layout);profiles,boundary=self.model(columns);i,_,j,_=wrf_layout.shape;return profiles.reshape(i,j,57,-1).permute(0,2,1,3),boundary.reshape(i,j,-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") |
|
|
