Советы по удалению вложенных циклов в коде numpy, включая условия и замаскированные данные - PullRequest
0 голосов
/ 12 июня 2019

Я пишу функцию для расчета смешанного слоя океана, но пока лучший способ, который я нашел, включает множество вложенных циклов, что кажется плохим и уже медленным в моих тестовых данных.Есть ли способ рефакторинга с учетом того, что код также включает условные проверки?

У меня есть рабочий, но медленный пример с использованием вложенных циклов, и я хотел бы ускорить это.Я начал с изменения, например, на «notmasked_edges», но уже не могу понять, что происходит в моей голове.

Это показывает трехмерную температуру и соленость (глубину, широту, долготу) и глубину 1D.координата (глубина).Возвращает 2D "смешанную глубину слоя" (широта, долгота)

def mld_sigmat(thetao, so, mid_depths):
    """To make the Sigma-T based MLD, defined here:
    'Experimental and diagnostic protocol for the physical component of the
    CMIP6 Ocean Model Intercomparison Project (OMIP)'
    """

    buoyancy_crit = 0.0003  # m/s2, as defined in paper above

    nk, nj, ni = thetao.shape

    mld = np.ma.masked_all(shape=(nj, ni))

    for jj in range(nj):
        for ii in range(ni):
            edges = np.ma.flatnotmasked_edges(thetao[:, jj, ii])
            if edges is None: continue

            # function_to_calculate_density takes 2 N-dim arrays and a
            # 1D depth coordinate which is the same length as the last
            # dimension of the N-dim arrays
            dens_local = np.ma.masked_all(shape=(nk))
            dens_displaced = np.ma.masked_all(shape=(nk))
            for kk in range(edges[0], edges[1]):
                dens_local[kk] = function_to_calculate_density(thetao[kk, jj, ii], this_so[kk, jj, ii], mid_depths[kk]) + 1e3
                dens_displaced[kk] = function_to_calculate_density(thetao[0, jj, ii], this_so[0, jj, ii], mid_depths[kk]) + 1e3

            buoyancy_change = -9.8 * (dens_displaced - dens_local) / dens_local

            indices = np.where(buoyancy_change > buoyancy_crit)[0]

            if len(indices) == 0:  # Full depth convection (criteria never met)
                mld[jj, ii] = mid_depths[edges[1]]
            elif indices[0] == 0:  # No convection
                mld[jj, ii] = mid_depths[edges[0]]
            else:  # MLD is somewhere in between. Take the depth 1 before and then add the extra bit on
                ddepth = mid_depths[indices[0]] - mid_depths[indices[0]-1]
                dbuoyancy = buoyancy_change[indices[0]] - buoyancy_change[indices[0]-1]
                ddepth_dbuoyancy = ddepth / dbuoyancy
                mld[jj, ii] = mid_depths[indices[0]-1] + ddepth_dbuoyancy * (buoyancy_change[indices[0]] - buoyancy_crit)

    return mld
Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...