python - 用 python 制作 3D blob 更快的方法？

Question

是否有更好的方法来制作 3D 密度函数？

def make_spot_3d(bright, spread, x0,y0,z0):
    # Create x and y indices
    x = np.linspace(-50, 50, 200)
    y = np.linspace(-50, 50, 200)
    z = np.linspace(-50, 50, 200)

    X, Y, Z = np.meshgrid(x, y, z)
    Intensity = np.uint16(bright*np.exp(-((X-x0)/spread)**2
                                        -((Y-y0)/spread)**2
                                        -((Z-z0)/spread)**2))

    return Intensity

该函数可以生成一个 3D numpy 数组，可以用 mayavi 绘制该数组。

但是，当该函数用于生成 Blob 簇(~100)时，如下所示:

Spots = np.asarray([make_spot_3d(100,2, *loc) for loc in locations])
cluster = np.sum(Spots, axis=0)

产量例如:

执行时间约为 1 分钟(CPU i5)；我打赌这可能会更快。

Answer 1

一个明显的改进是使用广播来评估“稀疏”网格而不是完整的网格上的强度函数，例如:

X, Y, Z = np.meshgrid(x, y, z, sparse=True)

这将我的机器上的运行时间减少了大约 4 倍:

%timeit make_spot_3d(1., 1., 0, 0, 0)
1 loops, best of 3: 1.56 s per loop 

%timeit make_spot_3d_ogrid(1., 1., 0, 0, 0)
1 loops, best of 3: 359 ms per loop

您可以通过对位置、分布和亮度进行矢量化计算来消除列表理解中涉及的开销，例如:

def make_spots(bright, spread, x0, y0, z0):

    # Create x and y indices
    x = np.linspace(-50, 50, 200)
    y = np.linspace(-50, 50, 200)
    z = np.linspace(-50, 50, 200)

    # this will broadcast out to an (nblobs, ny, nx, nz) array
    dx = x[None, None, :, None] - x0[:, None, None, None]
    dy = y[None, :, None, None] - y0[:, None, None, None]
    dz = z[None, None, None, :] - z0[:, None, None, None]
    spread = spread[:, None, None, None]
    bright = bright[:, None, None, None]

    # we can save time by performing the exponentiation over 2D arrays
    # before broadcasting out to 4D, since exp(a + b) == exp(a) * exp(b)
    s2 = spread * spread
    a = np.exp(-(dx * dx) / s2)
    b = np.exp(-(dy * dy) / s2)
    c = np.exp(-(dz * dz) / s2)

    intensity = bright * a * b * c

    return intensity.astype(np.uint16)

其中 bright、spread、x0、y0 和 z0 为 1D向量。这将生成一个 (nblobs, ny, nx, nz) 数组，然后您可以在第一个轴上求和。根据您生成的 blob 数量以及您评估它们的网格有多大，创建此中间数组可能会在内存方面变得相当昂贵。

另一种选择是初始化单个(ny, nx, nz)输出数组并就地计算总和:

def sum_spots_inplace(bright, spread, x0, y0, z0):

    # Create x and y indices
    x = np.linspace(-50, 50, 200)
    y = np.linspace(-50, 50, 200)
    z = np.linspace(-50, 50, 200)

    dx = x[None, None, :, None] - x0[:, None, None, None]
    dy = y[None, :, None, None] - y0[:, None, None, None]
    dz = z[None, None, None, :] - z0[:, None, None, None]
    spread = spread[:, None, None, None]
    bright = bright[:, None, None, None]

    s2 = spread * spread
    a = np.exp(-(dx * dx) / s2)
    b = np.exp(-(dy * dy) / s2)
    c = np.exp(-(dz * dz) / s2)

    out = np.zeros((200, 200, 200), dtype=np.uint16)

    for ii in xrange(bright.shape[0]):
        out += bright[ii] * a[ii] * b[ii] * c[ii]

    return out

这将需要更少的内存，但潜在的缺点是它需要在 Python 中循环。

让您了解相对性能:

def sum_spots_listcomp(bright, spread, x0, y0, z0):
    return np.sum([make_spot_3d(bright[ii], spread[ii], x0[ii], y0[ii], z0[ii])
                   for ii in xrange(len(bright))], axis=0)

def sum_spots_vec(bright, spread, x0, y0, z0):
    return make_spots(bright, spread, x0, y0, z0).sum(0)

# some fake data
bright = np.random.rand(10) * 100
spread = np.random.rand(10) * 100
x0 = (np.random.rand(10) - 0.5) * 50
y0 = (np.random.rand(10) - 0.5) * 50
z0 = (np.random.rand(10) - 0.5) * 50

%timeit sum_spots_listcomp(bright, spread, x0, y0, z0)
# 1 loops, best of 3: 16.6 s per loop

%timeit sum_spots_vec(bright, spread, x0, y0, z0)
# 1 loops, best of 3: 1.03 s per loop

%timeit sum_spots_inplace(bright, spread, x0, y0, z0)
# 1 loops, best of 3: 330 ms per loop