feat: add numeric and symbolic scripts

scripts/dps_hierarchy/dps_hierarchy.py

@@ -0,0 +1,110 @@
+"""
+dps_hierarchy.py
+
+General DPS (Doherty-Parrilo-Spedalieri) level-k symmetric-extension SDP,
+for a bipartite qudit state rho_AB with local dimension d, extending party
+B to k Bose-symmetric copies.
+
+Key design point (discussed at length in the chat this was extracted
+from): the SDP *variable* sigma is parametrized directly on
+A x Sym^k(C^d), dimension d * C(d+k-1,k) -- POLYNOMIAL in k. But the PPT
+constraint (sigma^{T_A} >= 0) must be checked on the full, unsymmetrized
+embedding A x B_1 x ... x B_k, dimension d^{k+1} -- EXPONENTIAL in k. So
+this construction saves on free parameters but NOT on the size of the
+PSD cone that actually drives SDP solve time. See the README for measured
+timings (k=2: ~27x27 cone, sub-second; k=3: ~81x81 cone, ~2-4 minutes
+with SCS in the original sandbox).
+
+Requires: numpy, cvxpy.
+"""
+import math
+from itertools import permutations
+import numpy as np
+import cvxpy as cp
+
+
+def sym_isometry(d, k):
+    """Isometry W, shape (d**k, dim Sym^k(C^d)), spanning the totally
+    symmetric subspace of (C^d)^{tensor k}. Built by brute-force averaging
+    over all k! permutations of the k tensor factors -- fine for k up to
+    ~6-7; for larger k this construction itself becomes the bottleneck,
+    independently of the SDP."""
+    n = d ** k
+    P = np.zeros((n, n))
+    for perm in permutations(range(k)):
+        M = np.zeros((n, n))
+        for idx in np.ndindex(*([d] * k)):
+            new_idx = tuple(idx[perm[i]] for i in range(k))
+            row = 0
+            col = 0
+            for i in range(k):
+                row = row * d + new_idx[i]
+                col = col * d + idx[i]
+            M[row, col] = 1
+        P += M
+    P /= math.factorial(k)
+    eigvals, eigvecs = np.linalg.eigh(P)
+    cols = [eigvecs[:, i] for i in range(n) if abs(eigvals[i] - 1) < 1e-9]
+    return np.column_stack(cols)
+
+
+def partial_trace_keep_first_copy(full_expr, d, k):
+    """full_expr indexed by (a, b_1, ..., b_k) with combined index
+    a*d**k + b_1*d**(k-1) + ... + b_k. Traces out b_2..b_k, keeping (a,b_1)
+    -- i.e. returns the marginal on A x (first copy of B)."""
+    rest_dim = d ** (k - 1)
+    rows = []
+    for a in range(d):
+        for b1 in range(d):
+            row = []
+            for ap in range(d):
+                for b1p in range(d):
+                    terms = [full_expr[(a * d + b1) * rest_dim + r,
+                                        (ap * d + b1p) * rest_dim + r]
+                             for r in range(rest_dim)]
+                    row.append(sum(terms))
+            rows.append(row)
+    return cp.bmat(rows)
+
+
+def partial_transpose_first_system(full_expr, d1, d2):
+    """Partial transpose on the first (d1-dim) system of a
+    (d1*d2) x (d1*d2) matrix. Has the same eigenvalues as transposing the
+    second system instead (standard fact: M^{T_A} and M^{T_B} always share
+    a spectrum, since M^{T_B} = (M^{T_A})^T)."""
+    rows = []
+    for i in range(d1):
+        for kk in range(d2):
+            row = []
+            for j in range(d1):
+                for l in range(d2):
+                    row.append(full_expr[j * d2 + kk, i * d2 + l])
+            rows.append(row)
+    return cp.bmat(rows)
+
+
+def build_dps_problem(d, k):
+    """Returns (prob, rho_param, sigma) for the level-k DPS feasibility
+    SDP. Set rho_param.value = <(d*d)x(d*d) target state>, then call
+    dps_feasible(...) or prob.solve(...) directly."""
+    W = sym_isometry(d, k)
+    dim_sym = W.shape[1]
+    Iso = np.kron(np.eye(d), W)  # d**(k+1)  x  (d * dim_sym)
+    sigma = cp.Variable((d * dim_sym, d * dim_sym), hermitian=True)
+    full = Iso @ sigma @ Iso.conj().T
+    ptrace = partial_trace_keep_first_copy(full, d, k)
+    pt = partial_transpose_first_system(full, d1=d, d2=d ** k)
+    rho_param = cp.Parameter((d * d, d * d), hermitian=True)
+    constraints = [sigma >> 0, cp.trace(sigma) == 1,
+                   ptrace == rho_param, pt >> 0]
+    prob = cp.Problem(cp.Minimize(0), constraints)
+    return prob, rho_param, sigma
+
+
+def dps_feasible(prob, rho_param, rho_target, solver=cp.SCS, **solve_kwargs):
+    """Solve the (already-built) DPS problem for a given target state and
+    return True iff a valid extension was found (i.e. rho_target is NOT
+    certified entangled at this level)."""
+    rho_param.value = rho_target
+    prob.solve(solver=solver, **solve_kwargs)
+    return prob.status in ("optimal", "optimal_inaccurate")