484 lines
19 KiB
Python
484 lines
19 KiB
Python
import torch
|
||
from torch.autograd import Function
|
||
# import modules
|
||
from dataclasses import dataclass
|
||
from tqdm.auto import tqdm
|
||
import numpy as np
|
||
from getdata import GetDataset
|
||
# deep learning modules
|
||
import scipy.sparse as sp
|
||
from scipy.sparse.linalg import spilu
|
||
from scipy.io import loadmat
|
||
from torch.autograd import Variable
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import torch.optim as optim
|
||
from torch.utils.data import Dataset, DataLoader
|
||
import datetime
|
||
import pandas as pd
|
||
# Plot modules
|
||
import matplotlib.pyplot as plt
|
||
from scipy.interpolate import griddata
|
||
from pathlib import Path
|
||
from scipy.io import savemat
|
||
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
class ILUApply(Function):
|
||
@staticmethod
|
||
def forward(ctx, r_torch, ilu):
|
||
"""
|
||
r_torch: torch complex tensor, shape (Mi,)
|
||
ilu: fixed SciPy spilu object
|
||
"""
|
||
ctx.ilu = ilu
|
||
r_np = r_torch.detach().cpu().numpy()
|
||
z_np = ilu.solve(r_np) # z = M^{-1} r
|
||
z = torch.from_numpy(z_np).to(r_torch.device).to(r_torch.dtype)
|
||
return z
|
||
|
||
@staticmethod
|
||
def backward(ctx, grad_out):
|
||
"""
|
||
grad_out: dL/dz
|
||
complex: grad_r = M^{-H} grad_out
|
||
"""
|
||
ilu = ctx.ilu
|
||
g_np = grad_out.detach().cpu().numpy()
|
||
gr_np = ilu.solve(g_np, trans='H') # conjugate-transpose solve
|
||
grad_r = torch.from_numpy(gr_np).to(grad_out.device).to(grad_out.dtype)
|
||
return grad_r, None
|
||
|
||
@dataclass
|
||
class PINNConfig:
|
||
# 训练（仅 PDE 残差时建议 batch_size 较小如 16~32，避免 branch 输出塌缩）
|
||
epochs: int = 1_000
|
||
batch_size: int = 32
|
||
learning_rate: float = 1e-2
|
||
step_size: int = 500 # StepLR: 每 step_size 轮衰减一次
|
||
gamma: float = 0.95 # StepLR 衰减系数
|
||
max_grad_norm: float = 1.0 # 梯度裁剪，稳定训练
|
||
print_every: int = 1 # 每多少轮打印一次
|
||
save_every: int = 1000 # 每多少轮保存一次 checkpoint
|
||
|
||
# 数据（n 为样本数目，如 deepOnet_data_A1_100 表示 100 个样本）
|
||
matpath: str = "deepOnet_data_A1_1558_8_2"
|
||
|
||
# 保存
|
||
save_dir: str = "./model_save"
|
||
results_dir: str = "./results"
|
||
load_file_name: str = "model_A1_size_1558_2"
|
||
save_file_name: str = "model_A1_size_1558_2"
|
||
|
||
# 设备/精度
|
||
device: str = "cuda" if torch.cuda.is_available() else "cpu"
|
||
dtype: torch.dtype = torch.float64 # 和 .mat 的 float64 对齐（double）
|
||
num_workers: int = 10
|
||
pin_memory: bool = True
|
||
|
||
class Modified_MLP_Block(nn.Module):
|
||
def __init__(self, input_dim, hidden_channel, output_dim, hidden_size=6):
|
||
super(Modified_MLP_Block, self).__init__()
|
||
self.activation = nn.Tanh()
|
||
self.encodeU = nn.Linear(input_dim, hidden_channel)
|
||
self.encodeV = nn.Linear(input_dim, hidden_channel)
|
||
self.In = nn.Linear(input_dim, hidden_channel)
|
||
|
||
self.hidden_layers = nn.ModuleList([
|
||
nn.Linear(hidden_channel, hidden_channel) for _ in range(hidden_size)
|
||
])
|
||
self.out = nn.Linear(hidden_channel, output_dim)
|
||
self._init_weights()
|
||
|
||
def _init_weights(self):
|
||
torch.manual_seed(123)
|
||
gain = nn.init.calculate_gain('tanh')
|
||
for m in self.modules():
|
||
if isinstance(m, nn.Linear):
|
||
nn.init.xavier_uniform_(m.weight, gain=gain)
|
||
nn.init.zeros_(m.bias)
|
||
|
||
def forward(self, x):
|
||
U = self.activation(self.encodeU(x))
|
||
V = self.activation(self.encodeV(x))
|
||
Hidden = self.activation(self.In(x))
|
||
|
||
for layer in self.hidden_layers:
|
||
Z = self.activation(layer(Hidden))
|
||
Hidden = (1 - Z) * U + Z * V
|
||
|
||
x = self.out(Hidden)
|
||
return x
|
||
|
||
|
||
|
||
def _branch_norm2d(channels):
|
||
"""Branch 内使用 InstanceNorm2d：按样本、按通道独立归一化，不跨样本混合，避免输出塌缩。"""
|
||
return nn.InstanceNorm2d(channels)
|
||
|
||
|
||
def add_spatial_coord_channels(epsilon_data):
|
||
"""
|
||
在 ε 图像上拼接空间坐标通道，使 CNN 能区分「同形不同位」的介质。
|
||
epsilon_data: (B, 1, H, W) -> 返回 (B, 3, H, W)，通道为 [ε, x_norm, y_norm]，归一化到 [0,1]。
|
||
"""
|
||
B, _, H, W = epsilon_data.shape
|
||
device, dtype = epsilon_data.device, epsilon_data.dtype
|
||
j = torch.linspace(0, 1, W, device=device, dtype=dtype).view(1, 1, 1, W).expand(B, 1, H, W)
|
||
i = torch.linspace(0, 1, H, device=device, dtype=dtype).view(1, 1, H, 1).expand(B, 1, H, W)
|
||
return torch.cat([epsilon_data, j, i], dim=1)
|
||
|
||
|
||
class CNN_Branch_Residual(nn.Module):
|
||
"""带残差连接的 CNN 分支（Branch 用 InstanceNorm2d；输入含 ε + 空间坐标通道以区分同形不同位）。"""
|
||
|
||
def __init__(self, in_channels=3, num_classes=128):
|
||
super(CNN_Branch_Residual, self).__init__()
|
||
|
||
# 初始卷积层
|
||
self.initial = nn.Sequential(
|
||
nn.Conv2d(in_channels, 32, kernel_size=3, padding=1, bias=False),
|
||
_branch_norm2d(32),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(32, 32, kernel_size=3, padding=1, bias=False),
|
||
_branch_norm2d(32),
|
||
nn.ReLU(inplace=True)
|
||
)
|
||
|
||
# 残差块
|
||
self.res_block1 = ResidualBlock(32, 64, stride=2, norm_layer=_branch_norm2d)
|
||
self.res_block2 = ResidualBlock(64, 128, stride=2, norm_layer=_branch_norm2d)
|
||
|
||
# 全局平均池化
|
||
self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
|
||
|
||
# 全连接层
|
||
self.fc = nn.Linear(128, num_classes)
|
||
|
||
def forward(self, x):
|
||
x = self.initial(x)
|
||
x = self.res_block1(x)
|
||
x = self.res_block2(x)
|
||
x = self.global_avg_pool(x)
|
||
x = x.view(x.size(0), -1)
|
||
x = self.fc(x)
|
||
return x
|
||
|
||
|
||
class ResidualBlock(nn.Module):
|
||
"""残差块（支持 norm_layer，Branch 中传入 GroupNorm）。"""
|
||
|
||
def __init__(self, in_channels, out_channels, stride=1, norm_layer=None):
|
||
super(ResidualBlock, self).__init__()
|
||
if norm_layer is None:
|
||
norm_layer = nn.BatchNorm2d
|
||
|
||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3,
|
||
stride=stride, padding=1, bias=False)
|
||
self.bn1 = norm_layer(out_channels)
|
||
|
||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
|
||
stride=1, padding=1, bias=False)
|
||
self.bn2 = norm_layer(out_channels)
|
||
|
||
self.relu = nn.ReLU(inplace=True)
|
||
|
||
# 下采样连接
|
||
self.downsample = None
|
||
if stride != 1 or in_channels != out_channels:
|
||
self.downsample = nn.Sequential(
|
||
nn.Conv2d(in_channels, out_channels, kernel_size=1,
|
||
stride=stride, bias=False),
|
||
norm_layer(out_channels)
|
||
)
|
||
|
||
def forward(self, x):
|
||
identity = x
|
||
|
||
out = self.conv1(x)
|
||
out = self.bn1(out)
|
||
out = self.relu(out)
|
||
|
||
out = self.conv2(out)
|
||
out = self.bn2(out)
|
||
|
||
if self.downsample is not None:
|
||
identity = self.downsample(x)
|
||
|
||
out += identity
|
||
out = self.relu(out)
|
||
|
||
return out
|
||
|
||
class DeepONet(nn.Module):
|
||
def __init__(self, branch_input_dim, trunk_input_dim, hidden_channel, output_dim):
|
||
super(DeepONet, self).__init__()
|
||
self.output_dim = output_dim
|
||
# branch 输入为 ε + 2 个空间坐标通道，共 3 通道，便于区分同形不同位的介质
|
||
self.branch_net = CNN_Branch_Residual(in_channels=3, num_classes=output_dim)
|
||
self.trunk_net = Modified_MLP_Block(trunk_input_dim, hidden_channel, output_dim)
|
||
|
||
def forward(self, branch_input, trunk_input):
|
||
branch_input = add_spatial_coord_channels(branch_input)
|
||
branch_out = self.branch_net(branch_input)
|
||
trunk_out = self.trunk_net(trunk_input)
|
||
|
||
B1, B2 = branch_out[:, :self.output_dim//2], branch_out[:, self.output_dim//2:]
|
||
T1, T2 = trunk_out[:, :, :self.output_dim//2], trunk_out[:, :, self.output_dim//2:]
|
||
#print("B1 shape:", B1.shape, "B2 shape:", B2.shape)
|
||
#print("T1 shape:", T1.shape, "T2 shape:", T2.shape)
|
||
s_re = torch.einsum('bi,bni->bn', B1, T1) #实部
|
||
s_im = torch.einsum('bi,bni->bn', B2, T2)
|
||
return s_re, s_im
|
||
|
||
|
||
class PINN_maxwell():
|
||
def __init__(self, model, config: PINNConfig):
|
||
self.cfg = config
|
||
|
||
self.device = torch.device(self.cfg.device)
|
||
self.model = model.to(self.device, dtype=self.cfg.dtype)
|
||
self.batch_size = self.cfg.batch_size
|
||
self.learning_rate = self.cfg.learning_rate
|
||
self.matpath = self.cfg.matpath
|
||
self.loss_fn = nn.MSELoss()
|
||
self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
|
||
self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=self.cfg.step_size, gamma=self.cfg.gamma)
|
||
self.losses = []
|
||
self.lamda = []
|
||
self.save_file_name = self.cfg.save_file_name
|
||
self.load_file_name = self.cfg.load_file_name
|
||
self.save_dir = Path(self.cfg.save_dir)
|
||
self.train_set, self.test_set = self.load_dataset()
|
||
self.train_loader = DataLoader(self.train_set, self.cfg.batch_size, shuffle=True)
|
||
self.test_loader = DataLoader(self.test_set, batch_size=len(self.test_set), shuffle=False)
|
||
self.ilu_cache = {}
|
||
self._build_ilu_cache()
|
||
|
||
|
||
def load_model(self):
|
||
self.model.load_state_dict(torch.load(self.save_dir / f'{self.load_file_name}.pth', map_location=self.device, weights_only=True))
|
||
|
||
def E_function(self, epsilon_data, coord_data):
|
||
epsilon_data = epsilon_data.to(self.device)
|
||
coord_data = coord_data.to(self.device)
|
||
return self.model(epsilon_data, coord_data)
|
||
|
||
def get_data_loss(self, epsilon_data, coord_data, E_true):
|
||
E_re_pred, E_im_pred = self.E_function(epsilon_data, coord_data)
|
||
E_re_true = E_true[:,:, 0]
|
||
E_im_true = E_true[:,:, 1]
|
||
data_loss = self.loss_fn(E_re_pred, E_re_true) + self.loss_fn(E_im_pred, E_im_true)
|
||
return data_loss
|
||
|
||
def get_fem_loss(self, indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len):
|
||
"""indices: (B,) 每个样本的全局索引，用于从 ilu_cache 取对应的 ILU。"""
|
||
Ere_pred, Eim_pred = self.E_function(epsilon_data, coord_data)
|
||
E = torch.complex(Ere_pred, Eim_pred) # shape: (B, Mmax)
|
||
|
||
B, Mmax = E.shape
|
||
Mi = coord_len.squeeze(-1).long().to(self.device) # shape: (B,)
|
||
|
||
arangeM = torch.arange(Mmax, device=self.device) # (Mmax,)
|
||
mask_x = arangeM[None, :] < Mi[:, None] # (B, Mmax)
|
||
|
||
x_flat = E[mask_x] # (sum Mi,)
|
||
b_flat = b.to(self.device)[mask_x].to(x_flat.dtype) # (sum Mi,)
|
||
|
||
sumMi = int(Mi.sum().item())
|
||
offsets = torch.cumsum(
|
||
torch.cat([torch.zeros(1, device=self.device, dtype=torch.long), Mi[:-1]]),
|
||
dim=0
|
||
) # (B,)
|
||
|
||
Ai = Ai.to(self.device).long()
|
||
Aj = Aj.to(self.device).long()
|
||
Av = Av.to(self.device).to(x_flat.dtype)
|
||
mask_nnz = (Ai > 0) & (Aj > 0) # (B, Kmax)，padding 位置为 0
|
||
|
||
rows = (Ai - 1 + offsets.unsqueeze(1)).masked_select(mask_nnz)
|
||
cols = (Aj - 1 + offsets.unsqueeze(1)).masked_select(mask_nnz)
|
||
vals = Av.masked_select(mask_nnz)
|
||
|
||
y = torch.zeros(sumMi, dtype=x_flat.dtype, device=self.device)
|
||
y.scatter_add_(0, rows, vals * x_flat.index_select(0, cols))
|
||
r = y - b_flat
|
||
|
||
# 按样本用缓存的 ILU 计算 z，每个样本一个 ILU
|
||
z_parts = []
|
||
for i in range(B):
|
||
start = int(offsets[i].item())
|
||
m = int(Mi[i].item())
|
||
r_i = r[start : start + m]
|
||
ilu = self.ilu_cache[int(indices[i].item())]
|
||
z_i = ILUApply.apply(r_i, ilu)
|
||
z_parts.append(z_i)
|
||
z = torch.cat(z_parts, dim=0)
|
||
|
||
loss = (z.abs() ** 2).mean()
|
||
return loss
|
||
|
||
|
||
@torch.no_grad()
|
||
def test_E_loss(self):
|
||
self.model.eval()
|
||
total_loss = 0.0
|
||
for indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len in self.test_loader:
|
||
indices = indices.to(self.device)
|
||
epsilon_data = epsilon_data.to(self.device)
|
||
coord_data = coord_data.to(self.device)
|
||
loss = self.get_fem_loss(indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len)
|
||
total_loss += loss.item()
|
||
self.model.train()
|
||
return total_loss / len(self.test_loader)
|
||
|
||
def train(self, epochs, print_every=100, save_every=10000):
|
||
self.losses.append(['epoch', 'fem_loss', 'test_loss'])
|
||
start_time = datetime.datetime.now()
|
||
for epoch in tqdm(range(epochs), desc='Training'):
|
||
self.model.train()
|
||
total_loss = 0.0
|
||
fem_loss = 0.0
|
||
|
||
for indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len in self.train_loader:
|
||
indices = indices.to(self.device)
|
||
epsilon_data = epsilon_data.to(self.device)
|
||
coord_data = coord_data.to(self.device)
|
||
self.optimizer.zero_grad()
|
||
|
||
fem_loss = self.get_fem_loss(indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len)
|
||
|
||
loss = fem_loss
|
||
|
||
loss.backward()
|
||
if self.cfg.max_grad_norm > 0:
|
||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.max_grad_norm)
|
||
self.optimizer.step()
|
||
|
||
total_loss += loss.item()
|
||
|
||
avg_total_loss = total_loss / len(self.train_loader)
|
||
avg_test_loss = self.test_E_loss()
|
||
self.losses.append([epoch, avg_total_loss, avg_test_loss])
|
||
self.scheduler.step()
|
||
|
||
if epoch % print_every == 0:
|
||
print(f'Epoch {epoch}, Total Loss: {avg_total_loss}, test Loss {avg_test_loss}')
|
||
if (epoch + 1) % save_every == 0:
|
||
ckpt_path = self.save_dir / f'{self.save_file_name}_epoch{epoch + 1}.pth'
|
||
torch.save(self.model.state_dict(), ckpt_path)
|
||
torch.save(self.model.state_dict(), self.save_file_name)
|
||
print("Current learning rate:", self.optimizer.param_groups[0]['lr'])
|
||
print("Training Time:", (datetime.datetime.now() - start_time).total_seconds(), "s")
|
||
|
||
def plot_loss(self):
|
||
data = np.array(self.losses[1:])
|
||
epochs = data[:, 0]
|
||
train_loss = data[:, 1]
|
||
test_loss = data[:, 2]
|
||
plt.figure(figsize=(10, 6))
|
||
plt.title('Training/Test Loss')
|
||
plt.semilogy(epochs, train_loss, label='train_loss')
|
||
plt.semilogy(epochs, test_loss, label='test_loss')
|
||
plt.xlabel('Epochs')
|
||
plt.ylabel('Loss')
|
||
plt.legend()
|
||
plt.grid()
|
||
save_path = 'results/loss_plot2_1.png'
|
||
plt.savefig(save_path)
|
||
plt.show()
|
||
|
||
def load_dataset(self):
|
||
"""从 MATLAB 划分好的 .mat 文件加载数据，90% 训练 10% 测试由 mat 内已分好。"""
|
||
data_set = loadmat(self.matpath)
|
||
|
||
Epsilon_train = data_set['Eplison_train']
|
||
X_train = data_set['X_train']
|
||
Ez_train = data_set['Ez_train']
|
||
|
||
Epsilon_test = data_set['Eplison_test']
|
||
X_test = data_set['X_test']
|
||
Ez_test = data_set['Ez_test']
|
||
|
||
coord_len_train = data_set['coord_len_train']
|
||
coord_len_test = data_set['coord_len_test']
|
||
|
||
Ai_train, Aj_train = data_set['Ai_train'], data_set['Aj_train']
|
||
Ai_test, Aj_test = data_set['Ai_test'], data_set['Aj_test']
|
||
Av_train, Av_test = data_set['Av_train'], data_set['Av_test']
|
||
b_train, b_test = data_set['b_train'], data_set['b_test']
|
||
|
||
n_train, n_test = len(Epsilon_train), len(Epsilon_test)
|
||
print(f"Train samples: {n_train}, Test samples: {n_test}")
|
||
print(f"Train shapes: ε {Epsilon_train.shape}, X {X_train.shape}, Ez {Ez_train.shape}")
|
||
|
||
Train_dataset = GetDataset(
|
||
Epsilon_train, X_train, Ez_train,
|
||
Ai_train, Aj_train, Av_train, b_train, coord_len_train,
|
||
index_offset=0
|
||
)
|
||
Test_dataset = GetDataset(
|
||
Epsilon_test, X_test, Ez_test,
|
||
Ai_test, Aj_test, Av_test, b_test, coord_len_test,
|
||
index_offset=n_train
|
||
)
|
||
return Train_dataset, Test_dataset
|
||
|
||
def _build_ilu_cache(self):
|
||
"""训练开始前为每个样本的稀疏矩阵 A 预计算 ILU 并缓存，训练过程中 A 不变，只算一次。"""
|
||
for idx in tqdm(range(len(self.train_set)), desc="Building ILU cache (train)"):
|
||
ai = self.train_set.Ai[idx].numpy()
|
||
aj = self.train_set.Aj[idx].numpy()
|
||
av = self.train_set.Av[idx].numpy()
|
||
Mi = int(self.train_set.coord_len[idx].item())
|
||
mask = (ai > 0) & (aj > 0)
|
||
rows = (ai[mask] - 1).astype(np.int64)
|
||
cols = (aj[mask] - 1).astype(np.int64)
|
||
vals = av[mask]
|
||
A = sp.coo_matrix((vals, (rows, cols)), shape=(Mi, Mi)).tocsc()
|
||
self.ilu_cache[idx] = spilu(A)
|
||
for idx in tqdm(range(len(self.test_set)), desc="Building ILU cache (test)"):
|
||
ai = self.test_set.Ai[idx].numpy()
|
||
aj = self.test_set.Aj[idx].numpy()
|
||
av = self.test_set.Av[idx].numpy()
|
||
Mi = int(self.test_set.coord_len[idx].item())
|
||
mask = (ai > 0) & (aj > 0)
|
||
rows = (ai[mask] - 1).astype(np.int64)
|
||
cols = (aj[mask] - 1).astype(np.int64)
|
||
vals = av[mask]
|
||
A = sp.coo_matrix((vals, (rows, cols)), shape=(Mi, Mi)).tocsc()
|
||
self.ilu_cache[len(self.train_set) + idx] = spilu(A)
|
||
print(f"ILU cache built: {len(self.ilu_cache)} samples.")
|
||
|
||
def saveE_pred(self):
|
||
for indices, epsilon_data, coord_data, *_ in self.train_loader:
|
||
epsilon_data = epsilon_data.to(self.device)
|
||
coord_data = coord_data.to(self.device)
|
||
E_real, E_imag = self.E_function(epsilon_data, coord_data)
|
||
E_pred = torch.complex(E_real, E_imag)
|
||
E_pred_np = E_pred.detach().cpu().numpy()
|
||
savemat('E_train_pred_size_1588.mat', {"E_pred": E_pred_np})
|
||
for indices, epsilon_data, coord_data, *_ in self.test_loader:
|
||
epsilon_data = epsilon_data.to(self.device)
|
||
coord_data = coord_data.to(self.device)
|
||
E_real, E_imag = self.E_function(epsilon_data, coord_data)
|
||
E_pred = torch.complex(E_real, E_imag)
|
||
E_pred_np = E_pred.detach().cpu().numpy()
|
||
savemat('E_test_pred_size_1588.mat', {"E_pred": E_pred_np})
|
||
|
||
|
||
if __name__ == "__main__":
|
||
cfg = PINNConfig(
|
||
)
|
||
model = DeepONet(branch_input_dim=1, trunk_input_dim=2, hidden_channel=128, output_dim=64)
|
||
model = model.double()
|
||
pinn = PINN_maxwell(model, cfg)
|
||
# pinn.load_model()
|
||
#pinn.test_fem_loss()
|
||
pinn.train(epochs=cfg.epochs, print_every=cfg.print_every, save_every=cfg.save_every)
|
||
pinn.plot_loss()
|
||
# pinn.saveE_pred()
|
||
|
||
|
||
|