| import json,torch |
| from pathlib import Path |
| from torch import nn |
| import torch.nn.functional as F |
| import yaml |
| FEATURES=('GRACE','WGHM','precipitation','ET','surface_runoff','baseflow','snowmelt','latitude','longitude') |
| def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text()) |
| def sample(i): |
| y,x=torch.meshgrid(torch.linspace(-1,1,32),torch.linspace(-1,1,32),indexing='ij');w=torch.sin(3*x+.1*i)*torch.cos(2*y);g=F.avg_pool2d(w[None,None],6,1,3)[0,0,:32,:32];g=F.interpolate(g[None,None],size=(32,32),mode='bilinear')[0,0];return torch.stack((g,w,torch.sin(x),torch.cos(y),w*.2,w*.1,w*.05,y,x)).float() |
| class Block(nn.Module): |
| def __init__(self,c):super().__init__();self.n=nn.Sequential(nn.Conv2d(c,c,3,padding=1),nn.ReLU(),nn.BatchNorm2d(c),nn.Conv2d(c,c,3,padding=1),nn.BatchNorm2d(c)) |
| def forward(self,x):return F.relu(x+self.n(x)) |
| class GRACESEDA(nn.Module): |
| def __init__(self,channels=9,hidden=16):super().__init__();self.e1=nn.Conv2d(channels,hidden,3,2,1);self.e2=nn.Conv2d(hidden,hidden*2,3,2,1);self.e3=nn.Conv2d(hidden*2,hidden*4,3,2,1);self.b=Block(hidden*4);self.d=nn.Sequential(nn.Conv2d(hidden*4,hidden*2,3,padding=1),nn.ReLU(),nn.Conv2d(hidden*2,hidden,3,padding=1),nn.ReLU(),nn.Conv2d(hidden,1,1));self.model_config={'channels':channels,'hidden':hidden} |
| def forward(self,x):z=F.relu(self.e1(x));z=F.relu(self.e2(z));z=self.b(F.relu(self.e3(z)));return F.interpolate(self.d(z),size=(32,32),mode='bilinear',align_corners=False)[:,0] |
| def loss_fn(p,x):g=x[:,0];w=x[:,1];ae=abs(p.mean((1,2))-g.mean((1,2)));pc=((p-p.mean((1,2),keepdim=True))*(w-w.mean((1,2),keepdim=True))).mean((1,2))/(p.std((1,2))*w.std((1,2))+1e-6);mae=abs(p-w).mean((1,2));return (ae+(1-pc)*mae).mean() |
| def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n') |
|
|
