Merge pull request #125 from FloraSauerbronn/TS-Diagram

Creating new function plot_ts for TS diagram

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
       files: requirements-dev.txt
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.11.2
+  rev: v1.13.0
   hooks:
   - id: mypy
     exclude: docs/source/conf.py
@@ -50,7 +50,7 @@ repos:
     - id: add-trailing-comma
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.6.9
+  rev: v0.7.1
   hooks:
     - id: ruff
       args: ["--fix", "--show-fixes"]
@@ -73,7 +73,7 @@ repos:
     - id: nb-strip-paths
 
 - repo: https://github.com/tox-dev/pyproject-fmt
-  rev: 2.2.4
+  rev: v2.4.3
   hooks:
     - id: pyproject-fmt
 

gliderpy/__init__.py

@@ -6,11 +6,12 @@
     __version__ = "unknown"
 
 from .fetchers import GliderDataFetcher
-from .plotting import plot_track, plot_transect
+from .plotting import plot_cast, plot_track, plot_transect, plot_ts
 
 __all__ = [
     "GliderDataFetcher",
     "plot_track",
     "plot_transect",
     "plot_cast",
+    "plot_ts",
 ]

gliderpy/plotting.py

@@ -7,8 +7,11 @@
 
 try:
     import cartopy.crs as ccrs
+    import gsw
     import matplotlib.dates as mdates
     import matplotlib.pyplot as plt
+    import numpy as np
+
 except ModuleNotFoundError:
     warnings.warn(
         "gliderpy requires matplotlib and cartopy for plotting.",
@@ -24,7 +27,7 @@
 
 
 @register_dataframe_method
-def plot_track(df: pd.DataFrame) -> tuple(plt.Figure, plt.Axes):
+def plot_track(df: pd.DataFrame) -> tuple[plt.Figure, plt.Axes]:
     """Plot a track of glider path coloured by temperature.
 
     :return: figures, axes
@@ -49,7 +52,7 @@ def plot_transect(
     var: str,
     ax: plt.Axes = None,
     **kw: dict,
-) -> tuple(plt.Figure, plt.Axes):
+) -> tuple[plt.Figure, plt.Axes]:
     """Make a scatter plot of depth vs time coloured by a user defined
     variable.
 
@@ -99,7 +102,7 @@ def plot_cast(
     var: str,
     ax: plt.Axes = None,
     color: str | None = None,
-) -> tuple:
+) -> tuple[plt.Figure, plt.Axes]:
     """Make a CTD profile plot of pressure vs property
     depending on what variable was chosen.
 
@@ -123,3 +126,62 @@ def plot_cast(
     ax.invert_yaxis()
 
     return fig, ax
+
+
+@register_dataframe_method
+def plot_ts(
+    df: pd.DataFrame,
+) -> tuple[plt.Figure, plt.Axes]:
+    """Make a TS diagram for all profiles in the DataFrame.
+
+    :return: figure, axes
+    """
+    df["sa"] = gsw.conversions.SA_from_SP(
+        df["salinity"],
+        df["pressure"],
+        df["longitude"],
+        df["latitude"],
+    )
+
+    df["ct"] = gsw.conversions.CT_from_t(
+        df["sa"],
+        df["temperature"],
+        df["pressure"],
+    )
+
+    g = df.groupby(["longitude", "latitude"])
+
+    fig, ax = plt.subplots(figsize=(10, 10))
+
+    for _name, group in g:
+        sc = plt.scatter(
+            group["sa"],
+            group["ct"],
+            c=group["pressure"],
+            cmap="plasma_r",
+            s=30,
+        )
+
+    plt.xlabel("Absolute Salinity(g/kg)")
+    plt.ylabel("Conservative Temperature (°C)")
+    cbar = plt.colorbar(sc)
+    cbar.set_label("Pressure (dbar)")
+
+    # Define salinity and temperature grids
+    salinity_grid = np.linspace(df["sa"].min() - 5, df["sa"].max() + 5, 100)
+    temperature_grid = np.linspace(df["ct"].min() - 5, df["ct"].max() + 5, 100)
+    sal, temp = np.meshgrid(salinity_grid, temperature_grid)
+    sigma = gsw.sigma0(sal, temp)
+
+    contours = plt.contour(
+        sal,
+        temp,
+        sigma,
+        levels=np.arange(20, 30, 1),
+        colors="grey",
+        linestyles="--",
+    )
+
+    plt.clabel(contours, inline=True, fmt="%1.1f", fontsize=8, colors="black")
+
+    return fig, ax

pyproject.toml

@@ -22,6 +22,7 @@ classifiers = [
   "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
 ]
 dynamic = [
   "dependencies",

requirements.txt

@@ -1,4 +1,5 @@
 erddapy
+gsw
 httpx
 pandas
 pandas-flavor

tests/test_plotting.py

@@ -11,7 +11,7 @@
 import pytest
 
 from gliderpy.fetchers import GliderDataFetcher
-from gliderpy.plotting import plot_cast, plot_track, plot_transect
+from gliderpy.plotting import plot_cast, plot_track, plot_transect, plot_ts
 
 root = Path(__file__).parent
 
@@ -80,3 +80,10 @@ def test_plot_cast(glider_data):
     """Test plot_cast accessor."""
     fig, ax = plot_cast(glider_data, 0, var="temperature", color="blue")
     return fig
+
+
+@pytest.mark.mpl_image_compare(baseline_dir=root.joinpath("baseline/"))
+def test_plot_ts(glider_data):
+    """Test plot_ts accessor."""
+    fig, ax = plot_ts(glider_data)
+    return fig