Merge branch 'qiskit-community:main' into update-sphinx-theme

qiskit-community · Oct 2, 2023 · 94bc31c · 94bc31c
2 parents b21fc9a + bf8510d
commit 94bc31c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/qiskit_research/protein_folding/protein_folding_result.py b/qiskit_research/protein_folding/protein_folding_result.py
@@ -21,7 +21,7 @@
 
 if _optionals.HAS_MATPLOTLIB:
     # pylint: disable=import-error,unused-import
-    from matplotlib.pyplot import figure
+    from matplotlib.figure import Figure
 
 
 class ProteinFoldingResult:
@@ -168,7 +168,7 @@ def save_xyz_file(
     @_optionals.HAS_MATPLOTLIB.require_in_call
     def get_figure(
         self, title: str = "Protein Structure", ticks: bool = False, grid: bool = False
-    ) -> "figure":
+    ) -> Figure:
         """
         Generates a figure of the molecule in 3D.
 

diff --git a/qiskit_research/protein_folding/utils/protein_plotter.py b/qiskit_research/protein_folding/utils/protein_plotter.py
@@ -11,18 +11,20 @@
 """An auxiliary class that plots aminoacids of a molecule
  in a ProteinFoldingResult."""
 from typing import Optional
+
 import numpy as np
 from qiskit.utils import optionals as _optionals
 
-
 from qiskit_research.protein_folding.utils.protein_shape_file_gen import (
     ProteinShapeFileGen,
 )
 
 if _optionals.HAS_MATPLOTLIB:
     # pylint: disable=import-error,unused-import
-    from matplotlib.pyplot import figure
     from matplotlib.axes import Axes
+    from matplotlib.figure import Figure
+    from mpl_toolkits.mplot3d.axes3d import Axes3D
+
 # pylint: disable=too-few-public-methods
 
 
@@ -53,7 +55,7 @@ def __init__(self, shape_gen: ProteinShapeFileGen) -> None:
         )
 
         self._fig = plt.figure()
-        self._ax_graph = self._fig.add_subplot(projection="3d")
+        self._ax_graph: Axes3D = self._fig.add_subplot(projection="3d")
 
     def _draw_main_chain(self):
         """
@@ -146,7 +148,7 @@ def _format_graph(
 
     def get_figure(
         self, title: str = "Protein Structure", ticks: bool = False, grid: bool = False
-    ) -> "figure":
+    ) -> Figure:
         """
         Plots the molecule in 3D.