""" 数值试井代理模型 - 单案例 Solver vs Surrogate 对比脚本 主要功能： 1. 使用 C++ 数值求解器生成真实试井曲线 2. 使用训练好的代理模型进行预测 3. 对比 Solver 与 Surrogate 的输出结果 4. 计算 RMSE / MAE / R2 等指标 5. 绘制压力、导数、斜率三部分对比图 6. 输出 JSON 分析结果 该脚本通常用于： - 检查代理模型是否学到真实物理行为 - 分析代理模型在哪些阶段误差较大 - 验证不同 schedule 对模型预测的影响 - 调试自动拟合(autofit)效果 """ # pylint: disable=import-error,wrong-import-position,wrong-import-order,line-too-long, # pylint: disable=too-many-locals,too-many-arguments,too-many-positional-arguments,invalid-name from __future__ import annotations import argparse import json import sys from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.append(str(ROOT)) import joblib import matplotlib.pyplot as plt import numpy as np import torch from src.common.config import Config from src.common.experiment_paths import ( config_for_stage, model_checkpoint_for_tag, normalize_tag, processed_path_for_tag, ) from src.data.curve_processing import clean_curve_for_dataset, is_valid_curve, resample_curve_to_features from src.data.param_features import param_feature_transform_from_meta, transform_param_features from src.data.params import Params, Schedule from src.data.runner_client import CppRunner, read_result_bin from src.data.schedule_features import build_schedule_model_vector from src.evaluation.autofit_objective import dual_log_objective from src.models.forward_surrogate import ForwardSurrogate # 默认单案例参数 # 用于没有传命令行参数时的快速测试 DEFAULT_SINGLE_CASE = { "config": "configs/data_gen_family_random.yaml", "tag": "family_random_50k", "no_schedule": False, "well_index": 0, "params": { "k": 0.025, "skin": 0.0, "wellboreC": 0.01, "phi": 0.0245, "h": 9.144, "Cf": 0.0004315, }, } def parse_args() -> argparse.Namespace: """解析单个样本对比所需的 processed 数据、模型、样本索引和输出目录。""" parser = argparse.ArgumentParser( description="Compare solver output and surrogate prediction on a single parameter set" ) parser.add_argument("--config", type=str, default=DEFAULT_SINGLE_CASE["config"], help="Config yaml path") parser.add_argument( "--stage", choices=["fixed_case", "case_neighborhood", "family_random", "family_random_v2_q"], default=None, help="Optional stage to infer config path", ) parser.add_argument("--processed", type=str, default=None, help="Processed dataset path") parser.add_argument("--model", type=str, default=None, help="Surrogate checkpoint path") parser.add_argument("--output-dir", type=str, default=None, help="Output directory") parser.add_argument( "--tag", type=str, default=DEFAULT_SINGLE_CASE["tag"], help="Experiment tag for processed/model lookup", ) parser.add_argument( "--no-schedule", action="store_true", help="When --model is omitted, infer the no-schedule checkpoint path", ) parser.add_argument("--k", type=float, default=DEFAULT_SINGLE_CASE["params"]["k"]) parser.add_argument("--skin", type=float, default=DEFAULT_SINGLE_CASE["params"]["skin"]) parser.add_argument("--wellboreC", type=float, default=DEFAULT_SINGLE_CASE["params"]["wellboreC"]) parser.add_argument("--phi", type=float, default=DEFAULT_SINGLE_CASE["params"]["phi"]) parser.add_argument("--h", type=float, default=DEFAULT_SINGLE_CASE["params"]["h"]) parser.add_argument("--Cf", type=float, default=DEFAULT_SINGLE_CASE["params"]["Cf"]) parser.add_argument( "--well-index", type=int, default=DEFAULT_SINGLE_CASE["well_index"], help="Well index for solver output", ) parser.add_argument( "--timeQ", type=str, default=None, help="Override schedule durations, e.g. '1000,1000,1000' or '0,1000,1000,1000'.", ) parser.add_argument( "--q", type=str, default=None, help="Override schedule rates, e.g. '200,100,0' or '0,200,100,0'.", ) parser.add_argument( "--section-index", type=int, default=None, help="Override tested section index. Defaults to the last section of the overridden schedule.", ) return parser.parse_args() def calc_metrics( y_true: np.ndarray, y_pred: np.ndarray, eps_range: float = 1e-3, eps_var: float = 1e-6, ) -> dict: """计算 RMSE、MAE、Bias、NRMSE、R2 等回归指标。""" # 残差 = 代理模型预测值 - 数值求解器真实值 err = y_pred - y_true mse = np.mean(err**2) rmse = float(np.sqrt(mse)) mae = float(np.mean(np.abs(err))) bias = float(np.mean(err)) value_range = float(np.max(y_true) - np.min(y_true)) ss_tot = float(np.sum((y_true - np.mean(y_true)) ** 2)) ss_res = float(np.sum(err**2)) # 曲线几乎为常数时，NRMSE 和 R2 的分母会过小，此时返回 NaN 更真实。 valid_nrmse = value_range > eps_range valid_r2 = ss_tot > eps_var nrmse = float(rmse / value_range) if valid_nrmse else np.nan r2 = float(1.0 - ss_res / ss_tot) if valid_r2 else np.nan return { "rmse": rmse, "mae": mae, "bias": bias, "abs_bias": float(abs(bias)), "nrmse": nrmse, "r2": r2, } def infer_curve_layout(meta: dict, curve_dim: int) -> dict: """从元数据读取曲线分段布局；旧数据没有布局时按压力/导数/斜率三等分回退。""" # curve_layout 描述了拼接曲线的结构布局 # 包括每一段 pressure/derivative/slope 的起止位置 curve_layout = meta.get("curve_layout") if curve_layout is not None: return curve_layout # 兼容早期 processed 文件：没有显式 layout 时仍按 pressure/derivative/slope 三段切分。 n_time_points = curve_dim // 3 return { "n_time_points": int(n_time_points), "parts": [ {"name": "log_pressure", "start": 0, "end": n_time_points}, {"name": "log_derivative", "start": n_time_points, "end": 2 * n_time_points}, {"name": "slope", "start": 2 * n_time_points, "end": 3 * n_time_points}, ], } def split_curve_by_layout(curve: np.ndarray, layout: dict) -> dict[str, np.ndarray]: """按照 curve_layout 将拼接曲线拆成 log_pressure、log_derivative 和 slope 三段。""" # parts 用于保存拆分后的不同曲线段 # 例如 log_pressure / derivative / slope parts: dict[str, np.ndarray] = {} for part in layout["parts"]: start = int(part["start"]) end = int(part["end"]) parts[str(part["name"])] = curve[start:end] return parts def resolve_cf(args: argparse.Namespace, cfg: Config) -> float: """解析压缩系数 Cf：优先使用命令行参数，否则从配置默认参数中读取。""" # 命令行优先级最高 # 如果用户显式提供 Cf，则直接使用 if args.Cf is not None: return float(args.Cf) fixed_cfg = cfg.raw.get("params", {}).get("fixed_params", {}).get("Cf", {}) or {} if bool(fixed_cfg.get("enabled", False)): return float(fixed_cfg["value"]) raise ValueError("Cf 未在命令行提供，且配置文件里也没有启用固定 Cf") def parse_float_list(raw: str, name: str) -> list[float]: """解析输入文本、命令行或配置值，转换为后续流程可直接使用的结构。""" # 用于解析类似 "1000,1000,1000" 的字符串输入 values: list[float] = [] for token in raw.replace(";", ",").replace(" ", ",").split(","): token = token.strip() if not token: continue values.append(float(token)) if not values: raise ValueError(f"{name} 不能为空") return values def resolve_case_schedule(cfg: Config, args: argparse.Namespace | None = None) -> Schedule: """解析单案例使用的流量制度；命令行提供 timeQ/q 时覆盖配置中的默认制度。""" if args is not None and (args.timeQ is not None or args.q is not None or args.section_index is not None): if args.timeQ is None or args.q is None: raise ValueError("覆盖流量制度时必须同时提供 --timeQ 和 --q") timeQ = parse_float_list(args.timeQ, "--timeQ") q = parse_float_list(args.q, "--q") if len(timeQ) != len(q): raise ValueError(f"--timeQ 和 --q 长度必须一致，当前分别为 {len(timeQ)} 和 {len(q)}") # 兼容报表写法：首行 0,0 只表示初始状态，不作为真实流量段传给求解器。 # 某些报表第一行仅用于表示初始状态 # 不是真实生产段，因此自动丢弃 if len(timeQ) >= 2 and timeQ[0] <= 0.0 and q[0] == 0.0: timeQ = timeQ[1:] q = q[1:] section_index = int(args.section_index) if args.section_index is not None else len(timeQ) schedule = Schedule(sectionIndex=section_index, timeQ=timeQ, q=q) if not schedule.validate(): raise ValueError( "命令行覆盖的流量制度非法。注意 compare 脚本里的 timeQ 表示每段持续时间，" "如果有初始行请写成 0,0；有效段必须 timeQ>0 且 q>=0。" ) return schedule schedule_cfg = cfg.raw["schedule"]["case_schedule"] timeQ = list(map(float, schedule_cfg["timeQ"])) q = list(map(float, schedule_cfg["q"])) policy = cfg.raw["schedule"].get("section_policy", {}) or {} mode = str(policy.get("mode", "fixed_last")).lower() n_sections = len(timeQ) if mode == "fixed_last": section_index = n_sections elif mode == "fixed_value": section_index = int(np.clip(int(policy.get("fixed_value", n_sections)), 1, n_sections)) else: section_index = int(np.clip(int(schedule_cfg.get("default_section_index", n_sections)), 1, n_sections)) schedule = Schedule(sectionIndex=section_index, timeQ=timeQ, q=q) if not schedule.validate(): raise ValueError("配置文件中的 case_schedule 非法，无法用于单样本对比") return schedule def resolve_paths(args: argparse.Namespace) -> tuple[Config, Path, Path, Path]: """根据命令行参数、实验标签和 use_schedule 开关解析配置、预处理数据、模型和输出路径。""" tag = normalize_tag(args.tag) config_path = args.config if config_path is None: config_path = str(config_for_stage(args.stage) or Path("configs/data_gen_family_random.yaml")) processed_path = Path(args.processed) if args.processed is not None else processed_path_for_tag(tag) model_path = ( Path(args.model) if args.model is not None else model_checkpoint_for_tag(tag, use_schedule=not args.no_schedule) ) if args.output_dir is not None: output_dir = Path(args.output_dir) else: suffix = "" if not args.no_schedule else "_no_schedule" output_dir = Path("results") / (f"single_compare_{tag}{suffix}" if tag else f"single_compare{suffix}") return Config(config_path), processed_path, model_path, output_dir def build_params_from_args(args: argparse.Namespace, cfg: Config, schedule: Schedule) -> Params: """根据命令行参数和默认值构造目标 Params 对象。""" # 构建完整物理参数对象 # 后续会直接送入数值求解器和代理模型 return Params( k=float(args.k), skin=float(args.skin), wellboreC=float(args.wellboreC), phi=float(args.phi), h=float(args.h), Cf=resolve_cf(args, cfg), schedule=schedule, ) def run_solver_and_extract_curve(cfg: Config, params: Params, well_index: int) -> tuple[np.ndarray, dict]: """调用 C++ 求解器运行一次正演，并把双对数输出重采样为模型曲线向量。""" # 创建 C++ 求解器客户端 # 实际底层会调用数值试井求解程序 runner = CppRunner(cfg=cfg) try: ok = runner.run_simulation(params, override_schedule=params.schedule, include_schedule=True) result = read_result_bin(runner.result_bin) if runner.result_bin.exists() else None # 某些求解器返回码可能失败但 result.bin 已写完；只要结果可读就继续对比。 if not ok and result is None: raise RuntimeError("求解器运行失败，未生成有效结果") if not ok and result is not None: print("Warning: runner_client 返回失败，但 result.bin 可读取，继续使用求解器结果。") if result is None: raise RuntimeError("无法读取 result.bin") if not result["loglog"]: raise RuntimeError("求解器未返回 loglog 数据") if well_index < 0 or well_index >= len(result["loglog"]): raise IndexError(f"well-index={well_index} 超出范围，可用井数={len(result['loglog'])}") loglog = result["loglog"][well_index] t = np.asarray(loglog["t"], dtype=np.float64) p = np.asarray(loglog["p"], dtype=np.float64) d = np.asarray(loglog["deriv"], dtype=np.float64) # 求解器原始点数不一定等于模型输出维度，先清洗再重采样到训练时的固定网格。 # 对原始求解器曲线进行清洗 # 去除非法点、异常点、重复点等 cleaned = clean_curve_for_dataset(cfg, t, p, d) if cleaned is None: raise RuntimeError("求解器返回曲线在清洗后无效") t_clean, p_clean, d_clean = cleaned valid, reason = is_valid_curve(cfg, t_clean, p_clean, d_clean) if not valid: raise RuntimeError(f"求解器曲线未通过有效性检查: {reason}") # 将原始不规则时间曲线重采样到固定特征网格 # 这是代理模型训练时使用的统一输入格式 curve_feat = resample_curve_to_features(cfg, t_clean, p_clean, d_clean) raw = { "t": t_clean.tolist(), "p": p_clean.tolist(), "d": d_clean.tolist(), "n_steps": int(result["nSteps"]), "n_wells": int(result["nWells"]), } return curve_feat, raw finally: runner.close() def build_schedule_vector(cfg: Config, schedule: Schedule) -> np.ndarray: """把 Schedule 编码成正演代理模型可直接接收的流量制度特征向量。""" return build_schedule_model_vector(cfg, schedule) def load_model(checkpoint_path: Path) -> tuple[ForwardSurrogate, bool, torch.device]: """加载模型检查点，按保存的维度和超参数重建网络并切换到评估模式。""" # 加载训练好的 PyTorch checkpoint checkpoint = torch.load(checkpoint_path, map_location="cpu") use_schedule = bool(checkpoint.get("use_schedule", True)) model = ForwardSurrogate( param_dim=int(checkpoint["param_dim"]), schedule_dim=int(checkpoint["schedule_dim"]), curve_dim=int(checkpoint["curve_dim"]), hidden_dim=int(checkpoint["hidden_dim"]), dropout=float(checkpoint["dropout"]), use_schedule=use_schedule, ) model.load_state_dict(checkpoint["model_state_dict"]) model.eval() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) return model, use_schedule, device def predict_surrogate_curve( processed: dict, model: ForwardSurrogate, device: torch.device, use_schedule: bool, params: Params, schedule: Schedule, cfg: Config, ) -> np.ndarray: """使用正演代理模型预测单个参数和流量制度对应的完整曲线。""" scaler_params = processed["scaler_params"] scaler_schedule = processed["scaler_schedule"] scaler_curve = processed["scaler_curve"] param_transform = param_feature_transform_from_meta(processed.get("meta", {})) # 构建参数特征向量 # 顺序必须与训练阶段保持一致 params_vec = np.asarray( [params.k, params.skin, params.wellboreC, params.phi, params.h, params.Cf], dtype=np.float32, ).reshape(1, -1) # 将流量制度 schedule 编码成模型输入向量 schedule_vec = build_schedule_vector(cfg, schedule).reshape(1, -1) # 输入特征必须使用训练时保存的 scaler 和参数变换，否则单案例预测会发生分布偏移。 params_x = scaler_params.transform(transform_param_features(params_vec, param_transform)).astype(np.float32) schedule_x = scaler_schedule.transform(schedule_vec).astype(np.float32) with torch.no_grad(): params_t = torch.tensor(params_x, dtype=torch.float32, device=device) if use_schedule: schedule_t = torch.tensor(schedule_x, dtype=torch.float32, device=device) pred_scaled = model(params_t, schedule_t).cpu().numpy() else: pred_scaled = model(params_t, None).cpu().numpy() return scaler_curve.inverse_transform(pred_scaled)[0].astype(np.float32) def plot_comparison( curve_true: np.ndarray, curve_pred: np.ndarray, curve_layout: dict, output_path: Path, params: Params, schedule: Schedule, model_path: Path, use_schedule: bool, ) -> dict: """绘制数值求解器与代理模型的单案例曲线对比图。""" true_parts = split_curve_by_layout(curve_true, curve_layout) pred_parts = split_curve_by_layout(curve_pred, curve_layout) # 整体曲线误差指标 overall = calc_metrics(curve_true, curve_pred) # 自动拟合目标函数 # 用于衡量双对数曲线形态是否一致 autofit = dual_log_objective(curve_true, curve_pred, curve_layout) part_names = ["log_pressure", "log_derivative", "slope"] title_map = { "log_pressure": "Log Pressure", "log_derivative": "Log |Derivative|", "slope": "Slope of Log Pressure vs Log Time", } r2_text = "nan" if np.isnan(overall["r2"]) else f"{overall['r2']:.4f}" fig, axes = plt.subplots(3, 2, figsize=(14, 12)) fig.suptitle( "Solver vs Surrogate\n" f"k={params.k:.6g}, skin={params.skin:.6g}, wellboreC={params.wellboreC:.6g}, " f"phi={params.phi:.6g}, h={params.h:.6g}, Cf={params.Cf:.6g}\n" f"sectionIndex={schedule.sectionIndex}, timeQ={schedule.timeQ}, q={schedule.q}\n" f"use_schedule={use_schedule}, model={model_path.name}, " f"Overall RMSE={overall['rmse']:.4f}, MAE={overall['mae']:.4f}, " f"AutoFitObj={autofit['dual_log_objective']:.4f}, " f"R2={r2_text}" ) summary = {"overall": overall, "autofit": autofit, "parts": {}} # 分别对 pressure / derivative / slope 三部分绘图 for row, name in enumerate(part_names): # 左列画曲线重合程度，右列画逐时间点残差，便于定位误差集中在哪个时期。 y_true = true_parts[name] y_pred = pred_parts[name] # 残差 = 代理模型预测值 - 数值求解器真实值 err = y_pred - y_true x = np.arange(len(y_true)) metrics = calc_metrics(y_true, y_pred) summary["parts"][name] = metrics nrmse_text = "nan" if np.isnan(metrics["nrmse"]) else f"{metrics['nrmse']:.4f}" r2_text = "nan" if np.isnan(metrics["r2"]) else f"{metrics['r2']:.4f}" ax_l = axes[row, 0] ax_l.plot(x, y_true, label="Solver", linewidth=2, alpha=0.85) ax_l.plot(x, y_pred, label="Surrogate", linewidth=2, alpha=0.85) ax_l.set_title( f"{title_map[name]} | RMSE={metrics['rmse']:.4f}, MAE={metrics['mae']:.4f}, " f"Bias={metrics['bias']:.4f}, NRMSE={nrmse_text}, R2={r2_text}" ) ax_l.set_xlabel("Resampled Time Index") ax_l.set_ylabel("Value") ax_l.grid(True, alpha=0.3) ax_l.legend() ax_r = axes[row, 1] ax_r.plot(x, err, linewidth=1.5, alpha=0.85) ax_r.axhline(0.0, linestyle="--", linewidth=1) ax_r.set_title(f"{title_map[name]} Error") ax_r.set_xlabel("Resampled Time Index") ax_r.set_ylabel("Surrogate - Solver") ax_r.grid(True, alpha=0.3) plt.tight_layout(rect=[0, 0, 1, 0.94]) plt.savefig(output_path, dpi=150, bbox_inches="tight") plt.close() return summary def main() -> None: """抽取一个测试样本，绘制代理模型曲线与真实数值求解曲线的对比图。""" args = parse_args() cfg, processed_path, model_path, output_dir = resolve_paths(args) output_dir.mkdir(parents=True, exist_ok=True) if not processed_path.exists(): raise FileNotFoundError(f"Processed 数据不存在: {processed_path}") if not model_path.exists(): raise FileNotFoundError(f"模型 checkpoint 不存在: {model_path}") processed = joblib.load(processed_path) curve_layout = infer_curve_layout(processed["meta"], int(processed["meta"]["curve_dim"])) # 先用同一组参数和制度跑真实求解器，再用代理模型预测同一输入，保证对比公平。 schedule = resolve_case_schedule(cfg, args) params = build_params_from_args(args, cfg, schedule) print("Running solver...") # 运行真实数值求解器 # 得到 ground truth 曲线 curve_solver, raw_solver = run_solver_and_extract_curve(cfg, params, args.well_index) print("Running surrogate...") model, use_schedule, device = load_model(model_path) # 使用代理模型预测同一组参数对应的曲线 curve_pred = predict_surrogate_curve( processed=processed, model=model, device=device, use_schedule=use_schedule, params=params, schedule=schedule, cfg=cfg, ) plot_path = output_dir / "single_case_comparison.png" summary = plot_comparison( curve_true=curve_solver, curve_pred=curve_pred, curve_layout=curve_layout, output_path=plot_path, params=params, schedule=schedule, model_path=model_path, use_schedule=use_schedule, ) # 汇总所有实验结果 # 最终会保存为 JSON 便于后续分析 summary_payload = { "config_path": str(cfg.path), "processed_path": str(processed_path), "model_path": str(model_path), "use_schedule": use_schedule, "device": str(device), "params": { "k": params.k, "skin": params.skin, "wellboreC": params.wellboreC, "phi": params.phi, "h": params.h, "Cf": params.Cf, }, "schedule": { "sectionIndex": schedule.sectionIndex, "timeQ": schedule.timeQ, "q": schedule.q, }, "metrics": summary, "solver_raw_loglog": raw_solver, } with open(output_dir / "single_case_summary.json", "w", encoding="utf-8") as f: json.dump(summary_payload, f, ensure_ascii=False, indent=2) overall = summary["overall"] autofit = summary["autofit"] r2_text = "nan" if np.isnan(overall["r2"]) else f"{overall['r2']:.6f}" print("\nSingle-case comparison complete.") print(f"Output dir: {output_dir}") print( f"Overall RMSE={overall['rmse']:.6f}, MAE={overall['mae']:.6f}, " f"Bias={overall['bias']:.6f}, " f"AutoFitObj={autofit['dual_log_objective']:.6f}, " f"R2={r2_text}" ) print(f"Plot saved: {plot_path}") print(f"Summary saved: {output_dir / 'single_case_summary.json'}") if __name__ == "__main__": main()