update sample_most_likely

qiskit_optimization/applications/optimization_application.py

@@ -11,13 +11,15 @@
 # that they have been altered from the originals.
 
 """An abstract class for optimization application classes."""
-from typing import Union, Dict
-from collections import OrderedDict
 from abc import ABC, abstractmethod
+from collections import OrderedDict
+from typing import Dict, Union
 
 import numpy as np
-
 from qiskit.opflow import StateFn
+from qiskit.quantum_info import Statevector
+from qiskit.result import QuasiDistribution
+
 from qiskit_optimization.algorithms import OptimizationResult
 from qiskit_optimization.problems.quadratic_program import QuadraticProgram
 
@@ -59,29 +61,54 @@ def _result_to_x(self, result: Union[OptimizationResult, np.ndarray]) -> np.ndar
         return x
 
     @staticmethod
-    def sample_most_likely(state_vector: Union[np.ndarray, Dict]) -> np.ndarray:
+    def sample_most_likely(
+        state_vector: Union[QuasiDistribution, Statevector, np.ndarray, Dict]
+    ) -> np.ndarray:
         """Compute the most likely binary string from state vector.
 
         Args:
-            state_vector: state vector or counts.
+            state_vector: state vector or counts or quasi-probabilities.
 
         Returns:
             binary string as numpy.ndarray of ints.
+
+        Raises:
+            ValueError: if state_vector is not QuasiDistribution, Statevector,
+                np.ndarray, or dict.
         """
-        if isinstance(state_vector, (OrderedDict, dict)):
+        if isinstance(state_vector, QuasiDistribution):
+            probabilities = state_vector.binary_probabilities()
+            binary_string = max(probabilities.items(), key=lambda kv: kv[1])[0]
+            x = np.asarray([int(y) for y in reversed(list(binary_string))])
+            return x
+        elif isinstance(state_vector, Statevector):
+            probabilities = state_vector.probabilities()
+            n = state_vector.num_qubits
+            k = np.argmax(np.abs(probabilities))
+            x = np.zeros(n)
+            for i in range(n):
+                x[i] = k % 2
+                k >>= 1
+            return x
+        elif isinstance(state_vector, (OrderedDict, dict)):
             # get the binary string with the largest count
-            binary_string = sorted(state_vector.items(), key=lambda kv: kv[1])[-1][0]
+            binary_string = max(state_vector.items(), key=lambda kv: kv[1])[0]
             x = np.asarray([int(y) for y in reversed(list(binary_string))])
             return x
         elif isinstance(state_vector, StateFn):
             binary_string = list(state_vector.sample().keys())[0]
             x = np.asarray([int(y) for y in reversed(list(binary_string))])
             return x
-        else:
+        elif isinstance(state_vector, np.ndarray):
             n = int(np.log2(state_vector.shape[0]))
             k = np.argmax(np.abs(state_vector))
             x = np.zeros(n)
             for i in range(n):
                 x[i] = k % 2
                 k >>= 1
             return x
+        else:
+            raise ValueError(
+                "state vector should be QuasiDistribution, Statevector, ndarray, or dict. "
+                f"But it is {type(state_vector)}."
+            )

test/applications/test_optimization_application.py

@@ -0,0 +1,44 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2022.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Test OptimizationApplication class"""
+
+import unittest
+from test.optimization_test_case import QiskitOptimizationTestCase
+
+import numpy as np
+from ddt import data, ddt
+from qiskit.opflow import StateFn
+from qiskit.result import QuasiDistribution
+
+from qiskit_optimization.applications import OptimizationApplication
+
+
+@ddt
+class TestOptimizationApplication(QiskitOptimizationTestCase):
+    """Test OptimizationApplication class"""
+
+    @data(
+        np.array([0, 0, 1, 0]),
+        StateFn([0, 0, 1, 0]),
+        {"10": 0.8, "01": 0.2},
+        QuasiDistribution({"10": 0.8, "01": 0.2}),
+    )
+    def test_sample_most_likely(self, state_vector):
+        """Test sample_most_likely"""
+
+        result = OptimizationApplication.sample_most_likely(state_vector)
+        np.testing.assert_allclose(result, [0, 1])
+
+
+if __name__ == "__main__":
+    unittest.main()