Ctd plot

gliderpy/plotting.py

@@ -9,6 +9,7 @@
     import cartopy.crs as ccrs
     import matplotlib.dates as mdates
     import matplotlib.pyplot as plt
+
 except ModuleNotFoundError:
     warnings.warn(
         "gliderpy requires matplotlib and cartopy for plotting.",
@@ -47,6 +48,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
@@ -57,7 +59,11 @@ def plot_transect(
     """
     cmap = kw.get("cmap", None)
 
-    fig, ax = plt.subplots(figsize=(17, 2))
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(17, 2))
+    else:
+        fig = ax.figure
+
     cs = ax.scatter(
         df.index,
         df["pressure"],
@@ -76,3 +82,43 @@ def plot_transect(
     cbar.ax.set_ylabel(var)
     ax.set_ylabel("pressure")
     return fig, ax
+
+
+@register_dataframe_method
+def plot_ctd(
+    df: pd.DataFrame,
+    var: str,
+    ax: plt.Axes = None,
+    color: str = None
+) -> tuple:
+    """Make a CTD profile plot of pressure vs property
+    depending on what variable was chosen.
+
+    :param var: variable to plot against pressure
+    :param ax: existing axis to plot on (default: None)
+    :param color: color for the plot line (default: None)
+    :return: figure, axes
+    """
+    g = df.groupby(["longitude", "latitude"])
+    profile = g.get_group((list(g.groups)[0]))
+
+    if ax is None:
+        fig, ax1 = plt.subplots(figsize=(5, 6))
+        ax1.plot(profile[var], -profile["pressure"], label=var, color=color)
+        ax1.set_ylabel('Pressure')
+        ax1.set_xlabel(var)
+        ax1.legend()
+        return fig, ax1
+    else:
+        fig = ax.get_figure()
+        ax2 = ax.twiny()  # Create a new twinned axis
+        ax2.plot(profile[var], -profile["pressure"], label=var, color=color)
+        ax2.set_xlabel(var)
+
+        # Handle legends
+        lines, labels = ax.get_legend_handles_labels()
+        lines2, labels2 = ax2.get_legend_handles_labels()
+        ax.legend(lines + lines2, labels + labels2, loc="lower center")
+
+        return fig, ax2
+

tests/test_plotting.py

@@ -1,11 +1,10 @@
 """Test transect."""
 
 from pathlib import Path
-
 import pytest
 
 from gliderpy.fetchers import GliderDataFetcher
-from gliderpy.plotting import plot_track, plot_transect
+from gliderpy.plotting import plot_track, plot_transect, plot_ctd
 
 root = Path(__file__).parent
 
@@ -36,3 +35,16 @@ def test_plot_transect():
 
     # Return the figure for pytest-mpl to compare
     return fig
+
+
+@pytest.mark.mpl_image_compare(baseline_dir=root.joinpath("baseline/"))
+def test_plot_ctd():
+    glider_grab = GliderDataFetcher()
+
+    glider_grab.fetcher.dataset_id = "whoi_406-20160902T1700"
+    df = glider_grab.to_pandas()
+    # Generate the plot
+    fig, ax = plot_ctd(df, 'temperature')
+
+    # Return the figure for pytest-mpl to compare
+    return fig