892 lines
33 KiB
Python
892 lines
33 KiB
Python
import os
|
||
import json
|
||
import time
|
||
import math
|
||
from dataclasses import dataclass, asdict
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
from tqdm.auto import tqdm
|
||
|
||
import scipy.sparse as sp
|
||
from scipy.sparse.linalg import spilu
|
||
from scipy.io import loadmat, savemat
|
||
|
||
import torch
|
||
from torch.autograd import Function
|
||
import torch.nn as nn
|
||
from torch.utils.data import DataLoader
|
||
|
||
import matplotlib.pyplot as plt
|
||
|
||
from getdata import GetDataset
|
||
|
||
|
||
# ============================================================
|
||
# Device
|
||
# ============================================================
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
|
||
# ============================================================
|
||
# Custom autograd for ILU preconditioner
|
||
# ============================================================
|
||
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 = torch.from_numpy(z_np).to(r_torch.device).to(r_torch.dtype)
|
||
return z
|
||
|
||
@staticmethod
|
||
def backward(ctx, grad_out):
|
||
ilu = ctx.ilu
|
||
g_np = grad_out.detach().cpu().numpy()
|
||
gr_np = ilu.solve(g_np, trans='H')
|
||
grad_r = torch.from_numpy(gr_np).to(grad_out.device).to(grad_out.dtype)
|
||
return grad_r, None
|
||
|
||
|
||
# ============================================================
|
||
# Config
|
||
# ============================================================
|
||
@dataclass
|
||
class PINNConfig:
|
||
# 数据
|
||
matpath: str = "deepOnet_data_A1_1558_8_2"
|
||
|
||
# 模型
|
||
ckpt_path: str = "./model_save/model_A1_size_1558.pth"
|
||
|
||
# 保存
|
||
results_dir: str = "./results"
|
||
|
||
# 设备/精度
|
||
device: str = "cuda" if torch.cuda.is_available() else "cpu"
|
||
dtype: torch.dtype = torch.float64
|
||
|
||
# DataLoader
|
||
batch_size: int = 64
|
||
num_workers: int = 4
|
||
pin_memory: bool = True
|
||
|
||
# 模型结构
|
||
trunk_input_dim: int = 2
|
||
hidden_channel: int = 128
|
||
output_dim: int = 64
|
||
|
||
# 推理
|
||
warmup_steps: int = 3
|
||
build_ilu_cache: bool = True
|
||
|
||
# 保存选项
|
||
save_npz: bool = True
|
||
save_mat: bool = True
|
||
save_csv: bool = True
|
||
save_plots: bool = True
|
||
|
||
|
||
# ============================================================
|
||
# Networks
|
||
# ============================================================
|
||
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):
|
||
return nn.InstanceNorm2d(channels)
|
||
|
||
|
||
def add_spatial_coord_channels(epsilon_data):
|
||
"""
|
||
epsilon_data: (B, 1, H, W)
|
||
return: (B, 3, H, W) -> [epsilon, x_norm, y_norm]
|
||
"""
|
||
B, _, H, W = epsilon_data.shape
|
||
dev = epsilon_data.device
|
||
dtype = epsilon_data.dtype
|
||
|
||
x = torch.linspace(0, 1, W, device=dev, dtype=dtype).view(1, 1, 1, W).expand(B, 1, H, W)
|
||
y = torch.linspace(0, 1, H, device=dev, dtype=dtype).view(1, 1, H, 1).expand(B, 1, H, W)
|
||
return torch.cat([epsilon_data, x, y], dim=1)
|
||
|
||
|
||
class ResidualBlock(nn.Module):
|
||
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 CNN_Branch_Residual(nn.Module):
|
||
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 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
|
||
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) # (B, output_dim)
|
||
trunk_out = self.trunk_net(trunk_input) # (B, N, output_dim)
|
||
|
||
B1 = branch_out[:, :self.output_dim // 2]
|
||
B2 = branch_out[:, self.output_dim // 2:]
|
||
T1 = trunk_out[:, :, :self.output_dim // 2]
|
||
T2 = trunk_out[:, :, self.output_dim // 2:]
|
||
|
||
s_re = torch.einsum('bi,bni->bn', B1, T1)
|
||
s_im = torch.einsum('bi,bni->bn', B2, T2)
|
||
|
||
return s_re, s_im
|
||
|
||
|
||
# ============================================================
|
||
# Inference / Evaluation
|
||
# ============================================================
|
||
class PINN_maxwell:
|
||
def __init__(self, model, config: PINNConfig):
|
||
self.cfg = config
|
||
self.device = torch.device(self.cfg.device)
|
||
|
||
self.results_dir = Path(self.cfg.results_dir)
|
||
self.results_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
self.model = model.to(self.device, dtype=self.cfg.dtype)
|
||
self.loss_fn = nn.MSELoss()
|
||
|
||
self.train_set, self.test_set = self.load_dataset()
|
||
|
||
pin_memory = bool(self.cfg.pin_memory and self.device.type == "cuda")
|
||
|
||
self.eval_train_loader = DataLoader(
|
||
self.train_set,
|
||
batch_size=self.cfg.batch_size,
|
||
shuffle=False,
|
||
num_workers=self.cfg.num_workers,
|
||
pin_memory=pin_memory
|
||
)
|
||
|
||
self.eval_test_loader = DataLoader(
|
||
self.test_set,
|
||
batch_size=self.cfg.batch_size,
|
||
shuffle=False,
|
||
num_workers=self.cfg.num_workers,
|
||
pin_memory=pin_memory
|
||
)
|
||
|
||
self.ilu_cache = {}
|
||
if self.cfg.build_ilu_cache:
|
||
self._build_ilu_cache()
|
||
|
||
# ========================================================
|
||
# Dataset
|
||
# ========================================================
|
||
def load_dataset(self):
|
||
"""
|
||
训练集和测试集已经在 mat 内分好，这里直接读取，不重新划分。
|
||
"""
|
||
data_set = loadmat(self.cfg.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 {Epsilon_train.shape}, X {X_train.shape}, Ez {Ez_train.shape}")
|
||
print(f"Test shapes: epsilon {Epsilon_test.shape}, X {X_test.shape}, Ez {Ez_test.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
|
||
|
||
# ========================================================
|
||
# Model IO
|
||
# ========================================================
|
||
def load_model(self):
|
||
ckpt_path = Path(self.cfg.ckpt_path)
|
||
if not ckpt_path.exists():
|
||
raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
|
||
|
||
print(f"Loading checkpoint from: {ckpt_path}")
|
||
state_dict = torch.load(ckpt_path, map_location=self.device)
|
||
self.model.load_state_dict(state_dict, strict=True)
|
||
self.model.to(self.device, dtype=self.cfg.dtype)
|
||
self.model.eval()
|
||
|
||
# ========================================================
|
||
# Utils
|
||
# ========================================================
|
||
def _sync_device(self):
|
||
if self.device.type == "cuda":
|
||
torch.cuda.synchronize(self.device)
|
||
|
||
def _build_valid_mask(self, coord_len, max_len):
|
||
Mi = coord_len.view(-1).long().to(self.device)
|
||
arange = torch.arange(max_len, device=self.device).unsqueeze(0)
|
||
mask = arange < Mi.unsqueeze(1)
|
||
return mask, Mi
|
||
|
||
def _move_E_true_to_device(self, E_true):
|
||
"""
|
||
保持真值原始类型迁移到 device：
|
||
- complex -> complex128/complex64
|
||
- real -> float64/float32
|
||
"""
|
||
if torch.is_complex(E_true):
|
||
target_dtype = torch.complex128 if self.cfg.dtype == torch.float64 else torch.complex64
|
||
return E_true.to(self.device, dtype=target_dtype, non_blocking=True)
|
||
return E_true.to(self.device, dtype=self.cfg.dtype, non_blocking=True)
|
||
|
||
def _parse_E_true(self, E_true):
|
||
"""
|
||
兼容：
|
||
1) complex tensor: (B, M)
|
||
2) two-channel real tensor: (B, M, 2)
|
||
3) real tensor: (B, M)
|
||
return:
|
||
E_re_true, E_im_true, E_true_complex
|
||
"""
|
||
if torch.is_complex(E_true):
|
||
E_true_complex = E_true
|
||
E_re_true = E_true.real.to(self.cfg.dtype)
|
||
E_im_true = E_true.imag.to(self.cfg.dtype)
|
||
return E_re_true, E_im_true, E_true_complex
|
||
|
||
if E_true.ndim == 3 and E_true.shape[-1] == 2:
|
||
E_re_true = E_true[:, :, 0].to(self.cfg.dtype)
|
||
E_im_true = E_true[:, :, 1].to(self.cfg.dtype)
|
||
E_true_complex = torch.complex(E_re_true, E_im_true)
|
||
return E_re_true, E_im_true, E_true_complex
|
||
|
||
if E_true.ndim == 2:
|
||
E_re_true = E_true.to(self.cfg.dtype)
|
||
E_im_true = torch.zeros_like(E_re_true)
|
||
E_true_complex = torch.complex(E_re_true, E_im_true)
|
||
return E_re_true, E_im_true, E_true_complex
|
||
|
||
raise ValueError(
|
||
f"Unsupported E_true format: shape={tuple(E_true.shape)}, dtype={E_true.dtype}"
|
||
)
|
||
|
||
# ========================================================
|
||
# Forward / Loss
|
||
# ========================================================
|
||
def E_function(self, epsilon_data, coord_data):
|
||
epsilon_data = epsilon_data.to(self.device, dtype=self.cfg.dtype)
|
||
coord_data = coord_data.to(self.device, dtype=self.cfg.dtype)
|
||
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_true = self._move_E_true_to_device(E_true)
|
||
E_re_true, E_im_true, _ = self._parse_E_true(E_true)
|
||
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):
|
||
"""
|
||
这里 E_true 其实不参与 fem_loss 计算，保留接口只是为了兼容你原来的调用方式。
|
||
"""
|
||
Ere_pred, Eim_pred = self.E_function(epsilon_data, coord_data)
|
||
E = torch.complex(Ere_pred, Eim_pred)
|
||
|
||
B, Mmax = E.shape
|
||
|
||
Mi = coord_len.squeeze(-1).long().to(self.device)
|
||
arangeM = torch.arange(Mmax, device=self.device)
|
||
mask_x = arangeM[None, :] < Mi[:, None]
|
||
|
||
x_flat = E[mask_x]
|
||
b_flat = b.to(self.device)[mask_x].to(x_flat.dtype)
|
||
|
||
sumMi = int(Mi.sum().item())
|
||
offsets = torch.cumsum(
|
||
torch.cat([
|
||
torch.zeros(1, device=self.device, dtype=torch.long),
|
||
Mi[:-1]
|
||
]),
|
||
dim=0
|
||
)
|
||
|
||
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)
|
||
|
||
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
|
||
|
||
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
|
||
|
||
# ========================================================
|
||
# ILU Cache
|
||
# ========================================================
|
||
def _build_ilu_cache(self):
|
||
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 _compute_sample_fem_metrics(self, global_idx, pred_complex_np, ai, aj, av, b_vec, Mi):
|
||
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()
|
||
x = pred_complex_np[:Mi]
|
||
b_valid = b_vec[:Mi]
|
||
|
||
r = A.dot(x) - b_valid
|
||
residual_l2 = float(np.linalg.norm(r))
|
||
relative_residual_l2 = float(residual_l2 / (np.linalg.norm(b_valid) + 1e-12))
|
||
|
||
if self.cfg.build_ilu_cache:
|
||
ilu = self.ilu_cache[int(global_idx)]
|
||
else:
|
||
ilu = spilu(A)
|
||
|
||
z = ilu.solve(r)
|
||
fem_precond_mse = float(np.mean(np.abs(z) ** 2))
|
||
return fem_precond_mse, residual_l2, relative_residual_l2
|
||
|
||
# ========================================================
|
||
# Inference / Evaluation
|
||
# ========================================================
|
||
@torch.no_grad()
|
||
def infer_and_save_split(self, loader, split_name="train"):
|
||
self.model.eval()
|
||
|
||
split_dir = self.results_dir / split_name
|
||
split_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# warmup
|
||
if self.device.type == "cuda" and self.cfg.warmup_steps > 0:
|
||
warmup_count = 0
|
||
for batch in loader:
|
||
if warmup_count >= self.cfg.warmup_steps:
|
||
break
|
||
_, epsilon_data, coord_data, *_ = batch
|
||
epsilon_data = epsilon_data.to(self.device, dtype=self.cfg.dtype, non_blocking=True)
|
||
coord_data = coord_data.to(self.device, dtype=self.cfg.dtype, non_blocking=True)
|
||
_ = self.model(epsilon_data, coord_data)
|
||
warmup_count += 1
|
||
self._sync_device()
|
||
|
||
all_indices = []
|
||
all_epsilon = []
|
||
all_coord = []
|
||
|
||
all_E_true_raw = []
|
||
all_E_true_real = []
|
||
all_E_true_imag = []
|
||
all_E_true_complex = []
|
||
|
||
all_E_pred_real = []
|
||
all_E_pred_imag = []
|
||
all_E_pred_complex = []
|
||
|
||
all_Ai = []
|
||
all_Aj = []
|
||
all_Av = []
|
||
all_b = []
|
||
all_coord_len = []
|
||
|
||
sample_metrics = []
|
||
|
||
total_forward_time = 0.0
|
||
total_samples = 0
|
||
total_valid_points = 0
|
||
|
||
sum_sq_err = 0.0
|
||
sum_abs_err = 0.0
|
||
sum_sq_true = 0.0
|
||
|
||
sum_sq_err_re = 0.0
|
||
sum_sq_err_im = 0.0
|
||
sum_abs_err_re = 0.0
|
||
sum_abs_err_im = 0.0
|
||
|
||
fem_loss_list = []
|
||
|
||
for batch in tqdm(loader, desc=f"Infer {split_name}"):
|
||
indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len = batch
|
||
|
||
indices = indices.to(self.device, non_blocking=True)
|
||
epsilon_data = epsilon_data.to(self.device, dtype=self.cfg.dtype, non_blocking=True)
|
||
coord_data = coord_data.to(self.device, dtype=self.cfg.dtype, non_blocking=True)
|
||
E_true = self._move_E_true_to_device(E_true)
|
||
|
||
B = epsilon_data.shape[0]
|
||
Mmax = coord_data.shape[1]
|
||
|
||
mask, Mi = self._build_valid_mask(coord_len, Mmax)
|
||
|
||
self._sync_device()
|
||
t0 = time.perf_counter()
|
||
E_re_pred, E_im_pred = self.E_function(epsilon_data, coord_data)
|
||
self._sync_device()
|
||
t1 = time.perf_counter()
|
||
|
||
forward_time = t1 - t0
|
||
total_forward_time += forward_time
|
||
total_samples += B
|
||
total_valid_points += int(mask.sum().item())
|
||
|
||
E_re_true, E_im_true, E_true_complex = self._parse_E_true(E_true)
|
||
E_pred_complex = torch.complex(E_re_pred, E_im_pred)
|
||
|
||
diff_complex = E_pred_complex - E_true_complex
|
||
diff_re = E_re_pred - E_re_true
|
||
diff_im = E_im_pred - E_im_true
|
||
|
||
sum_sq_err += float((diff_complex.abs()[mask] ** 2).sum().item())
|
||
sum_abs_err += float(diff_complex.abs()[mask].sum().item())
|
||
sum_sq_true += float((E_true_complex.abs()[mask] ** 2).sum().item())
|
||
|
||
sum_sq_err_re += float((diff_re[mask] ** 2).sum().item())
|
||
sum_sq_err_im += float((diff_im[mask] ** 2).sum().item())
|
||
sum_abs_err_re += float(diff_re.abs()[mask].sum().item())
|
||
sum_abs_err_im += float(diff_im.abs()[mask].sum().item())
|
||
|
||
fem_loss = self.get_fem_loss(indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len)
|
||
fem_loss_list.append(float(fem_loss.item()))
|
||
|
||
all_indices.append(indices.detach().cpu().numpy())
|
||
all_epsilon.append(epsilon_data.detach().cpu().numpy())
|
||
all_coord.append(coord_data.detach().cpu().numpy())
|
||
|
||
all_E_true_raw.append(E_true.detach().cpu().numpy())
|
||
all_E_true_real.append(E_re_true.detach().cpu().numpy())
|
||
all_E_true_imag.append(E_im_true.detach().cpu().numpy())
|
||
all_E_true_complex.append(E_true_complex.detach().cpu().numpy())
|
||
|
||
all_E_pred_real.append(E_re_pred.detach().cpu().numpy())
|
||
all_E_pred_imag.append(E_im_pred.detach().cpu().numpy())
|
||
all_E_pred_complex.append(E_pred_complex.detach().cpu().numpy())
|
||
|
||
all_Ai.append(Ai.detach().cpu().numpy())
|
||
all_Aj.append(Aj.detach().cpu().numpy())
|
||
all_Av.append(Av.detach().cpu().numpy())
|
||
all_b.append(b.detach().cpu().numpy())
|
||
all_coord_len.append(coord_len.detach().cpu().numpy())
|
||
|
||
pred_np = E_pred_complex.detach().cpu().numpy()
|
||
true_np = E_true_complex.detach().cpu().numpy()
|
||
Mi_np = Mi.detach().cpu().numpy()
|
||
|
||
Ai_np = Ai.detach().cpu().numpy()
|
||
Aj_np = Aj.detach().cpu().numpy()
|
||
Av_np = Av.detach().cpu().numpy()
|
||
b_np = b.detach().cpu().numpy()
|
||
|
||
forward_time_per_sample_ms = forward_time * 1000.0 / B
|
||
|
||
for i in range(B):
|
||
m = int(Mi_np[i])
|
||
|
||
pred_i = pred_np[i, :m]
|
||
true_i = true_np[i, :m]
|
||
diff_i = pred_i - true_i
|
||
|
||
mse_i = float(np.mean(np.abs(diff_i) ** 2))
|
||
rmse_i = float(np.sqrt(mse_i))
|
||
mae_i = float(np.mean(np.abs(diff_i)))
|
||
rel_l2_i = float(np.linalg.norm(diff_i) / (np.linalg.norm(true_i) + 1e-12))
|
||
|
||
mse_re_i = float(np.mean((pred_i.real - true_i.real) ** 2))
|
||
mse_im_i = float(np.mean((pred_i.imag - true_i.imag) ** 2))
|
||
mae_re_i = float(np.mean(np.abs(pred_i.real - true_i.real)))
|
||
mae_im_i = float(np.mean(np.abs(pred_i.imag - true_i.imag)))
|
||
|
||
fem_precond_mse_i = np.nan
|
||
residual_l2_i = np.nan
|
||
relative_residual_l2_i = np.nan
|
||
|
||
if self.cfg.build_ilu_cache:
|
||
fem_precond_mse_i, residual_l2_i, relative_residual_l2_i = \
|
||
self._compute_sample_fem_metrics(
|
||
global_idx=int(indices[i].item()),
|
||
pred_complex_np=pred_np[i],
|
||
ai=Ai_np[i],
|
||
aj=Aj_np[i],
|
||
av=Av_np[i],
|
||
b_vec=b_np[i],
|
||
Mi=m
|
||
)
|
||
|
||
sample_metrics.append({
|
||
"split": split_name,
|
||
"global_index": int(indices[i].item()),
|
||
"coord_len": m,
|
||
"mse_complex": mse_i,
|
||
"rmse_complex": rmse_i,
|
||
"mae_complex": mae_i,
|
||
"rel_l2": rel_l2_i,
|
||
"mse_real": mse_re_i,
|
||
"mse_imag": mse_im_i,
|
||
"mae_real": mae_re_i,
|
||
"mae_imag": mae_im_i,
|
||
"fem_precond_mse": fem_precond_mse_i,
|
||
"residual_l2": residual_l2_i,
|
||
"relative_residual_l2": relative_residual_l2_i,
|
||
"forward_time_ms": forward_time_per_sample_ms
|
||
})
|
||
|
||
mse_complex = sum_sq_err / max(total_valid_points, 1)
|
||
rmse_complex = math.sqrt(mse_complex)
|
||
mae_complex = sum_abs_err / max(total_valid_points, 1)
|
||
rel_l2_global = math.sqrt(sum_sq_err / (sum_sq_true + 1e-12))
|
||
|
||
mse_real = sum_sq_err_re / max(total_valid_points, 1)
|
||
mse_imag = sum_sq_err_im / max(total_valid_points, 1)
|
||
mae_real = sum_abs_err_re / max(total_valid_points, 1)
|
||
mae_imag = sum_abs_err_im / max(total_valid_points, 1)
|
||
|
||
mean_fem_loss = float(np.mean(fem_loss_list)) if len(fem_loss_list) > 0 else float("nan")
|
||
mean_forward_time_ms = total_forward_time * 1000.0 / max(total_samples, 1)
|
||
throughput = total_samples / max(total_forward_time, 1e-12)
|
||
|
||
sample_df = pd.DataFrame(sample_metrics)
|
||
|
||
summary = {
|
||
"split": split_name,
|
||
"num_samples": int(total_samples),
|
||
"num_valid_points": int(total_valid_points),
|
||
"mse_complex": float(mse_complex),
|
||
"rmse_complex": float(rmse_complex),
|
||
"mae_complex": float(mae_complex),
|
||
"rel_l2_global": float(rel_l2_global),
|
||
"mse_real": float(mse_real),
|
||
"mse_imag": float(mse_imag),
|
||
"mae_real": float(mae_real),
|
||
"mae_imag": float(mae_imag),
|
||
"mean_fem_loss": float(mean_fem_loss),
|
||
"total_forward_time_s": float(total_forward_time),
|
||
"mean_forward_time_ms_per_sample": float(mean_forward_time_ms),
|
||
"throughput_samples_per_s": float(throughput),
|
||
}
|
||
|
||
if len(sample_df) > 0:
|
||
for col in ["rel_l2", "fem_precond_mse", "residual_l2", "relative_residual_l2"]:
|
||
if col in sample_df.columns and sample_df[col].notna().any():
|
||
summary[f"mean_{col}"] = float(np.nanmean(sample_df[col].values))
|
||
summary[f"median_{col}"] = float(np.nanmedian(sample_df[col].values))
|
||
summary[f"max_{col}"] = float(np.nanmax(sample_df[col].values))
|
||
|
||
save_data = {
|
||
"indices": np.concatenate(all_indices, axis=0),
|
||
"epsilon_data": np.concatenate(all_epsilon, axis=0),
|
||
"coord_data": np.concatenate(all_coord, axis=0),
|
||
|
||
"E_true_raw": np.concatenate(all_E_true_raw, axis=0),
|
||
"E_true_real": np.concatenate(all_E_true_real, axis=0),
|
||
"E_true_imag": np.concatenate(all_E_true_imag, axis=0),
|
||
"E_true_complex": np.concatenate(all_E_true_complex, axis=0),
|
||
|
||
"E_pred_real": np.concatenate(all_E_pred_real, axis=0),
|
||
"E_pred_imag": np.concatenate(all_E_pred_imag, axis=0),
|
||
"E_pred_complex": np.concatenate(all_E_pred_complex, axis=0),
|
||
|
||
"Ai": np.concatenate(all_Ai, axis=0),
|
||
"Aj": np.concatenate(all_Aj, axis=0),
|
||
"Av": np.concatenate(all_Av, axis=0),
|
||
"b": np.concatenate(all_b, axis=0),
|
||
"coord_len": np.concatenate(all_coord_len, axis=0),
|
||
}
|
||
|
||
if self.cfg.save_npz:
|
||
np.savez_compressed(split_dir / f"{split_name}_all_data.npz", **save_data)
|
||
|
||
if self.cfg.save_mat:
|
||
mat_dict = dict(save_data)
|
||
for k, v in summary.items():
|
||
if isinstance(v, (int, float, np.number)):
|
||
mat_dict[f"summary_{k}"] = np.array([[v]])
|
||
elif isinstance(v, str):
|
||
mat_dict[f"summary_{k}"] = np.array([v], dtype=object)
|
||
|
||
if len(sample_df) > 0:
|
||
for col in sample_df.columns:
|
||
if pd.api.types.is_numeric_dtype(sample_df[col]):
|
||
mat_dict[f"sample_{col}"] = sample_df[col].to_numpy()
|
||
else:
|
||
mat_dict[f"sample_{col}"] = sample_df[col].astype(str).to_numpy(dtype=object)
|
||
|
||
savemat(split_dir / f"{split_name}_all_data.mat", mat_dict)
|
||
|
||
if self.cfg.save_csv:
|
||
sample_df.to_csv(split_dir / f"{split_name}_sample_metrics.csv", index=False, encoding="utf-8-sig")
|
||
|
||
with open(split_dir / f"{split_name}_summary.json", "w", encoding="utf-8") as f:
|
||
json.dump(summary, f, ensure_ascii=False, indent=2)
|
||
|
||
if self.cfg.save_plots and len(sample_df) > 0:
|
||
self._plot_split_metrics(split_name, sample_df, split_dir)
|
||
|
||
print(f"\n[{split_name}] Summary:")
|
||
for k, v in summary.items():
|
||
print(f"{k}: {v}")
|
||
|
||
return summary, sample_df, save_data
|
||
|
||
def _plot_split_metrics(self, split_name, sample_df, split_dir):
|
||
if "rel_l2" in sample_df.columns:
|
||
plt.figure(figsize=(8, 5))
|
||
plt.hist(sample_df["rel_l2"].dropna().values, bins=30)
|
||
plt.xlabel("Relative L2 Error")
|
||
plt.ylabel("Count")
|
||
plt.title(f"{split_name} - Relative L2 Error")
|
||
plt.grid(True, alpha=0.3)
|
||
plt.tight_layout()
|
||
plt.savefig(split_dir / f"{split_name}_rel_l2_hist.png", dpi=200)
|
||
plt.close()
|
||
|
||
if "forward_time_ms" in sample_df.columns:
|
||
plt.figure(figsize=(8, 5))
|
||
plt.hist(sample_df["forward_time_ms"].dropna().values, bins=30)
|
||
plt.xlabel("Forward Time per Sample (ms)")
|
||
plt.ylabel("Count")
|
||
plt.title(f"{split_name} - Forward Time")
|
||
plt.grid(True, alpha=0.3)
|
||
plt.tight_layout()
|
||
plt.savefig(split_dir / f"{split_name}_forward_time_hist.png", dpi=200)
|
||
plt.close()
|
||
|
||
if "fem_precond_mse" in sample_df.columns and sample_df["fem_precond_mse"].notna().any():
|
||
plt.figure(figsize=(8, 5))
|
||
plt.hist(sample_df["fem_precond_mse"].dropna().values, bins=30)
|
||
plt.xlabel("FEM Preconditioned MSE")
|
||
plt.ylabel("Count")
|
||
plt.title(f"{split_name} - FEM Preconditioned MSE")
|
||
plt.grid(True, alpha=0.3)
|
||
plt.tight_layout()
|
||
plt.savefig(split_dir / f"{split_name}_fem_precond_mse_hist.png", dpi=200)
|
||
plt.close()
|
||
|
||
def evaluate_and_infer_all(self):
|
||
self.results_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
train_summary, train_df, _ = self.infer_and_save_split(
|
||
self.eval_train_loader, split_name="train"
|
||
)
|
||
test_summary, test_df, _ = self.infer_and_save_split(
|
||
self.eval_test_loader, split_name="test"
|
||
)
|
||
|
||
summary_df = pd.DataFrame([train_summary, test_summary])
|
||
summary_df.to_csv(self.results_dir / "all_summary.csv", index=False, encoding="utf-8-sig")
|
||
|
||
cfg_dump = {}
|
||
for k, v in asdict(self.cfg).items():
|
||
if isinstance(v, torch.dtype):
|
||
cfg_dump[k] = str(v)
|
||
else:
|
||
cfg_dump[k] = v
|
||
|
||
with open(self.results_dir / "config.json", "w", encoding="utf-8") as f:
|
||
json.dump(cfg_dump, f, ensure_ascii=False, indent=2)
|
||
|
||
print("\n================ Overall Summary ================")
|
||
print(summary_df.to_string(index=False))
|
||
|
||
|
||
# ============================================================
|
||
# Main
|
||
# ============================================================
|
||
if __name__ == "__main__":
|
||
cfg = PINNConfig(
|
||
matpath="deepOnet_data_A1_1558_8_2", # 改成你的 .mat 路径
|
||
ckpt_path="./model_save/model_A1_size_1558_2_epoch1000.pth", # 改成你的模型权重路径
|
||
results_dir="./results",
|
||
batch_size=64,
|
||
num_workers=4,
|
||
device="cuda" if torch.cuda.is_available() else "cpu",
|
||
dtype=torch.float64,
|
||
trunk_input_dim=2,
|
||
hidden_channel=128,
|
||
output_dim=64,
|
||
build_ilu_cache=True,
|
||
)
|
||
|
||
model = DeepONet(
|
||
branch_input_dim=1,
|
||
trunk_input_dim=cfg.trunk_input_dim,
|
||
hidden_channel=cfg.hidden_channel,
|
||
output_dim=cfg.output_dim
|
||
).double()
|
||
|
||
pinn = PINN_maxwell(model, cfg)
|
||
|
||
# 只加载模型，不训练
|
||
pinn.load_model()
|
||
|
||
# 直接对 train/test 推理评估并保存
|
||
pinn.evaluate_and_infer_all()
|