sacflow/methods/sacflow_step.py

from __future__ import annotations
from typing import Dict, Tuple
import torch
import torch.nn.functional as F
from sacflow.utils.metrics import torch_soft_dice_loss, entropy_loss, confidence_and_margin, finite_difference_boundary
from sacflow.methods.task_space import centered_classifier_basis, random_basis, project_task_and_residual, project_to_residual
from sacflow.methods.source_memory import class_moments, moment_transport_residual
from sacflow.utils.misc import unwrap_model


def ce_loss_masked(logits, target, mask=None):
    loss = F.cross_entropy(logits, target.long(), reduction="none")
    if mask is not None:
        loss = loss * mask.float()
        return loss.sum() / (mask.float().sum() + 1e-6)
    return loss.mean()




def dice_loss_masked(logits, target, mask=None, num_classes=None, eps=1e-5):
    if num_classes is None:
        num_classes = logits.shape[1]
    probs = torch.softmax(logits, dim=1)
    target = target.clamp(0, num_classes - 1).long()
    onehot = torch.nn.functional.one_hot(target, num_classes).permute(0,4,1,2,3).float()
    if mask is not None:
        m = mask.float().unsqueeze(1)
        probs = probs * m
        onehot = onehot * m
    dims = tuple(range(2, logits.ndim))
    inter = (probs * onehot).sum(dims)
    denom = probs.sum(dims) + onehot.sum(dims)
    dice = (2 * inter + eps) / (denom + eps)
    return 1.0 - dice[:, 1:].mean()


def kl_masked(p, q, mask=None, eps=1e-8):
    # p,q probabilities [B,C,H,W,D]
    kl = (p * ((p+eps).log() - (q+eps).log())).sum(dim=1)
    if mask is not None:
        kl = kl * mask.float()
        return kl.sum() / (mask.float().sum() + 1e-6)
    return kl.mean()


def _stats_distance(mu_a, std_a, mu_b, std_b):
    # simple differentiable class/channel statistic distance
    return (mu_a - mu_b).abs().mean() + (std_a - std_b).abs().mean()


def build_task_basis(model, cfg, feat_dim, device, dtype):
    basis = cfg.get("sacflow", {}).get("basis", "centered_svd")
    W = model.final_classifier_weight().detach().to(device=device, dtype=dtype)
    if basis == "random":
        rank = max(1, min(W.shape[0]-1, feat_dim))
        return random_basis(feat_dim, rank, device, dtype)
    return centered_classifier_basis(W, foreground_only=cfg.get("sacflow", {}).get("foreground_only_basis", False))


def make_tau(B, device, dtype):
    return torch.rand(B, device=device, dtype=dtype)


def sacflow_forward_step(model, teacher, velocity_field, batch, memory, cfg):
    x = batch["image"]
    num_classes = cfg["data"]["num_classes"]
    # Use the wrapped model for the main forward when DDP is active; unwrap only for helper methods.
    base_model = unwrap_model(model)
    logits, feats = model(x, return_features=True)
    feat = feats["prelogit"]
    B, d, H, W, D = feat.shape
    with torch.no_grad():
        tlogits = teacher(x)
        tprobs = torch.softmax(tlogits, dim=1)
        if tprobs.shape[-3:] != (H,W,D):
            tprobs_f = F.interpolate(tprobs, size=(H,W,D), mode="trilinear", align_corners=False)
        else:
            tprobs_f = tprobs
        conf, margin, pseudo = confidence_and_margin(tprobs)
        mask = (conf > float(cfg["train"].get("pseudo_conf_threshold", 0.75))).float()
    Q = build_task_basis(base_model, cfg, d, feat.device, feat.dtype)
    F_task, Rt = project_task_and_residual(feat, Q)
    if cfg.get("sacflow", {}).get("flow_space", "residual") == "whole_feature":
        F_task = torch.zeros_like(feat)
        Rt = feat
        Q = torch.empty(d, 0, device=feat.device, dtype=feat.dtype)
    use_mem = cfg.get("sacflow", {}).get("use_compact_memory", True) and memory is not None
    whole_feature = cfg.get("sacflow", {}).get("flow_space", "residual") == "whole_feature"
    if use_mem and whole_feature and "feature_mu" in memory:
        src_mu = memory["feature_mu"].to(feat.device, feat.dtype)
        src_std = memory["feature_std"].to(feat.device, feat.dtype)
    elif use_mem and "residual_mu" in memory:
        src_mu = memory["residual_mu"].to(feat.device, feat.dtype)
        src_std = memory["residual_std"].to(feat.device, feat.dtype)
    else:
        # strict fallback: use target stats as weak source proxy; ablate this separately.
        # This makes R0 close to Rt and is intentionally weaker than compact-memory SACFlow.
        src_mu, src_std, _ = class_moments(Rt.detach(), tprobs_f.detach())
    R0, tgt_mu, tgt_std = moment_transport_residual(Rt.detach(), tprobs_f.detach(), src_mu, src_std)
    R1 = Rt
    tau = make_tau(B, feat.device, feat.dtype)
    tau_view = tau.view(B,1,1,1,1)
    sigma_tau = float(cfg.get("sacflow", {}).get("sigma_tau", 0.0) or 0.0)
    noise = torch.randn_like(Rt) if sigma_tau > 0 else torch.zeros_like(Rt)
    R_interp = (1 - tau_view) * R0 + tau_view * R1 + sigma_tau * tau_view * (1 - tau_view) * noise
    u = R1 - R0 + sigma_tau * (1 - 2*tau_view) * noise
    anchors = {"probs": tprobs_f.detach(), "boundary": finite_difference_boundary(tprobs_f.detach())}
    use_v = bool(cfg.get("sacflow", {}).get("use_velocity_field", True))
    if use_v and velocity_field is not None:
        v = velocity_field(R_interp, tau, anchors)
        if cfg.get("sacflow", {}).get("project_velocity_to_nullspace", True) and Q.numel() > 0:
            v = project_to_residual(v, Q)
        fm_loss = F.mse_loss(v, u.detach())
        path_mode = cfg.get("sacflow", {}).get("path_from_velocity", "one_step")
        if path_mode == "one_step":
            Rtau = R0 + tau_view * v
        else:
            Rtau = R_interp
    else:
        v = torch.zeros_like(Rt)
        fm_loss = torch.tensor(0.0, device=feat.device)
        Rtau = R_interp
    Ftau = F_task + Rtau
    # Classify path states through the wrapped model so DDP can track gradients.
    path_logits = model(prelogit_features=Ftau)
    # resize mask/pseudo if needed
    pseudo_f = pseudo
    mask_f = mask
    if pseudo.shape[-3:] != path_logits.shape[-3:]:
        pseudo_f = F.interpolate(pseudo[:, None].float(), size=path_logits.shape[-3:], mode="nearest")[:,0].long()
        mask_f = F.interpolate(mask[:, None].float(), size=path_logits.shape[-3:], mode="nearest")[:,0]
    path_probs = torch.softmax(path_logits, dim=1)
    task_loss = ce_loss_masked(path_logits, pseudo_f, mask_f)
    dice = dice_loss_masked(path_logits, pseudo_f, mask_f, num_classes)
    ent = entropy_loss(path_logits)
    if tprobs.shape[-3:] != path_logits.shape[-3:]:
        tprobs_path = F.interpolate(tprobs.detach(), size=path_logits.shape[-3:], mode="trilinear", align_corners=False)
    else:
        tprobs_path = tprobs.detach()
    task_kl = kl_masked(path_probs, tprobs_path, mask_f)
    b_pred = finite_difference_boundary(path_probs)
    b_ref = finite_difference_boundary(tprobs_path)
    boundary_loss = F.l1_loss(b_pred, b_ref)
    null_leak = torch.tensor(0.0, device=feat.device)
    if Q.numel() > 0:
        null_leak = (Rtau - R0 - project_to_residual(Rtau - R0, Q)).pow(2).mean()

    # Domain-progress proxy in residual-statistic space. A valid path state should
    # move monotonically from source-like residual stats toward target residual stats.
    rtau_mu, rtau_std, _ = class_moments(Rtau, tprobs_f.detach())
    dist_src = _stats_distance(rtau_mu, rtau_std, src_mu.detach(), src_std.detach())
    dist_tgt = _stats_distance(rtau_mu, rtau_std, tgt_mu.detach(), tgt_std.detach())
    rho = dist_src / (dist_src + dist_tgt + 1e-6)
    domain_progress_loss = (rho - tau.mean()).abs()

    losses_cfg = cfg.get("sacflow", {}).get("losses", {})
    val_cfg = cfg.get("sacflow", {}).get("validation", {})
    use_path_weights = bool(val_cfg.get("use_weights", True))
    if use_path_weights:
        alpha = float(val_cfg.get("task_alpha", 1.0))
        beta = float(val_cfg.get("anatomy_beta", 1.0))
        # Use detached scalar weights so the gate selects/weights path states but
        # does not create degenerate gradients that simply lower the weight.
        accepted_weight = torch.exp(-alpha * task_kl.detach() - beta * boundary_loss.detach()).clamp(
            min=float(val_cfg.get("min_weight", 0.05)), max=1.0
        )
    else:
        accepted_weight = torch.tensor(1.0, device=feat.device, dtype=feat.dtype)

    path_loss = accepted_weight * (
        float(losses_cfg.get("path_ce", 1.0)) * task_loss +
        float(losses_cfg.get("path_dice", 1.0)) * dice
    )
    loss = (
        float(losses_cfg.get("fm", 1.0)) * fm_loss +
        path_loss +
        float(losses_cfg.get("task_kl", 0.25)) * task_kl +
        float(losses_cfg.get("boundary", 0.05)) * boundary_loss +
        float(losses_cfg.get("null_leakage", 0.1)) * null_leak +
        float(losses_cfg.get("domain_progress", 0.05)) * domain_progress_loss +
        float(losses_cfg.get("entropy", 0.01)) * ent
    )
    logs = {
        "loss_total": loss.detach(),
        "loss_fm": fm_loss.detach(),
        "loss_path_ce": task_loss.detach(),
        "loss_path_dice": dice.detach(),
        "loss_task_kl": task_kl.detach(),
        "loss_boundary": boundary_loss.detach(),
        "loss_null_leak": null_leak.detach(),
        "loss_domain_progress": domain_progress_loss.detach(),
        "loss_entropy": ent.detach(),
        "sacflow_rho_mean": rho.detach(),
        "sacflow_tau_mean": tau.mean().detach(),
        "sacflow_velocity_mag": v.detach().abs().mean(),
        "sacflow_residual_gap": (R1-R0).detach().abs().mean(),
        "pseudo_conf_mean": conf.detach().mean(),
        "pseudo_accept_rate": mask.detach().mean(),
        "path_weight": accepted_weight.detach(),
    }
    return loss, logs