diff --git a/tests/contract.py b/tests/contract.py new file mode 100644 index 0000000..b78cfed --- /dev/null +++ b/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) diff --git a/tests/search_tree.py b/tests/search_tree.py new file mode 100644 index 0000000..783a02a --- /dev/null +++ b/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) diff --git a/tests/tree.pkl.bak b/tests/tree.pkl.bak new file mode 100644 index 0000000..a0ccdcd Binary files /dev/null and b/tests/tree.pkl.bak differ