2

I'm trying to improve a function used in a rasterizer in Python. Currently, I sequentially calculate the x-y coordinates of all pixels inside multiple bounding boxes. The bounding boxes will contain the minimum and maximum x-y coordinates for each box. I have a fairly good function that can be looped over and a partially vectorized version to complete this task. Any suggestions to improve the speed using Numpy?

Given:

  • an array of shape(N, 2) containing the minimum x & y coordinates of (N) bounding boxes
    • x_y_mins = [[62, 53], [47, 13], ...]
  • an array of shape(N, 2) containing the maximum x & y coordinates of (N) bounding boxes
    • x_y_maxs = [[63, 54], [54, 23], ...]

Construct:

  • an array of shape(M, 2) stacking all the M x-y coordinate tuples corresponding to the (m) pixel indices inside each of the (N) bounding boxes
    • pixel_coords = [[62, 53], [62, 54], [63, 53], [63, 54], [47, 13], [47, 14], ...]

Loop version:

import numpy as np
# x_y_mins: np.ndarray[np.int32] of shape [N, 2]
# x_y_maxs: np.ndarray[np.int32] of shape [N, 2]
x_y_mins = np.random.randint(0, 200, size=(60000, 2), dtype=np.int32)
x_y_maxs = x_y_mins + np.random.randint(1, 4, size=(60000, 2), dtype=np.int32)

all_pixel_coords_list = []
for x_y_min, x_y_max in zip(x_y_mins, x_y_maxs):

    x_coords = np.arange(x_y_mins[0], x_y_maxs[0])
    y_coords = np.arange(x_y_mins[1], x_y_maxs[1])
    x_grid, y_grid = np.ix_(x_coords, y_coords)

    pixel_coords = np.empty([x_y_maxs[0] - x_y_mins[0], x_y_maxs[1] - x_y_mins[1], 2])
    pixel_coords[..., 0] = x_grid
    pixel_coords[..., 1] = y_grid
    all_pixel_coords_list.append(pixel_coords.reshape(-1, 2))
all_pixel_coords = np.concatenate(all_pixel_coords_list, axis=0)

Alternatively, this can also be solved with np.mgrid / np.meshgrid / np.repeat with broadcasting. These are generally slightly slower

The above looped code is a faster solution to this question Using numpy to build an array of all combinations of two arrays and itself derived from Cartesian product of x and y array points into single array of 2D points

These questions attempt to solve a single x-y coordinate list from a single combination between one list of x coordinates and one list of y coordinates.

My issue is that these coordinates lists are then looped over slowly in Python and built into a single combined array after calling np.ix_/np.meshgrid 60,000 times. Hoping to vectorize everything and solve all coordinate lists from the concatenation of successive combinations of x & y coordinates all at once.

(mostly) Vectorized version:

import numpy as np
# x_y_mins: np.ndarray[np.int32] of shape [N, 2]
# x_y_maxs: np.ndarray[np.int32] of shape [N, 2]
x_y_mins = np.random.randint(0, 200, size=(60000, 2), dtype=np.int32)
x_y_maxs = x_y_mins + np.random.randint(1, 4, size=(60000, 2), dtype=np.int32)

x_lengths = x_y_maxs[:, 0] - x_y_mins[:, 0]
y_lengths = x_y_maxs[:, 1] - x_y_mins[:, 1]

x_coords = np.repeat(x_y_maxs[:, 0] - x_lengths.cumsum(), x_lengths) + np.arange(x_lengths.sum())
x_repeats = np.repeat(y_lengths, x_lengths)
x_grid = np.repeat(x_coords, x_repeats)

y_length_cumsum = y_lengths.cumsum()
y_coords = np.repeat(x_y_maxs[:, 1] - y_length_cumsum, y_lengths) + np.arange(y_lengths.sum())
y_length_indices = np.concatenate([[0], y_length_cumsum])
y_coords_splits = [slice(first, second) for first, second in zip(y_length_indices, y_length_indices[1:])]
y_grid = np.concatenate([extended for y_coords_slice, x_repeat in zip(y_coords_splits, x_lengths) for extended in [y_coords[y_coords_slice]] * x_repeat])

all_pixel_coords = np.stack((x_grid, y_grid), axis=1)
Kyle Vrooman
  • 33
  • 1
  • 3
  • This looks like a 2d version of a more common "how to take multiple slices from a 1d array?". Do you concatenate a list of slices: `[arr[x0:x1], arr[x2:x3], arr[x4:x5]...]`, or first concatenate `aranges`, and then index, or something "faster"? If the slices all have the same length, it's easy to "vectorize" the index generation. But if they vary in length it's is hard to improve on some sort of concatenate of a list comprehension. – hpaulj Sep 16 '21 at 20:04
  • Each of the ```aranges``` will have different lengths and thus the 60,000 iterations if stacked would make a jagged array. Agree that a standarized length would be much easier to vectorize / run opeations over the axes. Hoping for insights from the community on vectorizing "jagged" arrays. The approach I took was flattening everything to a 1-D array and keeping track of the appropriate indexes. – Kyle Vrooman Sep 17 '21 at 17:48

0 Answers0