MPSSimulator - Matrix Product State statevector simulator, JAX-backed.
Ported from the "TurboQuant TUREQ MPSSimulator v8.2 MatryoshkaFlash"
prototype (private research notebook, never published as part of the
dense-evolution package). Two real bugs were found and fixed by
independent verification against DenseSVSimulator before this module
existed in its current form:
-
The original applied Lloyd-Max quantization to the SVD singular
values on every truncation ("PolarQuantizer"). Measured a real ~0.5%
Total Variation Distance error against DenseSVSimulator on an
8-qubit entangling test circuit, with ZERO bond-dimension savings to
show for it. Dropped entirely -- this module keeps only the plain
adaptive SVD truncation (JSD-budget-driven bond dimension, the
author's own stopping criterion -- standard SVD truncation,
non-standard stopping metric).
-
get_top_k_probable_states (originally "_extract_top_k_paths") picked
a single "best" bond index via argmax at each step instead of
correctly summing over the bond dimension. Measured 0/8 correct
states against the exact contraction on the same test circuit,
values off by ~30x. Fixed by propagating the true partial-contraction
vector through each bond (matches exactly, to machine precision, on
every state it finds) -- but note it's a genuine greedy beam search,
not an exact top-k finder: recall of the true top states grows with
beam width k but isn't guaranteed complete for any fixed k.
Originally ported in plain numpy (matching the prototype), then
converted to jax.numpy so the core tensor contractions (einsum, SVD) run
on the same backend as the rest of dense_evolution instead of a second,
inconsistent numerics stack. Re-verified against DenseSVSimulator after
the conversion -- see test_mps.py.
Uses whatever jax_enable_x64 precision is currently active in the
process (does not toggle it itself) -- same convention as
DenseSVSimulator/Chunk, which rely on the caller (dashboard_core.py's
run_simulation) to set precision, since jax_enable_x64 is a process-wide
flag and toggling it locally would leak to unrelated code running later
in the same process.
For circuits with LOW entanglement (product states, GHZ/Bell-like chains,
shallow local circuits), the bond dimension stays small regardless of
qubit count, so this scales to hundreds of qubits where DenseSVSimulator
(or Chunk) cannot -- see get_probabilities_sampled and
get_top_k_probable_states, neither of which ever materializes a
(2**n,)-shaped array. For HIGHLY entangled circuits the bond dimension
grows and this degrades back toward the same exponential cost
DenseSVSimulator has -- it is not a universal replacement, it is
complementary.
MPSSimulator
MPSSimulator(
n_qubits: int,
max_bond: int = 64,
svd_cutoff: float = 1e-12,
jsd_budget: float = 1e-05,
)
Matrix Product State simulator with adaptive SVD-truncated bond
dimension (JSD-budget driven), no lossy post-truncation quantization.
JAX-backed core (einsum, SVD).
Parameters
n_qubits : int
max_bond : int -- hard cap on bond dimension chi
svd_cutoff : float -- singular values below this are dropped outright
jsd_budget : float -- max tolerated Jensen-Shannon distance between the
full and truncated singular-value distributions
at each cut; chi is grown by 1 until satisfied
or max_bond is hit.
Source code in dense_evolution/mps.py
| def __init__(
self,
n_qubits: int,
max_bond: int = 64,
svd_cutoff: float = 1e-12,
jsd_budget: float = 1e-5,
):
self.n = n_qubits
self.chi = max_bond
self.eps = svd_cutoff
self.jsd_budget = jsd_budget
self.gammas: List[jnp.ndarray] = []
self.lambdas: List[jnp.ndarray] = [jnp.ones(1)] * (n_qubits + 1)
self.truncation_errors: List[float] = []
self.jsd_per_bond: List[float] = []
self.entanglement_entropy = np.zeros(max(n_qubits - 1, 0))
self._bond_history: List[int] = []
# Counts truncations where max_bond was hit before jsd_budget could
# be satisfied -- the while loop below exits silently in that case,
# and avg_JSD (a mean over all steps) can look deceptively low even
# when the final contracted state is badly wrong (verified: TVD
# ~0.97 against DenseSVSimulator on an 8-qubit/15-layer entangling
# circuit with max_bond=2, while avg_JSD read 0.0534).
self.budget_violations: int = 0
# Cached compiled closure for run_circuit_jit -- built lazily on
# first use (self.n/self.chi/self.eps/self.jsd_budget are fixed
# for this instance's lifetime), never rebuilt per call. Same
# caching pattern as Chunk.__init__'s self._multi_chunk_runner.
self._mps_runner = None
for _ in range(n_qubits):
g = jnp.zeros((1, 2, 1), dtype=jnp.complex64 if not jax.config.jax_enable_x64 else jnp.complex128)
g = g.at[0, 0, 0].set(1.0)
self.gammas.append(g)
|
apply_gate_1q
apply_gate_1q(gate: ndarray, qubit: int) -> None
O(chi^2) -- updates only Gamma[qubit].
Source code in dense_evolution/mps.py
| def apply_gate_1q(self, gate: jnp.ndarray, qubit: int) -> None:
"""O(chi^2) -- updates only Gamma[qubit]."""
gate = jnp.asarray(gate)
self.gammas[qubit] = jnp.einsum("ij,ljr->lir", gate, self.gammas[qubit])
|
apply_gate_2q
apply_gate_2q(gate_2q: ndarray, q1: int, q2: int) -> None
2-qubit gate with adaptive SVD truncation. O(chi^3).
Source code in dense_evolution/mps.py
| def apply_gate_2q(self, gate_2q: jnp.ndarray, q1: int, q2: int) -> None:
"""2-qubit gate with adaptive SVD truncation. O(chi^3)."""
gate_2q = jnp.asarray(gate_2q)
if abs(q1 - q2) != 1:
self._apply_nonlocal_2q(gate_2q, q1, q2)
return
if q1 > q2:
q1, q2 = q2, q1
gate_2q = jnp.transpose(gate_2q, (1, 0, 3, 2))
g1 = self.gammas[q1]
g2 = self.gammas[q2]
lam = self.lambdas[q2]
theta = jnp.einsum("lik,k,kjr->lijr", g1, lam, g2)
chiL, d1, d2, chiR = theta.shape
theta_new = jnp.einsum("abcd,ecdf->eabf", gate_2q, theta)
theta_mat = theta_new.reshape(chiL * d1, d2 * chiR)
U_t, S_t, Vh_t, trunc_err, jsd_val = self._svd_truncate(theta_mat)
chi_new = len(S_t)
s_sq = S_t**2
p_dist = s_sq / (jnp.sum(s_sq) + 1e-20)
p_v = p_dist[p_dist > 1e-20]
ee = float(-jnp.sum(p_v * jnp.log2(p_v))) if len(p_v) > 1 else 0.0
if q1 < len(self.entanglement_entropy):
self.entanglement_entropy[q1] = ee
self.lambdas[q2] = S_t
self.gammas[q1] = U_t.reshape(chiL, d1, chi_new)
self.gammas[q2] = Vh_t.reshape(chi_new, d2, chiR)
self._bond_history.append(chi_new)
self.jsd_per_bond.append(jsd_val)
|
apply_ccx
apply_ccx(c1: int, c2: int, tgt: int) -> None
Toffoli via standard T-gate decomposition (all 1q/2q gates).
Source code in dense_evolution/mps.py
| def apply_ccx(self, c1: int, c2: int, tgt: int) -> None:
"""Toffoli via standard T-gate decomposition (all 1q/2q gates)."""
inv2 = 1.0 / np.sqrt(2.0)
h = inv2 * jnp.array([[1, 1], [1, -1]], dtype=complex)
t = jnp.array([[1, 0], [0, jnp.exp(1j * jnp.pi / 4)]], dtype=complex)
tdg = jnp.array([[1, 0], [0, jnp.exp(-1j * jnp.pi / 4)]], dtype=complex)
self.apply_gate_1q(h, tgt)
self.apply_cx(c2, tgt)
self.apply_gate_1q(tdg, tgt)
self.apply_cx(c1, tgt)
self.apply_gate_1q(t, tgt)
self.apply_cx(c2, tgt)
self.apply_gate_1q(tdg, tgt)
self.apply_cx(c1, tgt)
self.apply_gate_1q(t, c2)
self.apply_gate_1q(t, tgt)
self.apply_gate_1q(h, tgt)
self.apply_cx(c1, c2)
self.apply_gate_1q(t, c1)
self.apply_gate_1q(tdg, c2)
self.apply_cx(c1, c2)
|
get_probabilities_sampled
get_probabilities_sampled(
n_samples: int = 100000, seed: Optional[int] = None
) -> dict
Returns a {bitstring: empirical_probability} dict from n_samples
sequential draws -- the only entry point safe for n_qubits > 24.
Source code in dense_evolution/mps.py
| def get_probabilities_sampled(
self, n_samples: int = 100_000, seed: Optional[int] = None
) -> dict:
"""Returns a {bitstring: empirical_probability} dict from n_samples
sequential draws -- the only entry point safe for n_qubits > 24."""
from collections import Counter
rng = np.random.default_rng(seed)
counts: Counter = Counter()
for _ in range(n_samples):
bits = self._sample_bitstring(rng)
counts["".join(map(str, bits))] += 1
return {bitstr: c / n_samples for bitstr, c in counts.items()}
|
get_top_k_probable_states
get_top_k_probable_states(
k: int = 128,
) -> Tuple[np.ndarray, np.ndarray]
Greedy beam search (beam width k) for approximately-most-probable
basis states, without ever contracting to a full statevector.
Returns (indices, probabilities): indices are computational-basis
integers, probabilities are exact for the states found (not
approximated), sorted descending. Recall of the TRUE top states
improves with k but is not guaranteed for any fixed k -- see the
module docstring.
Source code in dense_evolution/mps.py
| def get_top_k_probable_states(self, k: int = 128) -> Tuple[np.ndarray, np.ndarray]:
"""Greedy beam search (beam width k) for approximately-most-probable
basis states, without ever contracting to a full statevector.
Returns (indices, probabilities): indices are computational-basis
integers, probabilities are exact for the states found (not
approximated), sorted descending. Recall of the TRUE top states
improves with k but is not guaranteed for any fixed k -- see the
module docstring."""
paths: List[Tuple[int, jnp.ndarray]] = [(0, jnp.array([1.0 + 0.0j]))]
for i in range(self.n):
candidates = []
gamma = self.gammas[i]
lam = self.lambdas[i + 1] if (i + 1) < len(self.lambdas) else jnp.ones(gamma.shape[2])
for idx_p, vec_p in paths:
for bit in (0, 1):
new_vec = jnp.einsum("l,lr->r", vec_p, gamma[:, bit, :]) * lam
weight = float(jnp.sum(jnp.abs(new_vec) ** 2))
candidates.append(((idx_p << 1) | bit, new_vec, weight))
candidates.sort(key=lambda c: c[2], reverse=True)
paths = [(idx, vec) for idx, vec, _ in candidates[:k]]
indices = np.array([p[0] for p in paths])
amplitudes = np.array([
complex(vec[0]) if len(vec) == 1 else complex(jnp.sum(vec))
for _, vec in paths
])
probabilities = np.abs(amplitudes) ** 2
order = np.argsort(-probabilities)
return indices[order], probabilities[order]
|
run_circuit_jit
run_circuit_jit(ops: List) -> None
Runs an entire circuit through a single jax.lax.scan-fused,
@jax.jit-compiled kernel instead of one eager Python call per gate
-- the eager path (apply_gate_1q/apply_gate_2q/_apply_nonlocal_2q,
all still available and unchanged) has zero @jax.jit anywhere and
pays a host-device sync on every 2-qubit gate's bond-dimension
search; measured 88.9s vs Qiskit Aer's 0.64s on a 60-qubit stress
circuit -- see README changelog for the real before/after number
this method produces on that same circuit.
Trade-off, explicit and intentional (not hidden): every gamma/
lambda is kept at a fixed max_bond-padded size for the rest of
this instance's lifetime after this call. Structurally correct
either way (zero-padding is mathematically transparent to every
other method here -- contract_to_statevector, get_top_k_probable_
states, etc. all still work correctly on the padded arrays,
verified), just not memory-minimal for genuinely low-entanglement
circuits, which is this module's whole point for very large qubit
counts. Use the eager methods directly instead when memory, not
speed, is the priority -- this is an addition, not a replacement.
ops: same convention as DenseSVSimulator.run_circuit_jit_beast_mode
-- list of (name, *args) tuples/lists. Unlike that method, SWAP is
never decomposed into 3xCX (kept as one real gate, see
_compile_mps_ops's docstring for why that matters here).
Source code in dense_evolution/mps.py
| def run_circuit_jit(self, ops: List) -> None:
"""Runs an entire circuit through a single jax.lax.scan-fused,
@jax.jit-compiled kernel instead of one eager Python call per gate
-- the eager path (apply_gate_1q/apply_gate_2q/_apply_nonlocal_2q,
all still available and unchanged) has zero @jax.jit anywhere and
pays a host-device sync on every 2-qubit gate's bond-dimension
search; measured 88.9s vs Qiskit Aer's 0.64s on a 60-qubit stress
circuit -- see README changelog for the real before/after number
this method produces on that same circuit.
Trade-off, explicit and intentional (not hidden): every gamma/
lambda is kept at a fixed max_bond-padded size for the rest of
this instance's lifetime after this call. Structurally correct
either way (zero-padding is mathematically transparent to every
other method here -- contract_to_statevector, get_top_k_probable_
states, etc. all still work correctly on the padded arrays,
verified), just not memory-minimal for genuinely low-entanglement
circuits, which is this module's whole point for very large qubit
counts. Use the eager methods directly instead when memory, not
speed, is the priority -- this is an addition, not a replacement.
ops: same convention as DenseSVSimulator.run_circuit_jit_beast_mode
-- list of (name, *args) tuples/lists. Unlike that method, SWAP is
never decomposed into 3xCX (kept as one real gate, see
_compile_mps_ops's docstring for why that matters here).
"""
compiled_rows = _compile_mps_ops(ops, self.n)
dtype = self.gammas[0].dtype
lambda_dtype = self.lambdas[0].dtype
if compiled_rows:
ops_array = jnp.array(compiled_rows, dtype=jnp.float64)
else:
ops_array = jnp.zeros((0, 5), dtype=jnp.float64)
if self._mps_runner is None:
self._mps_runner = _build_mps_runner(self.n, self.chi, self.eps, self.jsd_budget)
gammas_padded = jnp.stack([_pad_gamma(g, self.chi).astype(dtype) for g in self.gammas])
lambdas_padded = jnp.stack([_pad_lambda(l, self.chi).astype(lambda_dtype) for l in self.lambdas])
final_gammas, final_lambdas, diag = self._mps_runner(gammas_padded, lambdas_padded, ops_array)
self.gammas = [final_gammas[i] for i in range(self.n)]
self.lambdas = [final_lambdas[i] for i in range(self.n + 1)]
# Bookkeeping parity: jax.lax.scan's stacked per-step diagnostics
# (diag) replace the eager path's Python list.append()s inside the
# loop -- same final content, populated differently. Only 2-qubit
# steps (g_id >= 20, includes SWAP -- the eager _apply_nonlocal_2q
# path routes its SWAPs through apply_gate_2q too, so its history
# lists grow on those as well, not just the "real" gate) count.
if compiled_rows:
chi_history, jsd_history, trunc_err_history, entropy_history = (
np.asarray(diag[0]), np.asarray(diag[1]), np.asarray(diag[2]), np.asarray(diag[3]))
g_ids = np.asarray([row[0] for row in compiled_rows])
q1_ids = np.asarray([int(row[1]) for row in compiled_rows])
is_2q_mask = g_ids >= 20
for i in np.nonzero(is_2q_mask)[0]:
chi_new = int(chi_history[i])
jsd_val = float(jsd_history[i])
self._bond_history.append(chi_new)
self.jsd_per_bond.append(jsd_val)
self.truncation_errors.append(float(trunc_err_history[i]))
q1 = q1_ids[i]
if q1 < len(self.entanglement_entropy):
self.entanglement_entropy[q1] = float(entropy_history[i])
if jsd_val > self.jsd_budget:
if self.budget_violations == 0:
warnings.warn(
f"MPSSimulator: bond dimension capped at max_bond={self.chi}, "
f"jsd_budget={self.jsd_budget:.1e} not honored "
f"(jsd={jsd_val:.2e}) -- results may be unreliable, "
f"consider raising max_bond.",
UserWarning,
stacklevel=2,
)
self.budget_violations += 1
|