WIP: Add ax and figsize arg to plotting functions (#121)

* Add ax argument to plot_windrose

* Add figsize argument to plot functions

* pre-commmit fixes

* fix failing test

* add subplots test

* add baseline for new test

Co-authored-by: Filipe Fernandes <[email protected]>

Author: petebachant

tests/output/df/test_subplots.png

@@ -2,6 +2,7 @@
 
 
 import matplotlib
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import pytest
@@ -68,3 +69,11 @@ def test_windrose_np_plot_and_pd_plot():
     kind = "scatter"
     ax = plot_windrose(wd, ws, kind=kind, alpha=0.2)
     return ax.figure
+
+
+@pytest.mark.mpl_image_compare(baseline_dir="output/df")
+def test_subplots():
+    fig, (ax1, ax2) = plt.subplots(ncols=2, subplot_kw={"projection": "windrose"})
+    plot_windrose(df, ax=ax1)
+    plot_windrose(df, ax=ax2)
+    return fig

tests/test_windrose_pandas.py

@@ -90,14 +90,22 @@ def __init__(self, *args, **kwargs):
         self.clear()
 
     @staticmethod
-    def from_ax(ax=None, fig=None, rmax=None, rect=None, *args, **kwargs):
+    def from_ax(
+        ax=None,
+        fig=None,
+        rmax=None,
+        figsize=FIGSIZE_DEFAULT,
+        rect=None,
+        *args,
+        **kwargs,
+    ):
         """
         Return a WindroseAxes object for the figure `fig`.
         """
         if ax is None:
             if fig is None:
                 fig = plt.figure(
-                    figsize=FIGSIZE_DEFAULT,
+                    figsize=figsize,
                     dpi=DPI_DEFAULT,
                     facecolor="w",
                     edgecolor="w",
@@ -693,10 +701,10 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @staticmethod
-    def from_ax(ax=None, fig=None, *args, **kwargs):
+    def from_ax(ax=None, fig=None, figsize=FIGSIZE_DEFAULT, *args, **kwargs):
         if ax is None:
             if fig is None:
-                fig = plt.figure(figsize=FIGSIZE_DEFAULT, dpi=DPI_DEFAULT)
+                fig = plt.figure(figsize=figsize, dpi=DPI_DEFAULT)
             ax = WindAxes(fig, 1, 1, 1, *args, **kwargs)
             fig.add_axes(ax)
             return ax
@@ -788,36 +796,36 @@ def histogram(direction, var, bins, nsector, normed=False, blowto=False):
 
 
 @docstring.copy(WindroseAxes.contour)
-def wrcontour(direction, var, ax=None, rmax=None, **kwargs):
+def wrcontour(direction, var, ax=None, rmax=None, figsize=FIGSIZE_DEFAULT, **kwargs):
     """
     Draw contour probability density function and return Weibull
     distribution parameters.
     """
-    ax = WindroseAxes.from_ax(ax, rmax=rmax)
+    ax = WindroseAxes.from_ax(ax, rmax=rmax, figsize=figsize)
     ax.contour(direction, var, **kwargs)
     ax.set_legend()
     return ax
 
 
 @docstring.copy(WindroseAxes.contourf)
-def wrcontourf(direction, var, ax=None, rmax=None, **kwargs):
-    ax = WindroseAxes.from_ax(ax, rmax=rmax)
+def wrcontourf(direction, var, ax=None, rmax=None, figsize=FIGSIZE_DEFAULT, **kwargs):
+    ax = WindroseAxes.from_ax(ax, rmax=rmax, figsize=figsize)
     ax.contourf(direction, var, **kwargs)
     ax.set_legend()
     return ax
 
 
 @docstring.copy(WindroseAxes.box)
-def wrbox(direction, var, ax=None, rmax=None, **kwargs):
-    ax = WindroseAxes.from_ax(ax, rmax=rmax)
+def wrbox(direction, var, ax=None, rmax=None, figsize=FIGSIZE_DEFAULT, **kwargs):
+    ax = WindroseAxes.from_ax(ax, rmax=rmax, figsize=figsize)
     ax.box(direction, var, **kwargs)
     ax.set_legend()
     return ax
 
 
 @docstring.copy(WindroseAxes.bar)
-def wrbar(direction, var, ax=None, rmax=None, **kwargs):
-    ax = WindroseAxes.from_ax(ax, rmax=rmax)
+def wrbar(direction, var, ax=None, rmax=None, figsize=FIGSIZE_DEFAULT, **kwargs):
+    ax = WindroseAxes.from_ax(ax, rmax=rmax, figsize=figsize)
     ax.bar(direction, var, **kwargs)
     ax.set_legend()
     return ax
@@ -833,23 +841,26 @@ def wrpdf(
     Nbins=10,
     ax=None,
     rmax=None,
+    figsize=FIGSIZE_DEFAULT,
     *args,
     **kwargs,
 ):
     """
     Draw probability density function and return Weitbull distribution
     parameters
     """
-    ax = WindAxes.from_ax(ax)
+    ax = WindAxes.from_ax(ax, figsize=figsize)
     ax, params = ax.pdf(var, bins, Nx, bar_color, plot_color, Nbins, *args, **kwargs)
     return (ax, params)
 
 
-def wrscatter(direction, var, ax=None, rmax=None, *args, **kwargs):
+def wrscatter(
+    direction, var, ax=None, rmax=None, figsize=FIGSIZE_DEFAULT, *args, **kwargs
+):
     """
     Draw scatter plot
     """
-    ax = WindroseAxes.from_ax(ax, rmax=rmax)
+    ax = WindroseAxes.from_ax(ax, rmax=rmax, figsize=figsize)
     direction = -np.array(direction) + np.radians(90)
     ax.scatter(direction, var, *args, **kwargs)
     return ax
@@ -914,6 +925,7 @@ def plot_windrose(
     direction_name=DIR_DEFAULT,
     by=None,
     rmax=None,
+    ax=None,
     **kwargs,
 ):
     """Plot windrose from a pandas DataFrame or a numpy array."""
@@ -924,7 +936,9 @@ def plot_windrose(
         direction = df[direction_name].values
     else:
         direction = direction_or_df
-    return plot_windrose_np(direction, var, kind=kind, by=by, rmax=rmax, **kwargs)
+    return plot_windrose_np(
+        direction, var, kind=kind, by=by, rmax=rmax, ax=ax, **kwargs
+    )
 
 
 def plot_windrose_df(
@@ -934,16 +948,24 @@ def plot_windrose_df(
     direction_name=DIR_DEFAULT,
     by=None,
     rmax=None,
+    ax=None,
     **kwargs,
 ):
     """Plot windrose from a pandas DataFrame."""
     var = df[var_name].values
     direction = df[direction_name].values
-    return plot_windrose_np(direction, var, by=by, rmax=rmax, **kwargs)
+    return plot_windrose_np(direction, var, by=by, rmax=rmax, ax=ax, **kwargs)
 
 
 def plot_windrose_np(
-    direction, var, kind="contour", clean_flag=True, by=None, rmax=None, **kwargs
+    direction,
+    var,
+    kind="contour",
+    clean_flag=True,
+    by=None,
+    rmax=None,
+    ax=None,
+    **kwargs,
 ):
     """Plot windrose from a numpy array."""
     if kind in D_KIND_PLOT.keys():
@@ -957,7 +979,7 @@ def plot_windrose_np(
     if clean_flag:
         direction, var = clean(direction, var)
     if by is None:
-        ax = f_plot(direction=direction, var=var, rmax=rmax, **kwargs)
+        ax = f_plot(direction=direction, var=var, rmax=rmax, ax=ax, **kwargs)
         if kind not in ["pdf"]:
             ax.set_legend()
         return ax