File size: 4,335 Bytes
ff0fadf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import torch
import torch.nn as nn
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
class Model(nn.Module):
"""
DeepONet 模型。
采用双网络架构 (Branch Net 和 Trunk Net) 逼近非线性算子。
- Branch Net 处理输入函数 (如初始条件、边界条件)。
- Trunk Net 处理目标评估点的坐标 (物理位置)。
两者输出特征做内积后得到目标位置的物理场预测值。
"""
def __init__(self, args, device):
super(Model, self).__init__()
self.__name__ = "DeepONet"
self.args = args
# 位置处理
if args.unified_pos and args.geotype != "unstructured":
self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
trunk_in_dim = args.ref ** len(args.shapelist)
else:
trunk_in_dim = args.space_dim
# 分支网络输入维度（函数输入+时间嵌入）
branch_in_dim = args.fun_dim
if args.time_input:
branch_in_dim += args.n_hidden # 增加时间嵌入维度
# 分支网络（函数空间）
self.branch_net = StandardMLP(
input_dim=branch_in_dim,
output_dim=args.n_hidden,
hidden_dims=[args.n_hidden] * args.branch_depth,
activation=args.act,
use_bias=True
)
# 主干网络（物理空间）
self.trunk_net = StandardMLP(
input_dim=trunk_in_dim,
output_dim=args.n_hidden,
hidden_dims=[args.n_hidden] * args.trunk_depth,
activation=args.act,
use_bias=True
)
# 时间处理层
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),
)
# 输出层 - 支持多个输出通道
self.out_layer = nn.Linear(args.n_hidden, args.out_dim)
self.bias = nn.Parameter(torch.zeros(1, 1, args.out_dim))
def structured_forward(self, x, fx, T=None, geo=None):
B, N, _ = x.shape # x: [B, N, space_dim]
if self.args.unified_pos:
x = self.pos.repeat(B, 1, 1) # [B, N, d]
# 处理时间信息
if T is not None and self.args.time_input:
# 确保时间输入维度正确 [B, 1] -> [B]
T = T.view(-1)
# 生成时间嵌入 [B, D]
T_emb = timestep_embedding(T, self.args.n_hidden)
# 通过时间处理层 [B, D] -> [B, D]
T_emb = self.time_fc(T_emb)
# 扩展时间嵌入以匹配空间点 [B, D] -> [B, N, D]
T_emb_expanded = T_emb.unsqueeze(1).expand(-1, N, -1)
# 将时间嵌入连接到函数输入
if fx is not None:
# [B, N, fun_dim] + [B, N, D] = [B, N, fun_dim + D]
fx = torch.cat([fx, T_emb_expanded], dim=-1)
else:
fx = T_emb_expanded # [B, N, D]
# 分支网络处理
branch_feat = self.branch_net(fx) # [B, N, D]
# 主干网络处理
trunk_feat = self.trunk_net(x) # [B, N, D]
# 点积操作 + 输出层
inner = branch_feat * trunk_feat
out = self.out_layer(inner) + self.bias
return out
def unstructured_forward(self, x, fx, T=None):
B, N, _ = x.shape
if T is not None and self.args.time_input:
T = T.view(-1)
T_emb = timestep_embedding(T, self.args.n_hidden)
T_emb = self.time_fc(T_emb)
T_emb_expanded = T_emb.unsqueeze(1).expand(-1, N, -1)
if fx is not None:
fx = torch.cat([fx, T_emb_expanded], dim=-1)
else:
fx = T_emb_expanded
branch_feat = self.branch_net(fx)
trunk_feat = self.trunk_net(x)
inner = branch_feat * trunk_feat
out = self.out_layer(inner) + self.bias
return out
def forward(self, x, fx, T=None, geo=None):
if self.args.geotype == "unstructured":
return self.unstructured_forward(x, fx, T)
else:
return self.structured_forward(x, fx, T)
|