jaunatisblue
49b27a5840
Some checks failed
Build wheels / build (ubuntu-latest, 3.11) (push) Has been cancelled
Build wheels / build (ubuntu-latest, 3.12) (push) Has been cancelled
Build wheels / build (ubuntu-latest, 3.13) (push) Has been cancelled
Tests / check (push) Has been cancelled
Tests / build (ubuntu-latest, 3.11) (push) Has been cancelled
Tests / build (ubuntu-latest, 3.12) (push) Has been cancelled
Tests / build (ubuntu-latest, 3.13) (push) Has been cancelled
35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
"""Search contraction path and save."""
|
|
import time, os, pickle
|
|
from qibotn.parallel import parallel_path_search
|
|
from qibotn.observables import build_random_circuit
|
|
import qibo, quimb as qu
|
|
|
|
from mpi4py import MPI
|
|
|
|
NQUBITS, NLAYERS, WORKERS = 20, 10, 96
|
|
|
|
comm = MPI.COMM_WORLD
|
|
rank, size = comm.Get_rank(), comm.Get_size()
|
|
method = 'mpi' if size > 1 else 'processpool'
|
|
|
|
circuit = build_random_circuit(NQUBITS, NLAYERS)
|
|
qibo.set_backend("qibotn", platform="quimb")
|
|
backend = qibo.get_backend()
|
|
backend.configure_tn_simulation(ansatz="tn")
|
|
qc = backend._qibo_circuit_to_quimb(circuit, backend.circuit_ansatz)
|
|
tn = qc.local_expectation(qu.pauli('x') & qu.pauli('z'), (0, 1), rehearse='tn')
|
|
|
|
if rank == 0:
|
|
print(f"Searching {NQUBITS}q {NLAYERS}l, method={method}, ranks={size}, workers/rank={WORKERS}...")
|
|
t0 = time.time()
|
|
tree = parallel_path_search(tn, tn.outer_inds(), method=method,
|
|
total_repeats=1024, max_time=300, n_workers=WORKERS,trial_timeout=60)
|
|
t_search = time.time() - t0
|
|
|
|
if rank == 0:
|
|
os.makedirs('data', exist_ok=True)
|
|
path = f"data/tree_q{NQUBITS}_l{NLAYERS}.pkl"
|
|
with open(path, 'wb') as f:
|
|
pickle.dump(tree, f)
|
|
print(f"Search: {t_search:.2f}s Peak: {tree.max_size() * 16 / 1e9:.2f} GB Saved: {path}")
|