find_mld.py

import xarray as xr

def expfit(phi, dim):
    
    def exp_func(z, phi_bottom, delta_phi, z0):
        return phi_bottom - delta_phi * np.exp(- z / z0)
    
    mean_phi_max = phi.max(dim).mean()
    mean_phi_min = phi.min(dim).mean()
    delta_phi = mean_phi_max - mean_phi_min
    
    fit = phi.curvefit(phi[dim], 
                       exp_func,  
                       p0={'phi_bottom': mean_phi_max,
                           'delta_phi': delta_phi,
                           'z0': 1000})
    phi_fit = exp_func(phi[dim], 
               fit.curvefit_coefficients.sel(param='phi_bottom'),
               fit.curvefit_coefficients.sel(param='delta_phi'),
               fit.curvefit_coefficients.sel(param='z0'))
    return phi_fit


def mld_thres(phi, phi_step, dim='depth', ref_depth=10):
    """
    Find the Mixed Layer Depth (MLD) using a threshold criterium
    
    Paramaters
    ----------
    phi : array_like
        Vertical profiles, either density, salinity, temperature
    phi_step : float
        The step tin
    ref_depth : float, optional
        The depth at which the reference value is taken

    Returns
    -------
    mld_interp : array_like
        The mixed layer depth interpolated between depth levels
    """
    # Define the depth vector
    z = phi[dim] 
    # Drop the surface or first level
    z_sdrop = z.isel(**{dim: slice(1, None)})
    phi_sdrop = phi.isel(**{dim: slice(1, None)})
    
    # Compute regression coefficients for linear interpolation
    delta_z = z.diff(dim)
    delta_phi = phi.diff(dim)
    alpha = delta_z / delta_phi
    beta = z_sdrop - alpha * phi_sdrop
    
    # Get the reference denisty at the reference depth
    phi_ref = phi_sdrop.sel(**{dim: ref_depth})
    # Evaluate the density at the base of the mixed layer
    phi_mlb = phi_ref + phi_step
    
    # Mask profile where the density is larger or equals than the density at the base of the mixed layer and smaller or equals than the reference depth
    not_ml = (phi_sdrop >= phi_mlb) & (z_sdrop >= ref_depth)
    success = (not_ml.sum(dim) > 0)
    # Find the first depth level at which the
    #z_bcast = z_sdrop * xr.ones_like(phi_sdrop)
    ind_below_ml = (z_sdrop.where(not_ml).fillna(9e6).astype(int)
                                        .argmin(dim).compute())#.astype(int)#
                                         #.argmin(dim))
                                         #.compute())
    # Linear interpolation between the neighbouring poin to find the MLD
    #z_below_ml = z.isel(**{dim: ind_below_ml}).where(success)
    z_below_ml = z.sel(**{dim: z_sdrop[ind_below_ml]})#.where(success)
    #z_below_ml = z_below_ml.where(z_below_ml > ref_depth)
    mld_interp = (alpha.isel(**{dim: ind_below_ml}) * phi_mlb 
                  + beta.isel(**{dim: ind_below_ml}))
    #mld_interp = alpha.sel(**{dim: ind_below_ml}) * phi_mlb + beta.sel(**{dim: ind_below_ml})
    #return z_below_ml
    return mld_interp.drop(dim)    
    

def mld_rvar(phi, dim='depth'):
    # Get the depth vector and its length
    z = phi[dim]
    n = z.size

    # Define a running window of maximum size of the vector length
    running_window = phi.rolling(**{dim: n}, min_periods=1)

    # Compute the cumulative standard deviation (delta), 
    # maximum, minimum and amplitude (sigma)
    delta = running_window.std()
    phi_min = running_window.min()
    phi_max = running_window.max() 
    sigma = phi_max - phi_min

    # The ratio of the standard deviation and the amplitude yields
    # the khi quantity
    chi = delta / sigma
    # I found that detrending the khi quantity was necessary to avoid
    # detecting minima well below the mixed layer due to the variance 
    # decrease with depth
    # The retained solution for the moment is to remove a exponential curve
    # but this could be improved in the future
    trend = expfit(chi, dim)
    chi_dtr = (chi - trend).isel(**{dim: slice(1, None)})
    
    # Find the depth zn1 where khi is minimum
    k_zn1 = chi_dtr.fillna(9e6).argmin(dim=dim).compute()
    zn1 = z.isel(**{dim: k_zn1})
    
    # Find the first depth zn2 above zn1 that satisfies a small 
    # change in standard deviation
    delta_diff = delta.diff(dim=dim)                          
    r = delta_diff / delta_diff.sel(**{dim: z[k_zn1 + 1]})
    not_ml = (r <= 0.3) & (z <= zn1)
    success = (not_ml.sum(dim) > 0)
    k_zn2 = z.where(not_ml).fillna(-9e6).argmax(dim).compute()
    zn2 = z.isel(**{dim: k_zn2}).where(success)

    return zn2