Finding Mixed Layer Depth with numba, dask and xarray

find_mld_with_numba.py

def mld_thres(phi, phi_step, ref_depth=10, dim='depth'):
    
    # Get the depth vector
    z = phi[dim]
    # Find the index closest to the reference depth
    k_ref = int((np.abs(z - ref_depth)).argmin())
    
    from numba import guvectorize
    @guvectorize(['void(float64[:], float64[:], intp, float64, float64[:])'],
              '(n),(n),(),()->()', nopython=True)
    def mld_gufunc(phi, z, k_ref, phi_step, res):
        phi_ref = phi[k_ref]
        phi_mlb = phi_ref + phi_step
        for k in range(k_ref, phi.shape[0]):
            if phi[k] >= phi_mlb:
                k_mlb = k
                break
        # Make a linear interpolation to find the approximate MLD
        delta_z = z[k_mlb - 1] - z[k_mlb]
        alpha = (phi[k_mlb - 1] - phi[k_mlb]) / delta_z
        beta = z[k_mlb] - alpha * phi[k_mlb]
        res[0] = alpha * phi_mlb + beta
        
    mld = xr.apply_ufunc(mld_gufunc, phi, z, k_ref, phi_step, 
                         input_core_dims=[[dim], [dim], [], []],
                         dask='parallelized')
    return mld