sane.py

def get_sf_from_gotm(ds):
    """create shape functions following sane et al. ds is a netcdf dataset output from GOTM"""
    n_layers = 16
    batch = 100
    t_initial = 50
    t_max = t_initial + batch
    
    # choose a time segment to look at
    nuh = ds['nuh'][t_initial:t_max, :, 0, 0]
    mld_surf = ds['mld_surf'][t_initial:t_max, 0, 0]
    z = ds['z'][0, :, 0, 0]

    x_normalized = []
    for t in range(nuh.shape[0]):
        nuh_t = nuh[t]
        # Figure out the nearest depth index to the MLD
        mld_idx_t = np.argmin(np.abs(z - mld_surf[t]))
        surface_to_mld = nuh_t[0:mld_idx_t]
        max_val = np.max(surface_to_mld)
        normalized = surface_to_mld / max_val
        # average into 16 bins
        means = [np.mean(chunk) for chunk in
            np.array_split(normalized, n_layers)
        ]
        means = [0] + means + [0]
        x_normalized.append(means)

    sf = np.array(x_normalized)
    return sf