from typing import List, Dict, Optional import torch import numpy as np from torch import Tensor import torch.nn as nn import torch.nn.functional as F from ..base_model import AutoCfdModel from onescience.modules.fourier.fno_layers import SpectralConv2d torch.manual_seed(0) np.random.seed(0) class FnoBlock(nn.Module): """ FNO 块 (Fourier Neural Operator Block)。 包含一个 2D 频域卷积路径和一个 1x1 的空间卷积残差路径， 最后加上可选的激活函数。 """ def __init__( self, in_chan: int, out_chan: int, modes1: int, modes2: int, act_fn: Optional[nn.Module] = None, ): super().__init__() self.in_chan = in_chan self.out_chan = out_chan self.modes1 = modes1 self.modes2 = modes2 self.act_fn = act_fn self.conv0 = SpectralConv2d( in_channels=self.in_chan, out_channels=self.out_chan, modes1=self.modes1, modes2=self.modes2 ) self.w0 = nn.Conv2d(self.in_chan, self.out_chan, 1) def forward(self, x: Tensor) -> Tensor: x1 = self.conv0(x) x2 = self.w0(x) x = x1 + x2 if self.act_fn is not None: x = self.act_fn(x) return x class Fno2d(AutoCfdModel): def __init__( self, in_chan: int, out_chan: int, n_case_params: int, loss_fn: nn.Module, num_layers: int, modes1: int = 12, modes2: int = 12, hidden_dim: int = 20, padding: Optional[int] = None, ): super().__init__(loss_fn) """ The overall network. It contains 4 layers of the Fourier layer. 1. Lift the input to the desire channel dimension by self.fc0 . 2. 4 layers of the integral operators u' = (W + K)(u). W defined by self.w; K defined by self.conv . 3. Project from the channel space to the output space by self.fc1 and self.fc2 . """ self.in_chan = in_chan self.out_chan = out_chan self.n_case_params = n_case_params self.num_layers = num_layers self.modes1 = modes1 self.modes2 = modes2 self.hidden_dim = hidden_dim self.padding = padding # pad the domain if input is non-periodic self.act_fn = nn.GELU() # Channel projection into `hidden_dim` channels # +1 for mask, +2 for coordinates self.fc0 = nn.Conv2d( in_chan + 1 + 2 + n_case_params, self.hidden_dim, 1, 1, 0, ) # input channel is 12: the solution of the previous 10 # timesteps + 2 locations # (u(t-10, x, y), ..., u(t-1, x, y), x, y) # FNO blocks blocks = [] for _ in range(self.num_layers): blocks.append( FnoBlock( self.hidden_dim, self.hidden_dim, self.modes1, self.modes2, self.act_fn, ) ) self.blocks = nn.Sequential(*blocks) self.fc1 = nn.Conv2d(self.hidden_dim, 128, 1, 1, 0) self.fc2 = nn.Conv2d(128, self.out_chan, 1, 1, 0) def forward( self, inputs: Tensor, case_params: Tensor, mask: Optional[Tensor] = None, label: Optional[Tensor] = None, ) -> Dict: """ Args: - input: (b, c, h, w) - labels: (b, c, h, w) - mask: (b, h, w), a binary mask for indicating the geometry 1 for interior, 0 for obstacles. Returns: a tensor of shape (b, c, h, w), the solution at the next timestep """ batch_size, n_chan, height, width = inputs.shape if mask is None: # When there is no mask, we assume that there is no obstacles. mask = torch.ones((batch_size, 1, height, width)).to(inputs.device) else: if mask.dim() == 3: # (B, h, w) mask = mask.unsqueeze(1) # (B, 1, h, w) inputs = torch.cat([inputs, mask], dim=1) # (B, c + 1, h, w) # Physical properties props = case_params # (B, p) props = props.unsqueeze(-1).unsqueeze(-1) # (B, p, 1, 1) props = props.repeat(1, 1, height, width) # (B, p, H, W) # Append (x, y) coordinates to every location grid = self.get_coords(inputs.shape, inputs.device) # (b, 2, h, w) inputs = torch.cat( (inputs, grid, props), dim=1 ) # (b, c + 2 + 2, h, w) # Project channels inputs = self.fc0(inputs) # (b, hidden_dim, h, w) # x = x.permute(0, 3, 1, 2) # (b, c, h, w)? if self.padding is not None: # pad the domain if input is non-periodic inputs = F.pad(inputs, [0, self.padding, 0, self.padding]) inputs = self.blocks(inputs) # (b, hidden_dim, h, w) if self.padding is not None: # pad the domain if inputis non-periodic inputs = inputs[..., : -self.padding, : -self.padding] inputs = self.fc1(inputs) # (b, 128, h, w) inputs = self.act_fn(inputs) preds = self.fc2(inputs) # (b, c_out, h, w) # Masked locations are not prediction preds = preds * mask if label is not None: label = label * mask loss = self.loss_fn(preds=preds, labels=label) return dict( preds=preds, loss=loss, ) return dict(preds=preds) def get_coords(self, shape, device): """ Return a tensor of shape (b, 2, h, w) such that the element at [:, :, i, j] is the (x, y) coordinates at the grid location (i, j). """ bsz, c, size_x, size_y = shape grid_x = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float) grid_x = grid_x.reshape(1, 1, size_x, 1).repeat([bsz, 1, 1, size_y]) grid_y = torch.tensor(np.linspace(0, 1, size_y), dtype=torch.float) grid_y = grid_y.reshape(1, 1, 1, size_y).repeat([bsz, 1, size_x, 1]) coords = torch.cat([grid_x, grid_y], dim=1).to(device) # (b, 2, h, w) return coords def generate( self, inputs: Tensor, case_params: Tensor, mask: Optional[Tensor] = None, ) -> Tensor: outputs = self.forward( inputs=inputs, case_params=case_params, mask=mask ) # (b, c, h, w) preds = outputs["preds"] return preds def generate_many( self, inputs: Tensor, case_params: Tensor, mask: Tensor, steps: int ) -> List[Tensor]: """ Args: x (Tensor): (c, h, w) case_params (Tensor): (p) mask (Tensor): (h, w), 1 for interior, 0 for obstacles. Returns: output: (steps, c, h, w) """ assert len(inputs.shape) == len(case_params.shape) + 2 if inputs.dim() == 3: # Add a dimension for batch size of 1 inputs = inputs.unsqueeze(0) case_params = case_params.unsqueeze(0) mask = mask.unsqueeze(0) assert inputs.shape[0] == case_params.shape[0] == mask.shape[0] cur_frame = inputs # (b, c, h, w) preds = [] for _ in range(steps): cur_frame = self.generate( inputs=cur_frame, case_params=case_params, mask=mask ) preds.append(cur_frame) return preds