Fix TQATrainer returning incorrect optimised QAOA angles (#35)

eggerdj · web-flow · commit bc035086d3f9 · 2025-11-06T09:45:06.000+01:00
diff --git a/README.md b/README.md
@@ -156,23 +156,24 @@ This repository is still in development: new functionality is being added and th
 
 ## Version tracking
 
-| Version | Added functionality                    | Pull request |
-|---------|----------------------------------------|--------------|
-|       1 | Track system information               |           #8 |
-|       2 | Add history mix-in                     |          #11 |
-|       3 | Add fidelity bounds for MPS            |           #6 |
-|       4 | Add QAOA angles functions              |          #14 |
-|       5 | Switch to qaoa_ansatz                  |          #16 |
-|       6 | Add Pauli Propagation                  |          #15 |
-|       7 | Add problem class in train             |          #17 |
-|       8 | Improve train tests                    |          #22 |
-|       9 | Add SAT map pre-processing             |          #19 |
-|      10 | Add recursive transition states trainer|          #20 |
-|      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 |
-|      15 | Remove python 3.9 support              |          #31 |
+| Version | Added functionality                                          | Pull request |
+|---------|--------------------------------------------------------------|--------------|
+|       1 | Track system information                                     |           #8 |
+|       2 | Add history mix-in                                           |          #11 |
+|       3 | Add fidelity bounds for MPS                                  |           #6 |
+|       4 | Add QAOA angles functions                                    |          #14 |
+|       5 | Switch to qaoa_ansatz                                        |          #16 |
+|       6 | Add Pauli Propagation                                        |          #15 |
+|       7 | Add problem class in train                                   |          #17 |
+|       8 | Improve train tests                                          |          #22 |
+|       9 | Add SAT map pre-processing                                   |          #19 |
+|      10 | Add recursive transition states trainer                      |          #20 |
+|      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 |
+|      15 | Remove python 3.9 support                                    |          #31 |
+|      16 | Fix TQATrainer qaoa_angles_function and returned ParamResult |          #35 |
 
 ## IBM Public Repository Disclosure
 
diff --git a/qaoa_training_pipeline/training/param_result.py b/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": 15,
+            "qaoa_training_pipeline_version": 16,
         }
 
         # Convert, e.g., np.float to float
diff --git a/qaoa_training_pipeline/training/tqa_trainer.py b/qaoa_training_pipeline/training/tqa_trainer.py
@@ -10,20 +10,61 @@
 
 from time import time
 from typing import Any, Dict, Optional
+
 import matplotlib.pyplot as plt
 import numpy as np
-
 from qiskit import QuantumCircuit
 from qiskit.quantum_info import SparsePauliOp
 from scipy.optimize import minimize
 
 from qaoa_training_pipeline.evaluation import EVALUATORS
 from qaoa_training_pipeline.evaluation.base_evaluator import BaseEvaluator
-from qaoa_training_pipeline.training.history_mixin import HistoryMixin
 from qaoa_training_pipeline.training.base_trainer import BaseTrainer
+from qaoa_training_pipeline.training.functions import BaseAnglesFunction
+from qaoa_training_pipeline.training.history_mixin import HistoryMixin
 from qaoa_training_pipeline.training.param_result import ParamResult
 
 
+class TQATrainerFunction(BaseAnglesFunction):
+    """Wrapper function around TQATrainer.tqa_schedule.
+
+    This function allows for a default ``rep`` value to be set, which is passed
+    to ``tqa_schedule``. If no value has been set yet, and no value is provided
+    to :meth:`__call__`, then an error is raised.
+    """
+
+    def __init__(self, tqa_schedule_method: callable, reps: int | None = None) -> None:
+        """Create an instance of a TQATrainer QAOA angles function.
+
+        Args:
+            tqa_schedule_method: The :meth:`TQATrainer.tqa_schedule` method for
+                an instance of :class:`TQATrainer`.
+            reps: The default reps to use if not overridden by
+                :meth:`TQATrainer.train`. If None and ``reps`` is not provided
+                in :meth:`__call__`, an error is raised. Defaults to None.
+        """
+        super().__init__()
+        self._tqa_schedule = tqa_schedule_method
+        self.reps = reps
+
+    # pylint: disable=unused-argument
+    def __call__(self, x: list, reps: int | None = None) -> list:
+        if reps is None:
+            reps = self.reps
+        if reps is None:
+            raise ValueError(
+                f"reps must be provided to {self.__class__.__name__}(reps=...) or "
+                + "set with trainer.train(..., reps=...)"
+            )
+        return self._tqa_schedule(reps=reps, dt=x[0])
+
+    # pylint: disable=unused-argument
+    @classmethod
+    def from_config(cls, config: dict) -> None:
+        """Create a TQATrainer from a config dictionary."""
+        raise RuntimeError(f"{cls.__name__} cannot be constructed from a config.")
+
+
 class TQATrainer(BaseTrainer, HistoryMixin):
     """Trotterized Quantum Annealing parameter generation.
 
@@ -33,13 +74,24 @@ class TQATrainer(BaseTrainer, HistoryMixin):
     be used as an initial point generator which is independent of problem instance and
     does not do any optimization. However, we can also use this class to perform a
     SciPy optimization of the end point of the TQA schedule.
+
+
+    .. note::
+
+        :attr:`qaoa_angles_function` for :class:`TQATrainer` requires knowledge
+        of the number of repetitions. ``reps`` can be provided by calling
+        ``qaoa_angles_function(params, reps=reps)``. If no value for ``reps`` is
+        provided, the value for the most recent call to :meth:`train` will be
+        used. If :meth:`train` has not been called yet for the instance of
+        :class:`TQATrainer`, then an error is raised.
     """
 
     def __init__(
         self,
         evaluator: Optional[BaseEvaluator] = None,
         minimize_args: Optional[Dict[str, Any]] = None,
         energy_minimization: bool = False,
+        initial_dt: float = 0.75,
     ) -> None:
         """Initialize an instance.
 
@@ -51,11 +103,18 @@ def __init__(
             energy_minimization: Allows us to switch between minimizing the energy or maximizing
                 the energy. The default and assumed convention in this repository is to
                 maximize the energy.
+            initial_dt: Initial dt if not provided to :meth:`train`. Defaults to
+                ``0.75``.
         """
-        BaseTrainer.__init__(self, evaluator)
+        BaseTrainer.__init__(
+            self,
+            evaluator,
+            qaoa_angles_function=TQATrainerFunction(self.tqa_schedule, reps=None),
+        )
         HistoryMixin.__init__(self)
 
         self._minimize_args = {"method": "COBYLA", "options": {"maxiter": 20, "rhobeg": 0.1}}
+        self.initial_dt = initial_dt
 
         minimize_args = minimize_args or {}
         self._minimize_args.update(minimize_args)
@@ -76,9 +135,17 @@ def train(
         mixer: Optional[QuantumCircuit] = None,
         initial_state: Optional[QuantumCircuit] = None,
         ansatz_circuit: Optional[QuantumCircuit] = None,
+        initial_dt: float | None = None,
     ) -> ParamResult:
         """Train the QAOA parameters."""
         self.reset_history()
+        # Set the reps attribute on the angles function, if it supports one.
+        # This allow us to override it later with a function that doesn't
+        # require reps. We set reps here so that ParamResult.from_scipy_result
+        # correctly populates "optimized_qaoa_angles" and so we can call
+        # `trainer.qaoa_angles_function()` in scripts.
+        if hasattr(self._qaoa_angles_function, "reps"):
+            self._qaoa_angles_function.reps = reps
 
         def _energy(x):
             """Optimize the energy by minimizing the negative energy.
@@ -105,21 +172,15 @@ def _energy(x):
 
         start = time()
 
+        initial_dt = initial_dt or self.initial_dt
         if self.evaluator is None:
-            tqa_dt = 0.75
-            param_result = ParamResult(
-                self.tqa_schedule(reps, dt=tqa_dt), time() - start, self, None
-            )
+            param_result = ParamResult([initial_dt], time() - start, self, None)
         else:
-            params0 = [0.75]
+            params0 = [initial_dt]
             result = minimize(_energy, params0, **self._minimize_args)
             param_result = ParamResult.from_scipy_result(
                 result, params0, time() - start, self._sign, self
             )
-            param_result["optimized_params"] = self.tqa_schedule(
-                reps, dt=param_result["optimized_params"]
-            )
-
         param_result.add_history(self)
 
         return param_result
diff --git a/test/training/test_tqa.py b/test/training/test_tqa.py
@@ -8,40 +8,116 @@
 
 """Classes to test the TQA trainer."""
 
-from test import TrainingPipelineTestCase
-
 from qiskit.quantum_info import SparsePauliOp
 
 from qaoa_training_pipeline.evaluation.mps_evaluator import MPSEvaluator
+from qaoa_training_pipeline.training.param_result import ParamResult
 from qaoa_training_pipeline.training.tqa_trainer import TQATrainer
 
+# Disable import order for this line. Python has a stdlib test module, but this
+# is our own one. Therefore, it is imported with third-party libraries.
+from test import TrainingPipelineTestCase  # pylint: disable=wrong-import-order
+
 
 class TestTQA(TrainingPipelineTestCase):
     """Class to test the TQA trainer."""
 
     def test_no_optim(self):
         """ "Test that we can run without doing any optimization."""
-        result = TQATrainer().train(None, 3)
-
-        self.assertListEqual(result["optimized_params"], [0.875, 0.625, 0.375, 0.125, 0.375, 0.625])
+        reps = 3
+        trainer = TQATrainer()
+
+        with self.assertRaises(
+            ValueError,
+            msg="Calling qaoa_angles_function without reps=... "
+            + "on untrained TQATrainer should raise an error.",
+        ):
+            _ = trainer.qaoa_angles_function([0.2])
+
+        self.assertTrue(
+            len(trainer.qaoa_angles_function([0.2], reps=reps)) == 2 * reps,
+            msg="Calling qaoa_angles_function with reps=... on untrained "
+            + "TQATrainer should return list of angles.",
+        )
+
+        result = trainer.train(None, reps)
+
+        self.assertListEqual(
+            result["optimized_qaoa_angles"],
+            [0.875, 0.625, 0.375, 0.125, 0.375, 0.625],
+            msg="Number of QAOA angles is not as expected.",
+        )
+        self.assertListEqual(
+            result["optimized_params"],
+            [0.75],
+            msg="Optimized params with default argument should be [0.75]",
+        )
 
         # Check that history is not present.
         self.assertTrue(len(result["energy_history"]) == 0)
         self.assertTrue(len(result["parameter_history"]) == 0)
         self.assertTrue(len(result["energy_evaluation_time"]) == 0)
+        # Double check that the default number of reps for qaoa_angles_function
+        # is the same as the most recent run.
+        self.assertTrue(
+            len(trainer.qaoa_angles_function(result["optimized_params"])) == 2 * reps,
+            msg="Calling qaoa_angles_function without reps=... "
+            + "on trained TQATrainer should return list of angles.",
+        )
+
+        result = trainer.train(None, reps + 1)
+        self.assertTrue(
+            len(trainer.qaoa_angles_function(result["optimized_params"])) == 2 * (reps + 1),
+            msg="Calling qaoa_angles_function without reps=... "
+            + "on trained TQATrainer should return list of angles.",
+        )
 
     def test_optim(self):
         """Test that we can optimize the dt of the TQA schedule."""
         evaluator = MPSEvaluator()
 
+        reps = 4
         trainer = TQATrainer(evaluator)
 
+        with self.assertRaises(
+            ValueError,
+            msg="Calling qaoa_angles_function without reps=... "
+            + "on untrained TQATrainer should raise an error.",
+        ):
+            _ = trainer.qaoa_angles_function([0.2])
+
+        self.assertTrue(
+            len(trainer.qaoa_angles_function([0.2], reps=reps)) == 2 * reps,
+            msg="Calling qaoa_angles_function with reps=... "
+            + "on untrained TQATrainer should return list of angles.",
+        )
+
         cost_op = SparsePauliOp.from_list([("ZIIZ", -1), ("IZIZ", -1), ("IIZZ", -1)])
 
-        result = trainer.train(cost_op, reps=4)
+        result: ParamResult = trainer.train(cost_op, reps=reps)
 
         self.assertEqual(result["success"], "True")
-        self.assertEqual(len(result["optimized_params"]), 8)
+        self.assertEqual(
+            len(result["optimized_params"]),
+            1,
+            msg="There is only one parameter, dt, for TQATrainer.",
+        )
+        self.assertEqual(
+            len(result["optimized_qaoa_angles"]),
+            2 * reps,
+            msg="Number of QAOA angles is not as expected.",
+        )
+        self.assertTrue(
+            len(trainer.qaoa_angles_function(result["optimized_params"])) == 2 * reps,
+            msg="Calling qaoa_angles_function without reps=... "
+            + "on trained TQATrainer should return list of angles.",
+        )
+        self.assertListEqual(
+            result["optimized_qaoa_angles"],
+            trainer.qaoa_angles_function(result["optimized_params"]),
+            msg="Calling qaoa_angles_function without reps=... "
+            + "on trained TQATrainer should return the same angles.",
+        )
 
         # Check that history is present.
         self.assertTrue(len(result["energy_history"]) > 0)
@@ -55,7 +131,10 @@ def test_from_config(self):
         trainer = TQATrainer.from_config(config)
         self.assertIsNone(trainer.evaluator)
 
-        config = {"evaluator": "MPSEvaluator", "evaluator_init": {"bond_dim_circuit": 2}}
+        config = {
+            "evaluator": "MPSEvaluator",
+            "evaluator_init": {"bond_dim_circuit": 2},
+        }
 
         trainer = TQATrainer.from_config(config)
         self.assertTrue(isinstance(trainer.evaluator, MPSEvaluator))

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ def __init__(`
`47`	`47`	`"system": platform.system(),`
`48`	`48`	`"processor": platform.processor(),`
`49`	`49`	`"platform": platform.platform(),`
`50`		`- "qaoa_training_pipeline_version": 15,`
	`50`	`+ "qaoa_training_pipeline_version": 16,`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`# Convert, e.g., np.float to float`