[FEAT] add teos10 function for adding variables

glidertest/tools.py

@@ -37,7 +37,21 @@ def _necessary_variables_check(ds: xr.Dataset, vars: list):
     if missing_vars:
         msg = f"Required variables {list(missing_vars)} do not exist in the supplied dataset."
         raise KeyError(msg)
-
+
+
+def _calc_teos10_variables(ds):
+    """
+    Calculates TEOS 10 variables not present in the dataset
+    :param ds:
+    :return:
+    """
+    _necessary_variables_check(ds, ['DEPTH', 'LONGITUDE', 'LATITUDE', 'TEMP', 'PSAL'])
+    if 'DENSITY' not in ds.variables:
+        SA = gsw.SA_from_SP(ds.PSAL, ds.DEPTH, ds.LONGITUDE, ds.LATITUDE)
+        CT = gsw.CT_from_t(SA, ds.TEMP, ds.DEPTH)
+        ds['DENSITY'] = ('N_MEASUREMENTS', gsw.rho(SA, CT, ds.DEPTH).values)
+    return ds
+
 
 def compute_grid2d(x, y, v, xi=1, yi=1):
     """
@@ -199,19 +213,14 @@ def plot_basic_vars(ds: xr.Dataset, v_res=1, start_prof=0, end_prof=-1):
     Chiara Monforte
     """
     _necessary_variables_check(ds, ['PROFILE_NUMBER', 'DEPTH', 'TEMP', 'PSAL', 'LATITUDE', 'LONGITUDE'])
+    ds = _calc_teos10_variables(ds)
     p = 1
     z = v_res
     tempG, profG, depthG = compute_grid2d(ds.PROFILE_NUMBER, ds.DEPTH, ds.TEMP, p, z)
     salG, profG, depthG = compute_grid2d(ds.PROFILE_NUMBER, ds.DEPTH, ds.PSAL, p, z)
-
-    if 'DENSITY' not in ds.variables:
-        ds['DENSITY'] = (('N_MEASUREMENTS'), np.full(ds.dims['N_MEASUREMENTS'], np.nan))
-        SA = gsw.SA_from_SP(ds.PSAL, ds.DEPTH, ds.LONGITUDE, ds.LATITUDE)
-        CT = gsw.CT_from_t(SA, ds.TEMP, ds.DEPTH)
-        ds['DENSITY'] = gsw.rho(SA, CT, ds.DEPTH)
-
     denG, profG, depthG = compute_grid2d(ds.PROFILE_NUMBER, ds.DEPTH, ds.DENSITY, p, z)
 
+
     tempG = tempG[start_prof:end_prof, :]
     salG = salG[start_prof:end_prof, :]
     denG = denG[start_prof:end_prof, :]
@@ -681,19 +690,16 @@ def check_temporal_drift(ds: xr.Dataset, var: str, ax: plt.Axes = None, **kw: di
     else:
         fig = plt.gcf()
 
-    if var not in ds.variables:
-        print(f'{var} does not exist in the dataset. Make sure the spelling is correct or add this variable to your dataset')
-    else:
-        ax[0].scatter(mdates.date2num(ds.TIME), ds[var], s=10)
-        ax[0].xaxis.set_major_formatter(DateFormatter('%Y-%m-%d'))
-        ax[0].set(ylim=(np.nanpercentile(ds[var], 0.01), np.nanpercentile(ds[var], 99.99)), ylabel=var)
+    ax[0].scatter(mdates.date2num(ds.TIME), ds[var], s=10)
+    ax[0].xaxis.set_major_formatter(DateFormatter('%Y-%m-%d'))
+    ax[0].set(ylim=(np.nanpercentile(ds[var], 0.01), np.nanpercentile(ds[var], 99.99)), ylabel=var)
 
-        c = ax[1].scatter(ds[var], ds.DEPTH, c=mdates.date2num(ds.TIME), s=10)
-        ax[1].set(xlim=(np.nanpercentile(ds[var], 0.01), np.nanpercentile(ds[var], 99.99)), ylabel='Depth (m)', xlabel=var)
-        ax[1].invert_yaxis()
+    c = ax[1].scatter(ds[var], ds.DEPTH, c=mdates.date2num(ds.TIME), s=10)
+    ax[1].set(xlim=(np.nanpercentile(ds[var], 0.01), np.nanpercentile(ds[var], 99.99)), ylabel='Depth (m)', xlabel=var)
+    ax[1].invert_yaxis()
 
-        [a.grid() for a in ax]
-        plt.colorbar(c, format=DateFormatter('%b %d'))
+    [a.grid() for a in ax]
+    plt.colorbar(c, format=DateFormatter('%b %d'))
     return fig, ax
 
 
@@ -723,13 +729,14 @@ def check_monotony(da):
         return True
 
 
-def plot_prof_monotony(ds: xr.DataArray, ax: plt.Axes = None, **kw: dict, ) -> tuple({plt.Figure, plt.Axes}):
+def plot_profIncrease(ds: xr.Dataset, ax: plt.Axes = None, **kw: dict, ) -> tuple({plt.Figure, plt.Axes}):
+
     """
     This function can be used to plot the profile number and check for any possible issues with the profile index assigned.
 
     Parameters
     ----------
-    ds: xarray in OG1 format with at least PROFILE_NUMBER, TIME, DEPTH. Data should not be gridded
+    ds: xarray dataset in OG1 format with at least PROFILE_NUMBER, TIME, DEPTH. Data should not be gridded
     ax: axis to plot the data
 
     Returns 
@@ -1005,9 +1012,7 @@ def calc_DEPTH_Z(ds):
     ----------------
     Eleanor Frajka-Williams
     """
-    # Ensure the required variables are present
-    if 'PRES' not in ds.variables or 'LATITUDE' not in ds.variables or 'LONGITUDE' not in ds.variables:
-        raise ValueError("Dataset must contain 'PRES', 'LATITUDE', and 'LONGITUDE' variables.")
+    _necessary_variables_check(ds, ['PRES', 'LONGITUDE', 'LATITUDE'])
 
     # Initialize the new variable with the same dimensions as dive_num
     ds['DEPTH_Z'] = (['N_MEASUREMENTS'], np.full(ds.dims['N_MEASUREMENTS'], np.nan))

notebooks/demo.ipynb

@@ -221,7 +221,7 @@
    "outputs": [],
    "source": [
     "tools.check_monotony(ds.PROFILE_NUMBER)\n",
-    "tools.plot_prof_monotony(ds)"
+    "tools.plot_profIncrease(ds)"
    ]
   },
   {
@@ -460,7 +460,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,

tests/test_tools.py

@@ -4,10 +4,12 @@
 import math
 import numpy as np
 import matplotlib
-matplotlib.use('agg') # use agg backend to prevent creating plot windows during tests
+matplotlib.use('agg')  # use agg backend to prevent creating plot windows during tests
+
 
 def test_plots(start_prof=0, end_prof=100):
     ds = fetchers.load_sample_dataset()
+    ds = ds.drop_vars(['DENSITY'])
     fig, ax = tools.plot_basic_vars(ds,start_prof=start_prof, end_prof=end_prof)
     assert ax[0].get_ylabel() == 'Depth (m)'
     assert ax[0].get_xlabel() == f'Average Temperature [C] \nbetween profile {start_prof} and {end_prof}'
@@ -73,7 +75,7 @@ def test_temporal_drift(var='DOXY'):
 def test_profile_check():
     ds = fetchers.load_sample_dataset()
     tools.check_monotony(ds.PROFILE_NUMBER)
-    tools.plot_prof_monotony(ds)
+    tools.plot_profIncrease(ds)
 
 
 def test_check_monotony():