Самый быстрый метод билинейного взвешивания двумерного облака точек на сетке

Имея точку (0.8, 0.6) и интенсивность (3), можно выполнить билинейную обратную интерполяцию одной точки на сетку 2x2 с целочисленными индексами (0,0) -> (1,1).

Точки и сетка ориентированы так, что первая запись увеличивается вниз, а вторая запись увеличивается вправо.

В такой сетке вес только указанной выше координаты становится:

  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 в качестве выходных массивов. Мне нужно выполнить вышеуказанное для очень большого (десятки миллионов) числа координат, где координаты имеют значения от 0,0 до 4095,0. Я написал подпрограмму numpy ниже, и хотя она достаточно быстра (1,25 секунды) для 100_000 точек, я бы предпочел, чтобы она была быстрее, так как мне нужно вызывать ее несколько раз для 10_000_000 точек данных, которые у меня есть.

Я подумал о векторизации кода numpy вместо цикла for, но затем я генерирую в основном пустой массив 4096x4096 для каждой точки, которую затем суммирую. Для этого потребуется 1000 ТБ оперативной памяти.

Я также пробовал наивную реализацию в cupy, но поскольку я использую цикл for, он становится слишком медленным.

В моем коде я генерирую взвешенный массив 2x2 для каждой точки, умножаю массив на интенсивность и добавляю их к основным массивам путем нарезки. Есть ли способ лучше?

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')

person TomNorway    schedule 18.01.2021    source источник


Ответы (2)


Вы бы хотели использовать 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

Кроме этого, это будет невозможно с этим подходом. Поскольку индексация целочисленного массива не обновляется постепенно.

Например,

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.])
person armamut    schedule 18.01.2021
comment
Хороший! np.add.at был гениален! Я немного поискал, узнал о np.bincount и написал еще один ответ ниже. Пожалуйста, взгляните! - person TomNorway; 18.01.2021

Спасибо @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

Пустая реализация:

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

И реализация чашки:

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
person TomNorway    schedule 18.01.2021
comment
Хорошо сделано! Но я пытаюсь понять код. Результаты отличаются от вашего вопроса или я что-то упустил? - person armamut; 18.01.2021
comment
Ага, понятно! Я учусь новому каждый день! - person armamut; 18.01.2021