Add custom ansatz op to state vector evaluator (#29)

Author: eggerdj

README.md

@@ -171,6 +171,7 @@ This repository is still in development: new functionality is being added and th
 |      11 | Bug fix in train pre-processing data   |          #24 |
 |      12 | More data in result saving in train.py |          #26 |
 |      13 | Create PPEvaluator from configs        |          #25 |
+|      14 | Custom ansatz operator to state vector |          #29 |
 
 ## IBM Public Repository Disclosure
 

qaoa_training_pipeline/evaluation/statevector_evaluator.py

@@ -8,7 +8,7 @@
 
 """Statevector-based QAOA evaluator."""
 
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from qiskit import QuantumCircuit
 from qiskit.circuit.library import qaoa_ansatz
@@ -43,7 +43,7 @@ def evaluate(
         params: List[float],
         mixer: Optional[QuantumCircuit] = None,
         initial_state: Optional[QuantumCircuit] = None,
-        ansatz_circuit: Optional[QuantumCircuit] = None,
+        ansatz_circuit: Optional[Union[QuantumCircuit, SparsePauliOp]] = None,
     ) -> float:
         """Evaluate the expectation value of a cost operator given a set of parameters.
 
@@ -56,16 +56,27 @@ def evaluate(
             initial_state (Optional[QuantumCircuit], optional): the initial state of the QAOA.
         """
 
-        if ansatz_circuit is not None:
-            raise NotImplementedError("Custom ansatz circuits are not yet supported.")
+        if isinstance(ansatz_circuit, SparsePauliOp):
+            ansatz_op = ansatz_circuit
+        elif ansatz_circuit is None:
+            ansatz_op = cost_op
+        else:
+            raise NotImplementedError(
+                "Custom ansatz circuits in format"
+                f"{ansatz_circuit.__class__.__name__} are not yet supported."
+            )
 
         circuit = qaoa_ansatz(
-            cost_op,
+            ansatz_op,
             reps=len(params) // 2,
             mixer_operator=mixer,
             initial_state=initial_state,
         )
 
+        # Prevents the edge case where the cost op is the identity.
+        if len(circuit.parameters) != len(params):
+            raise ValueError("The QAOA Circuit does not have the correct number of parameters. ")
+
         result = self.primitive.run([(circuit, cost_op, params)]).result()
 
         return float(result[0].data.evs)

qaoa_training_pipeline/training/param_result.py

@@ -47,7 +47,7 @@ def __init__(
             "system": platform.system(),
             "processor": platform.processor(),
             "platform": platform.platform(),
-            "qaoa_training_pipeline_version": 13,
+            "qaoa_training_pipeline_version": 14,
         }
 
         # Convert, e.g., np.float to float

test/evaluation/test_statevector.py

@@ -48,3 +48,19 @@ def test_optimize(self):
         result = trainer.train(cost_op=self.cost_op, params0=[0.2, 0.3])
 
         self.assertGreaterEqual(len(result["energy_history"]), 3)
+
+    def test_custom_ansatz(self):
+        """Test that we can construct the ansatz from a different operator."""
+        ansatz_op = SparsePauliOp.from_list([("ZI", 1)])
+
+        angles = [1.2, 1.3]
+
+        energy1 = self.evaluator.evaluate(self.cost_op, params=angles, ansatz_circuit=ansatz_op)
+        energy2 = self.evaluator.evaluate(self.cost_op, params=angles)
+
+        self.assertTrue(abs(energy1 - energy2) > 0.1)
+
+        energy1 = self.evaluator.evaluate(self.cost_op, params=angles, ansatz_circuit=self.cost_op)
+        energy2 = self.evaluator.evaluate(self.cost_op, params=angles)
+
+        self.assertAlmostEqual(energy1, energy2, places=6)