versatile masking

docs/api.rst

@@ -117,7 +117,9 @@ General Utilities
    :toctree: ./_generated/
 
    utils.time_average_per_dive
-   utils.mask_to_depth_array
+   utils.mask_above_depths
+   utils.mask_below_depths
+   utils.mask_profile_depth
    utils.merge_dimensions
    utils.calc_glider_vert_velocity
    utils.calc_dive_phase

docs/physics.md

@@ -22,11 +22,12 @@ plt.show()
 
 
 ```python
-mld = gt.physics.mixed_layer_depth(dives, depth, dens0)
+import matplotlib.pyplot as plt
+mld = gt.physics.mixed_layer_depth(ds, 'density', verbose=False)
 mld_smoothed = mld.rolling(10, min_periods=3).mean()
 
-mld_mask = gt.physics.mixed_layer_depth(dives, depth, dens0, return_as_mask=True)
-mld_grid = gt.grid_data(x, y, mld_mask, verbose=False)
+mld_mask = gt.utils.mask_below_depth(ds, mld)
+mld_grid = gt.grid_data(ds.dives, ds.depth, mld_mask, verbose=False)
 
 fig, ax = plt.subplots(1, 2, figsize=[9, 3], dpi=100, sharey=True)
 
@@ -37,7 +38,7 @@ gt.plot(mld_grid, ax=ax[1])
 
 ax[0].set_ylabel('Depth (m)')
 [a.set_xlabel('Dives') for a in ax]
-xticks(rotation=0)
+plt.xticks(rotation=0)
 
 fig.tight_layout()
 ```

docs/whats-new.rst

@@ -24,14 +24,15 @@ What's New
     - Dark count corrections for optical sensors(:pull:'110'). By 'Isabelle Giddy <https://github.com/isgiddy>'_.
 
 
-v2022.xx (unreleased)
+v2023.xx (unreleased)
 ------------------------
 
-.. _whats-new.2022.xx:
+.. _whats-new.2023.xx:
 
 New Features
 ~~~~~~~~~~~~
 - added import for VOTO seaexplorer data (:pull:`170`) By `Martin Mohrmann <https://github.com/MartinMohrmann>`_.
+- added versatile, depth dependent masking (:pull:`172`) By `Martin Mohrmann <https://github.com/MartinMohrmann>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

glidertools/physics.py

@@ -29,9 +29,7 @@
     warnings.warn(message, category=GliderToolsWarning)
 
 
-def mixed_layer_depth(
-    ds, variable, thresh=0.01, ref_depth=10, return_as_mask=False, verbose=True
-):
+def mixed_layer_depth(ds, variable, thresh=0.01, ref_depth=10, verbose=True):
     """
     Calculates the MLD for ungridded glider array.
 
@@ -58,16 +56,12 @@ def mixed_layer_depth(
     mld = (
         ds[[variable, "depth"]]
         .groupby("dives")
-        .apply(mld_profile, variable, thresh, ref_depth, return_as_mask, verbose)
+        .apply(mld_profile, variable, thresh, ref_depth, verbose)
     )
-
-    if return_as_mask:
-        return np.concatenate([el for el in mld])
-    else:
-        return mld
+    return mld
 
 
-def mld_profile(df, variable, thresh, ref_depth, mask=False, verbose=True):
+def mld_profile(df, variable, thresh, ref_depth, verbose=True):
     exception = False
     df = df.dropna(subset=[variable, "depth"])
     if len(df) == 0:
@@ -105,10 +99,7 @@ def mld_profile(df, variable, thresh, ref_depth, mask=False, verbose=True):
             )
     if verbose and exception:
         warnings.warn(message, category=GliderToolsWarning)
-    if mask:
-        return depth <= mld
-    else:
-        return mld
+    return mld
 
 
 def potential_density(salt_PSU, temp_C, pres_db, lat, lon, pres_ref=0):

glidertools/utils.py

@@ -48,32 +48,62 @@ def time_average_per_dive(dives, time):
     return diveavg
 
 
-def mask_to_depth_array(dives, depth, var):
+def mask_above_depth(ds, depths):
     """
-    Use when function returns a boolean section (as a mask) and you would
-    like to return the depth of the positive mask (True) for each dive. This is
-    useful for cases like MLD which returns a mask. Note that this is for
-    ungridded data in "series" format.
+    Masks all data above depths.
+
+    Parameters:
+    -----------
+    df : xarray.Dataframe or pandas.Dataframe
+    mask_depths : float (for constant depth masking) or pandas.Series as
+        returned e.g. by the mixed_layer_depth function
+    """
+    return _mask_depth(ds, depths, above=True)
 
-    Parameters
-    ----------
-    dives : np.array, dtype=float, shape=[n, ]
-        discrete dive numbers (down = d.0; up = d.5) that matches depth and var
-        length
-    depth : np.array, dtype=float, shape=[n, ]
-        depth of each measurement
-    var : np.array, dtype=bool, shape=[n,]
-        mask array
 
+def mask_below_depth(ds, depths):
     """
+    Masks all data below depths.
 
-    from numpy import array, diff, r_
-    from pandas import Series
+    Parameters:
+    -----------
+    df : xarray.Dataframe or pandas.Dataframe
+    mask_depths : float (for constant depth masking) or pandas.Series as
+        returned e.g. by the mixed_layer_depth function
+    """
+    return _mask_depth(ds, depths, above=False)
+
+
+def mask_profile_depth(df, mask_depth, above):
+    """
+    masks either above or below mask_depth. If type(mask_depth)=np.nan,
+    the whole profile will be masked. Warning: This function is for a SINGLE
+    profile only, for masking a complete Glider Dataset please look for
+    utils.mask_above_depth and/or utils.mask_below_depth.
+
+    Parameters:
+    -----------
+    df : xarray.Dataframe or pandas.Dataframe
+    mask_depths : float (for constant depth masking) or pandas.Series as
+        returned e.g. by the mixed_layer_depth function
+    above : boolean
+        Mask either above mask_depth (True) or below (False)
+    """
+    if type(mask_depth) not in [int, float]:
+        # this case for calling from _mask_depth
+        mask_depth = mask_depth.loc[df.index[0]]
+    if above:
+        mask = df.depth > mask_depth
+    else:
+        mask = df.depth < mask_depth
+    return mask
 
-    i = r_[False, diff(var)].astype(bool)
-    idx_depth = Series(array(depth)[i], index=array(dives)[i])
 
-    return idx_depth
+def _mask_depth(ds, depths, above=True):
+    ds = ds.reset_coords().to_pandas().set_index("dives")
+    mask = ds.groupby("dives").apply(mask_profile_depth, depths, above)
+    # mask = mask if above else ~mask
+    return mask.values
 
 
 def merge_dimensions(df1, df2, interp_lim=3):

tests/test_physics.py

@@ -7,6 +7,7 @@
     potential_density,
     spice0,
 )
+from glidertools.utils import mask_above_depth, mask_below_depth
 
 
 filenames = "./tests/data/p542*.nc"
@@ -58,6 +59,15 @@ def test_mixed_layer_depth():
     assert mld.max() < 40
 
 
+def test_masking():
+    # We "know" that the mld for this dataset is >10m and <40m
+    mld = mixed_layer_depth(dat, "temp_raw")
+    mask = mask_above_depth(dat, mld)
+    assert dat.depth[mask].max() > 10
+    mask = mask_below_depth(dat, mld)
+    assert dat.depth[mask].max() < 40
+
+
 def test_potential_density():
     pot_den = potential_density(
         dat.salt_raw, dat.temp_raw, dat.pressure, dat.latitude, dat.longitude