dense_evolution/mitigation/renyi.py

"""Sandwiched Quantum Renyi Divergence for full density-matrix diagnostics

(Muller-Lennert, Reeb, Wolf, Wilde, "On quantum Renyi entropies: a new

generalization and some applications", arXiv:1306.3142, Definition 1).



D_alpha(rho||sigma) = 1/(alpha-1) * log2 Tr[(sigma^e rho sigma^e)^alpha],

e = (1-alpha)/(2*alpha), with the alpha->1 limit reducing to the standard

quantum relative entropy and alpha=1/2 reducing to a fidelity-based form.



Originated from a Colab proposal with a real bug in its case_general

branch: `tr_inner = jnp.maximum(tr_inner, 1.0)` floors the inner trace at

1.0 even when the true value is < 1 (the normal case for non-commuting

rho, sigma), silently forcing every result to log2(1)=0 -- confirmed

directly in the Colab's own printed output (alpha=1.5 gave exactly

0.000000 across an entire rotation sweep). A second, deeper bug survived

the floor-value fix alone: for alpha > 1, a trace below 1 is not a

numerical artifact to clamp away, it is the genuine signature of a

support mismatch (supp(rho) not contained in supp(sigma)), which the

divergence must report as +inf, not a finite (and wrong-signed) number --

verified by hand on two different pure states: Tr[Q^1.5] = 0.6759,

matching the closed-form prediction (|<sigma|rho>|^2)^alpha exactly, and

plugging that into the naive formula gives a finite NEGATIVE divergence,

worse than the original bug's silent zero since it looks plausible

instead of visibly wrong.



Fixed and validated in Dense-Evolution-Discovery, Experiment 29

(https://tatopenn-cell.github.io/Dense-Evolution-Discovery/sandwiched_renyi_density_matrix/):

against the alpha->1 relative-entropy limit (matches an independent numpy

reference to 4 decimal places), the commuting/diagonal case (reduces

exactly to the classical Renyi divergence), and the support-violation

+inf case (verified at alpha>1, confirmed the branch does not fire

spuriously at alpha<1).



Its originally proposed use case -- replacing the JSD-based truncation

criterion in dense_evolution.mps's bond-dimension search -- was

independently disproven: on the diagonal singular-value spectrum used

there, rho and sigma commute, so a non-commuting-aware divergence induces

the exact same truncation ordering as JSD (5 benchmark configurations,

byte-identical chi_used and truncation error every time) -- nothing for

it to add in that setting. Promoted here instead for the genuinely

non-commuting full-density-matrix diagnostic use case it WAS validated

against: alongside uhlmann_fidelity, tracking a Bell state degraded by

amplitude damping, where the two metrics' noise-sensitivity curves

visibly diverge from each other.

"""
import jax
import jax.numpy as jnp

__all__ = ["sandwiched_renyi_divergence", "sandwiched_renyi_divergence_jit"]

_EPS = 1e-12


def _case_half(rho, sigma):
    ev_r, ec_r = jnp.linalg.eigh(rho)
    safe_ev_r = jnp.where(ev_r > _EPS, ev_r, 0.0)
    sqrt_rho = (ec_r * jnp.sqrt(safe_ev_r)) @ jnp.conj(ec_r).T
    uhlmann_mat = sqrt_rho @ sigma @ sqrt_rho
    ev_u = jnp.linalg.eigvalsh(uhlmann_mat)
    safe_ev_u = jnp.where(ev_u > _EPS, ev_u, 0.0)
    fidelity = jnp.clip(jnp.sum(jnp.sqrt(safe_ev_u)), 0.0, 1.0)
    return -2.0 * jnp.log2(jnp.maximum(fidelity, _EPS))


def _case_one(rho, sigma):
    ev_r, ec_r = jnp.linalg.eigh(rho)
    safe_ev_r = jnp.where(ev_r > _EPS, ev_r, 1.0)
    log_rho = (ec_r * jnp.where(ev_r > _EPS, jnp.log2(safe_ev_r), 0.0)) @ jnp.conj(ec_r).T
    ev_s, ec_s = jnp.linalg.eigh(sigma)
    safe_ev_s = jnp.where(ev_s > _EPS, ev_s, 1.0)
    log_sigma = (ec_s * jnp.where(ev_s > _EPS, jnp.log2(safe_ev_s), 0.0)) @ jnp.conj(ec_s).T
    return jnp.trace(rho @ (log_rho - log_sigma)).real


def _case_general(rho, sigma, alpha):
    exponent = (1.0 - alpha) / (2.0 * alpha)
    ev_s, ec_s = jnp.linalg.eigh(sigma)
    mask_s = ev_s > _EPS
    safe_ev_s = jnp.where(mask_s, ev_s, 1.0)
    pow_ev_s = jnp.where(mask_s, safe_ev_s ** exponent, 0.0)
    sigma_pow = (ec_s * pow_ev_s) @ jnp.conj(ec_s).T

    int_m = sigma_pow @ rho @ sigma_pow
    v_int = jnp.linalg.eigvalsh(int_m)
    mask_int = v_int > _EPS
    safe_v_int = jnp.where(mask_int, v_int, 1.0)
    pow_v_int = jnp.where(mask_int, safe_v_int ** alpha, 0.0)

    tr_inner = jnp.sum(pow_v_int)

    # For alpha > 1, D_alpha is finite only when supp(rho) subset supp(sigma);
    # tr_inner < 1 there is the genuine signature of a support mismatch, not
    # a numerical artifact -- see the module docstring's bug-history note.
    is_support_violation = (alpha > 1.0) & (tr_inner < 1.0 - 1e-9)
    tr_inner_safe = jnp.maximum(tr_inner, _EPS)
    finite_result = ((1.0 / (alpha - 1.0)) * jnp.log2(tr_inner_safe)).real
    return jnp.where(is_support_violation, jnp.inf, finite_result)


def _sandwiched_renyi_divergence_core(rho: jnp.ndarray, sigma: jnp.ndarray, alpha: float = 1.5) -> jnp.ndarray:
    # alpha must be a jnp value, not a raw Python float, before it reaches
    # _case_general: jax.lax.cond traces BOTH branches regardless of which
    # one runs, and `1.0 / (alpha - 1.0)` at alpha=1.0 raises a Python
    # ZeroDivisionError if alpha is still a plain float at that point (the
    # case_one branch's own guard never gets a chance to skip it) -- under
    # jax.jit this conversion happens automatically (all arguments get
    # traced), which is why the original un-split Discovery version never
    # hit this; the eager (non-jit) core path here needs it explicitly.
    alpha = jnp.asarray(alpha, dtype=jnp.float64)
    is_half = jnp.isclose(alpha, 0.5)
    is_one = jnp.isclose(alpha, 1.0)
    return jax.lax.cond(
        is_half, lambda: _case_half(rho, sigma),
        lambda: jax.lax.cond(is_one, lambda: _case_one(rho, sigma), lambda: _case_general(rho, sigma, alpha)),
    )


def sandwiched_renyi_divergence(rho: jnp.ndarray, sigma: jnp.ndarray, alpha: float = 1.5) -> float:
    """Sandwiched quantum Renyi divergence D_alpha(rho||sigma), in bits

    (log2). `rho`, `sigma` are density matrices of the same dimension;

    `alpha` selects the order (0.5 -> fidelity-based, 1.0 -> standard

    relative entropy, both handled as exact closed-form limits rather

    than through the general formula's own alpha->0.5/1 numerical

    instability).



    Zero when rho == sigma at every alpha (verified). For alpha > 1,

    returns `+inf` when supp(rho) is not contained in supp(sigma) --

    e.g. two different pure (rank-1) states -- rather than a finite

    number; see the module docstring for why this is the mathematically

    correct behavior, not an edge-case failure.



    KNOWN LIMITATION at exactly alpha=1.0: the same support-violation

    check is NOT applied to the alpha=1 (relative-entropy) branch, which

    instead clips log(0)-type contributions to 0 rather than diverging --

    e.g. D_1(rho||sigma) for two different pure states returns 0.0, not

    +inf, even though the true relative entropy diverges there too.

    Experiment 29 validated the alpha=1 branch only against full-rank

    (depolarized) inputs specifically to sidestep this exactly-singular

    case (`scipy.linalg.logm` itself raises `LogmExactlySingularWarning`

    on singular inputs -- an inherent ill-conditioning of relative

    entropy near degenerate support, not unique to this implementation).

    Do not rely on alpha=1 to correctly flag a support mismatch; use

    alpha slightly above 1 (e.g. 1.001) if that matters for your use case.



    Validation-only, like `uhlmann_fidelity`: meant to grade a correction

    against a known reference state, not to feed into one (see

    `uhlmann_fidelity`'s docstring for the full "ideal state as oracle"

    argument, which applies here identically).

    """
    rho = jnp.asarray(rho, dtype=jnp.complex128)
    sigma = jnp.asarray(sigma, dtype=jnp.complex128)
    return float(_sandwiched_renyi_divergence_core(rho, sigma, alpha))


sandwiched_renyi_divergence_jit = jax.jit(_sandwiched_renyi_divergence_core)
"""`jax.jit`-compiled entry point for `sandwiched_renyi_divergence`. `rho`/

`sigma` must already be `complex128`. Returns a jnp scalar, not a Python

`float`."""