baseline测试

tests/contract.py

@@ -0,0 +1,27 @@
+import time
+import pickle
+
+
+def run(input="tree.pkl"):
+    with open(input, "rb") as f:
+        data = pickle.load(f)
+
+    sliced_tree = data["sliced_tree"]
+    arrays = data["arrays"]
+    n_slices = sliced_tree.nslices
+    print(f"Total slices: {n_slices}")
+
+    t0 = time.perf_counter()
+    total = sum(sliced_tree.contract_slice(arrays, i, backend='numpy',implementation='cotengra') for i in range(n_slices))
+    t1 = time.perf_counter()
+
+    print(f"Contract: {t1 - t0:.4f} s")
+    #print(f"Result: {total:.10f}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", type=str, default="tree.pkl.bak")
+    args = parser.parse_args()
+    run(args.input)

tests/search_tree.py

@@ -0,0 +1,56 @@
+import time
+import pickle
+import numpy as np
+import quimb.tensor as qtn
+import cotengra as ctg
+
+
+def build_qft(n_qubits):
+    circ = qtn.Circuit(n_qubits, dtype=np.complex128)
+    for i in range(n_qubits):
+        circ.apply_gate('H', i)
+        for j in range(i + 1, n_qubits):
+            circ.apply_gate('CPHASE', np.pi / 2 ** (j - i), i, j)
+    return circ
+
+
+def run(n_qubits, output="tree.pkl"):
+    print(f"Circuit: QFT {n_qubits} qubits")
+
+    circ = build_qft(n_qubits)
+    psi = circ.psi
+
+    Z = np.array([[1, 0], [0, -1]], dtype=np.complex128)
+    bra = psi.conj().reindex({f'k{i}': f'b{i}' for i in range(n_qubits)})
+    obs = qtn.Tensor(Z, inds=(f'k0', f'b0'))
+    net = psi & obs & bra
+
+    t0 = time.perf_counter()
+    opt = ctg.HyperOptimizer(
+        methods=['kahypar'],
+        max_repeats=32,
+        minimize='combo',
+        parallel=8,
+    )
+    tree = net.contraction_tree(optimize=opt)
+    t1 = time.perf_counter()
+    print(f"Path search: {t1 - t0:.4f} s")
+    print(tree)
+
+    sliced_tree = tree.slice(target_size=2**28)
+    print(f"Total slices: {sliced_tree.nslices}")
+
+    arrays = [t.data for t in net.tensors]
+
+    with open(output, "wb") as f:
+        pickle.dump({"sliced_tree": sliced_tree, "arrays": arrays}, f)
+    print(f"Saved to {output}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n_qubits", type=int, default=18)
+    parser.add_argument("--output", type=str, default="tree.pkl")
+    args = parser.parse_args()
+    run(args.n_qubits, args.output)