上传文件至 /

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

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

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

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

说明.md

@@ -0,0 +1,18 @@
+不同几何形状的散射体电磁响应求解
+
+
+
+支干网络输入介电分布；
+
+主干网络输入坐标；
+
+ax-b作为损失函数；
+
+
+
+模型训练完成后，使用cnn\_branch\_inference.py进行推理预测；
+
+
+
+
+