quantum-shadow-maps_v2/scripts/mixture_search.py

"""mixture_search.py -- searching (and partly proving) the biseparable supremum of
min(||M_AB||_*, ||M_AC||_*, ||M_AD||_*) via convex mixtures of pure product states
across different bipartitions.

min() of convex functions is NOT itself convex, so (unlike Phi_sym or a single
||M_S||_*) the supremum over biseparable states can genuinely lie ABOVE what any single
pure product state achieves, and can only be found by explicitly searching mixtures.

Contains:
  - exact closed-form derivation/verification for the 2-component Bell-pair mixture
    family rho(p) = p*(Bell_AB x Bell_CD) + (1-p)*(Bell_AC x Bell_BD):
        M_AB(p) = 5 - 4p,  M_AC(p) = 1 + 4p,  M_AD(p) = 3 + 2|2p-1|
    so min(...)(p) is maximized EXACTLY at p=1/2, value = 3 (proven by hand from the
    2x2-block eigenvalue structure of the mixed correlation tensor).
  - general K-component mixture optimizations (free internal state parameters + free
    softmax weights) that repeatedly rediscover this same value 3.0 as the best found,
    across 2-, 3-, 4- and 6-component mixture families.
"""
import numpy as np
from scipy.optimize import minimize
from core import state_2_2, state_1_3
from cluster import cluster_map, nuc

I2 = np.eye(2, dtype=complex)
X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.array([[1, 0], [0, -1]], dtype=complex)
_PAULI = [I2, X, Y, Z]


def _kron4(a, b, c, d):
    return np.kron(np.kron(a, b), np.kron(c, d))


_OPS = np.zeros((4, 4, 4, 4, 16, 16), dtype=complex)
for _i0 in range(4):
    for _i1 in range(4):
        for _i2 in range(4):
            for _i3 in range(4):
                _OPS[_i0, _i1, _i2, _i3] = _kron4(_PAULI[_i0], _PAULI[_i1], _PAULI[_i2], _PAULI[_i3])
_OPS_FLAT = _OPS.reshape(256, 16, 16)


def full_tensor_mixed_fast(rho):
    """Fast Pauli-tensor extraction for a general (possibly mixed) 16x16 density
    matrix, via a single vectorized einsum over all 256 precomputed Pauli operators."""
    vals = np.einsum('kij,ji->k', _OPS_FLAT, rho)
    return vals.real.reshape(4, 4, 4, 4)


def triple_mixed(rho):
    C = full_tensor_mixed_fast(rho)
    return nuc(cluster_map(C, 0, 1)), nuc(cluster_map(C, 0, 2)), nuc(cluster_map(C, 0, 3))


def bellpair_state(pairing):
    bell = np.array([1, 0, 0, 1]) / np.sqrt(2)
    (p1a, p1b), (p2a, p2b) = pairing
    psi = np.zeros(16, dtype=complex)
    for x in range(2):
        for y in range(2):
            for u in range(2):
                for v in range(2):
                    idx = [0, 0, 0, 0]
                    idx[p1a] = x
                    idx[p1b] = y
                    idx[p2a] = u
                    idx[p2b] = v
                    lin = idx[0] * 8 + idx[1] * 4 + idx[2] * 2 + idx[3]
                    psi[lin] = bell[x * 2 + y] * bell[u * 2 + v]
    return psi


# ---------------------------------------------------------------------
# Exact closed form for the 2-component Bell-pair mixture family
# ---------------------------------------------------------------------
def bell_mixture_exact_formula(p):
    M_AB = 5 - 4 * p
    M_AC = 1 + 4 * p
    M_AD = 3 + 2 * abs(2 * p - 1)
    return M_AB, M_AC, M_AD


def bell_mixture_numeric(p):
    rho1 = np.outer(bellpair_state(((0, 1), (2, 3))), np.conj(bellpair_state(((0, 1), (2, 3)))))
    rho2 = np.outer(bellpair_state(((0, 2), (1, 3))), np.conj(bellpair_state(((0, 2), (1, 3)))))
    rho = p * rho1 + (1 - p) * rho2
    return triple_mixed(rho)


# ---------------------------------------------------------------------
# General K-component mixture optimizations
# ---------------------------------------------------------------------
def neg_min_mixture_2comp(params):
    """2 components: pure states biseparable across AB|CD and AC|BD, full internal
    freedom, weight via sigmoid."""
    x1, x2 = params[0:16], params[16:32]
    w = 1 / (1 + np.exp(-params[32]))
    rho1 = np.outer(state_2_2(x1, (0, 1), (2, 3)), np.conj(state_2_2(x1, (0, 1), (2, 3))))
    rho2 = np.outer(state_2_2(x2, (0, 2), (1, 3)), np.conj(state_2_2(x2, (0, 2), (1, 3))))
    rho = w * rho1 + (1 - w) * rho2
    return -min(triple_mixed(rho))


def neg_min_mixture_3comp(params):
    """3 components: pure states biseparable across each of the three 2|2 cuts, full
    internal freedom, softmax weights."""
    x1, x2, x3 = params[0:16], params[16:32], params[32:48]
    w = np.exp(params[48:51] - np.max(params[48:51]))
    w = w / np.sum(w)
    psis = [state_2_2(x1, (0, 1), (2, 3)), state_2_2(x2, (0, 2), (1, 3)), state_2_2(x3, (0, 3), (1, 2))]
    rho = sum(wi * np.outer(p, np.conj(p)) for wi, p in zip(w, psis))
    return -min(triple_mixed(rho))


def neg_min_mixture_4slot(params):
    """4 slots: all three 2|2 cut types plus one 1|3 cut type, full internal freedom,
    softmax weights (optimizer is free to zero out unused slots)."""
    x0, x1, x2, x3 = params[0:16], params[16:32], params[32:48], params[48:66]
    w = np.exp(params[66:70] - np.max(params[66:70]))
    w = w / np.sum(w)
    psis = [state_2_2(x0, (0, 1), (2, 3)), state_2_2(x1, (0, 2), (1, 3)),
            state_2_2(x2, (0, 3), (1, 2)), state_1_3(x3, 0)]
    rho = sum(wi * np.outer(p, np.conj(p)) for wi, p in zip(w, psis))
    return -min(triple_mixed(rho))


def neg_min_mixture_6slot(params):
    """6 slots: AB|CD, AC|BD, AD|BC, AB|CD (2nd copy), AC|BD (2nd copy), A|BCD -- allows
    two independently-parametrized states of the SAME cut type to mix together too."""
    xs = [params[16 * k:16 * (k + 1)] for k in range(5)]
    x5 = params[80:98]
    w = np.exp(params[98:104] - np.max(params[98:104]))
    w = w / np.sum(w)
    psis = [state_2_2(xs[0], (0, 1), (2, 3)), state_2_2(xs[1], (0, 2), (1, 3)),
            state_2_2(xs[2], (0, 3), (1, 2)), state_2_2(xs[3], (0, 1), (2, 3)),
            state_2_2(xs[4], (0, 2), (1, 3)), state_1_3(x5, 0)]
    rho = sum(wi * np.outer(p, np.conj(p)) for wi, p in zip(w, psis))
    return -min(triple_mixed(rho))


def multistart(objective, nparams, n_restarts, seed0, label, maxiter=2000):
    best = -np.inf
    bx = None
    for i in range(n_restarts):
        rng = np.random.default_rng(seed0 + i)
        x0 = rng.normal(size=nparams) * 0.7
        res = minimize(objective, x0, method='Powell',
                        options={'maxiter': maxiter, 'xtol': 1e-8, 'ftol': 1e-10})
        v = -res.fun
        if v > best:
            best = v
            bx = res.x
    print(f'{label}: best min(M_AB,M_AC,M_AD) = {best:.6f}  ({n_restarts} restarts)')
    return best, bx


if __name__ == "__main__":
    print("=== Exact closed form vs numeric verification, Bell-pair mixture family ===")
    for p in [0.0, 0.25, 0.5, 0.75, 1.0]:
        formula = bell_mixture_exact_formula(p)
        numeric = bell_mixture_numeric(p)
        print(f"  p={p:.2f}  formula={tuple(round(x, 4) for x in formula)}  "
              f"numeric={tuple(round(x, 4) for x in numeric)}")

    print()
    print("=== General mixture optimizations (stress-testing the p=1/2 optimum, 3.0) ===")
    multistart(neg_min_mixture_2comp, 33, 6, 3000, "2-component (AB|CD + AC|BD, free params)")
    multistart(neg_min_mixture_3comp, 51, 3, 5000, "3-component (all three 2|2 cuts, free weights)")
    multistart(neg_min_mixture_4slot, 70, 2, 7000, "4-slot (three 2|2 + one 1|3, free weights)")
    multistart(neg_min_mixture_6slot, 104, 1, 8000, "6-slot (duplicated cut types)")