From a97c962f02d52b153ee0e081ff7c6b8062555434 Mon Sep 17 00:00:00 2001
From: Simon Cross <hodgestar@gmail.com>
Date: Thu, 22 Jul 2021 11:37:32 +0200
Subject: [PATCH] More carefully track the state of the matplotlib figure and
 axes used by Bloch sphere.

Previously .render(...) took fig and axes arguments and overwrote
self.fig and self.axes if the arguments were None. The actual values of
the arguments were completely discarded (!!). In addition, not supplying
a value for fig and axes would overwrite the values passed to the
constructor as Bloch(..., fig=fig, axes=axes).

The new behaviour is as follows:

* If an external figure is passed to the constructor it is retained.

* If no figure is supplied, .render() will create a figure if one
  has not been created yet, or if the figure has been closed, will
  recreate it.

* .render() will always call self.fig.canvas.update() to ensure that
  any existing figure that is already shown will be updated.

Previously it was possible to cause .render() to raise obscure
exceptions by, e.g., calling .render(), closing the figure, calling
.render() again. See, for example,
https://github.com/qutip/qutip/issues/1616.

The external figure tracking was backported from Qiskit's copy of
Bloch and was inspired by PR
https://github.com/Qiskit/qiskit-terra/pull/1091/.
---
 qutip/bloch.py | 36 ++++++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/qutip/bloch.py b/qutip/bloch.py
index e29908b976..d06bf5bd4d 100644
--- a/qutip/bloch.py
+++ b/qutip/bloch.py
@@ -111,7 +111,13 @@ class Bloch:
     def __init__(self, fig=None, axes=None, view=None, figsize=None,
                  background=False):
         # Figure and axes
+        self._ext_fig = False
+        if fig is not None:
+            self._ext_fig = True
         self.fig = fig
+        self._ext_axes = False
+        if axes is not None:
+            self._ext_axes = True
         self.axes = axes
         # Background axes, default = False
         self.background = background
@@ -178,8 +184,6 @@ def __init__(self, fig=None, axes=None, view=None, figsize=None,
         # Style of points, 'm' for multiple colors, 's' for single color
         self.point_style = []
 
-        # status of rendering
-        self._rendered = False
         # status of showing
         if fig is None:
             self._shown = False
@@ -405,7 +409,7 @@ def make_sphere(self):
         """
         Plots Bloch sphere and data sets.
         """
-        self.render(self.fig, self.axes)
+        self.render()
 
     def run_from_ipython(self):
         try:
@@ -414,20 +418,21 @@ def run_from_ipython(self):
         except NameError:
             return False
 
-    def render(self, fig=None, axes=None):
+    def render(self):
         """
         Render the Bloch sphere and its data sets in on given figure and axes.
         """
-        if self._rendered:
-            self.axes.clear()
-
-        self._rendered = True
-
-        # Figure instance for Bloch sphere plot
-        if not fig:
+        if not self._ext_fig:
+            if self.fig is not None and not plt.fignum_exists(self.fig.number):
+                # a previous call to .rander() created the figure, but the
+                # figure not longer exists so we trigger recreating it:
+                self.fig = None
+                self.axes = None
+
+        if not self._ext_fig and self.fig is None:
             self.fig = plt.figure(figsize=self.figsize)
 
-        if not axes:
+        if not self._ext_axes and self.axes is None:
             self.axes = _axes3D(self.fig, azim=self.view[0],
                                 elev=self.view[1])
 
@@ -437,6 +442,7 @@ def render(self, fig=None, axes=None):
             self.axes.set_ylim3d(-1.3, 1.3)
             self.axes.set_zlim3d(-1.3, 1.3)
         else:
+            self.axes.clear()
             self.plot_axes()
             self.axes.set_axis_off()
             self.axes.set_xlim3d(-0.7, 0.7)
@@ -454,6 +460,8 @@ def render(self, fig=None, axes=None):
         self.plot_front()
         self.plot_axes_labels()
         self.plot_annotations()
+        # Trigger an update of the Bloch sphere if it is already shown:
+        self.fig.canvas.draw()
 
     def plot_back(self):
         # back half of sphere
@@ -623,7 +631,7 @@ def show(self):
         """
         Display Bloch sphere and corresponding data sets.
         """
-        self.render(self.fig, self.axes)
+        self.render()
         if self.run_from_ipython():
             if self._shown:
                 display(self.fig)
@@ -653,7 +661,7 @@ def save(self, name=None, format='png', dirc=None, dpin=None):
         File containing plot of Bloch sphere.
 
         """
-        self.render(self.fig, self.axes)
+        self.render()
         # Conditional variable for first argument to savefig
         # that is set in subsequent if-elses
         complete_path = ""