src/remedi/optimization.py

"""Convex finite-action minimax regret and matched generation rules."""
from dataclasses import dataclass
import time
import numpy as np
from scipy.special import logsumexp

@dataclass
class Decision:
    probabilities: np.ndarray
    objective: float
    worst_regret: float
    constraint_violation: float
    seconds: float
    status: str

def validate_problem(mean, factor, diagonal, beta, tau, prior=None):
    mean = np.asarray(mean, dtype=float)
    factor = np.asarray(factor, dtype=float)
    diagonal = np.asarray(diagonal, dtype=float)
    n = len(mean)
    if n == 0 or mean.shape != (n,) or factor.ndim != 2 or factor.shape[0] != n or diagonal.shape != (n,):
        raise ValueError("Incompatible finite-action dimensions")
    if not all(np.isfinite(x).all() for x in [mean, factor, diagonal]):
        raise ValueError("Nonfinite loss or uncertainty values")
    if (diagonal < 0).any() or not np.isfinite(beta) or not np.isfinite(tau) or beta < 0 or tau <= 0:
        raise ValueError("Need nonnegative uncertainty and positive temperature")
    prior = np.ones(n)/n if prior is None else np.asarray(prior, dtype=float)
    if prior.shape != (n,) or not np.isfinite(prior).all() or (prior <= 0).any():
        raise ValueError("Every eligible action needs a positive finite prior")
    return mean, factor, diagonal, prior/prior.sum()

def regret_components(p, mean, factor, diagonal, beta):
    """All comparator constraints in O(N*S), including diagonal cross-terms."""
    q = factor.T @ p
    variance = np.sum((factor-q[None, :])**2, axis=1)
    variance += np.dot(diagonal, p*p)-2*diagonal*p+diagonal
    return np.dot(p, mean)-mean+beta*np.sqrt(np.maximum(variance, 0))

def softmin(values, tau, prior=None):
    values = np.asarray(values, float)
    if tau <= 0 or not np.isfinite(values).all():
        raise ValueError("Softmin needs finite values and positive temperature")
    prior = np.ones(len(values))/len(values) if prior is None else np.asarray(prior, float)
    if (prior <= 0).any():
        raise ValueError("Prior must have positive support")
    logits = np.log(prior)-values/tau
    return np.exp(logits-logsumexp(logits))

def solve_regret(mean, factor, diagonal, beta=1., tau=.05, prior=None,
                 tolerance=1e-6, constraint_generation=True, max_rounds=100, solver="CLARABEL"):
    """Optimize a distribution over N candidate molecule-dose pairs.

    mean is (N,), factor is (N, S), and diagonal is (N,). Their covariance
    is factor @ factor.T + diag(diagonal). beta sets the uncertainty radius
    and tau sets KL regularization. Return probabilities and solver diagnostics.
    """
    import cvxpy as cp
    start = time.perf_counter()
    mean, factor, diagonal, prior = validate_problem(mean, factor, diagonal, beta, tau, prior)
    n = len(mean)
    # At zero uncertainty, the optimizer is the prior-weighted Gibbs distribution.
    if beta == 0:
        p = softmin(mean, tau, prior)
        regret = float(p@mean-mean.min())
        objective = regret+tau*float(np.sum(p*np.log(np.maximum(p, 1e-300)/prior)))
        return Decision(p, objective, regret, 0., time.perf_counter()-start, "analytic")
    active = set(range(n)) if not constraint_generation else {int(np.argmin(mean))}
    probability, epigraph = cp.Variable(n), cp.Variable()
    p = None
    for _ in range(max_rounds):
        constraints = [probability >= 0, cp.sum(probability) == 1]
        for b in sorted(active):
            e = np.zeros(n); e[b] = 1
            # Compare the randomized intervention with the single action b.
            contrast = probability-e
            norm = cp.norm(cp.hstack([factor.T@contrast, cp.multiply(np.sqrt(diagonal), contrast)]), 2)
            constraints.append(epigraph >= contrast@mean+beta*norm)
        problem = cp.Problem(cp.Minimize(epigraph+tau*cp.sum(cp.kl_div(probability, prior))), constraints)
        options = {"tol_gap_abs": tolerance*.1, "tol_feas": tolerance*.1} if solver == "CLARABEL" else {}
        problem.solve(solver=solver, warm_start=True, **options)
        if problem.status not in {"optimal", "optimal_inaccurate"} or probability.value is None:
            raise RuntimeError(f"Cone solver failed: {problem.status}")
        p = np.maximum(np.asarray(probability.value).ravel(), 0)
        p /= p.sum()
        # Check every comparator, including those outside the active constraint set.
        values = regret_components(p, mean, factor, diagonal, beta)
        worst = int(np.argmax(values))
        violation = max(0., float(values[worst]-epigraph.value))
        if violation <= tolerance:
            status = problem.status
            break
        # Add the most violated comparator and solve the expanded program.
        active.add(worst)
    else:
        status = "iteration_limit"
    robust_regret = float(np.max(regret_components(p, mean, factor, diagonal, beta)))
    kl = float(np.sum(p*np.log(np.maximum(p, 1e-300)/prior)))
    return Decision(p, robust_regret+tau*kl, robust_regret, violation,
                    time.perf_counter()-start, status)

def finite_minimax(losses, tau=.05, prior=None, solver="CLARABEL"):
    import cvxpy as cp
    losses = np.asarray(losses, float)
    if losses.ndim != 2 or not np.isfinite(losses).all():
        raise ValueError("Scenarios must be a finite S-by-N matrix")
    n = losses.shape[1]
    prior = np.ones(n)/n if prior is None else np.asarray(prior, float)
    regrets = losses-losses.min(axis=1, keepdims=True)
    p, z = cp.Variable(n), cp.Variable()
    problem = cp.Problem(cp.Minimize(z+tau*cp.sum(cp.kl_div(p, prior))),
                         [p >= 0, cp.sum(p) == 1, regrets@p <= z])
    problem.solve(solver=solver)
    if p.value is None:
        raise RuntimeError(f"Finite minimax failed: {problem.status}")
    values = np.maximum(np.asarray(p.value), 0)
    return values/values.sum()

def generate_baseline(name, mean, factor, diagonal, losses, beta=1., tau=.05, seed=0):
    n = len(mean)
    if name == "uniform":
        return np.ones(n)/n
    if name == "mean":
        p = np.zeros(n); p[int(np.argmin(mean))] = 1; return p
    if name == "gibbs":
        return softmin(mean, tau)
    if name == "marginal":
        return softmin(mean+beta*np.sqrt((factor**2).sum(1)+diagonal), tau)
    if name == "thompson":
        # Exact empirical distribution of scenario-wise minimizers.
        return np.bincount(np.argmin(losses, axis=1), minlength=n)/len(losses)
    if name == "finite-minimax":
        return finite_minimax(losses, tau)
    if name == "absolute":
        import cvxpy as cp
        p = cp.Variable(n)
        norm = cp.norm(cp.hstack([factor.T@p, cp.multiply(np.sqrt(diagonal), p)]), 2)
        objective = cp.Minimize(mean@p+beta*norm+tau*cp.sum(cp.kl_div(p, np.ones(n)/n)))
        problem = cp.Problem(objective, [p >= 0, cp.sum(p) == 1]); problem.solve(solver="CLARABEL")
        if p.value is None: raise RuntimeError("Absolute-loss solver failed")
        values = np.maximum(p.value, 0); return values/values.sum()
    raise ValueError(f"Unknown baseline {name}")