Format with Black

src/qibotn/MPSContractionHelper.py

@@ -1,70 +1,75 @@
 from cuquantum import contract, contract_path, CircuitToEinsum, tensor
 
+
 class MPSContractionHelper:
     """
     A helper class to compute various quantities for a given MPS.
-    
-    Interleaved format is used to construct the input args for `cuquantum.contract`. 
+
+    Interleaved format is used to construct the input args for `cuquantum.contract`.
     A concrete example on how the modes are populated for a 7-site MPS is provided below:
-    
-          0     2     4     6     8    10     12    14        
+
+          0     2     4     6     8    10     12    14
     bra -----A-----B-----C-----D-----E-----F-----G-----
-             |     |     |     |     |     |     |     
-            1|    3|    5|    7|    9|   11|   13|     
-             |     |     |     |     |     |     |     
+             |     |     |     |     |     |     |
+            1|    3|    5|    7|    9|   11|   13|
+             |     |     |     |     |     |     |
     ket -----a-----b-----c-----d-----e-----f-----g-----
           15    16    17    18    19    20    21    22
-    
-    
+
+
     The follwing compute quantities are supported:
-    
+
         - the norm of the MPS.
         - the equivalent state vector from the MPS.
         - the expectation value for a given operator.
         - the equivalent state vector after multiplying an MPO to an MPS.
-    
-    Note that for the nth MPS tensor (rank-3), the modes of the tensor are expected to be `(i,p,j)` 
-    where i denotes the bonding mode with the (n-1)th tensor, p denotes the physical mode for the qubit and 
+
+    Note that for the nth MPS tensor (rank-3), the modes of the tensor are expected to be `(i,p,j)`
+    where i denotes the bonding mode with the (n-1)th tensor, p denotes the physical mode for the qubit and
     j denotes the bonding mode with the (n+1)th tensor.
-    
+
     Args:
         num_qubits: The number of qubits for the MPS.
     """
-    
+
     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)]
-        offset = 2*num_qubits+1
-        self.ket_modes = [(i+offset, 2*i+1, i+1+offset) 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
+        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):
         """
         Contract the corresponding tensor network to form the norm of the MPS.
 
         Args:
-            mps_tensors: A list of rank-3 ndarray-like tensor objects. 
-                The indices of the ith tensor are expected to be bonding index to the i-1 tensor, 
+            mps_tensors: A list of rank-3 ndarray-like tensor objects.
+                The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
                 the physical mode, and then the bonding index to the i+1th tensor.
-            options: Specify the contract and decompose options. 
+            options: Specify the contract and decompose options.
 
         Returns:
             The norm of the MPS.
         """
         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.append([]) # output
+            interleaved_inputs.extend(
+                [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):
         """
         Contract the corresponding tensor network to form the state vector representation of the MPS.
 
         Args:
-            mps_tensors: A list of rank-3 ndarray-like tensor objects. 
-                The indices of the ith tensor are expected to be bonding index to the i-1 tensor, 
+            mps_tensors: A list of rank-3 ndarray-like tensor objects.
+                The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
                 the physical mode, and then the bonding index to the i+1th tensor.
-            options: Specify the contract and decompose options. 
+            options: Specify the contract and decompose options.
 
         Returns:
             An ndarray-like object as the state vector.
@@ -73,28 +78,30 @@ def contract_state_vector(self, mps_tensors, options=None):
         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.
 
         Args:
-            mps_tensors: A list of rank-3 ndarray-like tensor objects. 
-                The indices of the ith tensor are expected to be bonding index to the i-1 tensor, 
+            mps_tensors: A list of rank-3 ndarray-like tensor objects.
+                The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
                 the physical mode, and then the bonding index to the i+1th tensor.
-            operator: A ndarray-like tensor object. 
-                The modes of the operator are expected to be output qubits followed by input qubits, e.g, 
-                ``A, B, a, b`` where `a, b` denotes the inputs and `A, B'` denotes the outputs. 
-            qubits: A sequence of integers specifying the qubits that the operator is acting on. 
-            options: Specify the contract and decompose options. 
+            operator: A ndarray-like tensor object.
+                The modes of the operator are expected to be output qubits followed by input qubits, e.g,
+                ``A, B, a, b`` where `a, b` denotes the inputs and `A, B'` denotes the outputs.
+            qubits: A sequence of integers specifying the qubits that the operator is acting on.
+            options: Specify the contract and decompose options.
             normalize: Whether to scale the expectation value by the normalization factor.
 
         Returns:
             An ndarray-like object as the state vector.
         """
-        
+
         interleaved_inputs = []
         extra_mode = 3 * self.num_qubits + 2
         operator_modes = [None] * len(qubits) + [self.bra_modes[q][1] for q in qubits]
@@ -105,19 +112,18 @@ def contract_expectation(self, mps_tensors, operator, qubits, options=None, norm
             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:
             norm = 1
         return self._contract(interleaved_inputs, options=options) / norm
-
-    def _contract(self, interleaved_inputs, options=None):
 
+    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,73 +2,80 @@
 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> 
+    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, 
-    i, 
-    **kwargs
-):
+
+def mps_site_right_swap(mps_tensors, 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))
-    mps_tensors[i:i+2] = (a, b)
+    a, _, b = contract_decompose(
+        "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, 
-    gate, 
-    qubits, 
-    **kwargs
-):
+
+def apply_gate(mps_tensors, gate, qubits, **kwargs):
     """
     Apply the gate operand to the MPS tensors in-place.
-    
+
     Args:
-        mps_tensors: A list of rank-3 ndarray-like tensor objects. 
-            The indices of the ith tensor are expected to be the bonding index to the i-1 tensor, 
+        mps_tensors: A list of rank-3 ndarray-like tensor objects.
+            The indices of the ith tensor are expected to be the bonding index to the i-1 tensor,
             the physical mode, and then the bonding index to the i+1th tensor.
-        gate: A ndarray-like tensor object representing the gate operand. 
-            The modes of the gate is expected to be output qubits followed by input qubits, e.g, 
-            ``A, B, a, b`` where ``a, b`` denotes the inputs and ``A, B`` denotes the outputs. 
+        gate: A ndarray-like tensor object representing the gate operand.
+            The modes of the gate is expected to be output qubits followed by input qubits, e.g,
+            ``A, B, a, b`` where ``a, b`` denotes the inputs and ``A, B`` denotes the outputs.
         qubits: A sequence of integers denoting the qubits that the gate is applied onto.
-        algorithm: The contract and decompose algorithm to use for gate application. 
+        algorithm: The contract and decompose algorithm to use for gate application.
             Can be either a `dict` or a `ContractDecomposeAlgorithm`.
-        options: Specify the contract and decompose options. 
-    
+        options: Specify the contract and decompose options.
+
     Returns:
         The updated MPS tensors.
     """
-    
+
     n_qubits = len(qubits)
     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)
-        elif i+1 == j:
+            return apply_gate(mps_tensors, gate.transpose(1, 0, 3, 2), (j, i), **kwargs)
+        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))
-            mps_tensors[i:i+2] = (a, b) # in-place update
+            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)
+            )
+            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 @@ def state_vector_operands(self):
         for key in qubits_frontier:
             out_list.append(qubits_frontier[key])
 
-        operand_exp_interleave = [x for y in zip(
-            operands, mode_labels) for x in y]
+        operand_exp_interleave = [x for y in zip(operands, mode_labels) for x in y]
         operand_exp_interleave.append(out_list)
         return operand_exp_interleave