nmWTAI-Platform/ML/nmWTAI-ML/src/evaluation/autofit_objective.py

# -*- coding: utf-8 -*-
"""自动拟合候选曲线目标函数。

本模块提供轻量级的曲线误差计算，用于比较目标曲线和代理模型预测曲线的匹配度。
目标函数与 C++ 真实求解器的当前误差公式保持一致，并分别计算 log_pressure 与
log_derivative 两条曲线的贡献，适合在参数筛选、PSO 候选排序和局部邻域验证中复用。
"""

from __future__ import annotations

import numpy as np




def split_curve_by_layout(curve: np.ndarray, layout: dict) -> dict[str, np.ndarray]:
    """按照 curve_layout 拆分曲线。"""
    parts: dict[str, np.ndarray] = {}
    for part in layout["parts"]:
        start = int(part["start"])
        end = int(part["end"])
        parts[str(part["name"])] = np.asarray(curve[start:end], dtype=np.float64)
    return parts


def calculate_curve_objective_1d(target: np.ndarray, pred: np.ndarray) -> float:
    """在自然对数曲线上计算与 C++ calculatePointError 等价的均方根误差。"""
    target = np.asarray(target, dtype=np.float64).reshape(-1)
    pred = np.asarray(pred, dtype=np.float64).reshape(-1)

    if target.size == 0 or pred.size != target.size:
        return float("inf")
    if not (np.isfinite(target).all() and np.isfinite(pred).all()):
        return float("inf")

    # 模型曲线已经是 ln(raw)。因此 C++ 中的 logError 就是两列之差；
    # raw 相对误差可稳定地化简为 1-exp(-|delta_log|)，无需先 exp 回原始尺度。
    log_error = np.abs(target - pred)
    relative_error = -np.expm1(-log_error)
    point_error = 0.7 * log_error + 0.3 * relative_error
    return float(np.sqrt(np.mean(point_error**2)))


def dual_log_objective(
    curve_target: np.ndarray,
    curve_pred: np.ndarray,
    curve_layout: dict,
    w_pressure: float = 0.5,
    w_derivative: float = 0.5,
) -> dict[str, float]:
    """分别计算压力和导数目标，并按权重合成双对数自动拟合目标。"""
    parts_target = split_curve_by_layout(curve_target, curve_layout)
    parts_pred = split_curve_by_layout(curve_pred, curve_layout)

    p_obj = calculate_curve_objective_1d(parts_target["log_pressure"], parts_pred["log_pressure"])
    d_obj = calculate_curve_objective_1d(parts_target["log_derivative"], parts_pred["log_derivative"])

    total_w = max(float(w_pressure) + float(w_derivative), 1e-12)
    combined = (float(w_pressure) * p_obj + float(w_derivative) * d_obj) / total_w

    return {
        "log_pressure_objective": float(p_obj),
        "log_derivative_objective": float(d_obj),
        "dual_log_objective": float(combined),
    }


def prediction_curve_time_from_meta(meta: dict, curve_layout: dict) -> np.ndarray:
    """读取预处理模型数据中保存的固定物理时间网格。"""
    mode = str(meta.get("curve_time_mode", "missing"))
    if mode != "fixed":
        raise ValueError(
            "surrogate scoring requires curve_time_mode='fixed'; "
            f"got curve_time_mode={mode!r}"
        )

    raw_time = meta.get("prediction_curve_time")
    if raw_time is None:
        raise ValueError("curve_time_mode is fixed but prediction_curve_time is missing")

    prediction_time = np.asarray(raw_time, dtype=np.float64).reshape(-1)
    pressure_part = next(
        (part for part in curve_layout["parts"] if str(part["name"]) == "log_pressure"),
        None,
    )
    if pressure_part is None:
        raise ValueError("curve_layout has no log_pressure part")
    expected_size = int(pressure_part["end"]) - int(pressure_part["start"])

    if prediction_time.size != expected_size:
        raise ValueError(
            f"prediction_curve_time has {prediction_time.size} points; expected {expected_size}"
        )
    if not np.isfinite(prediction_time).all() or np.any(prediction_time <= 0.0):
        raise ValueError("prediction_curve_time must contain positive finite values")
    if np.any(np.diff(prediction_time) <= 0.0):
        raise ValueError("prediction_curve_time must be strictly increasing")
    return prediction_time


def _prepare_timed_raw_curve(
    time: np.ndarray,
    pressure: np.ndarray,
    derivative: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """清洗并排序一条原始压力/导数曲线，不在此处进行重采样。"""
    time = np.asarray(time, dtype=np.float64).reshape(-1)
    pressure = np.asarray(pressure, dtype=np.float64).reshape(-1)
    derivative = np.asarray(derivative, dtype=np.float64).reshape(-1)
    if time.size != pressure.size or time.size != derivative.size:
        raise ValueError("time, pressure and derivative lengths do not match")

    valid = (
        np.isfinite(time)
        & np.isfinite(pressure)
        & np.isfinite(derivative)
        & (time > 0.0)
    )
    time = time[valid]
    pressure = pressure[valid]
    derivative = derivative[valid]
    if time.size < 3:
        raise ValueError("timed curve has fewer than three valid points")

    order = np.argsort(time, kind="stable")
    time = time[order]
    pressure = pressure[order]
    derivative = derivative[order]
    keep = np.ones(time.size, dtype=bool)
    keep[1:] = time[1:] > time[:-1]
    time = time[keep]
    pressure = pressure[keep]
    derivative = derivative[keep]
    if time.size < 3:
        raise ValueError("timed curve has fewer than three unique time points")
    return time, pressure, derivative


def _calculate_raw_curve_objective_1d(target: np.ndarray, pred: np.ndarray) -> float:
    """在原始数值上复现 nmCalculationAutoFitPSO::calculatePointError。"""
    target = np.asarray(target, dtype=np.float64).reshape(-1)
    pred = np.asarray(pred, dtype=np.float64).reshape(-1)
    if target.size == 0 or pred.size != target.size:
        return float("inf")
    if not (np.isfinite(target).all() and np.isfinite(pred).all()):
        return float("inf")

    target_abs = np.maximum(np.abs(target), 1.0e-10)
    pred_abs = np.maximum(np.abs(pred), 1.0e-10)
    log_error = np.abs(np.log(target_abs) - np.log(pred_abs))
    relative_error = np.abs(target - pred) / np.maximum(
        np.maximum(np.abs(target), np.abs(pred)),
        1.0e-10,
    )
    point_error = 0.7 * log_error + 0.3 * relative_error
    return float(np.sqrt(np.mean(point_error**2)))


def timed_raw_dual_objective(
    target_time: np.ndarray,
    target_pressure: np.ndarray,
    target_derivative: np.ndarray,
    pred_time: np.ndarray,
    pred_pressure: np.ndarray,
    pred_derivative: np.ndarray,
    n_common_points: int = 50,
    w_pressure: float = 0.5,
    w_derivative: float = 0.5,
) -> dict[str, float]:
    """按照 C++ 目标函数的规则，在同一物理时间网格上比较两条原始曲线。"""
    if int(n_common_points) < 2:
        raise ValueError("n_common_points must be at least two")

    target_time, target_pressure, target_derivative = _prepare_timed_raw_curve(
        target_time,
        target_pressure,
        target_derivative,
    )
    pred_time, pred_pressure, pred_derivative = _prepare_timed_raw_curve(
        pred_time,
        pred_pressure,
        pred_derivative,
    )

    overlap_start = max(float(target_time[0]), float(pred_time[0]))
    overlap_end = min(float(target_time[-1]), float(pred_time[-1]))
    if overlap_start <= 0.0 or overlap_start >= overlap_end:
        return {
            "log_pressure_objective": float("inf"),
            "log_derivative_objective": float("inf"),
            "dual_log_objective": float("inf"),
        }

    common_time = np.geomspace(overlap_start, overlap_end, int(n_common_points))
    target_pressure_common = np.interp(common_time, target_time, target_pressure)
    target_derivative_common = np.interp(common_time, target_time, target_derivative)
    pred_pressure_common = np.interp(common_time, pred_time, pred_pressure)
    pred_derivative_common = np.interp(common_time, pred_time, pred_derivative)

    p_obj = _calculate_raw_curve_objective_1d(
        target_pressure_common,
        pred_pressure_common,
    )
    d_obj = _calculate_raw_curve_objective_1d(
        target_derivative_common,
        pred_derivative_common,
    )
    total_w = max(float(w_pressure) + float(w_derivative), 1.0e-12)
    combined = (float(w_pressure) * p_obj + float(w_derivative) * d_obj) / total_w
    return {
        "log_pressure_objective": float(p_obj),
        "log_derivative_objective": float(d_obj),
        "dual_log_objective": float(combined),
    }


def timed_dual_log_objective(
    target_time: np.ndarray,
    target_pressure: np.ndarray,
    target_derivative: np.ndarray,
    pred_time: np.ndarray,
    pred_curve: np.ndarray,
    curve_layout: dict,
    n_common_points: int = 50,
    w_pressure: float = 0.5,
    w_derivative: float = 0.5,
) -> dict[str, float]:
    """按照 C++ 时间对齐规则，将预测对数曲线与原始目标曲线进行比较。"""
    pred_parts = split_curve_by_layout(pred_curve, curve_layout)
    pred_log_pressure = pred_parts["log_pressure"]
    pred_log_derivative = pred_parts["log_derivative"]
    pred_time = np.asarray(pred_time, dtype=np.float64).reshape(-1)
    if pred_time.size != pred_log_pressure.size or pred_time.size != pred_log_derivative.size:
        raise ValueError("pred_time length does not match predicted curve parts")

    target_time, target_pressure, target_derivative = _prepare_timed_raw_curve(
        target_time,
        target_pressure,
        target_derivative,
    )
    # 不使用超出目标曲线实际时间范围的固定网格输出做插值；这些位置在训练时已由掩码排除。
    covered = pred_time <= target_time[-1]
    if int(np.sum(covered)) < 3:
        raise ValueError("target range contains fewer than three prediction time points")
    pred_time = pred_time[covered]
    pred_log_pressure = pred_log_pressure[covered]
    pred_log_derivative = pred_log_derivative[covered]

    with np.errstate(over="ignore", invalid="ignore"):
        pred_pressure = np.exp(pred_log_pressure)
        pred_derivative = np.exp(pred_log_derivative)

    target_end = float(target_time[-1])
    if pred_time[-1] < target_end:
        dt = float(pred_time[-1] - pred_time[-2])
        if dt <= 0.0:
            raise ValueError("prediction time grid has a non-positive final interval")
        fraction = (target_end - float(pred_time[-1])) / dt
        pressure_end = pred_pressure[-1] + fraction * (pred_pressure[-1] - pred_pressure[-2])
        derivative_end = pred_derivative[-1] + fraction * (
            pred_derivative[-1] - pred_derivative[-2]
        )
        pred_time = np.append(pred_time, target_end)
        pred_pressure = np.append(pred_pressure, max(float(pressure_end), 1.0e-300))
        pred_derivative = np.append(pred_derivative, max(float(derivative_end), 1.0e-300))

    return timed_raw_dual_objective(
        target_time=target_time,
        target_pressure=target_pressure,
        target_derivative=target_derivative,
        pred_time=pred_time,
        pred_pressure=pred_pressure,
        pred_derivative=pred_derivative,
        n_common_points=n_common_points,
        w_pressure=w_pressure,
        w_derivative=w_derivative,
    )