Enable efficient contraction order simulation

benchmarks/efficient_contraction_order_statevec.py

@@ -0,0 +1,127 @@
+"""
+Efficient contraction order statevector simulation of MBQC patterns
+===================================================================
+
+Here we benchmark our efficient contraction order statevector simulator for MBQC,
+which uses the TensorNetwork backend.
+
+The methods and modules we use are the followings:
+    1. :meth:`graphix.pattern.Pattern.simulate_pattern`
+        Pattern simulation with Statevector backend.
+        :meth:`graphix.pattern.Pattern.minimize_space` locally minimizes the space of the pattern.
+    2. :mod:`graphix.sim.tensornet`
+        Pattern simulation with TensorNetwork backend.
+        This enables the efficient contraction order simulation of the pattern.
+        Here we use the `cotengra` optimizer for the contraction order optimization.
+"""
+
+# %%
+# Firstly, let us import relevant modules:
+
+from copy import deepcopy
+from time import perf_counter
+
+import cotengra as ctg
+import quimb as qu
+from numpy.random import PCG64, Generator
+
+from graphix.random_circuit import get_rand_circuit
+
+# %%
+# Next, set global seed and number of thread workers
+GLOBAL_SEED = 2
+qu.core._NUM_THREAD_WORKERS = 1
+
+# %%
+# We then run simulations.
+# Let's benchmark the simulation time of the statevector simulator and the efficient contraction order simulator.
+
+n_qubit_list = list(range(2, 31))
+
+sv_sim = []
+eco_sim = []
+eco_sim_wo_ctg = []
+max_space_ls = []
+
+for n_qubit in n_qubit_list:
+    rng = Generator(PCG64(GLOBAL_SEED))
+    circuit = get_rand_circuit(n_qubit, n_qubit, rng)
+    pattern = circuit.transpile().pattern
+    pattern.standardize()
+    pat_original = deepcopy(pattern)
+    start = perf_counter()
+    pat_original.minimize_space()
+    max_sp = pat_original.max_space()
+    max_space_ls.append(max_sp)
+    sv = pat_original.simulate_pattern(max_qubit_num=max_sp)
+    end = perf_counter()
+    sv_sim.append(end - start)
+    del sv
+
+    # efficient contraction order simulation (eco-sim): tn
+    tn = pattern.simulate_pattern("tensornetwork")
+    output_inds = [tn._dangling[str(index)] for index in tn.default_output_nodes]
+
+    start = perf_counter()
+    tn_sv = tn.to_statevector(
+        backend="numpy",
+        skip=False,
+        optimize=ctg.HyperOptimizer(minimize="combo", max_time=600, progbar=True),
+    )
+    end = perf_counter()
+    eco_sim.append(end - start)
+    del tn_sv
+
+    # eco-sim: tn without cotengra optimizer
+    start = perf_counter()
+    tn_sv = tn.to_statevector(
+        backend="numpy",
+        skip=False,
+        optimize=None,
+    )
+    end = perf_counter()
+    eco_sim_wo_ctg.append(end - start)
+    del tn_sv, tn
+
+# %%
+# Finally, we save the results to a text file.
+with open("sqrqcresults.txt", "w") as f:
+    f.write("n_qubit, sv_sim, eco_sim, eco_sim_wo_ctg, max_space_ls\n")
+    for i in range(len(n_qubit_list)):
+        f.write(f"{n_qubit_list[i]}, {sv_sim[i]}, {eco_sim[i]}, {eco_sim_wo_ctg[i]}, {max_space_ls[i]}\n")
+
+# %%
+# We can now plot the simulation time results.
+import matplotlib.pyplot as plt
+import numpy as np
+
+data = np.loadtxt("sqrqcresults.txt", delimiter=",", skiprows=1)
+n_qubits = data[:, 0].astype(int)
+sv_sim = data[:, 1].astype(float)
+eco_sim = data[:, 2].astype(float)
+eco_sim_wo_ctg = data[:, 3].astype(float)
+max_sp = data[:, 4].astype(int)
+
+fig, ax1 = plt.subplots(figsize=(8, 5))
+color = "tab:red"
+ax1.set_xlabel("Original Circuit Size [qubit]")
+ax1.set_ylabel("Simulation time [sec]")
+ax1.set_yscale("log")
+ax1.scatter(n_qubits, sv_sim, marker="x", label="MBQC Statevector (minimizing sp)")
+ax1.scatter(n_qubits, eco_sim, marker="x", label="MBQC TN base (with cotengra)")
+ax1.scatter(n_qubits, eco_sim_wo_ctg, marker="x", label="MBQC TN base (without cotengra)")
+ax1.tick_params(axis="y")
+ax1.legend(loc="upper left")
+ax1.set_title("Simulation time (Square RQC)")
+plt.grid(True, which="Major")
+
+ax2 = ax1.twinx()
+color = "tab:blue"
+ax2.set_ylabel("Max Space [qubit]")
+ax2.plot(n_qubits, max_sp, color=color, linestyle="--", label="Max Space")
+ax2.tick_params(axis="y")
+ax2.legend(loc="lower right")
+
+plt.rcParams["svg.fonttype"] = "none"
+plt.savefig("simulation_time_wo_p.png")
+plt.close()

examples/MBQCvqe.py

@@ -22,9 +22,11 @@
 import networkx as nx
 import numpy as np
 from scipy.optimize import minimize
+
 from graphix import Circuit
 from graphix.simulator import PatternSimulator
 
+
 # %%
 # Define the Hamiltonian for the VQE problem (Example: H = Z0Z1 + X0 + X1)
 def create_hamiltonian():
@@ -33,6 +35,7 @@ def create_hamiltonian():
     H = np.kron(Z, Z) + np.kron(X, np.eye(2)) + np.kron(np.eye(2), X)
     return H
 
+
 # %%
 # Function to build the VQE circuit
 def build_vqe_circuit(n_qubits, params):
@@ -45,6 +48,7 @@ def build_vqe_circuit(n_qubits, params):
         circuit.cnot(i, i + 1)
     return circuit
 
+
 # %%
 class MBQCVQE:
     def __init__(self, n_qubits, hamiltonian):
@@ -85,6 +89,7 @@ def compute_energy(self, params):
         energy = tn.expectation_value(self.hamiltonian, qubit_indices=range(self.n_qubits))
         return energy
 
+
 # %%
 # Set parameters for VQE
 n_qubits = 2
@@ -99,13 +104,14 @@ def compute_energy(self, params):
 def cost_function(params):
     return mbqc_vqe.compute_energy(params)
 
+
 # %%
 # Random initial parameters
 initial_params = np.random.rand(n_qubits * 3)
 
 # %%
 # Perform the optimization using COBYLA
-result = minimize(cost_function, initial_params, method='COBYLA', options={'maxiter': 100})
+result = minimize(cost_function, initial_params, method="COBYLA", options={"maxiter": 100})
 
 print(f"Optimized parameters: {result.x}")
 print(f"Optimized energy: {result.fun}")

tests/random_circuit.py → graphix/random_circuit.py

@@ -84,6 +84,21 @@ def get_rand_circuit(
     use_rzz: bool = False,
     use_ccx: bool = False,
 ) -> Circuit:
+    """Generate a random circuit.
+
+    Parameters
+    ----------
+    nqubits : int
+        Number of qubits in the circuit. Must be at least 2.
+    depth : int
+        Number of layers in the circuit.
+    rng : Generator
+        Random number generator.
+    use_rzz : bool, optional
+        Whether to include RZZ gates, by default False.
+    use_ccx : bool, optional
+        Whether to include CCX gates, by default False.
+    """
     circuit = Circuit(nqubits)
     gate_choice = (
         functools.partial(circuit.ry, angle=np.pi / 4),

graphix/sim/statevec.py

@@ -1,29 +1,30 @@
 from copy import deepcopy
 
 import numpy as np
+from numpy.typing import NDArray
 
 import graphix.sim.base_backend
-from graphix.clifford import CLIFFORD, CLIFFORD_CONJ, CLIFFORD_MUL
+from graphix.clifford import CLIFFORD, CLIFFORD_CONJ
 from graphix.ops import Ops
 
 
 class StatevectorBackend(graphix.sim.base_backend.Backend):
     """MBQC simulator with statevector method."""
 
-    def __init__(self, pattern, max_qubit_num=20, pr_calc=True):
+    def __init__(self, pattern, pr_calc=True, max_qubit_num=None):
         """
         Parameters
         -----------
         pattern : :class:`graphix.pattern.Pattern` object
             MBQC pattern to be simulated.
         backend : str, 'statevector'
             optional argument for simulation.
-        max_qubit_num : int
-            optional argument specifying the maximum number of qubits
-            to be stored in the statevector at a time.
         pr_calc : bool
             whether or not to compute the probability distribution before choosing the measurement result.
             if False, measurements yield results 0/1 with 50% probabilities each.
+        max_qubit_num : int, optional
+            optional argument specifying the maximum number of qubits
+            to be stored in the statevector at a time.
         """
         # check that pattern has output nodes configured
         # assert len(pattern.output_nodes) > 0
@@ -34,9 +35,7 @@
         self.Nqubit = 0
         self.to_trace = []
         self.to_trace_loc = []
-        self.max_qubit_num = max_qubit_num
-        if pattern.max_space() > max_qubit_num:
-            raise ValueError("Pattern.max_space is larger than max_qubit_num. Increase max_qubit_num and try again")
+        self.max_qubit_num = _validate_max_qubit_num(max_qubit_num, pattern.max_space())
         super().__init__(pr_calc)
 
     def qubit_dim(self):
@@ -182,21 +181,19 @@
 class Statevec:
     """Simple statevector simulator"""
 
-    def __init__(self, nqubit=1, plus_states=True):
+    def __init__(self, nqubit=1, psi=None, plus_states=True):
         """Initialize statevector
 
         Parameters
         ----------
         nqubit : int, optional:
             number of qubits. Defaults to 1.
+        psi : numpy.ndarray, optional
+            statevector. Defaults to None.
         plus_states : bool, optional
             whether or not to start all qubits in + state or 0 state. Defaults to +
         """
-        if plus_states:
-            self.psi = np.ones((2,) * nqubit) / 2 ** (nqubit / 2)
-        else:
-            self.psi = np.zeros((2,) * nqubit)
-            self.psi[(0,) * nqubit] = 1
+        self.psi = _initial_state(nqubit, psi, plus_states)
 
     def __repr__(self):
         return f"Statevec, data={self.psi}, shape={self.dims()}"
@@ -225,8 +222,13 @@
             target qubits' indices
         """
         op_dim = int(np.log2(len(op)))
-        # TODO shape = (2,)* 2 * op_dim
-        shape = [2 for _ in range(2 * op_dim)]
+        shape = (
+            [
+                2,
+            ]
+            * 2
+            * op_dim
+        )
         op_tensor = op.reshape(shape)
         self.psi = np.tensordot(
             op_tensor,
@@ -409,6 +411,48 @@
         return np.dot(st2.psi.flatten().conjugate(), st1.psi.flatten())
 
 
-def _get_statevec_norm(psi):
+def _get_statevec_norm(psi) -> float:
     """returns norm of the state"""
     return np.sqrt(np.sum(psi.flatten().conj() * psi.flatten()))
+
+
+def _initial_state(nqubit: int = 1, psi: NDArray = None, plus_states: bool = True) -> NDArray:
+    """Create initial state
+
+    Parameters
+    ----------
+    nqubit : int, optional:
+        number of qubits. Defaults to 1.
+    psi : numpy.ndarray, optional
+        statevector. Defaults to None.
+    plus_states : bool, optional
+        whether or not to start all qubits in + state or 0 state. Defaults to +
+
+    Returns
+    -------
+    numpy.ndarray
+        statevector
+    """
+    if isinstance(psi, np.ndarray):
+        return psi
+    if not isinstance(nqubit, int):
+        raise TypeError("nqubit must be an integer")
+    if nqubit < 0:
+        raise ValueError("nqubit must be a non-negative integer")
+    if plus_states:
+        return np.ones((2,) * nqubit) / 2 ** (nqubit / 2)
+    psi = np.zeros((2,) * nqubit)
+    psi[(0,) * nqubit] = 1
+    return psi
+
+
+def _validate_max_qubit_num(max_qubit_num: int | None, max_space: int) -> int | None:
+    if max_qubit_num is None:
+        return
+    if not isinstance(max_qubit_num, int):
+        raise ValueError("max_qubit_num must be an integer")
+    if max_qubit_num < 1:
+        raise ValueError("max_qubit_num must be a positive integer")
+    if max_qubit_num < max_space:
+        raise ValueError("Pattern.max_space is larger than max_qubit_num. Increase max_qubit_num and try again")
+    return max_qubit_num