查找numpy数组中最接近meshgrids的元素的索引

cunj1qz1  于 2023-06-29  发布在  其他
关注(0)|答案(1)|浏览(146)

假设我有以下网格omegaphiyphix

import numpy as np

phi=np.linspace(-0.75,0.75,num=4)
omega=np.linspace(-100,100,num=5)
omega,phiy,phix=np.meshgrid(omega,phi,phi,sparse=True,indexing='ij')

rand=np.random.rand(5,4)

然后我想得到下面的C

def f(j,k):
   return np.argmin(np.abs(phi[j]**2+np.sin(phi[k])-phi))

C=np.empty((5,4,4))
for i in range(5):
  for j in range(4):
    for k in range(4):
       C[i,j,k]=rand[i,f(j,k)]

有没有更紧凑更好的方法来得到它,使用一些meshgrids的工具?

dfddblmv

dfddblmv1#

你可以通过消除for循环并只使用向量化操作来计算你的f函数来更快地做到这一点:

# f_mat[i,j] contains f(i,j)
f_mat = np.abs((phi[None].T**2 + np.sin(phi))[...,None] - phi).argmin(-1)
C = rand[:,f_mat]

相关问题