將二維點雲雙線性加權到網格上的最快方法

Question

給定一個點(0.8, 0.6)和強度(3) ，可以將單個點雙線性“反向插值”到具有 integer 索引(0,0) -> (1,1)的 2x2 網格上。

點和網格的方向是第一個條目向下增加，第二個條目向右增加。

在這樣的網格上，上面坐標的權重變為：

  0.08 |  0.12  
----------------
  0.32 |  0.48

我們可以乘以與坐標相關的強度，得到一個 2x2 網格雙線性加權強度：

  0.24 |  0.36 
----------------
  0.96 |  1.44

並且可以像這樣繪制：

對於幾個點，可以將它們加權到同一個數組上（下面的完整代碼）：

points = np.array([(0.8, 0.6), (2.2, 2.6),(5, 1), (3, 4.2), (8.5, 8.2)])
intens = np.array([3, 3, 1, 1, 2])
image, weight = bilinear(points, intens)

對於我的工作，我需要weights和intensity*weights為 output arrays。 我需要在非常大（數百萬）的坐標上執行上述操作，其中坐標的值從 0.0 到 4095.0。 我在下面編寫了一個 numpy 例程，雖然它在 100_000 個點上相當快（1.25 秒），但我更希望它更快，因為我需要在我擁有的 10_000_000 個數據點上多次調用它。

我考慮對 numpy 代碼進行矢量化而不是 for 循環，但隨后我為每個點生成一個 4096x4096 的大部分空數組，然后我將對其求和。 那將需要 1000 TB 的內存。

我也嘗試過在 cupy 中實現一個簡單的實現，但是由於我使用了 for 循環，它變得太慢了。

在我的代碼中，我為每個點生成一個 2x2 加權數組，將數組乘以強度，然后通過切片將它們添加到主 arrays 中。 有沒有更好的辦法？

import numpy as np

def bilinear(points, intensity):
    """Bilinear weighting of points onto a grid.
    Extent of grid given by min and max of points in each dimension
    points should have shape (N, 2)
    intensity should have shape (N,)
    """
    floor = np.floor(points)
    ceil = floor + 1
    floored_indices = np.array(floor, dtype=int)
    low0, low1 = floored_indices.min(0)
    high0, high1 = floored_indices.max(0)
    floored_indices = floored_indices - (low0, low1)
    shape = (high0 - low0 + 2, high1-low1 + 2)
    
    weights_arr = np.zeros(shape, dtype=float)
    int_arr = np.zeros(shape, dtype=float)
    
    upper_diff = ceil - points
    lower_diff = points - floor
    w1 = np.prod((upper_diff), axis=1)
    w2 = upper_diff[:,0]*lower_diff[:,1]
    w3 = lower_diff[:,0]*upper_diff[:,1]
    w4 = np.prod((lower_diff), axis=1)
    
    for i, index in enumerate(floored_indices):
        s = np.s_[index[0]:index[0]+2, index[1]:index[1]+2]
        weights = np.array([[w1[i], w2[i]], [w3[i], w4[i]]])
        weights_arr[s] += weights
        int_arr[s] += intensity[i]*weights
    return int_arr, weights_arr

rng = np.random.default_rng()
N_points = 10_000 #  use 10_000 so it is quick
image_shape = (256, 256) # Use 256 so it isn't so big

points = rng.random((N_points, 2)) * image_shape
intensity = rng.random(N_points)

image, weight = bilinear(points, intensity)

為了測試代碼，我還提供了以下繪圖代碼 - 僅用於少量（~10）點，否則散點將覆蓋整個圖像。

import matplotlib.pyplot as plt
floor = np.floor(points) - 0.5
lower, left = floor.min(0)
upper, right = (floor).max(0) + 2
extent = (left, right, upper, lower)

fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(6,3))
ax1.scatter(*points[:,::-1].T, c='red')
im1 = ax1.imshow(weight, clim=(image.min(), image.max()), extent=extent)
ax1.set(title='Weight', xlim=(left - 1, right + 1), ylim = (upper + 1, lower - 1))
colorbar(im1)

ax2.scatter(*points[:,::-1].T , c='red')
im2 = ax2.imshow(image, extent=extent)

ax2.set(title='Weight x Intensity', xlim=(left - 1, right + 1), ylim = (upper + 1, lower - 1))
colorbar(im2)
plt.tight_layout()
plt.show()

# If labeling the first point
# ax1.text(*points[0].T, f"({points[0,0]}, {points[0,1]})", va='bottom', ha='center', color='red')
# ax2.text(*points[0].T, f"({points[0,0]}, {points[0,1]}, {intens[0]})", va='bottom', ha='center', color='red')

Answer 1

您想使用np.add.at 。 見， https://numpy.org/doc/stable/reference/generated/numpy.ufunc.at.html

def bilinear_2(points, intensity):
    # Create empty matrices, starting from 0 to p.max
    w = np.zeros((points[:, 0].max().astype(int) + 2, points[:, 1].max().astype(int) + 2))
    i = np.zeros_like(w)
    
    # Calc weights
    floor = np.floor(points)
    ceil = floor + 1
    upper_diff = ceil - points
    lower_diff = points - floor
    w1 = upper_diff[:, 0] * upper_diff[:, 1]
    w2 = upper_diff[:, 0] * lower_diff[:, 1]
    w3 = lower_diff[:, 0] * upper_diff[:, 1]
    w4 = lower_diff[:, 0] * lower_diff[:, 1]
    
    # Get indices
    ix, iy = floor[:, 0].astype(int), floor[:, 1].astype(int)

    # Use np.add.at. See, https://numpy.org/doc/stable/reference/generated/numpy.ufunc.at.html
    np.add.at(w, (ix, iy), w1)
    np.add.at(w, (ix, iy+1), w2)
    np.add.at(w, (ix+1, iy), w3)
    np.add.at(w, (ix+1, iy+1), w4)

    np.add.at(i, (ix, iy), w1 * intensity)
    np.add.at(i, (ix, iy+1), w2 * intensity)
    np.add.at(i, (ix+1, iy), w3 * intensity)
    np.add.at(i, (ix+1, iy+1), w4 * intensity)
    
    # Clip (to accomodate image size to be the same as your bilinear function)
    iix, iiy = points[:, 0].min().astype(int), points[:, 1].min().astype(int)
    i, w = i[iix:, iiy:], w[iix:, iiy:]
    return i, w

# At 10_000 samples:
%time image, weight = bilinear(points, intensity)
%time image_2, weight_2 = bilinear_2(points, intensity)
>>>
CPU times: user 178 ms, sys: 3.73 ms, total: 182 ms
Wall time: 185 ms
CPU times: user 9.63 ms, sys: 601 µs, total: 10.2 ms
Wall time: 10 ms

# These tests passes
np.testing.assert_allclose(weight, weight_2)
np.testing.assert_allclose(image, image_2)

# At 100K samples
N_points = 100_000
image_shape = (256, 256)
points = rng.random((N_points, 2)) * image_shape
intensity = rng.random(N_points)

%time image_2, weight_2 = bilinear_2(points, intensity)

CPU times: user 115 ms, sys: 66 ms, total: 181 ms
Wall time: 181 ms


# At 10M samples
N_points = 10_000_000
image_shape = (256, 256)
points = rng.random((N_points, 2)) * image_shape
intensity = rng.random(N_points)

%time image_2, weight_2 = bilinear_2(points, intensity)

CPU times: user 8.23 s, sys: 656 ms, total: 8.88 s
Wall time: 9.31 s

除此之外，這種方法是不可能的。 因為 integer 數組索引不會增量更新。

例如，

a = np.zeros(5)
a[np.array((1,1,2))] += 1
a
>>> array([0., 1., 1., 0., 0.])

但;

a = np.zeros(5)
np.add.at(a, ([1,1,2]), 1)
a
>>> array([0., 2., 1., 0., 0.])

Answer 2

感謝@armamut 的好答案，它啟發了我看了一下，我發現np.bincount ，它也在 cupy 中實現。 事實證明 bincount 實現速度更快，而 cupy 實現速度非常快，后者可能會進一步改進。 因為我不得不處理幾個元組才能讓它工作。

# Timings
points = np.random.random((10_000_000, 2)) * (256, 256)
intens = np.random.random((10_000_000))

pcupy = cp.asarray(points)
icupy = cp.asarray(intens)
%time bilinear_bincount_cupy(pcupy, icupy)
%time bilinear_bincount_numpy(points, intens)
%time bilinear_2(points, intens)
Wall time: 456 ms
Wall time: 2.57 s
Wall time: 5.37 s

numpy 實現：

def bilinear_bincount_numpy(points, intensities):
    """Bilinear weighting of points onto a grid.
    Extent of grid given by min and max of points in each dimension
    points should have shape (N, 2)
    intensity should have shape (N,)
    """
    floor = np.floor(points)
    ceil = floor + 1
    floored_indices = np.array(floor, dtype=int)
    low0, low1 = floored_indices.min(0)
    high0, high1 = floored_indices.max(0)
    floored_indices = floored_indices - (low0, low1)
    shape = (high0 - low0 + 2, high1-low1 + 2)

    upper_diff = ceil - points
    lower_diff = points - floor

    w1 = np.prod((upper_diff), axis=1)
    w2 = upper_diff[:,0]*lower_diff[:,1]
    w3 = lower_diff[:,0]*upper_diff[:,1]
    w4 = np.prod((lower_diff), axis=1)

    shifts = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
    indices = floored_indices[:, None] + shifts
    indices = (indices * (shape[1], 1)).sum(-1)
    weights = np.array([w1, w2, w3, w4]).T

    weight_bins = np.bincount(indices.flatten(), weights=weights.flatten())
    intens_bins = np.bincount(indices.flatten(), weights=(intensities[:, None]*weights).flatten())

    all_weight_bins = np.zeros(np.prod(shape))
    all_intens_bins = np.zeros(np.prod(shape))

    all_weight_bins[:len(weight_bins)] = weight_bins
    all_intens_bins[:len(weight_bins)] = intens_bins

    weight_image = all_weight_bins.reshape(shape)
    intens_image = all_intens_bins.reshape(shape)
    return intens_image, weight_image

和 cupy 的實現：

def bilinear_bincount_cupy(points, intensities):
    """Bilinear weighting of points onto a grid.
    Extent of grid given by min and max of points in each dimension
    points should be a cupy array of shape (N, 2)
    intensity should be a cupy array of shape (N,)
    """
    floor = cp.floor(points)
    ceil = floor + 1
    floored_indices = cp.array(floor, dtype=int)
    low0, low1 = floored_indices.min(0)
    high0, high1 = floored_indices.max(0)
    floored_indices = floored_indices - cp.array([low0, low1])
    shape = cp.array([high0 - low0 + 2, high1-low1 + 2])

    upper_diff = ceil - points
    lower_diff = points - floor

    w1 = upper_diff[:, 0] * upper_diff[:, 1]
    w2 = upper_diff[:, 0] * lower_diff[:, 1]
    w3 = lower_diff[:, 0] * upper_diff[:, 1]
    w4 = lower_diff[:, 0] * lower_diff[:, 1]

    shifts = cp.array([[0, 0], [0, 1], [1, 0], [1, 1]])
    indices = floored_indices[:, None] + shifts
    indices = (indices * cp.array([shape[1].item(), 1])).sum(-1)
    weights = cp.array([w1, w2, w3, w4]).T

    # These bins only fill up to the highest index - not to shape[0]*shape[1]
    weight_bins = cp.bincount(indices.flatten(), weights=weights.flatten())
    intens_bins = cp.bincount(indices.flatten(), weights=(intensities[:, None]*weights).flatten())
    
    # So we create a zeros array that is big enough
    all_weight_bins = cp.zeros(cp.prod(shape).item())
    all_intens_bins = cp.zeros_like(all_weight_bins)
    # And fill it here
    all_weight_bins[:len(weight_bins)] = weight_bins
    all_intens_bins[:len(weight_bins)] = intens_bins

    weight_image = all_weight_bins.reshape(shape.get())
    intens_image = all_intens_bins.reshape(shape.get())
    return intens_image, weight_image