density_calibration.py

"""Affine density calibration + gradient consistency (WP-13.1 B).

Extracted from hybrid_loss so the loss module stays thin. Torch helpers
normalize inputs once to (B, 1, H, W).
"""

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np
import torch


def _masked_mse_np(res: np.ndarray, w: Optional[np.ndarray]) -> float:
    if w is None:
        return float(np.mean(res * res))
    return float(np.sum(w * res * res) / (np.sum(w) + 1e-8))


def _masked_mse_torch(res: torch.Tensor, w: Optional[torch.Tensor]) -> torch.Tensor:
    if w is None:
        return (res * res).mean()
    return (w * res * res).sum() / (w.sum() + 1e-8)


def as_b1hw(t: torch.Tensor) -> torch.Tensor:
    """Normalize a spatial tensor to (B, 1, H, W)."""
    if t.dim() == 2:
        return t.unsqueeze(0).unsqueeze(0)
    if t.dim() == 3:
        # (C, H, W) or (1, H, W) → (1, C, H, W); for masks C=1
        return t.unsqueeze(0)
    if t.dim() == 4:
        return t
    raise ValueError(f"expected 2D/3D/4D tensor, got shape {tuple(t.shape)}")


def affine_fit_numpy(
    d_pred: np.ndarray,
    d_obs: np.ndarray,
    w: Optional[np.ndarray],
) -> Tuple[float, float, float, np.ndarray]:
    """Fit D_obs ≈ a·D_pred + b on VALID pixels (detached coefficients).

    Returns (residual_mse, a, b, residual_map) where residual_map = D_obs − (a·D_pred+b).
    a clamped to [0, 10]; negative a → 0. Flat D_pred → a=0, b=mean(D_obs).
    """
    if w is not None:
        m = w > 0.5
        if int(m.sum()) < 2:
            m = np.ones(d_pred.shape, dtype=bool)
        x = d_pred[m].astype(np.float64)
        y = d_obs[m].astype(np.float64)
    else:
        x = d_pred.ravel().astype(np.float64)
        y = d_obs.ravel().astype(np.float64)

    mean_x = float(np.mean(x))
    mean_y = float(np.mean(y))
    var_x = float(np.mean((x - mean_x) ** 2))
    if var_x < 1e-8:
        a = 0.0
        b = mean_y
    else:
        cov_xy = float(np.mean((x - mean_x) * (y - mean_y)))
        a = cov_xy / var_x
        if a < 0.0:
            a = 0.0
        a = min(a, 10.0)
        b = mean_y - a * mean_x

    cal = (a * d_pred.astype(np.float64) + b).astype(np.float32)
    res = (d_obs.astype(np.float32) - cal).astype(np.float32)
    mse = _masked_mse_np(res, w)
    return float(mse), float(a), float(b), res


def pearson_r_valid(
    d_pred: np.ndarray,
    d_obs: np.ndarray,
    w: Optional[np.ndarray],
) -> float:
    """Pearson r(D_pred, D_obs) on VALID pixels (diagnostic)."""
    if w is not None:
        m = w > 0.5
        if int(m.sum()) < 2:
            return 0.0
        x = d_pred[m].astype(np.float64)
        y = d_obs[m].astype(np.float64)
    else:
        x = d_pred.ravel().astype(np.float64)
        y = d_obs.ravel().astype(np.float64)
    if x.size < 2:
        return 0.0
    sx = float(np.std(x))
    sy = float(np.std(y))
    if sx < 1e-12 or sy < 1e-12:
        return 0.0
    return float(np.corrcoef(x, y)[0, 1])


def gradient_consistency_numpy(
    d_cal: np.ndarray,
    d_obs: np.ndarray,
    w: Optional[np.ndarray],
) -> float:
    """L_grad = mean|∇x| + mean|∇y| with VALID-pair forward diffs."""
    # Accept (H,W) or squeeze trailing singleton channel
    if d_cal.ndim > 2:
        d_cal = np.squeeze(d_cal)
        d_obs = np.squeeze(d_obs)
        if w is not None:
            w = np.squeeze(w)
    h, ww = d_cal.shape[-2], d_cal.shape[-1]
    if h < 2 or ww < 2:
        return 0.0

    if w is None:
        valid = np.ones((h, ww), dtype=bool)
    else:
        valid = w > 0.5

    pair_x = valid[:, :-1] & valid[:, 1:]
    gx_cal = d_cal[:, 1:] - d_cal[:, :-1]
    gx_obs = d_obs[:, 1:] - d_obs[:, :-1]
    pair_y = valid[:-1, :] & valid[1:, :]
    gy_cal = d_cal[1:, :] - d_cal[:-1, :]
    gy_obs = d_obs[1:, :] - d_obs[:-1, :]

    terms = []
    n_x = int(pair_x.sum())
    if n_x > 0:
        terms.append(float(np.mean(np.abs(gx_cal[pair_x] - gx_obs[pair_x]))))
    n_y = int(pair_y.sum())
    if n_y > 0:
        terms.append(float(np.mean(np.abs(gy_cal[pair_y] - gy_obs[pair_y]))))
    if not terms:
        return 0.0
    if len(terms) == 2:
        return terms[0] + terms[1]
    return terms[0] * 2.0


def affine_fit_torch(
    d_pred: torch.Tensor,
    d_target: torch.Tensor,
    w: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, float, float]:
    """Detached a,b; residual = d_target − (a·d_pred+b). float32 for MPS.

    Normalizes all inputs to (B,1,H,W) once (fixes mixed 2D-mask + 4D-pred).
    Residual is always (B,1,H,W) and differentiable in d_pred.
    """
    d_pred_n = as_b1hw(d_pred)
    d_target_n = as_b1hw(d_target)
    w_n = as_b1hw(w) if w is not None else None

    with torch.no_grad():
        if w_n is not None:
            m = w_n > 0.5
            if int(m.sum().item()) < 2:
                x = d_pred_n.reshape(-1).float()
                y = d_target_n.reshape(-1).float()
            else:
                x = d_pred_n[m].float()
                y = d_target_n[m].float()
        else:
            x = d_pred_n.reshape(-1).float()
            y = d_target_n.reshape(-1).float()
        mx = x.mean()
        my = y.mean()
        var_x = ((x - mx) ** 2).mean()
        if float(var_x.item()) < 1e-8:
            a_f = 0.0
            b_f = float(my.item())
        else:
            cov = ((x - mx) * (y - my)).mean()
            a_f = float((cov / var_x).item())
            if a_f < 0.0:
                a_f = 0.0
            a_f = min(a_f, 10.0)
            b_f = float(my.item()) - a_f * float(mx.item())

    a_t = torch.tensor(a_f, device=d_pred.device, dtype=d_pred.dtype)
    b_t = torch.tensor(b_f, device=d_pred.device, dtype=d_pred.dtype)
    res = d_target_n - (a_t * d_pred_n + b_t)
    return res, a_f, b_f


def residual_torch(
    d_pred: torch.Tensor,
    d_target: torch.Tensor,
    w: Optional[torch.Tensor],
    calibration: str,
) -> Tuple[torch.Tensor, float, float]:
    """Single residual helper for forward_tensor / physics_loss (WP-13.1 B).

    Returns (residual, a, b) with residual always (B,1,H,W).
    Pointwise when calibration is not affine.
    """
    if calibration == "affine":
        return affine_fit_torch(d_pred, d_target, w)
    d_pred_n = as_b1hw(d_pred)
    d_target_n = as_b1hw(d_target)
    return d_pred_n - d_target_n, 1.0, 0.0


def gradient_consistency_torch(
    d_cal: torch.Tensor,
    d_obs: torch.Tensor,
    w: Optional[torch.Tensor],
) -> torch.Tensor:
    """Differentiable L_grad; all inputs normalized once to (B,1,H,W)."""
    d_cal = as_b1hw(d_cal)
    d_obs = as_b1hw(d_obs)
    w4 = as_b1hw(w) if w is not None else None

    gx_cal = d_cal[..., :, 1:] - d_cal[..., :, :-1]
    gx_obs = d_obs[..., :, 1:] - d_obs[..., :, :-1]
    gy_cal = d_cal[..., 1:, :] - d_cal[..., :-1, :]
    gy_obs = d_obs[..., 1:, :] - d_obs[..., :-1, :]

    if w4 is not None:
        pair_x = (w4[..., :, :-1] > 0.5) & (w4[..., :, 1:] > 0.5)
        pair_y = (w4[..., :-1, :] > 0.5) & (w4[..., 1:, :] > 0.5)
        nx = pair_x.to(d_cal.dtype).sum().clamp(min=1.0)
        ny = pair_y.to(d_cal.dtype).sum().clamp(min=1.0)
        mx = (torch.abs(gx_cal - gx_obs) * pair_x.to(d_cal.dtype)).sum() / nx
        my = (torch.abs(gy_cal - gy_obs) * pair_y.to(d_cal.dtype)).sum() / ny
    else:
        mx = torch.abs(gx_cal - gx_obs).mean()
        my = torch.abs(gy_cal - gy_obs).mean()
    return mx + my