added plot_probabilities method in result class

openqaoa/optimizers/result.py

@@ -143,3 +143,53 @@ def plot_cost(self, figsize=(10, 8), label='Cost', linestyle='--', color='b', ax
         ax.set_title('Cost history')
 
         return
+
+    def plot_probabilities(self, figsize = (10,8),label='Probability distribution',color='tab:blue', ax=None):
+
+        """
+        Helper function to plot the probabilities corresponding to each basis states (with prob != 0) obtained from the optimized result
+
+        Parameters
+        ----------
+        figsize: `tuple`
+            The size of the figure to be plotted. Defaults to (10,8).
+        label: `str`
+            The label of the cost line, defaults to 'Probability distribution'.
+        color: `str`
+            The color of the line. Defaults to 'tab:blue'.
+        ax: 'matplotlib.axes._subplots.AxesSubplot'
+            Axis on which to plot the graph. Deafults to None
+        """
+
+        outcome = self.optimized['optimized measurement outcomes']
+
+        # converting to counts dictionary if outcome is statevector
+        if type(outcome) == type(np.array([])):
+            outcome = self.get_counts(outcome)
+            # setting norm to 1 since it might differ slightly for statevectors due to numerical preicision
+            norm = np.float64(1)
+        else: 
+            # needed to be able to divide the tuple by 'norm'
+            norm = np.float64(sum(outcome.values())) 
+
+        # sorting dictionary by increasing binary value (i.e. keys) and unzipping the states and counts
+        outcome_list = sorted(outcome.items())
+        states, counts = zip(*outcome_list)
+
+        # normalizing to obtain probabilities
+        probs = counts/norm
+
+        # formatting labels
+        labels = [r'$\left|{}\right>$'.format(state) for state in states]
+
+        if ax is None:
+            fig, ax = plt.subplots(figsize=figsize)
+
+        ax.bar(labels,probs, color=color)
+        ax.set_xlabel('Eigen-State')
+        ax.set_ylabel('Probability')
+        ax.set_title(label)
+        ax.tick_params(axis='x', labelrotation = 70)
+        ax.grid(True, axis='y', linestyle='--')
+
+        return

tests/test_result_object.py

@@ -8,6 +8,7 @@
 from openqaoa.optimizers.result import most_probable_bitstring
 from openqaoa.utilities import qaoa_probabilities
 from openqaoa.workflows.optimizer import QAOA  
+from openqaoa.devices import create_device
 
 class TestingLoggerClass(unittest.TestCase):
 
@@ -84,6 +85,34 @@ def test_plot_cost(self):
 
         q.results.plot_cost()
 
+    def test_plot_probabilities(self):
+
+        # Create the problem
+        g = nx.circulant_graph(6, [1])
+        vc = MinimumVertexCover(g, field =1.0, penalty=10).get_qubo_problem()
+
+        # first for state-vector based simulator:
+        q_sv = QAOA()
+        q_sv.set_circuit_properties(p=3, init_type='ramp')
+        q_sv.compile(vc)
+
+        q_sv.optimize()
+
+        q_sv.results.plot_probabilities()
+
+        # then for shot based simulator:
+        q_shot = QAOA()
+        q_shot_dev = create_device(location='local',name='qiskit.shot_simulator')
+        q_shot.set_device(q_shot_dev)
+
+        q_shot.compile(vc)
+        q_shot.optimize()
+
+        q_shot.results.plot_probabilities()
+
+
+
+
     def test_get_counts(self):
 
         # measurement outcome from a statevector_simulator backend