循环遍历任意维度的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

有没有更聪明的方法可以做到这一点，所以我不必写两种不同的情况？如果我有一个可以缩放到任意尺寸（即轴数）的版本，那就太好了。

Answer 1

类似的事情应该可以完成：

uniqueEvent

或者，如果我理解正确，

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

Answer 2

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

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

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

记住：尝试避免NumPy数组的for循环，因为它效率很低。