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