REFACTOR: Extend consistency on plot methods (#4723)

Co-authored-by: Devin <[email protected]>

_unittest/test_12_PostProcessing.py

@@ -2,7 +2,9 @@
 import sys
 
 from _unittest.conftest import config
+from matplotlib.figure import Figure
 import pytest
+from pyvista.plotting.plotter import Plotter
 
 from pyaedt import Circuit
 from pyaedt import Icepak
@@ -597,18 +599,18 @@ def test_70_far_field_data(self):
         assert ffdata.origin == [0, 0, 1]
 
         img1 = os.path.join(self.local_scratch.path, "ff_2d1.jpg")
-        ffdata.plot_2d_cut(primary_sweep="Theta", secondary_sweep_value="all", image_path=img1)
+        ffdata.plot_2d_cut(primary_sweep="Theta", secondary_sweep_value="all", image_path=img1, show=False)
         assert os.path.exists(img1)
         img2 = os.path.join(self.local_scratch.path, "ff_2d2.jpg")
-        ffdata.plot_2d_cut(secondary_sweep_value=[0, 1], image_path=img2)
+        ffdata.plot_2d_cut(secondary_sweep_value=[0, 1], image_path=img2, show=False)
         assert os.path.exists(img2)
         img3 = os.path.join(self.local_scratch.path, "ff_2d2.jpg")
-        ffdata.plot_2d_cut(image_path=img3)
+        ffdata.plot_2d_cut(image_path=img3, show=False)
         assert os.path.exists(img3)
         curve_2d = ffdata.plot_2d_cut(show=False)
-        assert len(curve_2d[0]) == 3
+        assert isinstance(curve_2d, Figure)
         data = ffdata.polar_plot_3d(show=False)
-        assert len(data) == 3
+        assert isinstance(data, Figure)
 
         img4 = os.path.join(self.local_scratch.path, "ff_3d1.jpg")
         ffdata.polar_plot_3d_pyvista(
@@ -623,7 +625,7 @@ def test_70_far_field_data(self):
         data_pyvista = ffdata.polar_plot_3d_pyvista(
             quantity="RealizedGain", show=False, background=[255, 0, 0], show_geometry=False, convert_to_db=True
         )
-        assert data_pyvista
+        assert isinstance(data_pyvista, Plotter)
 
     @pytest.mark.skipif(is_linux or sys.version_info < (3, 8), reason="FarFieldSolution not supported by IronPython")
     def test_71_antenna_plot(self, field_test):

pyaedt/generic/plot.py

@@ -385,7 +385,7 @@ def plot_polar_chart(
 
 @pyaedt_function_handler()
 @update_plot_settings
-def plot_3d_chart(plot_data, size=(2000, 1000), xlabel="", ylabel="", title="", snapshot_path=None):
+def plot_3d_chart(plot_data, size=(2000, 1000), xlabel="", ylabel="", title="", snapshot_path=None, show=True):
     """Create a Matplotlib 3D plot based on a list of data.
 
     Parameters
@@ -439,8 +439,11 @@ def plot_3d_chart(plot_data, size=(2000, 1000), xlabel="", ylabel="", title="",
 
 @pyaedt_function_handler()
 @update_plot_settings
-def plot_2d_chart(plot_data, size=(2000, 1000), show_legend=True, xlabel="", ylabel="", title="", snapshot_path=None):
-    """Create a Matplotlib plot based on a list of data.
+def plot_2d_chart(
+    plot_data, size=(2000, 1000), show_legend=True, xlabel="", ylabel="", title="", snapshot_path=None, show=True
+):
+    """Create a Matplotlib figure based on a list of data.
+
     Parameters
     ----------
     plot_data : list of list
@@ -470,10 +473,6 @@ def plot_2d_chart(plot_data, size=(2000, 1000), show_legend=True, xlabel="", yla
     fig, ax = plt.subplots(figsize=figsize)
     label_id = 1
     for plo_obj in plot_data:
-        if len(plo_obj) == 3:
-            label = plo_obj[2]
-        else:
-            label = "Trace " + str(label_id)
         if isinstance(plo_obj[0], np.ndarray):
             x = plo_obj[0]
             y = plo_obj[1]
@@ -544,11 +543,10 @@ def plot_matplotlib(
     show : bool, optional
         Whether to show the plot or return the matplotlib object. Default is `True`.
 
-
     Returns
     -------
     :class:`matplotlib.pyplot.Figure`
-        Matplotlib Figure object.
+        Matplotlib figure object.
     """
     dpi = 100.0
     figsize = (size[0] / dpi, size[1] / dpi)

pyaedt/modules/solutions.py

@@ -819,8 +819,9 @@ def plot(
         title="",
         snapshot_path=None,
         is_polar=False,
+        show=True,
     ):
-        """Create a matplotlib plot based on a list of data.
+        """Create a matplotlib figure based on a list of data.
 
         Parameters
         ----------
@@ -852,8 +853,8 @@ def plot(
 
         Returns
         -------
-        :class:`matplotlib.plt`
-            Matplotlib fig object.
+        :class:`matplotlib.pyplot.Figure`
+            Matplotlib figure object.
         """
         if is_ironpython:  # pragma: no cover
             return False
@@ -893,9 +894,9 @@ def plot(
         if len(data_plot) > 15:
             show_legend = False
         if is_polar:
-            return plot_polar_chart(data_plot, size, show_legend, x_label, y_label, title, snapshot_path)
+            return plot_polar_chart(data_plot, size, show_legend, x_label, y_label, title, snapshot_path, show=show)
         else:
-            return plot_2d_chart(data_plot, size, show_legend, x_label, y_label, title, snapshot_path)
+            return plot_2d_chart(data_plot, size, show_legend, x_label, y_label, title, snapshot_path, show=show)
 
     @pyaedt_function_handler(xlabel="x_label", ylabel="y_label", math_formula="formula")
     def plot_3d(
@@ -909,8 +910,9 @@ def plot_3d(
         formula=None,
         size=(2000, 1000),
         snapshot_path=None,
+        show=True,
     ):
-        """Create a matplotlib 3d plot based on a list of data.
+        """Create a matplotlib 3D figure based on a list of data.
 
         Parameters
         ----------
@@ -938,8 +940,8 @@ def plot_3d(
 
         Returns
         -------
-        :class:`matplotlib.plt`
-            Matplotlib fig object.
+        :class:`matplotlib.figure.Figure`
+            Matplotlib figure object.
         """
         if is_ironpython:
             return False  # pragma: no cover
@@ -985,7 +987,7 @@ def plot_3d(
             y_label = y_axis
         if not title:
             title = "Simulation Results Plot"
-        return plot_3d_chart(data_plot, size, x_label, y_label, title, snapshot_path)
+        return plot_3d_chart(data_plot, size, x_label, y_label, title, snapshot_path, show=show)
 
     @pyaedt_function_handler()
     def ifft(self, curve_header="NearE", u_axis="_u", v_axis="_v", window=False):
@@ -1640,6 +1642,7 @@ def plot_farfield_contour(
         Returns
         -------
         :class:`matplotlib.pyplot.Figure`
+            Matplotlib figure object.
 
         Examples
         --------
@@ -1736,6 +1739,7 @@ def plot_2d_cut(
             image_path=None,
             show=True,
             is_polar=False,
+            show_legend=True,
             **kwargs
     ):
         # fmt: on
@@ -1769,14 +1773,13 @@ def plot_2d_cut(
             If ``False``, the Matplotlib instance of the plot is shown.
         is_polar : bool, optional
             Whether this plot is a polar plot. The default is ``True``.
+        show_legend : bool, optional
+            Whether to display the legend or not. The default is ``True``.
 
         Returns
         -------
-        :class:`matplotlib.plt`
-            Whether to show the plotted curve.
-            If ``show=True``, a Matplotlib figure instance of the plot is returned.
-            If ``show=False``, the plotted curve is returned.
-
+        :class:`matplotlib.pyplot.Figure`
+            Matplotlib figure object.
 
         Examples
         --------
@@ -1849,30 +1852,29 @@ def plot_2d_cut(
                 return False
             curves.append([x, y, "{}={}".format(y_key, data[y_key][theta_idx])])
 
-        if show:
-            show_legend = True
-            if len(curves) > 15:
-                show_legend = False
-            if is_polar:
-                return plot_polar_chart(
-                    curves,
-                    xlabel=x_key,
-                    ylabel=quantity,
-                    title=title,
-                    snapshot_path=image_path,
-                    show_legend=show_legend,
-                )
-            else:
-                return plot_2d_chart(
-                    curves,
-                    xlabel=x_key,
-                    ylabel=quantity,
-                    title=title,
-                    snapshot_path=image_path,
-                    show_legend=show_legend,
-                )
+        # FIXME: See if we need to keep this check on the curves length
+        # if len(curves) > 15:
+        #     show_legend = False
+        if is_polar:
+            return plot_polar_chart(
+                curves,
+                xlabel=x_key,
+                ylabel=quantity,
+                title=title,
+                snapshot_path=image_path,
+                show_legend=show_legend,
+                show=show,
+            )
         else:
-            return curves
+            return plot_2d_chart(
+                curves,
+                xlabel=x_key,
+                ylabel=quantity,
+                title=title,
+                snapshot_path=image_path,
+                show_legend=show_legend,
+                show=show,
+            )
 
     # fmt: off
     @pyaedt_function_handler(farfield_quantity="quantity",
@@ -1913,14 +1915,12 @@ def polar_plot_3d(
             Full path for the image file. The default is ``None``, in which case a file is not exported.
         show : bool, optional
             Whether to show the plot. The default is ``True``.
-            If ``False``, the Matplotlib instance of the plot is shown.
+            If ``False``, the Matplotlib instance of the plot is not shown.
 
         Returns
         -------
-        :class:`matplotlib.plt`
-            Whether to show the plotted curve.
-            If ``show=True``, a Matplotlib figure instance of the plot is returned.
-            If ``show=False``, the plotted curve is returned.
+        :class:`matplotlib.pyplot.Figure`
+            Matplotlib figure object.
 
         Examples
         --------
@@ -1968,10 +1968,7 @@ def polar_plot_3d(
         x = r * np.sin(theta_grid) * np.cos(phi_grid)
         y = r * np.sin(theta_grid) * np.sin(phi_grid)
         z = r * np.cos(theta_grid)
-        if show:  # pragma: no cover
-            plot_3d_chart([x, y, z], xlabel="Theta", ylabel="Phi", title=title, snapshot_path=image_path)
-        else:
-            return x, y, z
+        return plot_3d_chart([x, y, z], xlabel="Theta", ylabel="Phi", title=title, snapshot_path=image_path, show=show)
 
     # fmt: off
     @pyaedt_function_handler(farfield_quantity="quantity", export_image_path="image_path")
@@ -2028,8 +2025,7 @@ def polar_plot_3d_pyvista(
         Returns
         -------
         bool or :class:`Pyvista.Plotter`
-            ``True`` when successful. The :class:`Pyvista.Plotter` is returned when ``show`` and
-            ``export_image_path`` are ``False``.
+            ``False`` when the method fails, Pyvista plotter object otherwise.
 
         Examples
         --------
@@ -2208,10 +2204,8 @@ def scale(value=1):
 
         if image_path:
             p.show(screenshot=image_path)
-            return True
-        elif show:  # pragma: no cover
+        if show:  # pragma: no cover
             p.show()
-            return True
         return p
 
     @pyaedt_function_handler()