import torch import math import torch.nn as nn import numpy as np import torch.nn.functional as F from onescience.modules.decoder.unet_decoder import ( UNetDecoder1D, UNetDecoder2D, UNetDecoder3D, ) from onescience.modules.encoder.unet_encoder import ( UNetEncoder1D, UNetEncoder2D, UNetEncoder3D, ) from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI from onescience.modules.head.unet_head import UNetHead1D, UNetHead2D, UNetHead3D from onescience.modules.mlp.MLP import StandardMLP from onescience.modules.embedding import timestep_embedding, unified_pos_embedding EncoderList = [None, UNetEncoder1D, UNetEncoder2D, UNetEncoder3D] DecoderList = [None, UNetDecoder1D, UNetDecoder2D, UNetDecoder3D] HeadList = [None, UNetHead1D, UNetHead2D, UNetHead3D] class Model(nn.Module): """ 多尺度物理场 U-Net 模型。 该模型支持结构化网格（1D/2D/3D）和非结构化网格（通过 GeoFNO 的几何投影）的物理场预测。 利用编码器和解码器实现多尺度特征提取与融合。 Args: args: 包含模型配置的参数命名空间 (如 task, geotype, n_hidden 等)。 device: 运行设备。 bilinear (bool, optional): U-Net 上采样是否使用双线性插值。默认值: True。 s1 (int, optional): 非结构化网格投影的潜在空间高度。默认值: 96。 s2 (int, optional): 非结构化网格投影的潜在空间宽度。默认值: 96。 形状: 输入 x: 坐标张量 (B, N, space_dim)。 输入 fx: 物理场特征张量 (B, N, fun_dim)。 输入 T: 可选的时间步张量。 输出: (B, N, out_dim) """ def __init__(self, args, device, bilinear=True, s1=96, s2=96): super(Model, self).__init__() self.__name__ = "U-Net" self.args = args self.bilinear = bilinear normtype = "bn" if args.task == "steady" else "in" # ========================================== # 1. 位置编码与预处理 (Embedding) # ========================================== if args.unified_pos and args.geotype != "unstructured": # structured mesh self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device) input_dim = args.fun_dim + args.ref ** len(args.shapelist) else: input_dim = args.fun_dim + args.space_dim self.preprocess = StandardMLP( input_dim=input_dim, hidden_dims=[args.n_hidden * 2], output_dim=args.n_hidden, activation=args.act, ) if args.time_input: self.time_fc = nn.Sequential( nn.Linear(args.n_hidden, args.n_hidden), nn.SiLU(), nn.Linear(args.n_hidden, args.n_hidden), ) # ========================================== # 2. 几何投影 (Geometry Projection) # ========================================== if self.args.geotype == "unstructured": self.fftproject_in = GeoSpectralConv2d( in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes, s1=s1, s2=s2 ) self.fftproject_out = GeoSpectralConv2d( in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes, s1=s1, s2=s2 ) self.iphi = IPHI() patch_size = [(size + (16 - size % 16) % 16) // 16 for size in [s1, s2]] self.padding = [(16 - size % 16) % 16 for size in [s1, s2]] else: patch_size = [(size + (16 - size % 16) % 16) // 16 for size in args.shapelist] self.padding = [(16 - size % 16) % 16 for size in args.shapelist] # ========================================== # 3. U-Net 核心 (Multiscale modules) # ========================================== dim = len(patch_size) num_stages = 4 self.encoder = EncoderList[dim]( in_channels=args.n_hidden, base_channels=args.n_hidden, num_stages=num_stages, bilinear=bilinear, normtype=normtype ) self.decoder = DecoderList[dim]( base_channels=args.n_hidden, num_stages=num_stages, bilinear=bilinear, normtype=normtype ) self.outc = HeadList[dim]( in_channels=args.n_hidden, out_channels=args.n_hidden ) # 最终投影 self.fc1 = nn.Linear(args.n_hidden, args.n_hidden) self.fc2 = nn.Linear(args.n_hidden, args.out_dim) def multiscale(self, x): """完全解耦的 U-Net 前向计算""" # 提取多尺度特征 features = self.encoder(x) # 融合上采样 x = self.decoder(features) # 预测头输出 return self.outc(x) def structured_geo(self, x, fx, T=None): B, N, _ = x.shape if self.args.unified_pos: x = self.pos.repeat(x.shape[0], 1, 1) if fx is not None: fx = torch.cat((x, fx), -1) fx = self.preprocess(fx) else: fx = self.preprocess(x) if T is not None: Time_emb = timestep_embedding(T, self.args.n_hidden).repeat(1, x.shape[1], 1) Time_emb = self.time_fc(Time_emb) fx = fx + Time_emb x = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist) # Padding if not all(item == 0 for item in self.padding): if len(self.args.shapelist) == 2: x = F.pad(x, [0, self.padding[1], 0, self.padding[0]]) elif len(self.args.shapelist) == 3: x = F.pad(x, [0, self.padding[2], 0, self.padding[1], 0, self.padding[0]]) # 极简调用多尺度 U-Net x = self.multiscale(x) # Unpadding if not all(item == 0 for item in self.padding): if len(self.args.shapelist) == 2: x = x[..., : -self.padding[0], : -self.padding[1]] elif len(self.args.shapelist) == 3: x = x[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]] x = x.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1) x = self.fc1(x) x = F.gelu(x) x = self.fc2(x) return x def unstructured_geo(self, x, fx, T=None): original_pos = x if fx is not None: fx = torch.cat((x, fx), -1) fx = self.preprocess(fx) else: fx = self.preprocess(x) if T is not None: Time_emb = timestep_embedding(T, self.args.n_hidden).repeat(1, x.shape[1], 1) Time_emb = self.time_fc(Time_emb) fx = fx + Time_emb # GeoFNO: 坐标映射与入场变换 x = self.fftproject_in( fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None ) # 极简调用多尺度 U-Net x = self.multiscale(x) # GeoFNO: 出场逆变换 x = self.fftproject_out( x, x_out=original_pos, iphi=self.iphi, code=None ).permute(0, 2, 1) x = self.fc1(x) x = F.gelu(x) x = self.fc2(x) return x def forward(self, x, fx, T=None, geo=None): if self.args.geotype == "unstructured": return self.unstructured_geo(x, fx, T) else: return self.structured_geo(x, fx, T)