只有在多维数组的Python中值大于零时才添加值

Question

我有这样一个数组: tmp.shape = (128, 64, 64)

我正在计算128轴上的所有零,如下所示:

nonzeros = np.count_nonzero(tmp, axis=0) // shape = (64, 64)

我有一个阵列 c.shape = (64, 64)

现在我想沿着128轴将c的值添加到tmp,但前提是tmp的值> 0:

for i in range(tmp.shape[0]):
    axis1 = tmp[i,:,:]
    tmp[i, :, :] += (c / nonzeros) // only if tmp[i, :, :] > 0

这可能在短期内完成吗？或者我必须使用多个循环？我希望任何人都可以建议没有其他循环的解决方案

这样的东西不起作用:

tmp[i, tmp > 0.0, tmp > 0.0] += (c / nonzeros)

IndexError:数组的索引太多了

长版

for i in range(tmp.shape[0]):
    for x in range(tmp.shape[1]):
        for y in range(tmp.shape[2]):
            pixel = tmp[i, x, y]
            if pixel != 0:
                pixel += (c[x,y] / nonzeros[x,y])

Answer 1

好吧，您基本上是将广播添加c/nonzeros到 tmp 数组中，除了 tmp 元素为零的地方。因此，一种方法是预先存储掩码0s，添加进去c/nonzeros，最后使用掩码来重置tmp元素。

因此，实施将是 -

mask = tmp==0
tmp+= c/nonzeros
tmp[mask] = 0

运行时测试

方法 -

# @DSM's soln
def fast(tmp, c, nonzeros):
    return tmp + np.where(tmp > 0, c/nonzeros, 0)

# Proposed in this post
def fast2(tmp, c, nonzeros):
    mask = tmp==0
    tmp+= c/nonzeros
    tmp[mask] = 0
    return tmp

时间安排 -

In [341]: # Setup inputs
     ...: M,N = 128,64
     ...: tmp = np.random.randint(0,10,(M,N,N)).astype(float)
     ...: c = np.random.rand(N,N)*100
     ...: nonzeros = np.count_nonzero(tmp, axis=0)
     ...: 
     ...: # Make copies for testing as input would be edited with the approaches
     ...: tmp1 = tmp.copy()
     ...: tmp2 = tmp.copy()
     ...: 

In [342]: %timeit fast(tmp1, c, nonzeros)
100 loops, best of 3: 2.22 ms per loop

In [343]: %timeit fast2(tmp2, c, nonzeros)
1000 loops, best of 3: 1.61 ms per loop

更短的替代方案

如果您正在寻找紧凑的代码，这里是另一个使用掩码non-0s进行广播元素级乘法并c/nonzeros添加到tmp，从而获得单行解决方案，如下所示 -

tmp += (tmp!=0)*(c/nonzeros)

注意：为了避免除以0，我们可以nonzeros使用0s除 之外的任何内容进行编辑0，1然后使用发布的方法，如下所示 -

nonzeros = np.where(nonzeros > 0, nonzeros, 1)