scripts/pixelmodel_robustness/metrics.py

"""Weight and image quality metrics."""

from __future__ import annotations

import math
from collections.abc import Mapping, Sequence

import numpy as np


def finite_metrics(values: np.ndarray) -> dict[str, int | float]:
    """NaN/Inf/subnormalを修復せずに数える。"""
    array = np.asarray(values)
    finite = np.isfinite(array)
    subnormal = np.isfinite(array) & (array != 0) & (np.abs(array) < np.finfo(array.dtype).tiny) if np.issubdtype(array.dtype, np.floating) else np.zeros(array.shape, dtype=bool)
    total = int(array.size)
    return {
        "total_count": total,
        "finite_only_count": int(finite.sum()),
        "finite_rate": float(finite.mean()) if total else 1.0,
        "nan_count": int(np.isnan(array).sum()),
        "posinf_count": int(np.isposinf(array).sum()),
        "neginf_count": int(np.isneginf(array).sum()),
        "subnormal_count": int(subnormal.sum()),
        "subnormal_rate": float(subnormal.mean()) if total else 0.0,
    }


def channel_change_metrics(reference: np.ndarray, candidate: np.ndarray) -> dict[str, int | float]:
    """RGBのbyte/bit/pixel変更数とrateを計算する。"""
    left = np.asarray(reference, dtype=np.uint8)
    right = np.asarray(candidate, dtype=np.uint8)
    if left.shape != right.shape or left.ndim != 3 or left.shape[-1] != 3:
        raise ValueError("reference and candidate must be equal-shaped RGB arrays")
    changed = left != right
    bits = np.unpackbits(np.bitwise_xor(left, right), axis=-1).reshape(*left.shape[:2], 24)
    result: dict[str, int | float] = {
        "changed_pixel_count": int(changed.any(axis=2).sum()),
        "changed_pixel_rate": float(changed.any(axis=2).mean()),
        "changed_bit_count": int(bits.sum()),
        "changed_bit_rate": float(bits.mean()),
    }
    for index, name in enumerate(("red", "green", "blue")):
        result[f"{name}_changed_byte_count"] = int(changed[..., index].sum())
        result[f"{name}_changed_byte_rate"] = float(changed[..., index].mean())
        result[f"{name}_changed_bit_count"] = int(bits[..., index * 8:(index + 1) * 8].sum())
        result[f"{name}_changed_bit_rate"] = float(bits[..., index * 8:(index + 1) * 8].mean())
    return result


def _finite_pair(reference: np.ndarray, candidate: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    mask = np.isfinite(reference) & np.isfinite(candidate)
    return reference[mask].astype(np.float64), candidate[mask].astype(np.float64)


def _metrics(reference: np.ndarray, candidate: np.ndarray) -> dict[str, float]:
    ref, cand = _finite_pair(reference, candidate)
    if not len(ref):
        return {"l2": float("nan"), "mse": float("nan"), "cosine": float("nan"), "max_abs": float("nan"), "changed_percent": 100.0}
    diff = cand - ref
    denom = np.linalg.norm(ref) * np.linalg.norm(cand)
    return {"l2": float(np.linalg.norm(diff)), "mse": float(np.mean(diff * diff)), "cosine": float(np.dot(ref, cand) / denom) if denom else 1.0, "max_abs": float(np.max(np.abs(diff))), "changed_percent": float(np.mean(diff != 0) * 100)}


def normalized_rmse(reference: np.ndarray, candidate: np.ndarray, eps: float = 1e-12) -> float:
    """有限要素のRMSEをreferenceのRMSで正規化する。"""
    ref, cand = _finite_pair(np.asarray(reference), np.asarray(candidate))
    if not len(ref):
        return float("nan")
    return float(np.sqrt(np.mean((cand - ref) ** 2)) / (np.sqrt(np.mean(ref ** 2)) + eps))


def ulp_distance(reference: np.ndarray, candidate: np.ndarray) -> np.ndarray:
    """有限float32のbit patternを全順序へ写像してULP距離を返す。"""
    ref = np.asarray(reference, dtype=np.float32)
    cand = np.asarray(candidate, dtype=np.float32)
    mask = np.isfinite(ref) & np.isfinite(cand)
    ref_bits = ref.view(np.int32).astype(np.int64)
    cand_bits = cand.view(np.int32).astype(np.int64)
    ref_order = np.where(ref_bits < 0, 0x80000000 - ref_bits, ref_bits)
    cand_order = np.where(cand_bits < 0, 0x80000000 - cand_bits, cand_bits)
    return np.abs(ref_order[mask] - cand_order[mask])


def weight_metrics(reference: np.ndarray, candidate: np.ndarray, layers: Mapping[str, tuple[int, int]]) -> dict[str, object]:
    """NaN/Infを補完せずにglobal/layer weight metricsを計算する。"""
    ref = np.asarray(reference).reshape(-1)
    cand = np.asarray(candidate).reshape(-1)
    base = _metrics(ref, cand)
    finite = finite_metrics(np.asarray(candidate))
    ulp = ulp_distance(reference, candidate)
    nonfinite_count = int((~np.isfinite(cand)).sum())
    result: dict[str, object] = {
        "finite": bool(np.isfinite(cand).all()),
        "nonfinite_count": int((~np.isfinite(ref)).sum() + nonfinite_count),
        "reference_nonfinite_count": int((~np.isfinite(ref)).sum()),
        "candidate_nonfinite_count": nonfinite_count,
        "global_l2": base["l2"], "global_mse": base["mse"], "global_cosine": base["cosine"],
        "global_max_abs": base["max_abs"], "global_changed_percent": float(np.mean(ref != cand) * 100),
        "normalized_rmse": normalized_rmse(ref, cand),
        "ulp_p50": float(np.percentile(ulp, 50)) if len(ulp) else float("nan"),
        "ulp_p95": float(np.percentile(ulp, 95)) if len(ulp) else float("nan"),
        **finite,
        "layers": [],
    }
    result["layers"] = []
    for name, (start, end) in layers.items():
        raw_ref = np.asarray(reference).reshape(-1)[start:end]
        raw_cand = np.asarray(candidate).reshape(-1)[start:end]
        layer_ref, layer_cand = _finite_pair(raw_ref, raw_cand)
        result["layers"].append({"layer": name, **_metrics(layer_ref, layer_cand), "finite": bool(np.isfinite(raw_cand).all()), "nonfinite_count": int((~np.isfinite(raw_cand)).sum()), "normalized_rmse": normalized_rmse(raw_ref, raw_cand)})
    return result


def _ssim(reference: np.ndarray, candidate: np.ndarray) -> float:
    from skimage.metrics import structural_similarity
    size = min(reference.shape[:2])
    win_size = min(7, size if size % 2 else size - 1)
    return float(structural_similarity(reference, candidate, channel_axis=2, data_range=255, win_size=win_size))


def image_metrics(reference: np.ndarray, candidate: np.ndarray, reference_size: int, file_size: int) -> dict[str, float | int]:
    """RGB画像のPSNR/SSIMとstorage metricsを返す。"""
    ref = np.asarray(reference, dtype=np.float64)
    cand = np.asarray(candidate, dtype=np.float64)
    mse = float(np.mean((ref - cand) ** 2))
    psnr = float("inf") if mse == 0 else float(20 * math.log10(255.0 / math.sqrt(mse)))
    return {"image_mse": mse, "psnr": psnr, "ssim": max(-1.0, min(1.0, _ssim(ref, cand))), "file_size_bytes": int(file_size), "compression_ratio": float(reference_size / file_size) if file_size else float("inf")}


def bootstrap_ci(values: Sequence[float], seed: int = 0, samples: int = 2000) -> tuple[float, float]:
    """paired値のpercentile bootstrap 95% CIを計算する。"""
    array = np.asarray(values, dtype=float)
    if array.size == 0:
        return float("nan"), float("nan")
    rng = np.random.default_rng(seed)
    means = np.array([rng.choice(array, array.size, replace=True).mean() for _ in range(samples)])
    return float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5))