Move axis creating to plotting module

This consolidates most of the plotting code to the plotting module so that it may be easier to support additional plotting backends (i.e. ploytly, bokeh)
rafmudaf · Apr 23, 2024 · a5cc358 · a5cc358
1 parent 24dbb46
commit a5cc358
Show file tree

Hide file tree

Showing 8 changed files with 6 additions and 55 deletions.
diff --git a/examples/torque2024_1turbine.ipynb b/examples/torque2024_1turbine.ipynb
@@ -715,7 +715,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax[2],\n",
     "    min_speed=min_speed,\n",
     "    max_speed=max_speed,\n",
     "    cmap='PuOr',\n",
@@ -920,7 +919,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax[1],\n",
     "    min_speed=min_speed,\n",
     "    max_speed=max_speed,\n",
     "    cmap='PuOr',\n",

diff --git a/examples/torque2024_4turbine.ipynb b/examples/torque2024_4turbine.ipynb
@@ -699,7 +699,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax[0],\n",
     "    min_speed=min_speed,\n",
     "    max_speed=max_speed,\n",
     "    cmap='PuOr',\n",
@@ -896,7 +895,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax[1],\n",
     "    min_speed=min_speed,\n",
     "    max_speed=max_speed,\n",
     "    cmap='PuOr',\n",

diff --git a/examples/wesc2023_demo.ipynb b/examples/wesc2023_demo.ipynb
@@ -211,7 +211,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax[0],\n",
     "    min_speed=min_speed,\n",
     "    max_speed=max_speed,\n",
     "    cmap='PuOr',\n",
@@ -227,7 +226,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax[1],\n",
     "    min_speed=min_speed,\n",
     "    max_speed=max_speed,\n",
     "    cmap='PuOr',\n",
@@ -407,7 +405,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax,\n",
     "    color_bar=True,\n",
     "    clevels=100,\n",
     "    cmap='PuOr',\n",
@@ -422,7 +419,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax,\n",
     "    color_bar=True,\n",
     "    clevels=100,\n",
     "    cmap='PuOr',\n",
@@ -533,7 +529,6 @@
     "max_speed = np.max(abs_diff)\n",
     "plot_plane(\n",
     "    diff,\n",
-    "    ax=ax[2],\n",
     "    min_speed=min_speed,\n",
     "    max_speed=max_speed,\n",
     "    cmap='PuOr',\n",

diff --git a/wcomp/base_interface.py b/wcomp/base_interface.py
@@ -133,10 +133,8 @@ class should be used to store the data. Then, the {py:meth}`wcomp.plotting.plot_
                 profile = WakeProfile(z, u)
 
                 # Plot the profile
-                ax = plt.gca()      # Get the current pyplot axis to place the plot
                 plot_profile(
                     profile,
-                    ax=ax,
                     color=self.LINE_PLOT_COLOR,
                     marker=self.LINE_PLOT_MARKER,
                     linestyle=self.LINE_PLOT_LINESTYLE,
@@ -187,10 +185,8 @@ class should be used to store the data. Then, the {py:meth}`wcomp.plotting.plot_
                 profile = WakeProfile(x, u)
 
                 # Plot the profile
-                ax = plt.gca()      # Get the current pyplot axis to place the plot
                 plot_profile(
                     profile,
-                    ax=ax,
                     color=self.LINE_PLOT_COLOR,
                     marker=self.LINE_PLOT_MARKER,
                     linestyle=self.LINE_PLOT_LINESTYLE,
@@ -242,10 +238,8 @@ class should be used to store the data. Then, the {py:meth}`wcomp.plotting.plot_
                 profile = WakeProfile(y, u)
 
                 # Plot the profile
-                ax = plt.gca()      # Get the current pyplot axis to place the plot
                 plot_profile(
                     profile,
-                    ax=ax,
                     color=self.LINE_PLOT_COLOR,
                     marker=self.LINE_PLOT_MARKER,
                     linestyle=self.LINE_PLOT_LINESTYLE,
@@ -298,10 +292,8 @@ class should be used to store the data. Then, the {py:meth}`wcomp.plotting.plot_
                 )
 
                 # Plot the plane
-                ax = plt.gca()      # Get the current pyplot axis to place the plot
                 plot_plane(
                     plane,
-                    ax=ax,
                     color=self.LINE_PLOT_COLOR,
                     marker=self.LINE_PLOT_MARKER,
                     linestyle=self.LINE_PLOT_LINESTYLE,
@@ -354,10 +346,8 @@ class should be used to store the data. Then, the {py:meth}`wcomp.plotting.plot_
                 )
 
                 # Plot the plane
-                ax = plt.gca()      # Get the current pyplot axis to place the plot
                 plot_plane(
                     plane,
-                    ax=ax,
                     color=self.LINE_PLOT_COLOR,
                     marker=self.LINE_PLOT_MARKER,
                     linestyle=self.LINE_PLOT_LINESTYLE,

diff --git a/wcomp/floris_interface.py b/wcomp/floris_interface.py
@@ -2,7 +2,6 @@
 import copy
 from pathlib import Path
 
-import matplotlib.pyplot as plt
 import numpy as np
 from floris.tools import FlorisInterface
 from windIO.utils.yml_utils import load_yaml
@@ -277,7 +276,6 @@ def vertical_profile_plot(
         y_coordinate: float,
         zmax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         cut_plane = self.fi.calculate_y_plane(
             crossstream_dist=y_coordinate,
@@ -294,7 +292,6 @@ def vertical_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='x',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -312,7 +309,6 @@ def streamwise_profile_plot(
         xmin: float,
         xmax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         cut_plane = self.fi.calculate_y_plane(
             crossstream_dist=y_coordinate,
@@ -329,7 +325,6 @@ def streamwise_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='x',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -347,7 +342,6 @@ def xsection_profile_plot(
         ymin: float,
         ymax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         cut_plane = self.fi.calculate_horizontal_plane(
             height=self.hub_height,
@@ -364,7 +358,6 @@ def xsection_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='y',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -401,7 +394,6 @@ def horizontal_contour(self, wind_direction: float) -> WakePlane:
         plane = WakePlane(x, y, u, "z")
         plot_plane(
             plane,
-            ax=plt.gca(),
             # cmap='Blues_r',
             # color_bar=True,
             clevels=100
@@ -432,7 +424,6 @@ def xsection_contour(self, wind_direction: float, x_coordinate: float) -> WakePl
         plane = WakePlane(y, z, u, "x")
         plot_plane(
             plane,
-            ax=plt.gca(),
             # cmap='Blues_r',
             # color_bar=True,
             clevels=100

diff --git a/wcomp/foxes_interface.py b/wcomp/foxes_interface.py
@@ -1,10 +1,8 @@
 
 from pathlib import Path
 
-import matplotlib.pyplot as plt
 import foxes.constants as FC
 import foxes.variables as FV
-import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import foxes
@@ -442,7 +440,6 @@ def vertical_profile_plot(
         y_coordinate: float,
         zmax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         # create points of interest, shape (n_states, n_points, 3):
         points = np.zeros((1, self.N_POINTS_1D, 3))
@@ -457,7 +454,6 @@ def vertical_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='x',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -475,7 +471,6 @@ def streamwise_profile_plot(
         xmin: float,
         xmax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         # create points of interest, shape (n_states, n_points, 3):
         points = np.zeros((1, self.N_POINTS_1D, 3))
@@ -491,7 +486,6 @@ def streamwise_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='x',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -509,7 +503,6 @@ def xsection_profile_plot(
         ymin: float,
         ymax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         # create points of interest, shape (n_states, n_points, 3):
         points = np.zeros((1, self.N_POINTS_1D, 3))
@@ -527,7 +520,6 @@ def xsection_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='y',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -568,7 +560,6 @@ def horizontal_contour(self, wind_direction: float) -> WakePlane:
         plane = WakePlane(x, y, u, "z")
         plot_plane(
             plane,
-            ax=plt.gca(),
             # cmap='Blues_r',
             # color_bar=True,
             clevels=100
@@ -602,7 +593,6 @@ def xsection_contour(self, wind_direction: float, x_coordinate: float) -> WakePl
         plane = WakePlane(y, z, u, "x")
         plot_plane(
             plane,
-            ax=plt.gca(),
             # cmap='Blues_r',
             # color_bar=True,
             clevels=100

diff --git a/wcomp/plotting.py b/wcomp/plotting.py
@@ -1,22 +1,20 @@
 
-import matplotlib
 import matplotlib.pyplot as plt
 
 from .output_struct import WakePlane, WakeProfile
 
-# Method to add plots to an axis
-# plot_1d(ax, series)
 
 def plot_profile(
     wake_profile: WakeProfile,
-    ax=None,
     # direction='x',
     # component='u',
     title="",
     **kwargs
 ):
-    if not ax:
+    if len(plt.get_fignums()) == 0:
         fig, ax = plt.subplots()
+    else:
+        ax = plt.gca()
 
     ax.plot(
         wake_profile.x1,
@@ -27,7 +25,6 @@ def plot_profile(
 
 def plot_plane(
     wake_plane: WakePlane,
-    ax: matplotlib.axes.Axes = None,
     min_speed: float = None,
     max_speed: float = None,
     cmap: str = "coolwarm",
@@ -66,9 +63,10 @@ def plot_plane(
             See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.tricontourf.html
             for available parameters.
     """
-
-    if not ax:
+    if len(plt.get_fignums()) == 0:
         fig, ax = plt.subplots()
+    else:
+        ax = plt.gca()
 
     if title:
         ax.set_title(title)

diff --git a/wcomp/pywake_interface.py b/wcomp/pywake_interface.py
@@ -1,7 +1,6 @@
 
 from pathlib import Path
 
-import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
 from py_wake import HorizontalGrid, XZGrid, YZGrid
@@ -221,7 +220,6 @@ def vertical_profile_plot(
         y_coordinate: float,
         zmax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         wake_data = self.sim_res.flow_map(
             XZGrid(
@@ -242,7 +240,6 @@ def vertical_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             color=self.LINE_PLOT_COLOR,
             marker=self.LINE_PLOT_MARKER,
             linestyle=self.LINE_PLOT_LINESTYLE,
@@ -258,7 +255,6 @@ def streamwise_profile_plot(
         xmin: float,
         xmax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         wake_data = self.sim_res.flow_map(
             XZGrid(
@@ -281,7 +277,6 @@ def streamwise_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='x',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -299,7 +294,6 @@ def xsection_profile_plot(
         ymin: float,
         ymax: float
     ) -> WakeProfile:
-        ax = plt.gca()
 
         wake_data = self.sim_res.flow_map(
             YZGrid(
@@ -320,7 +314,6 @@ def xsection_profile_plot(
         )
         plot_profile(
             profile,
-            ax=ax,
             # direction='x',
             # component='u',
             color=self.LINE_PLOT_COLOR,
@@ -353,7 +346,6 @@ def horizontal_contour(self, wind_direction: float) -> WakePlane:
         plane = WakePlane(x, y, u, "z")
         plot_plane(
             plane,
-            ax=plt.gca(),
             # cmap='Blues_r',
             # color_bar=True,
             clevels=100
@@ -380,7 +372,6 @@ def xsection_contour(self, wind_direction: float, x_coordinate: float) -> WakePl
         plane = WakePlane(y, z, u, "x")
         plot_plane(
             plane,
-            ax=plt.gca(),
             # cmap='Blues_r',
             # color_bar=True,
             clevels=100