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()