| import json,torch |
| from pathlib import Path |
| from torch import nn |
| import yaml |
| FIELDS=('Z1000','Z500','Z250','THICK300_700','T2M','T850','TCWV','WS10','OLR') |
| def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text()) |
| def state(i,s=0): |
| y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing='ij');c=torch.arange(9)[:,None,None];a=torch.sin((c+1)*x+.04*(i+s))*torch.cos((c%4+1)*y);sst=torch.cos(x+.01*(i+s))*torch.cos(y);return a.float(),sst[None].float() |
| class Net(nn.Module): |
| def __init__(self,ci,co,h):super().__init__();self.n=nn.Sequential(nn.Conv2d(ci,h,3,padding=1),nn.GELU(),nn.Conv2d(h,h,3,padding=1),nn.GELU(),nn.Conv2d(h,co,1)) |
| def forward(self,x):return self.n(x) |
| class DLESyM(nn.Module): |
| def __init__(self,atmosphere_channels=9,hidden=16):super().__init__();self.atm=Net(10,9,hidden);self.ocean=Net(4,1,hidden//2);self.precip=Net(9,1,hidden//2);self.model_config={'atmosphere_channels':atmosphere_channels,'hidden':hidden} |
| def atmosphere(self,a,sst):return a+self.atm(torch.cat((a,sst),1)) |
| def cycle(self,a,sst): |
| seq=[] |
| for _ in range(16):a=self.atmosphere(a,sst);seq.append(a) |
| f=torch.stack(seq);forcing=torch.stack((f[:,:,7].mean(0),f[:,:,0].mean(0),f[:,:,8].mean(0)),1);sst=sst+self.ocean(torch.cat((sst,forcing),1));return a,sst,self.precip(a) |
| def forward(self,a,sst):return self.cycle(a,sst) |
| def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n') |
|
|
