Luke Skywalker
Luke Skywalker

Reputation: 1521

Numpy broadcasting on multiple arrays

I have a basis for a plane in 3 dimensions: (u, v).

I would like to obtain all linear combinations of this basis to basically go through my whole plane:

for i in [0, 512[ and j in [0, 512[, get all (i * u + j * v).

I need this to be fast so for loops are not really an option. How can I do that with numpy broadcasting?

Had a look at https://docs.scipy.org/doc/numpy-1.13.0/user/basics.broadcasting.html and my impression is that it's not possible to do...

Tried:

# This is an orthonormal basis but there is no guarantee it is
u = np.array([1, 0, 0])
v = np.array([0, 1, 0])
tmp = np.arange(512)
factors = itertools.combinations(tmp, 2)
pixels = factors[0] * u + factors[1] + v

But obviously it does not work.

Is there a solution to this problem? And if yes then how?

Upvotes: 2

Views: 343

Answers (2)

blue note
blue note

Reputation: 29071

A few edits to your code

factors = itertools.product(tmp, tmp)
factors = list(zip(*factors))
factors = np.array(factors)
pixels = factors[0][:, None] * u + factors[1][:, None] + v

First, a math error: you want the Cartesian product, not the combinations.

Now, the actual syntax error: itertools produces a list of [(i1, j1), (i2, j2)..], not a numpy array. So, for your final line to work you have to

  1. zip(*) to convert the list to (i1, i2), (j1, j2) format
  2. make it a numpy array
  3. in the factors[0], factors[1] vectors, which are 1D, add [:, None] to convert the to columns, so that broadcasting works.

Upvotes: 1

w-m
w-m

Reputation: 11232

Multiply (u, v) with a 2D index grid:

ind = np.indices((512, 512))
pixels = ind[0, ..., np.newaxis] * u + ind[1, ..., np.newaxis] * v

>>> %timeit ind = np.indices((512, 512)); pixels = ind[0, ..., np.newaxis] * u + ind[1, ..., np.newaxis] * v
8.06 ms ± 69.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Multiply u with a 1D index range, multiply v with a 1D index range, broadcast and combine to 2D:

i512 = np.arange(512)[:, np.newaxis]
pixels = (i512 * u)[:, np.newaxis, :] + (i512 * v)[np.newaxis, :, :]

>>> %timeit i512 = np.arange(512)[:, np.newaxis]; pixels = (i512 * u)[:, np.newaxis, :] + (i512 * v)[np.newaxis, :, :]
4.06 ms ± 58.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Upvotes: 2

Related Questions