Allow user to specify Pauli string pattern for expecation calculation…

… [skip CI]

src/qibotn/backends/gpu.py

@@ -12,12 +12,35 @@ def __init__(self, runcard):
         super().__init__()
         import cuquantum  # pylint: disable=import-error
         from cuquantum import cutensornet as cutn  # pylint: disable=import-error
-
+
+        self.pauli_string_pattern = "XXXZ"
         if runcard is not None:
             self.MPI_enabled = runcard.get("MPI_enabled", False)
             self.MPS_enabled = runcard.get("MPS_enabled", False)
             self.NCCL_enabled = runcard.get("NCCL_enabled", False)
-            self.expectation_enabled = runcard.get("expectation_enabled", False)
+
+            expectation_enabled_value = runcard.get('expectation_enabled')
+
+            if expectation_enabled_value is True:
+                self.expectation_enabled = True
+
+                print("expectation_enabled is",self.expectation_enabled)
+            elif expectation_enabled_value is False:
+                self.expectation_enabled = False
+
+                print("expectation_enabled is",self.expectation_enabled)
+            elif isinstance(expectation_enabled_value, dict):
+                self.expectation_enabled = True
+                expectation_enabled_dict = runcard.get('expectation_enabled', {})
+
+                self.pauli_string_pattern = expectation_enabled_dict.get('pauli_string_pattern', None)
+
+                print("expectation_enabled is a dictionary",self.expectation_enabled,self.pauli_string_pattern )
+            else:
+                raise TypeError("expectation_enabled has an unexpected type")
+
+
+
         else:
             self.MPI_enabled = False
             self.MPS_enabled = False
@@ -144,7 +167,7 @@ def execute_circuit(
             if initial_state is not None:
                 raise_error(NotImplementedError, "QiboTN cannot support initial state.")
 
-            state = eval.expectation_pauli_tn(circuit, self.dtype)
+            state = eval.expectation_pauli_tn(circuit, self.dtype, self.pauli_string_pattern)
 
         elif (
             self.MPI_enabled == True
@@ -155,7 +178,7 @@ def execute_circuit(
             if initial_state is not None:
                 raise_error(NotImplementedError, "QiboTN cannot support initial state.")
 
-            state, rank = eval.expectation_pauli_tn_MPI(circuit, self.dtype, 32)
+            state, rank = eval.expectation_pauli_tn_MPI(circuit, self.dtype, self.pauli_string_pattern, 32)
 
             if rank > 0:
                 state = np.array(0)
@@ -169,7 +192,7 @@ def execute_circuit(
             if initial_state is not None:
                 raise_error(NotImplementedError, "QiboTN cannot support initial state.")
 
-            state, rank = eval.expectation_pauli_tn_nccl(circuit, self.dtype, 32)
+            state, rank = eval.expectation_pauli_tn_nccl(circuit, self.dtype, self.pauli_string_pattern, 32)
 
             if rank > 0:
                 state = np.array(0)

src/qibotn/eval.py

@@ -14,10 +14,10 @@ def dense_vector_tn(qibo_circ, datatype):
     return contract(*myconvertor.state_vector_operands())
 
 
-def expectation_pauli_tn(qibo_circ, datatype):
+def expectation_pauli_tn(qibo_circ, datatype, pauli_string):
     myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
     return contract(
-        *myconvertor.expectation_operands(PauliStringGen(qibo_circ.nqubits))
+        *myconvertor.expectation_operands(PauliStringGen(qibo_circ.nqubits, pauli_string))
     )
 
 
@@ -203,7 +203,7 @@ def dense_vector_tn_nccl(qibo_circ, datatype, n_samples=8):
     return result, rank
 
 
-def expectation_pauli_tn_nccl(qibo_circ, datatype, n_samples=8):
+def expectation_pauli_tn_nccl(qibo_circ, datatype, pauli_string, n_samples=8):
     from mpi4py import MPI  # this line initializes MPI
     import socket
     from cuquantum import Network
@@ -231,7 +231,7 @@ def expectation_pauli_tn_nccl(qibo_circ, datatype, n_samples=8):
     myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
     # mem_avail = cp.cuda.Device().mem_info[0]
     # print("Mem avail: aft convetor",mem_avail, "rank =",rank)
-    operands = myconvertor.expectation_operands(PauliStringGen(qibo_circ.nqubits))
+    operands = myconvertor.expectation_operands(PauliStringGen(qibo_circ.nqubits, pauli_string))
 
     # mem_avail = cp.cuda.Device().mem_info[0]
     # print("Mem avail: aft operand interleave",mem_avail, "rank =",rank)
@@ -290,7 +290,7 @@ def expectation_pauli_tn_nccl(qibo_circ, datatype, n_samples=8):
     return result, rank
 
 
-def expectation_pauli_tn_MPI(qibo_circ, datatype, n_samples=8):
+def expectation_pauli_tn_MPI(qibo_circ, datatype, pauli_string, n_samples=8):
     from mpi4py import MPI  # this line initializes MPI
     import socket
     from cuquantum import Network
@@ -310,7 +310,7 @@ def expectation_pauli_tn_MPI(qibo_circ, datatype, n_samples=8):
     myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
     # mem_avail = cp.cuda.Device().mem_info[0]
     # print("Mem avail: aft convetor",mem_avail, "rank =",rank)
-    operands = myconvertor.expectation_operands(PauliStringGen(qibo_circ.nqubits))
+    operands = myconvertor.expectation_operands(PauliStringGen(qibo_circ.nqubits, pauli_string))
     # mem_avail = cp.cuda.Device().mem_info[0]
     # print("Mem avail: aft operand interleave",mem_avail, "rank =",rank)
 
@@ -378,17 +378,17 @@ def dense_vector_mps(qibo_circ, gate_algo, datatype):
     )
 
 
-def PauliStringGen(nqubits):
+def PauliStringGen(nqubits, pauli_string):
     if nqubits <= 0:
         return "Invalid input. N should be a positive integer."
 
-    # characters = 'IXYZ'
-    characters = "XXXZ"
+    characters = pauli_string
+    #characters = "XXXZ"
 
     result = ""
 
     for i in range(nqubits):
         char_to_add = characters[i % len(characters)]
         result += char_to_add
-
+    print("pauli string", result)
     return result