Add new ax feature

gliderpy/plotting.py

@@ -46,6 +46,7 @@ def plot_track(df: pd.DataFrame) -> tuple(plt.Figure, plt.Axes):
 def plot_transect(
     df: pd.DataFrame,
     var: str,
+    ax: plt.Axes = None,
     **kw: dict,
 ) -> tuple(plt.Figure, plt.Axes):
     """Make a scatter plot of depth vs time coloured by a user defined
@@ -56,7 +57,18 @@ def plot_transect(
     """
     cmap = kw.get("cmap", None)
 
-    fig, ax = plt.subplots(figsize=(17, 2))
+    fignums = plt.get_fignums()
+    if ax is None and not fignums:
+        fig, ax = plt.subplots(figsize=(17, 2))
+    elif ax:
+        fig = ax.get_figure()
+    else:
+        ax = plt.gca()
+        fig = plt.gcf()
+
+    if not ax.yaxis_inverted():
+        ax.invert_yaxis()
+
     cs = ax.scatter(
         df.index,
         df["pressure"],
@@ -67,13 +79,15 @@ def plot_transect(
         cmap=cmap,
     )
 
-    ax.invert_yaxis()
     xfmt = mdates.DateFormatter("%H:%Mh\n%d-%b")
     ax.xaxis.set_major_formatter(xfmt)
 
     cbar = fig.colorbar(cs, orientation="vertical", extend="both")
     cbar.ax.set_ylabel(var)
     ax.set_ylabel("pressure")
+
+    ax.set_ylim(ax.get_ylim()[0], 0)
+
     return fig, ax
 
 

tests/test_plotting.py

@@ -1,36 +1,44 @@
-"""Test transect."""
-
+import cmcrameri
 import pytest
-import matplotlib as mpl
+import matplotlib.pyplot as plt
 from pathlib import Path
-from gliderpy.plotting import plot_track
-from gliderpy.plotting import plot_transect
+from gliderpy.plotting import plot_track, plot_transect
 from gliderpy.fetchers import GliderDataFetcher
 
 root = Path(__file__).parent
 
-@pytest.mark.mpl_image_compare(baseline_dir=root.joinpath("baseline/"))
-def test_plot_track():
+@pytest.fixture
+def glider_data():
     glider_grab = GliderDataFetcher()
-
     glider_grab.fetcher.dataset_id = "whoi_406-20160902T1700"
     df = glider_grab.to_pandas()
-    # Generate the plot
-    fig, ax = plot_track(df)
+    return df
 
+@pytest.mark.mpl_image_compare(baseline_dir=root.joinpath("baseline/"))
+def test_plot_track(glider_data):
+    # Generate the plot
+    fig, ax = plot_track(glider_data)
     # Return the figure for pytest-mpl to compare
     return fig
 
-
-
 @pytest.mark.mpl_image_compare(baseline_dir=root.joinpath("baseline/"))
-def test_plot_transect():
-    glider_grab = GliderDataFetcher()
-
-    glider_grab.fetcher.dataset_id = "whoi_406-20160902T1700"
-    df = glider_grab.to_pandas()
+def test_plot_transect(glider_data):
     # Generate the plot
-    fig, ax = plot_transect(df, 'temperature')
+    fig, ax = plot_transect(glider_data, 'temperature')
+    # Return the figure for pytest-mpl to compare
+    return fig
 
+@pytest.mark.mpl_image_compare(baseline_dir=root.joinpath("baseline/"))
+def test_plot_transect_multiple_figures(glider_data):
+    # Generate the plot with multiple figures
+    fig, (ax0, ax1) = plt.subplots(figsize=(15, 9), nrows=2, sharex=True, sharey=True)
+    glider_data.plot_transect(var="temperature", ax=ax0)
+    glider_data.plot_transect(var="salinity", ax=ax1, cmap=cmcrameri.cm.davos)
     # Return the figure for pytest-mpl to compare
     return fig
+
+def test_plot_transect_size(glider_data):
+    # Generate the plot with a specific size
+    fig, ax = plt.subplots(figsize=(15, 9))
+    glider_data.plot_transect(var="temperature", ax=ax)
+    assert fig.get_size_inches() == pytest.approx([15., 9.])