Spaces:
Paused
Paused
| """Geometry-conditioned field network with spatial Fourier features.""" | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from step_surrogate import features,TASKS as ORIGINAL_TASKS | |
| TASKS={**ORIGINAL_TASKS,'mechanical':ORIGINAL_TASKS['mechanical'][:4]} | |
| class DenseNet(nn.Module): | |
| def __init__(self,query_dim): | |
| super().__init__();self.encoder=nn.Sequential(nn.Linear(6,64),nn.SiLU(),nn.Linear(64,96),nn.SiLU());self.register_buffer('frequencies',torch.tensor([1.,2.,4.,8.])*np.pi) | |
| self.body=nn.Sequential(nn.Linear(query_dim+192+24+1,192),nn.SiLU(),nn.Linear(192,192),nn.SiLU()) | |
| self.heads=nn.ModuleDict({k:nn.Sequential(nn.Linear(192,96),nn.SiLU(),nn.Linear(96,len(v))) for k,v in TASKS.items()}) | |
| for task in ('filling','cooling'): | |
| self.register_buffer(task+'_sf_mean',torch.tensor(0.));self.register_buffer(task+'_sf_std',torch.tensor(1.)) | |
| def encode(self,cloud): | |
| a=self.encoder(cloud);return torch.cat([a.max(0).values,a.mean(0)]) | |
| def forward(self,cloud,query,time_s,task,encoded=None): | |
| code=self.encode(cloud) if encoded is None else encoded;t=torch.log1p(time_s.clamp(min=0)).reshape(-1,1)/10 | |
| phase=(query[:,:3,None]*self.frequencies).flatten(1);pos=torch.cat([phase.sin(),phase.cos()],1) | |
| out=self.heads[task](self.body(torch.cat([query,code.expand(len(query),-1),pos,t],1))) | |
| if task in ('filling','cooling'): | |
| out=out.clone();out[:,1]=(out[:,1].sigmoid()-getattr(self,task+'_sf_mean'))/getattr(self,task+'_sf_std') | |
| return out | |