Я пишу функцию для расчета смешанного слоя океана, но пока лучший способ, который я нашел, включает множество вложенных циклов, что кажется плохим и уже медленным в моих тестовых данных.Есть ли способ рефакторинга с учетом того, что код также включает условные проверки?
У меня есть рабочий, но медленный пример с использованием вложенных циклов, и я хотел бы ускорить это.Я начал с изменения, например, на «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