循环遍历任意维度的 numpy 数组中的元素

Question

我有如下代码：

def infball_proj(mu, beta):
    newmu = np.zeros(mu.shape)
    if len(mu.shape) == 2:
        for i in range(mu.shape[0]):
            for j in range(mu.shape[1]):
                if np.abs(mu[i,j]) > beta:
                    newmu[i,j] = np.sign(mu[i,j]) * beta
                else:
                    newmu[i,j] = mu[i,j]
        return newmu
    elif len(mu.shape) == 1:
        for i in range(mu.shape[0]):
            if np.abs(mu[i]) > beta:
                newmu[i] = np.sign(mu[i]) * beta
            else:
                newmu[i] = mu[i]
        return newmu

有没有更聪明的方法来做到这一点，所以我不必编写 2 个不同的案例？ 如果我可以有一个可以缩放到任意维度（即轴数）的版本，那就太好了。

Answer 1

这样的事情应该可以完成这项工作：

newmu = np.where(np.abs(mu) > beta, np.sign(mu) * beta, mu)

或者，如果我的逻辑正确，

newmu = np.minimum(np.abs(mu), beta) * np.sign(mu)

Answer 2

mu[np.abs(mu)>beta] = np.sign(mu[np.abs(mu)>beta]) * beta

np.abs(mu)>beta将创建一个布尔数组，然后可用于布尔索引。

LHS mu[np.abs(mu)>beta]将返回由布尔索引选择的元素的视图，并且可以分配给您想要的值，即 RHS。

记住：尽量避免 NumPy 数组的 for 循环，因为它非常低效。