commit 9032d2196cd1b9226b72f74b2d5f0de28b2f5b79 Author: caowei <2642367431@qq.com> Date: Tue Jun 23 16:43:02 2026 +0800 上传文件至 / diff --git a/cnn_branch_inference.py b/cnn_branch_inference.py new file mode 100644 index 0000000..3334016 --- /dev/null +++ b/cnn_branch_inference.py @@ -0,0 +1,891 @@ +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() diff --git a/cnn_branch_test2.py b/cnn_branch_test2.py new file mode 100644 index 0000000..4489d57 --- /dev/null +++ b/cnn_branch_test2.py @@ -0,0 +1,483 @@ +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() + + + diff --git a/getdata.py b/getdata.py new file mode 100644 index 0000000..609cd89 --- /dev/null +++ b/getdata.py @@ -0,0 +1,66 @@ +from scipy.interpolate import griddata +import torch +import numpy as np +from torch.utils.data import Dataset, DataLoader +from scipy.io import loadmat +class GetDataset(Dataset): + def __init__(self, epsilon=None, coord=None, Ez=None, Ai=None, Aj=None, Av=None, b=None, coord_len=None, index_offset=0): + super().__init__() + self.index_offset = index_offset + self.epsilon = torch.as_tensor(epsilon, dtype=torch.float64) + self.coord = torch.as_tensor(coord, dtype=torch.float64) + self.Ez = torch.as_tensor(Ez, dtype=torch.complex128) + self.Ai = torch.as_tensor(Ai, dtype=torch.int64) # (B, nnz) + self.Aj = torch.as_tensor(Aj, dtype=torch.int64) + self.Av = torch.as_tensor(Av, dtype=torch.complex128) + self.b = torch.as_tensor(b, dtype=torch.complex128) # (B, M) + self.coord_len = torch.as_tensor(coord_len, dtype=torch.int64) # (B, M) + def __getitem__(self, index): + global_index = self.index_offset + index + epsilon = self.epsilon[index] # (M,) + coord = self.coord[index] # (M, 2) + ez = self.Ez[index] # (M, 2) + + ai = self.Ai[index] # (nnz,) + aj = self.Aj[index] # (nnz,) + av = self.Av[index] + + b = self.b[index] # (M,) + coord_len = self.coord_len[index] # (1,) + return global_index, epsilon, coord, ez, ai, aj, av, b, coord_len + + def __len__(self): + return len(self.epsilon) + +# Usage in main +if __name__ == "__main__": + # Load the data + data_set = loadmat('deepOnet_data_A1') + + Epsilon_train = data_set['Eplison_train'] # (390, 64, 64) + X_train = data_set['X_train'] # (4096, 2) + Ez_train = data_set['Ez_train'] # (390, 4096, 2) + + Epsilon_test = data_set['Eplison_test'] # (98, 64, 64) + X_test = data_set['X_test'] # (4096, 2) + Ez_test = data_set['Ez_test'] # (98, 4096, 2) + + 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'] + + print(f"Train shapes: ε {Epsilon_train.shape}, X {X_train.shape}, Ez {Ez_train.shape}") + + # Prepare the dataset instances + Train_dataset = GetDataset(Epsilon_train, X_train, Ez_train, Ai_train, Aj_train, Av_train, b_train, coord_len_train) + Test_dataset = GetDataset(Epsilon_test, X_test, Ez_test, Ai_test, Aj_test, Av_test, b_test, coord_len_test) + + # Example DataLoader（返回首项为 global_index） + loader = DataLoader(Train_dataset, batch_size=4, shuffle=True) + for indices, epsilon_data, coord_data, E_true, Ai, Aj, Av, b, coord_len in loader: + print(f'indices: {indices}, B_eps shape: {epsilon_data.shape}, T_xy shape: {coord_data.shape}, Ez shape: {E_true.shape}') + break diff --git a/pido_test.sh b/pido_test.sh new file mode 100644 index 0000000..ce588c2 --- /dev/null +++ b/pido_test.sh @@ -0,0 +1,22 @@ +#!/bin/bash +#SBATCH --job-name=run_node06 # 作业名称 +#SBATCH --output=pidn_output_%j # 标准输出文件（%j 会被替换为作业 ID） +#SBATCH --error=pidn_error_%j.txt # 标准错误文件（%j 会被替换为作业 ID） +#SBATCH --time=100:00:00 # 运行时间限制 +#SBATCH --partition=gpu # 请求 GPU 分区 +#SBATCH --cpus-per-task=10 # 节点请求的cpu核心数 +#SBATCH --mem=128G # 请求内存大小 +#SBATCH --gres=gpu:1 # 请求 4 个 GPU 资源 +#SBATCH --nodelist=node06 # 指定使用 node06 节点 + + +# 激活虚拟环境 +eval "$(/public/apps/miniconda3/bin/conda shell.bash hook)" +conda activate deepnet + +# 切换到存放 Python 脚本的目录 +cd /public/home/cw/deepOnet_ax_b_complex/modelA1-5 + + +# 执行 Python 脚本 +python cnn_branch_test2.py \ No newline at end of file diff --git a/说明.md b/说明.md new file mode 100644 index 0000000..a486ec7 --- /dev/null +++ b/说明.md @@ -0,0 +1,18 @@ +不同几何形状的散射体电磁响应求解 + + + +支干网络输入介电分布； + +主干网络输入坐标； + +ax-b作为损失函数； + + + +模型训练完成后，使用cnn\_branch\_inference.py进行推理预测； + + + + +