Format with Black

src/qibotn/MPSContractionHelper.py

@@ -1,5 +1,6 @@
 from cuquantum import contract, contract_path, CircuitToEinsum, tensor
 
+
 class MPSContractionHelper:
     """
     A helper class to compute various quantities for a given MPS.
@@ -33,9 +34,11 @@ class MPSContractionHelper:
 
     def __init__(self, num_qubits):
         self.num_qubits = num_qubits
-        self.bra_modes = [(2*i, 2*i+1, 2*i+2) for i in range(num_qubits)]
+        self.bra_modes = [(2 * i, 2 * i + 1, 2 * i + 2) for i in range(num_qubits)]
-        offset = 2*num_qubits+1
+        offset = 2 * num_qubits + 1
-        self.ket_modes = [(i+offset, 2*i+1, i+1+offset) for i in range(num_qubits)]
+        self.ket_modes = [
+            (i + offset, 2 * i + 1, i + 1 + offset) for i in range(num_qubits)
+        ]
 
     def contract_norm(self, mps_tensors, options=None):
         """
@@ -52,8 +55,10 @@ class MPSContractionHelper:
         """
         interleaved_inputs = []
         for i, o in enumerate(mps_tensors):
-            interleaved_inputs.extend([o, self.bra_modes[i], o.conj(), self.ket_modes[i]])
+            interleaved_inputs.extend(
-        interleaved_inputs.append([]) # output
+                [o, self.bra_modes[i], o.conj(), self.ket_modes[i]]
+            )
+        interleaved_inputs.append([])  # output
         return self._contract(interleaved_inputs, options=options).real
 
     def contract_state_vector(self, mps_tensors, options=None):
@@ -73,10 +78,12 @@ class MPSContractionHelper:
         for i, o in enumerate(mps_tensors):
             interleaved_inputs.extend([o, self.bra_modes[i]])
         output_modes = tuple([bra_modes[1] for bra_modes in self.bra_modes])
-        interleaved_inputs.append(output_modes) # output
+        interleaved_inputs.append(output_modes)  # output
         return self._contract(interleaved_inputs, options=options)
 
-    def contract_expectation(self, mps_tensors, operator, qubits, options=None, normalize=False):
+    def contract_expectation(
+        self, mps_tensors, operator, qubits, options=None, normalize=False
+    ):
         """
         Contract the corresponding tensor network to form the state vector representation of the MPS.
 
@@ -105,11 +112,11 @@ class MPSContractionHelper:
             if i in qubits:
                 k_modes = (k_modes[0], extra_mode, k_modes[2])
                 q = qubits.index(i)
-                operator_modes[q] = extra_mode # output modes
+                operator_modes[q] = extra_mode  # output modes
                 extra_mode += 1
             interleaved_inputs.extend([o.conj(), k_modes])
         interleaved_inputs.extend([operator, tuple(operator_modes)])
-        interleaved_inputs.append([]) # output
+        interleaved_inputs.append([])  # output
         if normalize:
             norm = self.contract_norm(mps_tensors, options=options)
         else:
@@ -117,7 +124,6 @@ class MPSContractionHelper:
         return self._contract(interleaved_inputs, options=options) / norm
 
     def _contract(self, interleaved_inputs, options=None):
-
         path = contract_path(*interleaved_inputs, options=options)[0]
 
-        return contract(*interleaved_inputs, options=options, optimize={'path':path})
+        return contract(*interleaved_inputs, options=options, optimize={"path": path})

src/qibotn/MPSUtils.py

@@ -2,33 +2,32 @@ import cupy as cp
 from cuquantum.cutensornet.experimental import contract_decompose
 from cuquantum import contract
 
+
 def initial(num_qubits, dtype):
     """
     Generate the MPS with an initial state of |00...00>
     """
-    state_tensor = cp.asarray([1, 0], dtype=dtype).reshape(1,2,1)
+    state_tensor = cp.asarray([1, 0], dtype=dtype).reshape(1, 2, 1)
     mps_tensors = [state_tensor] * num_qubits
     return mps_tensors
 
-def mps_site_right_swap(
+
-    mps_tensors, 
+def mps_site_right_swap(mps_tensors, i, **kwargs):
-    i, 
-    **kwargs
-):
     """
     Perform the swap operation between the ith and i+1th MPS tensors.
     """
     # contraction followed by QR decomposition
-    a, _, b = contract_decompose('ipj,jqk->iqj,jpk', *mps_tensors[i:i+2], algorithm=kwargs.get('algorithm',None), options=kwargs.get('options',None))
+    a, _, b = contract_decompose(
-    mps_tensors[i:i+2] = (a, b)
+        "ipj,jqk->iqj,jpk",
+        *mps_tensors[i : i + 2],
+        algorithm=kwargs.get("algorithm", None),
+        options=kwargs.get("options", None)
+    )
+    mps_tensors[i : i + 2] = (a, b)
     return mps_tensors
 
-def apply_gate(
+
-    mps_tensors, 
+def apply_gate(mps_tensors, gate, qubits, **kwargs):
-    gate, 
-    qubits, 
-    **kwargs
-):
     """
     Apply the gate operand to the MPS tensors in-place.
 
@@ -52,23 +51,31 @@ def apply_gate(
     if n_qubits == 1:
         # single-qubit gate
         i = qubits[0]
-        mps_tensors[i] = contract('ipj,qp->iqj', mps_tensors[i], gate, options=kwargs.get('options',None)) # in-place update
+        mps_tensors[i] = contract(
+            "ipj,qp->iqj", mps_tensors[i], gate, options=kwargs.get("options", None)
+        )  # in-place update
     elif n_qubits == 2:
         # two-qubit gate
         i, j = qubits
         if i > j:
             # swap qubits order
-            return apply_gate(mps_tensors, gate.transpose(1,0,3,2), (j, i), **kwargs)
+            return apply_gate(mps_tensors, gate.transpose(1, 0, 3, 2), (j, i), **kwargs)
-        elif i+1 == j:
+        elif i + 1 == j:
             # two adjacent qubits
-            a, _, b = contract_decompose('ipj,jqk,rspq->irj,jsk', *mps_tensors[i:i+2], gate, algorithm=kwargs.get('algorithm',None), options=kwargs.get('options',None))
+            a, _, b = contract_decompose(
-            mps_tensors[i:i+2] = (a, b) # in-place update
+                "ipj,jqk,rspq->irj,jsk",
+                *mps_tensors[i : i + 2],
+                gate,
+                algorithm=kwargs.get("algorithm", None),
+                options=kwargs.get("options", None)
+            )
+            mps_tensors[i : i + 2] = (a, b)  # in-place update
         else:
             # non-adjacent two-qubit gate
             # step 1: swap i with i+1
             mps_site_right_swap(mps_tensors, i, **kwargs)
             # step 2: apply gate to (i+1, j) pair. This amounts to a recursive swap until the two qubits are adjacent
-            apply_gate(mps_tensors, gate, (i+1, j), **kwargs) 
+            apply_gate(mps_tensors, gate, (i + 1, j), **kwargs)
             # step 3: swap back i and i+1
             mps_site_right_swap(mps_tensors, i, **kwargs)
     else:

src/qibotn/QiboCircuitConvertor.py

@@ -44,8 +44,7 @@ class QiboCircuitToEinsum:
         for key in qubits_frontier:
             out_list.append(qubits_frontier[key])
 
-        operand_exp_interleave = [x for y in zip(
+        operand_exp_interleave = [x for y in zip(operands, mode_labels) for x in y]
-            operands, mode_labels) for x in y]
         operand_exp_interleave.append(out_list)
         return operand_exp_interleave