diff --git a/benchmarks/.gitignore b/benchmarks/.gitignore new file mode 100644 index 00000000..fbca2253 --- /dev/null +++ b/benchmarks/.gitignore @@ -0,0 +1 @@ +results/ diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 00000000..b72896e7 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,53 @@ +# Benchmarks + +`bench_suite.py` is the profiling and correctness harness behind the +optimisation work. Every workload is deterministic (fixed seeds) and returns a +*fingerprint*; `baseline.json` holds the fingerprints of the unoptimised code, +and a change is only accepted if the fingerprints still match. + +| Workload | What it exercises | Contract | +| --- | --- | --- | +| `surface_bitflip` | 5x5 surface code, bit-flip noise, 6 shots, χ = 64, natural qubit order | verdicts exact, posterior entries within 1e-2 | +| `css_optimised` | the same code under `qubit_order_strategy="Optimised"` (reverse Cuthill-McKee) | as above | +| `shor_depolarising` | Shor code, depolarising noise, 40 shots, χ = 128: the dense-readout path | as above | +| `classical_ldpc` | random (3,4) LDPC code, XOR constraints + dephasing-DMRG readout | overlaps exact to 1e-10 | +| `dmrg_ground_state` | DMRG on a 24-site transverse-field Ising chain | energy exact to 1e-10 | +| `dem_d3` | circuit-level detector error model, d=3 r=3 p=0.8%, eight busiest syndromes, χ = 32 | verdicts exact, class masses within 1e-2 | +| `dem_d5` | the same at d=5 r=5 p=0.5%, busiest syndrome, χ = 32 | as above | + +"Exact" means the value must agree to 1e-10; χ-truncated posterior entries get +1e-2 because a different but equally valid SVD gauge in a near-degenerate +spectrum changes which directions the truncation keeps (verdicts and +converged results are unaffected). A NaN never matches. + +## Running + +```bash +# time every workload and compare against the committed baseline +python benchmarks/bench_suite.py --check + +# one workload, with a cProfile dump and a text top-30 in benchmarks/results/ +python benchmarks/bench_suite.py --workload dem_d5 --profile + +# record a new baseline for a workload (only after the change is validated +# some other way: exact enumeration, agreement at converged chi) +python benchmarks/bench_suite.py --workload NAME --write-baseline +``` + +`--check` and `--write-baseline` are mutually exclusive, and a targeted +`--write-baseline` touches only the workloads it ran. `results/` is +gitignored; `summary.json` there keeps the last wall time per workload. + +## Writing a baseline from the reference code + +The baseline must come from code *without* the optimisations under test. +With a clean checkout of `main` next to this one: + +```bash +PYTHONPATH=/path/to/mdopt-main python benchmarks/bench_suite.py --workload NAME --write-baseline +PYTHONPATH=/path/to/this-checkout python benchmarks/bench_suite.py --workload NAME --check +``` + +`python -c` puts the current directory first on `sys.path`, so verify which +package a run imports (`mdopt.__file__`) from a neutral directory before +trusting a measurement. diff --git a/benchmarks/baseline.json b/benchmarks/baseline.json new file mode 100644 index 00000000..58d35064 --- /dev/null +++ b/benchmarks/baseline.json @@ -0,0 +1,416 @@ +{ + "classical_ldpc": [ + [ + 1.0, + 0.0, + 1.0, + 1.0 + ], + [ + 1.0, + 0.0, + 1.0, + 1.0 + ], + [ + 1.0, + 0.0, + 1.0, + 1.0 + ] + ], + "dmrg_ground_state": [ + -30.1997123268 + ], + "shor_depolarising": [ + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.6592642527, + 0.2556768373, + 0.6592642527, + 0.2556768373 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 0.0, + 0.2556768373, + 0.6592642527, + 0.2556768373, + 0.6592642527 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.5620575695, + 0.4290586074, + 0.5620575695, + 0.4290586074 + ], + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 0.0, + 0.46361496, + 0.0568776701, + 0.87983707, + 0.0878227179 + ], + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.87983707, + 0.0878227179, + 0.46361496, + 0.0568776701 + ], + [ + 0.0, + 0.2556768373, + 0.6592642527, + 0.2556768373, + 0.6592642527 + ], + [ + 1.0, + 0.9025425599, + 0.3321136927, + 0.2538275719, + 0.103387555 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + [ + 1.0, + 0.7551868377, + 0.4540849479, + 0.4025415357, + 0.2479112992 + ], + [ + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ] + ], + "surface_bitflip": [ + [ + 1.0, + 0.9999884886, + 0.0030496957, + 0.0037043265, + 0.0 + ], + [ + 1.0, + 0.9997212094, + 0.023448308, + 0.0027713539, + 0.0 + ], + [ + 1.0, + 0.998177969, + 0.0603082581, + 0.0019119129, + 2.86641e-05 + ], + [ + 1.0, + 0.9999879784, + 0.0030496899, + 0.0038395868, + 0.0 + ], + [ + 1.0, + 0.9999119531, + 0.0128654032, + 0.0032507742, + 0.0 + ], + [ + 1.0, + 0.9935503223, + 0.1133808331, + 0.0015948984, + 0.0 + ] + ], + "css_optimised": [ + [ + 1.0, + 0.993071558, + 0.1173987948, + 0.0051380829, + 6.20336e-05 + ], + [ + 1.0, + 0.9997980704, + 0.0188844941, + 0.0068698093, + 0.0 + ] + ], + "dem_d3": [ + [ + 1.0, + 0.3828748949, + 0.6171251051 + ], + [ + 0.0, + 0.9344395896, + 0.0655604104 + ], + [ + 0.0, + 0.5329658082, + 0.4670341918 + ], + [ + 0.0, + 0.6376500198, + 0.3623499802 + ], + [ + 0.0, + 0.9925403359, + 0.0074596641 + ], + [ + 1.0, + 0.1061289231, + 0.8938710769 + ], + [ + 1.0, + 0.20108188, + 0.79891812 + ], + [ + 0.0, + 0.9979570562, + 0.0020429438 + ] + ], + "dem_d5": [ + [ + 1.0, + 0.0131212224, + 0.9868787776 + ] + ] +} \ No newline at end of file diff --git a/benchmarks/bench_suite.py b/benchmarks/bench_suite.py new file mode 100644 index 00000000..279a483b --- /dev/null +++ b/benchmarks/bench_suite.py @@ -0,0 +1,360 @@ +"""Profiling benchmark suite for mdopt's hot paths. + +Each workload is deterministic (fixed seeds), sized to run in tens of seconds, +and returns a correctness fingerprint. The fingerprints are the contract for +the optimisation work: ``--check`` compares them against the committed +``benchmarks/baseline.json`` -- exact values (energies, verdicts, overlaps) +must match to 1e-10, and chi-truncated posterior entries must stay within +1e-2: a different but equally valid SVD gauge in a near-degenerate spectrum +changes which directions chi_max keeps, and that moves small class masses at +this level while leaving verdicts and converged results untouched (old and +new decoders agree to 6e-14 at chi=1e5). A real behaviour change moves +verdicts or exact values. ``--write-baseline`` +records a new baseline after a change that is validated some other way +(exact-enumeration tests, agreement at converged chi). + +Run: python benchmarks/bench_suite.py [--profile] [--check] [--workload NAME] +Profiles land in benchmarks/results/.pstats plus a text top-30. +""" + +import argparse +import cProfile +import io +import json +import math +import pstats +import time +from pathlib import Path + +import numpy as np +import qecstruct as qec +import stim + +# Imported once here, not inside the workloads: function-local imports made +# the first-run wall time and profile depend on invocation order (a workload +# run alone paid cold-import cost that a full sorted suite had already paid). +from mdopt.contractor.contractor import mps_mpo_contract +from mdopt.decoding.dem import decode_dem, dem_to_problem +from mdopt.examples.ising.ising import IsingMPO +from mdopt.decoding.decoding import ( + apply_bitflip_bias, + apply_constraints, + decode_css, + decode_message, + generate_pauli_error_string, + linear_code_constraint_sites, + linear_code_prepare_message, +) +from mdopt.mps.utils import ( + create_custom_product_state, + create_simple_product_state, + inner_product, +) +from mdopt.optimiser.dmrg import DMRG +from mdopt.optimiser.utils import SWAP, XOR_BULK, XOR_LEFT, XOR_RIGHT + +HERE = Path(__file__).parent +RESULTS = HERE / "results" + + +def wl_surface_bitflip(): + """Code-capacity surface-code decode: the quantum_surface workload.""" + code = qec.hypergraph_product(qec.repetition_code(5), qec.repetition_code(5)) + rng = np.random.default_rng(51) + outputs = [] + for _ in range(6): + error = generate_pauli_error_string( + len(code), 0.05, rng=rng, error_model="Bitflip" + ) + dense, success = decode_css( + code, + error, + chi_max=64, + bias_type="Bitflip", + bias_prob=0.05, + renormalise=True, + silent=True, + contraction_strategy="Optimised", + ) + # The full posterior, not just the verdict: a wrong posterior + # with an unmoved argmax must still move the fingerprint. + outputs.append([float(success)] + [round(float(x), 10) for x in dense]) + return outputs + + +def wl_css_optimised(): + """Surface-code decode under the RCM qubit ordering (optimise_qubit_order).""" + code = qec.hypergraph_product(qec.repetition_code(5), qec.repetition_code(5)) + rng = np.random.default_rng(52) + outputs = [] + for _ in range(2): + error = generate_pauli_error_string( + len(code), 0.05, rng=rng, error_model="Bitflip" + ) + dense, success = decode_css( + code, + error, + chi_max=64, + bias_type="Bitflip", + bias_prob=0.05, + renormalise=True, + silent=True, + contraction_strategy="Optimised", + qubit_order_strategy="Optimised", + ) + outputs.append([float(success)] + [round(float(x), 10) for x in dense]) + return outputs + + +def _dem_case(task, distance, rounds, p, seed, num_sampled, num_keep): + """A circuit-level DEM plus its busiest sampled syndromes. + + The busiest syndromes (most detection events) carry nontrivial posteriors + and drive the slowest contractions, so they are the ones worth timing and + fingerprinting; a quiet syndrome decodes to (1, 1e-10) and would not move. + """ + circuit = stim.Circuit.generated( + f"surface_code:rotated_memory_{task}", + distance=distance, + rounds=rounds, + after_clifford_depolarization=p, + before_measure_flip_probability=p, + after_reset_flip_probability=p, + ) + problem = dem_to_problem( + circuit.detector_error_model(decompose_errors=False, flatten_loops=True) + ) + sampler = circuit.compile_detector_sampler(seed=seed) + detections, _ = sampler.sample(num_sampled, separate_observables=True) + order = np.argsort(-detections.sum(axis=1), kind="stable")[:num_keep] + return problem, detections[order].astype(int) + + +def _dem_rows(problem, syndromes, chi_max): + outputs = [] + for syndrome in syndromes: + masses, flips = decode_dem(problem, syndrome, chi_max=chi_max) + posterior = masses / masses.sum() + # [verdict, *normalised class masses]: the verdict is exact, the + # masses are chi-truncated posterior entries. + outputs.append([float(flips[0])] + [round(float(x), 10) for x in posterior]) + return outputs + + +def wl_dem_d3(): + """Circuit-level DEM decode, d=3 r=3 p=0.8% memory-X (the Fig. 1d cell).""" + problem, syndromes = _dem_case("x", 3, 3, 0.008, seed=3, num_sampled=64, num_keep=8) + return _dem_rows(problem, syndromes, chi_max=32) + + +def wl_dem_d5(): + """Circuit-level DEM decode, d=5 r=5 p=0.5% memory-Z (the campaign cell).""" + problem, syndromes = _dem_case("z", 5, 5, 0.005, seed=5, num_sampled=32, num_keep=1) + return _dem_rows(problem, syndromes, chi_max=32) + + +def wl_shor_depolarising(): + """Small-code depolarising decode: dense readout path end to end.""" + code = qec.shor_code() + rng = np.random.default_rng(7) + outputs = [] + for _ in range(40): + error = generate_pauli_error_string(len(code), 0.1, rng=rng) + dense, success = decode_css( + code, + error, + chi_max=128, + bias_type="Depolarising", + bias_prob=0.1, + renormalise=True, + silent=True, + ) + # The full posterior, not just the verdict: a wrong posterior + # with an unmoved argmax must still move the fingerprint. + outputs.append([float(success)] + [round(float(x), 10) for x in dense]) + return outputs + + +def wl_classical_ldpc(): + """Classical LDPC pipeline: constraints + Dephasing DMRG readout.""" + outputs = [] + for seed in (11, 12, 13): + code = qec.random_regular_code(48, 36, 3, 4, qec.Rng(seed)) + first, second = linear_code_prepare_message( + code, 0.1, error_model=qec.BinarySymmetricChannel, seed=seed + ) + sites = linear_code_constraint_sites(code) + start = create_custom_product_state(first, form="Right-canonical") + state = create_custom_product_state(second, form="Right-canonical") + received = state.copy() + state = apply_bitflip_bias(mps=state, sites_to_bias="All", prob_bias_list=0.1) + state = apply_constraints( + state, + sites, + [XOR_LEFT, XOR_BULK, SWAP, XOR_RIGHT], + chi_max=64, + renormalise=True, + strategy="Optimised", + silent=True, + ) + _, overlap = decode_message( + message=state, + codeword=start, + num_runs=1, + chi_max_dmrg=64, + silent=True, + ) + # The DMRG verdict alone is blind to the constrained state: it stays + # 1.0 whenever the MAP codeword is unchanged. Observables of the state + # the hot path actually produces come first: its overlaps with the + # transmitted codeword and with the received (biased) message, and + # the Schmidt spectrum at the middle bond (a corrupted contraction + # leaves weight outside the codeword, which shows up here). + middle = state.num_sites // 2 + _, spectra = state.copy().move_orth_centre( + middle, return_singular_values=True, renormalise=True + ) + schmidt = sorted( + (float(v) for v in np.asarray(spectra[-1]).ravel()), reverse=True + ) + outputs.append( + [ + round(float(abs(inner_product(start, state))), 10), + round(float(abs(inner_product(received, state))), 10), + float(overlap), + *[round(v, 10) for v in schmidt[:4]], + ] + ) + return outputs + + +def wl_dmrg_ground_state(): + """Plain DMRG on a transverse-field Ising chain (optimiser hot path).""" + num_sites = 24 + mpo = IsingMPO(num_sites=num_sites, h_magnetic=1.0).hamiltonian_mpo() + mps = create_simple_product_state(num_sites, which="+") + engine = DMRG(mps, mpo, chi_max=48, cut=1e-12, mode="SA", silent=True) + engine.run(2) + # The energy depends on the optimised state everywhere the norm does not: + # renormalised bond updates make norm() ~ 1.0 for any state, correct or + # not, so it cannot serve as the correctness fingerprint. + ground = engine.mps + h_ground = mps_mpo_contract(ground, mpo, chi_max=int(1e4), renormalise=False) + energy = float(np.real(inner_product(ground, h_ground))) + return [round(energy, 10)] + + +BASELINE = HERE / "baseline.json" +# Workloads whose fingerprint rows are [verdict, *posterior entries at chi_max]. +POSTERIOR_WORKLOADS = { + "surface_bitflip", + "shor_depolarising", + "css_optimised", + "dem_d3", + "dem_d5", +} + +WORKLOADS = { + "surface_bitflip": wl_surface_bitflip, + "css_optimised": wl_css_optimised, + "shor_depolarising": wl_shor_depolarising, + "classical_ldpc": wl_classical_ldpc, + "dmrg_ground_state": wl_dmrg_ground_state, + "dem_d3": wl_dem_d3, + "dem_d5": wl_dem_d5, +} + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--profile", action="store_true") + parser.add_argument("--workload", choices=sorted(WORKLOADS), default=None) + # Mutually exclusive: writing the baseline first and then checking against + # it would compare every fingerprint with itself and always pass. + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--check", action="store_true", help="compare against baseline.json" + ) + mode.add_argument( + "--write-baseline", + action="store_true", + help="record the fingerprints of the workloads run by this invocation", + ) + args = parser.parse_args() + RESULTS.mkdir(exist_ok=True) + + names = [args.workload] if args.workload else sorted(WORKLOADS) + # Merge into the previous summary so a --workload run does not discard the + # other entries; profiled runs are marked, since cProfile inflates wall time. + summary = {} + if (RESULTS / "summary.json").exists(): + summary = json.loads((RESULTS / "summary.json").read_text()) + for name in names: + func = WORKLOADS[name] + started = time.perf_counter() + if args.profile: + profiler = cProfile.Profile() + fingerprint = profiler.runcall(func) + wall = time.perf_counter() - started + profiler.dump_stats(RESULTS / f"{name}.pstats") + stream = io.StringIO() + stats = pstats.Stats(profiler, stream=stream) + stats.sort_stats("cumulative").print_stats(30) + (RESULTS / f"{name}.top30.txt").write_text(stream.getvalue()) + else: + fingerprint = func() + wall = time.perf_counter() - started + summary[name] = { + "wall_s": round(wall, 3), + "fingerprint": fingerprint, + "profiled": bool(args.profile), + } + print(f"{name:>20}: {wall:7.2f} s fingerprint={fingerprint}", flush=True) + (RESULTS / "summary.json").write_text(json.dumps(summary, indent=2)) + if args.write_baseline: + # Update only the workloads this invocation ran: a targeted run must + # neither shrink the committed baseline to the selected workload nor + # promote stale cached fingerprints of the others. + baseline = json.loads(BASELINE.read_text()) if BASELINE.exists() else {} + for name in names: + baseline[name] = summary[name]["fingerprint"] + BASELINE.write_text(json.dumps(baseline, indent=2)) + print(f"baseline written: {BASELINE}") + if args.check: + baseline = json.loads(BASELINE.read_text()) + failures = [] + # Only the workloads this invocation ran; summary also carries cached + # entries from earlier runs, which a targeted --check must not judge. + for name in names: + failures += _compare(name, summary[name]["fingerprint"], baseline[name]) + if failures: + print("FINGERPRINT MISMATCH:\n " + "\n ".join(failures)) + raise SystemExit(1) + print("fingerprints match baseline") + + +def _compare(name, got, want, path=""): + """Exact for scalars/verdicts; 1e-2 for chi-truncated posterior entries.""" + tolerance = 1e-2 if name in POSTERIOR_WORKLOADS else 1e-10 + if isinstance(want, list): + if not isinstance(got, list) or len(got) != len(want): + return [f"{name}{path}: shape changed"] + return [ + f + for i, (g, w) in enumerate(zip(got, want)) + for f in _compare(name, g, w, f"{path}[{i}]") + ] + # the leading verdict of a posterior row is exact; the entries are not + if name in POSTERIOR_WORKLOADS and path.endswith("[0]") and path.count("[") == 2: + tolerance = 1e-10 + got_value, want_value = float(got), float(want) + # A NaN would otherwise pass: abs(nan) > tolerance is False. + if not (math.isfinite(got_value) and math.isfinite(want_value)): + return [f"{name}{path}: non-finite value {got} (baseline {want})"] + if abs(got_value - want_value) > tolerance: + return [f"{name}{path}: {got} vs baseline {want}"] + return [] + + +if __name__ == "__main__": + main() diff --git a/examples/decoding/dem_campaign/nishimori.py b/examples/decoding/dem_campaign/nishimori.py index 643b021c..34860a23 100644 --- a/examples/decoding/dem_campaign/nishimori.py +++ b/examples/decoding/dem_campaign/nishimori.py @@ -65,14 +65,19 @@ def run(distance, p, shots, chi=128, seed=0): syndrome = h_z @ mech % 2 truth = int(obs_vec @ mech % 2) t0 = time.perf_counter() - _, flips = decode_dem(problem, syndrome, chi_max=chi) - rec = { - "i": i, - "truth": truth, - "map": int(flips[0]), - "mwpm": int(matcher.decode(syndrome)[0]) % 2, - "t": round(time.perf_counter() - t0, 4), - } + rec = {"i": i, "truth": truth} + try: + _, flips = decode_dem(problem, syndrome, chi_max=chi) + rec["map"] = int(flips[0]) + except ArithmeticError as exc: + # A collapsed or negative class-mass vector is a truncation + # artefact of this chi: record it and score the shot as a + # failure rather than abort the cell (and every cell after + # it, since the resume would replay the same shot). + rec["map"] = 1 - truth + rec["artefact"] = str(exc)[:80] + rec["mwpm"] = int(matcher.decode(syndrome)[0]) % 2 + rec["t"] = round(time.perf_counter() - t0, 4) sink.write(json.dumps(rec) + "\n") if (i + 1) % 500 == 0: sink.flush() diff --git a/mdopt/contractor/contractor.py b/mdopt/contractor/contractor.py index 030d7439..25f297a9 100644 --- a/mdopt/contractor/contractor.py +++ b/mdopt/contractor/contractor.py @@ -5,13 +5,89 @@ from typing import Union, List, Tuple, cast import numpy as np -from opt_einsum import contract +from opt_einsum import contract_expression from mdopt.backend import array as A from mdopt.mps.canonical import CanonicalMPS from mdopt.mps.explicit import ExplicitMPS from mdopt.utils.utils import split_two_site_tensor +_EXPRESSIONS: dict = {} + + +def _contract_cached(subscripts, path, backend, *tensors): + """Evaluate a fixed einsum through a reusable opt_einsum expression. + + The sweep in :func:`mps_mpo_contract` evaluates the same few einsums + thousands of times per decode; ``contract`` re-parses the subscripts and + rebuilds path metadata on every call even when ``optimize`` is explicit. + An expression with an explicit path is shape-independent, so it is + cached per (subscripts, path) only -- keying on operand shapes made + truncation's data-dependent bond dimensions miss 7-14% of calls on + large codes -- and built from whatever shapes the first call carries. + """ + key = (subscripts, path) + expression = _EXPRESSIONS.get(key) + if expression is None: + expression = contract_expression( + subscripts, *(tensor.shape for tensor in tensors), optimize=list(path) + ) + _EXPRESSIONS[key] = expression + return expression(*tensors, backend=backend) + + +def _zip_first(left, right, mpo_left, mpo_right, backend): + """The zip-up's opening two-site tensor, ``ijk, klm, nojp, oqlr -> iprqm``. + + On the NumPy backend the three pairwise contractions of the cached + opt_einsum path are issued directly: the expression machinery added + roughly a third of the contraction's own cost per call on the small + tensors of a decode, and the operands are exactly the tensordots + opt_einsum itself would issue, so the result is the same to rounding. + """ + if backend != "numpy": + return _contract_cached( + "ijk, klm, nojp, oqlr -> iprqm", + ((0, 1), (1, 2), (0, 1)), + backend, + left, + right, + mpo_left, + mpo_right, + ) + pair = np.tensordot(left, right, axes=(2, 0)) # i j l m + ops = np.tensordot(mpo_left, mpo_right, axes=(1, 0)) # n j p q l r + out = np.tensordot(pair, ops, axes=([1, 2], [1, 4])) # i m n p q r + # n is the MPO's open left virtual leg (dimension 1): summed, as in the + # einsum where it is absent from the output. + return out.sum(axis=2).transpose(0, 2, 4, 3, 1) # i p r q m + + +def _zip_step(centre, right, mpo_tensor, backend): + """One zip-up sweep step, ``ijkl, lmn, komp -> ijpon`` (see _zip_first). + + On the NumPy backend the MPO tensor is contracted with the right MPS + tensor first, arranged so that the final tensordot's output already has + the ``i j p o n`` layout: the caller's reshape to a matrix is then a view + instead of a copy of the two-site tensor, and only the small + ``mpo x right`` intermediate is transposed. Same operands, same sums; + the result agrees with the einsum to rounding. + """ + if backend != "numpy": + return _contract_cached( + "ijkl, lmn, komp -> ijpon", + ((0, 1), (0, 1)), + backend, + centre, + right, + mpo_tensor, + ) + # k o m p -> k p o m, so that the free legs come out as p, o. + ops = np.tensordot( + mpo_tensor.transpose(0, 3, 1, 2), right, axes=(3, 1) + ) # k p o l n + return np.tensordot(centre, ops, axes=([2, 3], [0, 3])) # i j p o n + def apply_one_site_operator(tensor: np.ndarray, operator: np.ndarray) -> np.ndarray: """ @@ -57,8 +133,12 @@ def apply_one_site_operator(tensor: np.ndarray, operator: np.ndarray) -> np.ndar f"while the one given has {operator.ndim}." ) - tensor_updated = contract( - "ijk, jl -> ilk", tensor, operator, optimize=[(0, 1)], backend=backend + tensor_updated = _contract_cached( + "ijk, jl -> ilk", + ((0, 1),), + backend, + tensor, + operator, ) return A.to_device(np.asarray(tensor_updated)) @@ -139,17 +219,35 @@ def apply_two_site_unitary( b1_scaled = b_1 * (lam[:, None, None]) # with lambda_0 - t_with = contract( - "ijk, klm -> ijlm", b1_scaled, b_2, optimize=[(0, 1)], backend=backend + t_with = _contract_cached( + "ijk, klm -> ijlm", + ((0, 1),), + backend, + b1_scaled, + b_2, ) - t_with = contract( - "ijkl, jkmn -> imnl", t_with, unitary, optimize=[(0, 1)], backend=backend + t_with = _contract_cached( + "ijkl, jkmn -> imnl", + ((0, 1),), + backend, + t_with, + unitary, ) # without lambda_0 (for back-substitution) - t_wo = contract("ijk, klm -> ijlm", b_1, b_2, optimize=[(0, 1)], backend=backend) - t_wo = contract( - "ijkl, jkmn -> imnl", t_wo, unitary, optimize=[(0, 1)], backend=backend + t_wo = _contract_cached( + "ijk, klm -> ijlm", + ((0, 1),), + backend, + b_1, + b_2, + ) + t_wo = _contract_cached( + "ijkl, jkmn -> imnl", + ((0, 1),), + backend, + t_wo, + unitary, ) # split and back-substitute @@ -160,12 +258,12 @@ def apply_two_site_unitary( renormalise=False, return_truncation_error=True, ) - b_1_updated = contract( + b_1_updated = _contract_cached( "ijkl, mkl -> ijm", + ((0, 1),), + backend, t_wo, np.conjugate(b_2_updated), - optimize=[(0, 1)], - backend=backend, ) if A.GPU: @@ -251,7 +349,12 @@ def mps_mpo_contract( mps = mps.mixed_canonical(start_site) assert isinstance(mps, CanonicalMPS) if mps.orth_centre != start_site: - mps = cast(CanonicalMPS, mps.move_orth_centre(start_site, renormalise=False)) + # inplace: this function owns `mps` by now (copied above unless the + # caller asked for inplace, in which case it handed ownership over). + mps = cast( + CanonicalMPS, + mps.move_orth_centre(start_site, renormalise=False, inplace=True), + ) for i, tensor in enumerate(mpo): if tensor.ndim != 4: @@ -273,14 +376,12 @@ def mps_mpo_contract( orth_centre_index = start_site - two_site_mps_mpo_tensor = contract( - "ijk, klm, nojp, oqlr -> iprqm", + two_site_mps_mpo_tensor = _zip_first( mps.tensors[start_site], mps.tensors[start_site + 1], mpo[0], mpo[1], - optimize=[(0, 1), (1, 2), (0, 1)], - backend=backend, + backend, ).reshape( ( mps.tensors[start_site].shape[0], @@ -292,17 +393,16 @@ def mps_mpo_contract( # Sweep across the MPO for i in range(len(mpo) - 2): - mps.tensors[orth_centre_index], singular_values, b_r, _ = ( + mps.tensors[orth_centre_index], singular_values, b_r = ( split_two_site_tensor( two_site_mps_mpo_tensor, chi_max=chi_max, cut=cut, renormalise=renormalise, - return_truncation_error=True, ) ) - with A.stream(): - if A.GPU: + if A.GPU: + with A.stream(): mps.tensors[orth_centre_index] = A.to_device( mps.tensors[orth_centre_index] ) @@ -310,8 +410,7 @@ def mps_mpo_contract( singular_values = A.to_device(np.asarray(singular_values)) orth_centre_index += 1 - if isinstance(mps, CanonicalMPS): - mps.orth_centre = orth_centre_index + mps.orth_centre = orth_centre_index # Replace diag(s) @ b_r with broadcast multiply (no diag allocation) mps.tensors[orth_centre_index] = ( @@ -325,13 +424,11 @@ def mps_mpo_contract( ) ) - two_site_mps_mpo_tensor = contract( - "ijkl, lmn, komp -> ijpon", + two_site_mps_mpo_tensor = _zip_step( mps.tensors[orth_centre_index], mps.tensors[orth_centre_index + 1], mpo[i + 2], - optimize=[(0, 1), (0, 1)], - backend=backend, + backend, ).reshape( ( len(singular_values), @@ -342,15 +439,14 @@ def mps_mpo_contract( ) # Final split and update last tensor - mps.tensors[orth_centre_index], singular_values, b_r, _ = split_two_site_tensor( + mps.tensors[orth_centre_index], singular_values, b_r = split_two_site_tensor( two_site_mps_mpo_tensor, chi_max=chi_max, cut=cut, renormalise=renormalise, - return_truncation_error=True, ) - with A.stream(): - if A.GPU: + if A.GPU: + with A.stream(): mps.tensors[orth_centre_index] = A.to_device( mps.tensors[orth_centre_index] ) diff --git a/mdopt/examples/decoding/quantum_bivariate_bicycle.py b/mdopt/examples/decoding/quantum_bivariate_bicycle.py index 22113215..a9b091d3 100644 --- a/mdopt/examples/decoding/quantum_bivariate_bicycle.py +++ b/mdopt/examples/decoding/quantum_bivariate_bicycle.py @@ -5,6 +5,7 @@ import pickle import logging import argparse +from functools import partial from multiprocessing import Pool import numpy as np @@ -132,6 +133,19 @@ def parse_arguments(): required=True, help="Singular values smaller than that will be discarded in the SVD.", ) + parser.add_argument( + "--qubit_order_strategy", + type=str, + default="Natural", + choices=["Natural", "Optimised"], + help=( + "Qubit order along the MPS chain: the code's natural order, or the " + "reverse Cuthill-McKee order, which lowers the bond dimension the " + "contraction needs (measured: 5x faster on the surface code at " + "chi=64, and convergence at chi=64 instead of 128 on the [[72,12,6]] " + "bivariate-bicycle code)." + ), + ) return parser.parse_args() @@ -176,6 +190,7 @@ def run_single_experiment( silent, tolerance, cut, + qubit_order_strategy="Natural", ): """Run a single experiment.""" bb_code = create_bb_code(order_x, order_y, poly_a, poly_b) @@ -193,6 +208,7 @@ def run_single_experiment( contraction_strategy="Optimised", tolerance=tolerance, cut=cut, + qubit_order_strategy=qubit_order_strategy, ) except Exception as e: logging.error(f"Error during decoding: {e}", exc_info=True) @@ -210,6 +226,7 @@ def run_single_experiment( contraction_strategy="Optimised", tolerance=tolerance, cut=cut, + qubit_order_strategy=qubit_order_strategy, ) logging.info("Decoding finished with multiply_by_stabiliser=True.") except Exception as ex: @@ -247,6 +264,7 @@ def run_experiment( num_processes=1, tolerance=1e-8, cut=1e-8, + qubit_order_strategy="Natural", ): """Run the experiment consisting of multiple single experiments in parallel.""" logging.info( @@ -255,6 +273,8 @@ def run_experiment( f" TOLERANCE={tolerance}, CUT={cut}, ERROR_MODEL={error_model}, SEED={seed}" ) + worker = partial(run_single_experiment, qubit_order_strategy=qubit_order_strategy) + args = [ ( order_x, @@ -273,7 +293,7 @@ def run_experiment( ] with Pool(num_processes) as pool: - results = pool.starmap(run_single_experiment, args) + results = pool.starmap(worker, args) logging.info( f"Finished {num_experiments} experiments for ORDERS={order_x, order_y}," @@ -296,6 +316,7 @@ def run_experiment( "tolerance": tolerance, "cut": cut, "polynomials": [poly_a, poly_b], + "qubit_order_strategy": qubit_order_strategy, } @@ -311,10 +332,17 @@ def save_experiment_data( seed, tolerance, cut, + qubit_order_strategy="Natural", ): """Save the experiment data.""" error_model = error_model.replace(" ", "") - file_key = f"latticesize{order_x*order_y}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_seed{seed}.pkl" + # The ordering changes the decoder's results, so it is recorded in the data + # itself; the filename tag is derived from the stored value. + qubit_order_strategy = data.get("qubit_order_strategy", qubit_order_strategy) + order_tag = ( + "" if qubit_order_strategy == "Natural" else f"_order{qubit_order_strategy}" + ) + file_key = f"latticesize{order_x*order_y}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_seed{seed}{order_tag}.pkl" with open(file_key, "wb") as pickle_file: pickle.dump(data, pickle_file) logging.info( @@ -352,6 +380,7 @@ def main(): args.num_processes, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) save_experiment_data( experiment_data, @@ -365,6 +394,7 @@ def main(): args.seed, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) diff --git a/mdopt/examples/decoding/quantum_csp.py b/mdopt/examples/decoding/quantum_csp.py index 6035c154..5a5a66e8 100644 --- a/mdopt/examples/decoding/quantum_csp.py +++ b/mdopt/examples/decoding/quantum_csp.py @@ -8,6 +8,7 @@ import pickle import logging import argparse +from functools import partial from multiprocessing import Pool import numpy as np @@ -170,6 +171,19 @@ def parse_arguments(): required=True, help="Singular values smaller than that will be discarded in the SVD.", ) + parser.add_argument( + "--qubit_order_strategy", + type=str, + default="Natural", + choices=["Natural", "Optimised"], + help=( + "Qubit order along the MPS chain: the code's natural order, or the " + "reverse Cuthill-McKee order, which lowers the bond dimension the " + "contraction needs (measured: 5x faster on the surface code at " + "chi=64, and convergence at chi=64 instead of 128 on the [[72,12,6]] " + "bivariate-bicycle code)." + ), + ) return parser.parse_args() @@ -238,6 +252,7 @@ def run_single_experiment( tolerance, cut, seed=None, + qubit_order_strategy="Natural", ): """Run a single experiment (with a few random stabiliser gauges).""" csp_code = get_csp_code(num_qubits, batch, code_id) @@ -260,6 +275,7 @@ def _decode(multiply_by_stabiliser: bool): contraction_strategy="Optimised", tolerance=tolerance, cut=cut, + qubit_order_strategy=qubit_order_strategy, rng=rng, ) @@ -331,6 +347,7 @@ def run_experiment( num_processes=1, tolerance=0, cut=0, + qubit_order_strategy="Natural", ): """Run the experiment consisting of multiple single experiments in parallel.""" logging.info( @@ -358,6 +375,8 @@ def run_experiment( # as they were. shot_seeds = gauge_seed_sequences(seed, num_experiments) + worker = partial(run_single_experiment, qubit_order_strategy=qubit_order_strategy) + args = [ ( num_qubits, @@ -376,7 +395,7 @@ def run_experiment( ] with Pool(num_processes) as pool: - results = pool.starmap(run_single_experiment, args) + results = pool.starmap(worker, args) logging.info( f"Finished {num_experiments} experiments for NUM_QUBITS={num_qubits}," @@ -399,6 +418,7 @@ def run_experiment( "cut": cut, "batch": batch, "code_id": code_id, + "qubit_order_strategy": qubit_order_strategy, } @@ -415,10 +435,17 @@ def save_experiment_data( seed, tolerance, cut, + qubit_order_strategy="Natural", ): """Save the experiment data.""" error_model = error_model.replace(" ", "") - file_key = f"latticesize{num_qubits}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_batch{batch}_codeid{code_id}_seed{seed}.pkl" + # The ordering changes the decoder's results, so it is recorded in the data + # itself; the filename tag is derived from the stored value. + qubit_order_strategy = data.get("qubit_order_strategy", qubit_order_strategy) + order_tag = ( + "" if qubit_order_strategy == "Natural" else f"_order{qubit_order_strategy}" + ) + file_key = f"latticesize{num_qubits}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_batch{batch}_codeid{code_id}_seed{seed}{order_tag}.pkl" with open(file_key, "wb") as pickle_file: pickle.dump(data, pickle_file) logging.info( @@ -454,6 +481,7 @@ def main(): args.num_processes, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) save_experiment_data( experiment_data, @@ -468,6 +496,7 @@ def main(): args.seed, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) diff --git a/mdopt/examples/decoding/quantum_hypergraph_product.py b/mdopt/examples/decoding/quantum_hypergraph_product.py index e92eaacc..062088b6 100644 --- a/mdopt/examples/decoding/quantum_hypergraph_product.py +++ b/mdopt/examples/decoding/quantum_hypergraph_product.py @@ -5,6 +5,7 @@ import pickle import logging import argparse +from functools import partial from multiprocessing import Pool import numpy as np @@ -124,6 +125,19 @@ def parse_arguments(): required=True, help="Singular values smaller than that will be discarded in the SVD.", ) + parser.add_argument( + "--qubit_order_strategy", + type=str, + default="Natural", + choices=["Natural", "Optimised"], + help=( + "Qubit order along the MPS chain: the code's natural order, or the " + "reverse Cuthill-McKee order, which lowers the bond dimension the " + "contraction needs (measured: 5x faster on the surface code at " + "chi=64, and convergence at chi=64 instead of 128 on the [[72,12,6]] " + "bivariate-bicycle code)." + ), + ) return parser.parse_args() @@ -156,7 +170,16 @@ def generate_errors(system_size, error_rate, num_experiments, error_model, seed) def run_single_experiment( - system_size, chi_max, error, bias_prob, error_model, silent, tolerance, cut, seed + system_size, + chi_max, + error, + bias_prob, + error_model, + silent, + tolerance, + cut, + seed, + qubit_order_strategy="Natural", ): """Run a single experiment.""" check_degree, bit_degree = 4, 3 @@ -181,6 +204,7 @@ def run_single_experiment( contraction_strategy="Optimised", tolerance=tolerance, cut=cut, + qubit_order_strategy=qubit_order_strategy, ) except Exception as e: logging.error(f"Error during decoding: {e}", exc_info=True) @@ -198,6 +222,7 @@ def run_single_experiment( contraction_strategy="Optimised", tolerance=tolerance, cut=cut, + qubit_order_strategy=qubit_order_strategy, ) except Exception as ex: logging.error( @@ -231,6 +256,7 @@ def run_experiment( num_processes=1, tolerance=1e-8, cut=1e-8, + qubit_order_strategy="Natural", ): """Run the experiment consisting of multiple single experiments in parallel.""" logging.info( @@ -252,6 +278,8 @@ def run_experiment( len(qhgp_code) - qhgp_code.num_x_stabs() - qhgp_code.num_z_stabs(), ) + worker = partial(run_single_experiment, qubit_order_strategy=qubit_order_strategy) + args = [ ( system_size, @@ -268,7 +296,7 @@ def run_experiment( ] with Pool(num_processes) as pool: - results = pool.starmap(run_single_experiment, args) + results = pool.starmap(worker, args) logging.info( f"Starting {num_experiments} experiments for SYSTEM_SIZE={system_size}," @@ -292,6 +320,7 @@ def run_experiment( "tolerance": tolerance, "cut": cut, "code_parameters": code_parameters, + "qubit_order_strategy": qubit_order_strategy, } @@ -306,10 +335,17 @@ def save_experiment_data( seed, tolerance, cut, + qubit_order_strategy="Natural", ): """Save the experiment data.""" error_model = error_model.replace(" ", "") - file_key = f"latticesize{system_size}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_seed{seed}.pkl" + # The ordering changes the decoder's results, so it is recorded in the data + # itself; the filename tag is derived from the stored value. + qubit_order_strategy = data.get("qubit_order_strategy", qubit_order_strategy) + order_tag = ( + "" if qubit_order_strategy == "Natural" else f"_order{qubit_order_strategy}" + ) + file_key = f"latticesize{system_size}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_seed{seed}{order_tag}.pkl" with open(file_key, "wb") as pickle_file: pickle.dump(data, pickle_file) logging.info( @@ -341,6 +377,7 @@ def main(): args.num_processes, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) save_experiment_data( experiment_data, @@ -353,6 +390,7 @@ def main(): args.seed, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) diff --git a/mdopt/examples/decoding/quantum_surface.py b/mdopt/examples/decoding/quantum_surface.py index aa26aa87..dbf98d7f 100644 --- a/mdopt/examples/decoding/quantum_surface.py +++ b/mdopt/examples/decoding/quantum_surface.py @@ -5,6 +5,7 @@ import pickle import logging import argparse +from functools import partial from multiprocessing import Pool import numpy as np @@ -123,6 +124,19 @@ def parse_arguments(): required=True, help="Singular values smaller than that will be discarded in the SVD.", ) + parser.add_argument( + "--qubit_order_strategy", + type=str, + default="Natural", + choices=["Natural", "Optimised"], + help=( + "Qubit order along the MPS chain: the code's natural order, or the " + "reverse Cuthill-McKee order, which lowers the bond dimension the " + "contraction needs (measured: 5x faster on the surface code at " + "chi=64, and convergence at chi=64 instead of 128 on the [[72,12,6]] " + "bivariate-bicycle code)." + ), + ) return parser.parse_args() @@ -233,6 +247,7 @@ def run_experiment( num_processes=1, tolerance=1e-8, cut=1e-8, + qubit_order_strategy="Natural", ): """Run the experiment consisting of multiple single experiments in parallel.""" logging.info( @@ -241,6 +256,8 @@ def run_experiment( f" TOLERANCE={tolerance}, CUT={cut}, ERROR_MODEL={error_model}, SEED={seed}" ) + worker = partial(run_single_experiment, qubit_order_strategy=qubit_order_strategy) + args = [ ( lattice_size, @@ -256,10 +273,10 @@ def run_experiment( ] if num_processes == 1: - results = [run_single_experiment(*arg) for arg in args] + results = [worker(*arg) for arg in args] else: with Pool(num_processes) as pool: - results = pool.starmap(run_single_experiment, args) + results = pool.starmap(worker, args) logging.info( f"Finished {num_experiments} experiments for LATTICE_SIZE={lattice_size}," @@ -281,6 +298,7 @@ def run_experiment( "seed": seed, "tolerance": tolerance, "cut": cut, + "qubit_order_strategy": qubit_order_strategy, } @@ -295,10 +313,17 @@ def save_experiment_data( seed, tolerance, cut, + qubit_order_strategy="Natural", ): """Save the experiment data.""" error_model = error_model.replace(" ", "") - file_key = f"latticesize{lattice_size}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_seed{seed}.pkl" + # The ordering changes the decoder's results, so it is recorded in the data + # itself; the filename tag is derived from the stored value. + qubit_order_strategy = data.get("qubit_order_strategy", qubit_order_strategy) + order_tag = ( + "" if qubit_order_strategy == "Natural" else f"_order{qubit_order_strategy}" + ) + file_key = f"latticesize{lattice_size}_bonddim{chi_max}_errorrate{error_rate}_errormodel{error_model}_bias_prob{bias_prob}_numexperiments{num_experiments}_tolerance{tolerance}_cut{cut}_seed{seed}{order_tag}.pkl" with open(file_key, "wb") as pickle_file: pickle.dump(data, pickle_file) logging.info( @@ -330,6 +355,7 @@ def main(): args.num_processes, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) save_experiment_data( experiment_data, @@ -342,6 +368,7 @@ def main(): args.seed, args.tolerance, args.cut, + qubit_order_strategy=args.qubit_order_strategy, ) diff --git a/mdopt/mps/canonical.py b/mdopt/mps/canonical.py index 4a0f30ad..c2d5a6d0 100644 --- a/mdopt/mps/canonical.py +++ b/mdopt/mps/canonical.py @@ -152,7 +152,8 @@ def reverse(self) -> "CanonicalMPS": Returns a reversed version of the current MPS. """ reversed_tensors = [np.transpose(tensor) for tensor in reversed(self.tensors)] - if self.orth_centre: + # `is not None`, not truthiness: a centre at site 0 is a centre. + if self.orth_centre is not None: reversed_orth_centre = (self.num_sites - 1) - self.orth_centre return CanonicalMPS( reversed_tensors, reversed_orth_centre, self.tolerance, self.chi_max @@ -347,7 +348,8 @@ def move_orth_centre( final_pos: int, return_singular_values: bool = False, renormalise: bool = True, - ) -> Union["CanonicalMPS", Tuple["CanonicalMPS", List[list]]]: + inplace: bool = False, + ) -> Union["CanonicalMPS", Tuple["CanonicalMPS", List[np.ndarray]]]: """ Moves the orthogonality centre from its current position to ``final_pos``. @@ -363,6 +365,11 @@ def move_orth_centre( Whether to return the singular values obtained at each involved bond. renormalise : bool Whether to renormalise singular values during each SVD. + inplace : bool + Whether a rightward move may overwrite this instance's tensors + instead of deep-copying them first (a leftward move builds a new + object either way and leaves this instance unmoved). The returned + object is the one to use in both cases. Raises ------ @@ -377,7 +384,7 @@ def move_orth_centre( f"from 0 to {self.num_sites-1}, given {final_pos}." ) - singular_values = [] + singular_values: List[np.ndarray] = [] if self.orth_centre is None: self.orth_centre = self.check_orth_centre() # type: ignore @@ -386,7 +393,7 @@ def move_orth_centre( if self.orth_centre < final_pos: begin, final = self.orth_centre, final_pos - mps = self.copy() + mps = self if inplace else self.copy() elif self.orth_centre > final_pos: mps = self.reverse() begin = cast(int, mps.orth_centre) @@ -394,7 +401,62 @@ def move_orth_centre( else: return self + # A move factors the centre alone, (chi_l*d, chi_r), instead of the + # two-site tensor whenever the right neighbour is an isometry: the + # centre's singular values then ARE the two-site tensor's (theta = + # C B with B B^dag = 1 gives theta theta^dag = C C^dag), so the + # renormalised spectrum, the cut and the chi_max count all come out + # identical, and the move is the two-site SVD's answer at a fraction + # of the cost. That covers the DMRG sweeps' one-step moves and the + # explicit-form conversion, which ask for renormalised spectra. + for i in range(begin, final): + centre = mps.tensors[i] + chi_l, phys, chi_r = centre.shape + # A collapsed bond (dimension 0) has nothing to factor; the SVD + # branch carries that degenerate shape through unchanged. + # chi_r <= chi_max: with an isometric neighbour the revealed rank + # is at most chi_r, so this integer check rules out any move that + # would have to truncate before the factorisation is attempted. + if 0 < chi_r <= self.chi_max and chi_l * phys > 0: + # The single-site spectrum is the bond's Schmidt spectrum + # only if the right neighbour is an isometry. The bias + # appliers break that on every site they touch, so it is + # checked (a chi^2 * d * chi Gram product, cheaper than the + # SVD it enables) rather than assumed: with an isometric + # neighbour this move is exactly the two-site SVD's answer + # (same cut, same chi_max); otherwise the two-site SVD below + # runs, exactly as before. After one traversal the chain is + # canonical and every later move takes the fast path. + neighbour = mps.tensors[i + 1] + flat = neighbour.reshape(neighbour.shape[0], -1) + # An inexact dtype: integer or boolean tensors are valid input + # (the SVD promotes them), and the in-place subtraction below + # must not fail on them. No copy for float or complex. + gram_dtype = np.result_type(flat.dtype, 1.0) + flat_inexact = np.asarray(flat, dtype=gram_dtype) + gram = flat_inexact @ flat_inexact.conj().T + # ||G - I||_F <= 1e-12, the criterion find_orth_centre uses. + # The Frobenius norm bounds the spectral deviation of the + # neighbour from an isometry; a componentwise maximum does not + # (for a 400-dimensional Gram matrix it can be 400x smaller). + gram[np.diag_indices_from(gram)] -= 1.0 + if np.linalg.norm(gram) <= 1e-12: + u_l, s_bond, v_h, _ = svd( + centre.reshape(chi_l * phys, chi_r), + chi_max=self.chi_max, + renormalise=renormalise, + ) + singular_values.append(s_bond) + keep = len(s_bond) + mps.tensors[i] = u_l.reshape(chi_l, phys, keep) + # diag(s) @ (v_h . B): same idiom as the SVD branch below. + mps.tensors[i + 1] = ( + np.tensordot(v_h, neighbour, axes=(1, 0)) + * s_bond[:, None, None] + ) + mps.orth_centre = i + 1 + continue two_site_tensor = mps.two_site_tensor_next(i) u_l, singular_values_bond, v_r, _ = split_two_site_tensor( two_site_tensor, diff --git a/mdopt/mps/utils.py b/mdopt/mps/utils.py index e9555e86..216fb884 100644 --- a/mdopt/mps/utils.py +++ b/mdopt/mps/utils.py @@ -252,8 +252,8 @@ def inner_product( def _renormalise_after_truncation( - singular_values: List[float], truncation_error: Optional[float] -) -> List[float]: + singular_values: np.ndarray, truncation_error: Optional[float] +) -> np.ndarray: """Renormalises a Schmidt spectrum after ``chi_max`` truncation. :class:`ExplicitMPS` requires a unit-norm spectrum at every bond. Truncation @@ -263,11 +263,11 @@ def _renormalise_after_truncation( the constraint. """ if not truncation_error: - return singular_values + return np.asarray(singular_values, dtype=float) norm = float(np.linalg.norm(np.asarray(singular_values, dtype=float))) if norm == 0: - return singular_values - return list(np.asarray(singular_values, dtype=float) / norm) + return np.asarray(singular_values, dtype=float) + return np.asarray(singular_values, dtype=float) / norm def mps_from_dense( @@ -333,7 +333,7 @@ def mps_from_dense( ) tensors: list[np.ndarray] = [] - singular_values: list[list] = [] + singular_values: list = [] state_vector = state_vector.reshape((-1, phys_dim)) diff --git a/mdopt/optimiser/utils.py b/mdopt/optimiser/utils.py index 25c817f8..9b30f944 100644 --- a/mdopt/optimiser/utils.py +++ b/mdopt/optimiser/utils.py @@ -213,6 +213,23 @@ def mpo(self) -> List[np.ndarray]: return mpo +def _product_state_orth_centre(mps: CanonicalMPS, tolerance: float = 1e-12) -> int: + """The orthogonality centre :func:`find_orth_centre` assigns a product state. + + For a chain of bond dimension 1 every tensor is (1, d, 1); it is a left + and a right isometry exactly when its vector has unit norm. The scan then + reports the non-isometric sites as centres and the convention keeps the + first of them, or site 0 when all sites are isometric; a biased chain + whose first sites are still normalised (the decoders' logical prefix) + therefore gets its first biased site, not site 0. + """ + for site, tensor in enumerate(mps.tensors): + gram = float(np.vdot(tensor, tensor).real) + if not np.isclose(gram, 1.0, atol=tolerance, rtol=0.0): + return site + return 0 + + def apply_constraints( mps: CanonicalMPS, strings: List[List[List[int]]], @@ -282,6 +299,12 @@ def apply_constraints( if dense: mps_dense = mps.dense(flatten=True) + # One private copy up front, then every zip-up and move works in place: + # the contractor used to deep-copy the whole chain once per constraint, + # which on a 1700-site DEM chain was a measurable share of the decode. + if strings and not dense: + mps = mps.copy() + for string in tqdm(strings, disable=silent): string = ConstraintString(logical_tensors, string) mpo = string.mpo() @@ -297,6 +320,13 @@ def apply_constraints( continue # Ensure orthogonality centre is set and moved once per string + if mps.orth_centre is None and all(d == 1 for d in mps.bond_dimensions): + # A product state (every bond of dimension 1) needs no isometry + # scan: a (1, d, 1) tensor is an isometry exactly when its + # vector has unit norm, so the scan's answer -- the first + # non-isometric site, or site 0 when every site is isometric -- + # comes from the site norms alone. Same tolerance as the scan. + mps.orth_centre = _product_state_orth_centre(mps) if mps.orth_centre is None: orth_centres, flags_left, flags_right = find_orth_centre( mps, return_orth_flags=True @@ -324,7 +354,10 @@ def apply_constraints( mps.orth_centre = orth_centres[0] mps = mps.move_orth_centre( - final_pos=start_site, renormalise=False, return_singular_values=False + final_pos=start_site, + renormalise=False, + return_singular_values=False, + inplace=True, ) # type: ignore # Contract MPO string into the MPS (uses contractor that preserves dtype & avoids diag()) @@ -335,7 +368,7 @@ def apply_constraints( chi_max=chi_max, cut=cut, renormalise=False, - inplace=False, + inplace=True, ) if renormalise: diff --git a/mdopt/utils/utils.py b/mdopt/utils/utils.py index d4872942..667aaa90 100644 --- a/mdopt/utils/utils.py +++ b/mdopt/utils/utils.py @@ -6,24 +6,26 @@ import scipy from opt_einsum import contract -# --- Backend shim: prefer your GPU/array backend if available, else NumPy --- -try: - # expected to export a NumPy-like API (e.g., NumPy or CuPy) - from mdopt.backend import array as xp # type: ignore -except (ImportError, ModuleNotFoundError): - import numpy as xp # type: ignore +# The backend module exposes a NumPy-like API (NumPy, or CuPy when +# MDOPT_BACKEND=cupy and a CUDA device is usable) plus host/device transfer. +from mdopt.backend import array as xp def _to_numpy(a): - """Convert backend arrays (e.g., CuPy) to NumPy without copying if possible.""" - try: - import cupy as cp # type: ignore + """Bring a backend array to the host as a NumPy array (no copy on NumPy). - if isinstance(a, cp.ndarray): - return cp.asnumpy(a) - except Exception: - pass - return np.asarray(a) + Resolved through the backend's own transfer, which honours MDOPT_BACKEND + and the CUDA device probe; a per-call ``import cupy`` here once cost + ~18% of a decoding run. + """ + host = xp.to_host(a) + try: + return np.asarray(host) + except TypeError: + # A device array reached us while the NumPy backend is selected + # (CuPy installed, MDOPT_BACKEND unset); CuPy refuses the implicit + # conversion, so ask it explicitly. + return np.asarray(host.get()) def svd( @@ -32,10 +34,13 @@ def svd( chi_max: float = int(1e4), renormalise: bool = False, return_truncation_error: bool = False, -) -> Tuple[np.ndarray, List[float], np.ndarray, Optional[float]]: +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[float]]: """ Performs Singular Value Decomposition with different features. + The full decomposition is always taken; the comment at the backend call + records why strongly rectangular input is not reduced by QR/LQ first. + Parameters ---------- mat : np.ndarray @@ -55,8 +60,9 @@ def svd( ------- u_l : np.ndarray Unitary matrix having left singular vectors as columns. - singular_values : list - The singular values, sorted in non-increasing order. + singular_values : np.ndarray + The singular values kept after the cut and ``chi_max``, sorted in + non-increasing order, as a real one-dimensional array. v_r : np.ndarray Unitary matrix having right singular vectors as rows. truncation_error : Optional[float] @@ -79,7 +85,26 @@ def svd( for attempt in ("xp", "gesdd", "gesvd", "jitter"): try: if attempt == "xp": - u_l, s, v_h = xp.linalg.svd(a, full_matrices=False) # returns U, S, Vh + # The full decomposition, always. A QR/LQ pre-reduction of + # strongly rectangular input (factor A = QR, decompose the + # small factor, multiply Q back onto the kept singular + # vectors) was tried on PR #543 and removed: + # - with NumPy/SciPy wheels linked against Apple's Accelerate + # (the macOS 14+ arm64 wheels) it corrupted memory on the + # decoders' rank-deficient matrices: numpy.linalg.qr died + # with SIGBUS, and a [[72,12,6]] decode at chi_max=400 + # returned wrong verdicts on 24 of 24 shots while every + # unit test and benchmark fingerprint still passed; + # - with OpenBLAS it is memory-safe but not equivalent under + # truncation: it moved a chi_max=64 surface-code posterior + # entry by 0.017, beyond the benchmark suite's tolerance; + # - it saved 0-7% on the benchmark workloads. + # Reintroducing it needs `benchmarks/bench_suite.py --check` + # and tests/decoding/test_convergence.py to pass. + # No finiteness pre-scan: a non-finite input makes the SVD + # raise LinAlgError into the fallbacks below, and the scan + # would be a device sync on the GPU backend. + u_l, s, v_h = xp.linalg.svd(a, full_matrices=False) elif attempt == "gesdd": u_l, s, v_h = scipy.linalg.svd( _to_numpy(a), @@ -112,10 +137,10 @@ def svd( else: raise RuntimeError(f"All SVD methods failed. Last error: {last_exception}") - # Convert to NumPy for downstream consistency with the current codebase - u_l = _to_numpy(u_l) + # The spectrum comes to the host first: the truncation count is decided + # here, and the singular vectors are sliced while still on the backend, + # then converted. s = _to_numpy(s).astype(float, copy=False) # singular values are real non-negative - v_h = _to_numpy(v_h) # Truncate by cut and chi_max # int(chi_max) first would raise OverflowError on the chi_max=np.inf @@ -126,6 +151,9 @@ def svd( u_l = u_l[:, :max_num] s = s[:max_num] v_h = v_h[:max_num, :] + # Convert to NumPy for downstream consistency with the current codebase + u_l = _to_numpy(u_l) + v_h = _to_numpy(v_h) if renormalise and s.size > 0: norm = float(np.linalg.norm(s)) diff --git a/tests/decoding/test_convergence.py b/tests/decoding/test_convergence.py new file mode 100644 index 00000000..25c5ed98 --- /dev/null +++ b/tests/decoding/test_convergence.py @@ -0,0 +1,57 @@ +"""Slow decoding regression tests in the regime where truncation decides. + +These run the decoders at production-size bond dimensions on instances whose +verdict is known to be sensitive to the numerical details of the truncation +(a wrong singular-vector basis, a corrupted factorisation, a changed +tie-break) even when every cheap unit test and every fingerprint in +``benchmarks/`` passes. They take tens of minutes each and are skipped +unless ``MDOPT_RUN_SLOW=1`` is set; run them before merging any change to +the SVD, the orthogonality-centre moves or the contractor. +""" + +import os + +import numpy as np +import pytest + +from mdopt.decoding.decoding import create_bb_code, decode_css + +pytestmark = pytest.mark.skipif( + os.environ.get("MDOPT_RUN_SLOW") != "1", + reason="slow decoding regression test; set MDOPT_RUN_SLOW=1 to run", +) + + +def test_bb_72_12_6_natural_order_single_error_converges_in_chi(): + """A single Z error on the [[72,12,6]] code decodes to the identity at + chi_max=400 in the natural qubit order, and the verdict is the same at + chi_max=128. + + This instance exposed the QR pre-reduction of ``svd`` (since removed): + with it, the decode returned a flat or wrongly peaked + posterior at chi_max=400 on NumPy/Accelerate builds -- 24 of 24 shots + -- while every unit test and benchmark fingerprint still passed. Each + chi_max=400 decode takes 30-60 minutes on a laptop. + """ + code = create_bb_code(6, 6, "x**3 + y + y**2", "y**3 + x + x**2") + error = "I" * 58 + "Z" + "I" * 13 + assert len(error) == len(code) + peaks = {} + for chi_max in (128, 400): + posterior, success = decode_css( + code, + error, + chi_max=chi_max, + bias_type="Depolarising", + bias_prob=0.01, + renormalise=True, + silent=True, + contraction_strategy="Optimised", + qubit_order_strategy="Natural", + ) + posterior = np.asarray(posterior, dtype=float).ravel() + assert success == 1.0, f"chi_max={chi_max}: the identity class lost" + assert int(np.argmax(posterior)) == 0 + peaks[chi_max] = float(posterior.max()) + # A converged posterior for a single low-weight error is a delta. + assert peaks[chi_max] > 0.99, f"chi_max={chi_max}: peak {peaks[chi_max]}" diff --git a/tests/decoding/test_decoders.py b/tests/decoding/test_decoders.py index 820bc816..bd6ad3cf 100644 --- a/tests/decoding/test_decoders.py +++ b/tests/decoding/test_decoders.py @@ -508,11 +508,27 @@ def warnings_for(chi_max): ) return [r for r in caplog.records if "Negative logical amplitude" in r.message] - # The symplectic rewiring relocated where truncation bites on these seeded - # instances: the artefact now appears at chi_max=2 rather than 4. - assert warnings_for(2), "an aggressively truncated run should be flagged" + # Whether a given seeded decode produces a negative amplitude at a given + # chi_max is BLAS-dependent numerical noise (it differs between Accelerate + # and OpenBLAS and migrates under behaviour-preserving SVD changes), so the + # real decodes pin only the converged side. The emission itself is asserted + # deterministically on a readout that is known to carry a negative entry. assert not warnings_for(64), "a converged run should not be flagged" + from mdopt.mps.explicit import ExplicitMPS + + readout = np.array([0.9, -0.2, 0.1, 0.05]) + caplog.clear() + with ( + caplog.at_level(logging.WARNING), + patch.object(CanonicalMPS, "dense", return_value=readout), + patch.object(ExplicitMPS, "dense", return_value=readout), + ): + decode_css(code, errors[0], chi_max=4, silent=False) + assert any( + "Negative logical amplitude" in r.message for r in caplog.records + ), "a posterior with a negative amplitude must be flagged" + def test_max_product_readout_is_optimal_and_certified(): """Beam search should settle the readout without needing DMRG. diff --git a/tests/decoding/test_surface_ler_regression.py b/tests/decoding/test_surface_ler_regression.py new file mode 100644 index 00000000..f4007983 --- /dev/null +++ b/tests/decoding/test_surface_ler_regression.py @@ -0,0 +1,73 @@ +"""Logical-error-rate regression on the distance-5 surface code. + +A fixed seed makes the error sample deterministic, so the number of decoding +failures is a property of the code, not a Monte Carlo estimate, and "the +logical error rate must not increase" is an exact check. The bond dimension +is deliberately small: at ``chi_max=16`` the decoder is not converged, so +the verdicts depend on how the truncation is carried out, and a change to +the SVD, the orthogonality-centre moves or the contractor that alters the +truncation moves them. The pinned values were produced by ``main`` and by +the optimised branch alike (c5d2773 / 0ba5f0b, 2026-09-11). +""" + +import numpy as np +import qecstruct as qec + +from mdopt.decoding.decoding import decode_css, generate_pauli_error_string + +SEED = 2026 +SHOTS = 40 +CHI_MAX = 16 + + +def _surface_code(lattice_size): + rep = qec.repetition_code(lattice_size) + return qec.hypergraph_product(rep, rep) + + +def _verdicts(code, error_rate): + rng = np.random.default_rng(SEED) + verdicts = [] + for _ in range(SHOTS): + error = generate_pauli_error_string( + len(code), error_rate, error_model="Depolarising", rng=rng + ) + _, success = decode_css( + code, + error, + chi_max=CHI_MAX, + bias_type="Depolarising", + bias_prob=error_rate, + renormalise=True, + silent=True, + contraction_strategy="Optimised", + qubit_order_strategy="Natural", + ) + verdicts.append(int(success)) + return "".join(map(str, verdicts)) + + +def test_surface_5_low_error_rate_decodes_every_shot(): + """At p=0.05 (well below threshold) none of the 40 shots fails.""" + verdicts = _verdicts(_surface_code(5), 0.05) + failures = verdicts.count("0") + assert failures == 0, f"{failures} of {SHOTS} shots failed: {verdicts}" + + +def test_surface_5_truncated_ler_does_not_increase(): + """At p=0.08 and chi_max=16 the pinned sample has 4 failures. + + The count must not grow. The pattern is pinned too: a different pattern + with the same or a lower count means the truncation changed, which is + worth a review even when it looks like an improvement; update the pinned + string deliberately after that review. + """ + pinned = "1111111111111111011011101111111111011111" + verdicts = _verdicts(_surface_code(5), 0.08) + failures = verdicts.count("0") + assert failures <= pinned.count( + "0" + ), f"{failures} failures, pinned {pinned.count('0')}: {verdicts}" + assert ( + verdicts == pinned + ), f"verdict pattern changed:\n got {verdicts}\n pinned {pinned}" diff --git a/tests/mps/test_canonical.py b/tests/mps/test_canonical.py index b58e4cbd..0aa0f0c4 100644 --- a/tests/mps/test_canonical.py +++ b/tests/mps/test_canonical.py @@ -886,3 +886,219 @@ def test_marginal_does_not_produce_nans_when_the_centre_underflows(): marginalised = mps.marginal(sites_to_marginalise=list(range(10)), renormalise=True) assert np.all(np.isfinite(marginalised.dense(flatten=True))) + + +def _two_site_reference_move(mps, final_pos, renormalise): + """The two-site-SVD move spelled out: the reference for the one-site path. + + move_orth_centre no longer has a flag that forces the two-site branch + (renormalised moves and returned spectra take the one-site path too), so + an old-vs-new comparison has to build the reference explicitly. + """ + if mps.orth_centre == final_pos: + return mps.copy() + leftwards = mps.orth_centre > final_pos + work = mps.reverse() if leftwards else mps.copy() + begin = work.orth_centre + final = (mps.num_sites - 1 - final_pos) if leftwards else final_pos + for i in range(begin, final): + u_l, s_bond, v_r, _ = split_two_site_tensor( + work.two_site_tensor_next(i), + chi_max=mps.chi_max, + renormalise=renormalise, + strategy="svd", + return_truncation_error=True, + ) + work.tensors[i] = u_l + work.tensors[i + 1] = v_r * s_bond[:, None, None] + work.orth_centre = i + 1 + return work.reverse() if leftwards else work + + +def test_move_orth_centre_carries_a_collapsed_bond_through(): + """A bond of dimension 0 (a truncation that emptied the spectrum) must + move through the centre-only path without raising. + + Three example notebooks hit this in CI: a zero-width centre has nothing + to factor, so such a bond must take the two-site SVD branch, which + carries the degenerate shape through unchanged. + """ + from mdopt.mps.canonical import CanonicalMPS + + tensors = [ + np.zeros((1, 2, 0)), + np.zeros((0, 2, 1)), + np.zeros((1, 2, 1)), + ] + mps = CanonicalMPS(tensors, orth_centre=0, chi_max=4) + moved = mps.move_orth_centre(2, renormalise=False) + assert [t.shape for t in moved.tensors][0][2] == 0 + back = moved.move_orth_centre(0, renormalise=False) + assert len(back) == 3 + + +def test_move_orth_centre_collapses_a_sub_cut_spectrum_like_the_svd_path(): + """A centre whose whole spectrum sits below the 1e-12 cut must collapse + to a zero-width bond on the centre-only path exactly as on the two-site + SVD path, rather than propagating a sub-cut direction.""" + from mdopt.mps.canonical import CanonicalMPS + + def tiny_centre_mps(): + tensors = [ + np.array([[[1.0, 0.0], [0.0, 1.0]]]).reshape(1, 2, 2) * 5e-13, + np.eye(2).reshape(2, 2, 1), + np.array([1.0, 0.0]).reshape(1, 2, 1), + ] + return CanonicalMPS(tensors, orth_centre=0, chi_max=4) + + via_one_site = tiny_centre_mps().move_orth_centre(2, renormalise=False) + via_two_site = _two_site_reference_move(tiny_centre_mps(), 2, renormalise=False) + assert via_one_site.tensors[0].shape[2] == 0 + assert via_one_site.tensors[0].shape[2] == via_two_site.tensors[0].shape[2] + + +def test_move_orth_centre_qr_path_matches_svd_path_on_full_rank_states(): + """The one-site move must reproduce the two-site SVD move exactly. + + Real and complex random states, moves in both directions, with a + chi_max below the full Schmidt rank so the finite truncation is + exercised too: dense() and every bond dimension must agree between the + one-site path and the two-site reference spelled out above. + """ + from mdopt.mps.utils import mps_from_dense + + for seed, complex_case in ((11, False), (12, True), (13, False), (14, True)): + rng = np.random.default_rng(seed) + vec = rng.standard_normal(2**8) + if complex_case: + vec = vec + 1j * rng.standard_normal(2**8) + vec = vec / np.linalg.norm(vec) + for chi_max in (int(1e4), 3): + base = mps_from_dense(vec, form="Right-canonical") + base.chi_max = chi_max + for targets in ((7, 0), (0, 7), (4, 1, 6)): + fast, slow = base.copy(), base.copy() + for target in targets: + fast = fast.move_orth_centre(target, renormalise=False) + slow = _two_site_reference_move(slow, target, renormalise=False) + assert list(fast.bond_dimensions) == list(slow.bond_dimensions), ( + seed, + chi_max, + targets, + ) + assert np.allclose( + fast.dense(flatten=True), + slow.dense(flatten=True), + rtol=0.0, + atol=1e-11, + ), (seed, chi_max, targets) + + +def test_move_orth_centre_matches_svd_path_on_non_canonical_chains(): + """Behind a biased pair the neighbours are not isometries, so a move + that must truncate at chi_max has to take the two-site SVD; the fast + path may only handle moves that truncate nothing.""" + from mdopt.mps.utils import mps_from_dense + from mdopt.examples.decoding.decoding import apply_depolarising_bias + + rng = np.random.default_rng(5) + vec = rng.standard_normal(2**8) + vec = vec / np.linalg.norm(vec) + for chi_max in (2, 4, int(1e4)): + base = mps_from_dense(vec, form="Right-canonical", chi_max=chi_max) + biased = apply_depolarising_bias( + base, sites_to_bias=[0, 2, 4, 6], prob_bias_list=0.3 + ) + fast = biased.copy().move_orth_centre(0, renormalise=False) + slow = _two_site_reference_move(biased.copy(), 0, renormalise=False) + assert list(fast.bond_dimensions) == list(slow.bond_dimensions), chi_max + assert np.allclose( + fast.dense(flatten=True), slow.dense(flatten=True), rtol=0.0, atol=1e-10 + ), chi_max + + +def test_move_orth_centre_prunes_on_discarded_amplitude_not_on_the_centre_spectrum(): + """Review counterexamples: a sub-cut centre direction amplified by a large + neighbour entry must be KEPT (its amplitude is macroscopic), and a + direction a null neighbour row makes worthless may go. Both cases must + agree with the two-site SVD path on the represented state.""" + from mdopt.mps.canonical import CanonicalMPS + + def chain(neighbour): + centre = np.zeros((1, 2, 2)) + centre[0, 0, 0], centre[0, 1, 1] = 1.0, 5e-13 + last = np.array([1.0, 0.0]).reshape(1, 2, 1) + return CanonicalMPS([centre, neighbour, last], orth_centre=0, chi_max=4) + + amplified = np.zeros((2, 2, 1)) + amplified[0, 0, 0], amplified[1, 1, 0] = 1.0, 1e13 + null_row = np.zeros((2, 2, 1)) + null_row[0, 0, 0] = 1.0 + for neighbour in (amplified, null_row): + fast = chain(neighbour).move_orth_centre(2, renormalise=False) + slow = _two_site_reference_move(chain(neighbour), 2, renormalise=False) + assert np.allclose( + fast.dense(flatten=True), slow.dense(flatten=True), rtol=0.0, atol=1e-10 + ) + kept = chain(amplified).move_orth_centre(1, renormalise=False) + assert kept.tensors[0].shape[2] == 2, "the amplified direction carries amplitude 5" + + +def test_one_site_move_matches_two_site_move_with_renormalisation(): + """The centre-only move also covers renormalised moves and returned spectra. + + With an isometric neighbour B, theta = C B has theta theta^dag = C C^dag, + so the spectrum, its renormalisation, the cut and the chi_max count are + those of the centre alone. Both DMRG sweeps (renormalise=True) and the + explicit-form conversion (return_singular_values=True) go through it now; + the reference below is the two-site formula spelled out. + """ + from mdopt.utils.utils import split_two_site_tensor + + rng = np.random.default_rng(11) + for dtype in (float, complex): + vector = rng.standard_normal(2**7) + if dtype is complex: + vector = vector + 1j * rng.standard_normal(2**7) + vector /= np.linalg.norm(vector) + mps = mps_from_dense(vector, form="Right-canonical", chi_max=5) + assert isinstance(mps, CanonicalMPS) + mps.chi_max = 5 + + reference = mps.copy() + spectra = [] + for site in range(0, 4): + theta = reference.two_site_tensor_next(site) + u_l, s_bond, v_r, _ = split_two_site_tensor( + theta, chi_max=5, renormalise=True, return_truncation_error=True + ) + spectra.append(s_bond) + reference.tensors[site] = u_l + reference.tensors[site + 1] = v_r * s_bond[:, None, None] + reference.orth_centre = site + 1 + + moved, singular_values = mps.move_orth_centre( + 4, return_singular_values=True, renormalise=True + ) + assert moved.orth_centre == 4 + assert len(singular_values) == len(spectra) + for got, want in zip(singular_values, spectra): + assert np.allclose(got, want, rtol=0.0, atol=1e-12) + assert np.allclose(moved.dense(), reference.dense(), rtol=0.0, atol=1e-12) + assert moved.bond_dimensions == reference.bond_dimensions + + +def test_move_orth_centre_accepts_integer_and_boolean_tensors(): + """Exact-arithmetic product states are valid input; the isometry gate must + promote its Gram matrix rather than fail on an in-place float subtraction.""" + from mdopt.mps.canonical import CanonicalMPS + + for dtype in (int, bool): + tensors = [np.array([1, 0], dtype=dtype).reshape(1, 2, 1) for _ in range(4)] + mps = CanonicalMPS(tensors, orth_centre=0, chi_max=4) + moved = mps.move_orth_centre(3, renormalise=False) + assert moved.orth_centre == 3 + assert np.allclose(moved.dense(flatten=True)[0], 1.0) + back = moved.move_orth_centre(0, renormalise=True) + assert back.orth_centre == 0 + assert list(back.bond_dimensions) == [1, 1, 1] diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index 8e322f66..f22c4280 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -1,5 +1,9 @@ """Tests for the ``mdopt.utils.utils`` module.""" +import os +import subprocess +import sys +import textwrap import pytest import scipy import numpy as np @@ -630,3 +634,127 @@ def test_qr_accepts_infinite_chi_max(rng): q_small, r_small, _ = qr(mat, cut=1e-16, chi_max=3) assert q_small.shape[1] == 3 assert r_small.shape[0] == 3 + + +def test_svd_rectangular_inputs_match_direct_svd(): + """Rectangular inputs in both orientations, at and around the 2:1 aspect + ratio the removed QR/LQ pre-reduction used to take, must match a direct + numpy SVD: singular values, reconstruction, orthogonality, and chi_max + truncation. + """ + rng = np.random.default_rng(20240903) + shapes_and_paths = [ + ((16, 32), "wide 2:1"), + ((16, 31), "wide just under 2:1"), + ((32, 16), "tall 2:1"), + ((31, 16), "tall just under 2:1"), + ((16, 16), "square"), + ] + for complex_case in (False, True): + for shape, label in shapes_and_paths: + mat = rng.standard_normal(shape) + if complex_case: + mat = mat + 1j * rng.standard_normal(shape) + u_l, s, v_h, _ = svd(mat, cut=0.0, chi_max=int(1e4)) + u_ref, s_ref, v_ref = np.linalg.svd(mat, full_matrices=False) + assert np.allclose(s, s_ref, rtol=0.0, atol=1e-11), (label, complex_case) + assert np.allclose((u_l * s) @ v_h, mat, rtol=0.0, atol=1e-11), ( + label, + complex_case, + ) + eye = np.eye(u_l.shape[1]) + assert np.allclose(u_l.conj().T @ u_l, eye, rtol=0.0, atol=1e-12), label + assert np.allclose(v_h @ v_h.conj().T, eye, rtol=0.0, atol=1e-12), label + + chi = 5 + u_t, s_t, v_t, err = svd( + mat, cut=0.0, chi_max=chi, return_truncation_error=True + ) + assert s_t.shape == (chi,) + assert np.allclose(s_t, s_ref[:chi], rtol=0.0, atol=1e-11), label + # rtol=0 so the bound is genuinely absolute; the discarded + # spectrum's norm-square is O(100) here, where the default + # relative tolerance would hide errors of order 1e-3. + assert np.isclose( + err, float(np.linalg.norm(s_ref[chi:]) ** 2), rtol=0.0, atol=1e-9 + ), label + + +def test_svd_nonfinite_input_takes_the_fallback_chain(): + """A non-finite input must make the whole call raise. + + There is no finiteness pre-scan: the SVD of a NaN matrix raises, which + sends the call through the fallback chain, and the jitter attempt cannot + rescue a NaN either.""" + mat = np.full((8, 32), np.nan) + with pytest.raises(RuntimeError, match="All SVD methods failed"): + svd(mat) + + +def _graded_rank_deficient(rng, rows=636, cols=304, rank=237): + """A matrix of the decoders' centre-tensor shape with a spectrum spanning + sixteen orders of magnitude and a numerically null tail.""" + spectrum = np.concatenate([np.logspace(0, -16, rank), np.zeros(cols - rank)]) + left, _ = np.linalg.qr(rng.normal(size=(rows, cols))) + right, _ = np.linalg.qr(rng.normal(size=(cols, cols))) + return (left * spectrum) @ right + + +def test_svd_reconstructs_graded_rank_deficient_matrices(rng): + """``svd`` is exact on rank-deficient, wide-spectrum input of the shape the + decoders' centre tensors take. A small deterministic check for the default + suite; the allocator-churn stress version is opt-in and runs in a + subprocess (see below).""" + for _ in range(3): + mat = _graded_rank_deficient(rng) + u_l, s, v_h, _ = svd(mat, cut=1e-17, chi_max=400) + assert np.isfinite(u_l).all() and np.isfinite(v_h).all() + assert np.abs((u_l * s) @ v_h - mat).max() < 1e-10 + + +@pytest.mark.skipif( + os.environ.get("MDOPT_RUN_SLOW") != "1", + reason="allocator-churn stress test; set MDOPT_RUN_SLOW=1 to run", +) +def test_svd_under_allocator_churn_in_a_subprocess(): + """``svd`` stays exact call after call on graded rank-deficient input while + the allocator churns. + + NumPy's ``linalg.qr`` on the Accelerate framework (macOS arm64 wheels) + intermittently returned a factorisation whose product was not the input, + and died with SIGBUS, on matrices of this kind, and Accelerate's SVD + tripped malloc's heap check on the decoders' matrices. A native crash of + that sort would terminate pytest, so the loop runs in a subprocess and a + signal becomes an ordinary test failure. It is a guard, not a certain + detector (the fault is heap-state dependent); + ``tests/decoding/test_convergence.py`` is the deterministic check. + """ + script = textwrap.dedent(""" + import numpy as np + from mdopt.utils.utils import svd + + rng = np.random.default_rng(2026) + rows, cols, rank = 636, 304, 237 + spectrum = np.concatenate([np.logspace(0, -16, rank), np.zeros(cols - rank)]) + junk = [] + for _ in range(150): + left, _ = np.linalg.qr(rng.normal(size=(rows, cols))) + right, _ = np.linalg.qr(rng.normal(size=(cols, cols))) + mat = (left * spectrum) @ right + junk.append(rng.normal(size=rng.integers(1, 200_000))) + junk = junk[-10:] + u_l, s, v_h, _ = svd(mat, cut=1e-17, chi_max=400) + assert np.isfinite(u_l).all() and np.isfinite(v_h).all() + assert np.abs((u_l * s) @ v_h - mat).max() < 1e-10 + """) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=1800, + check=False, + ) + assert result.returncode == 0, ( + f"exit code {result.returncode} (a negative code is the signal that " + f"killed the process): {result.stderr[-2000:]}" + )