Is there a faster implementation of the following code?

Question

I have a one-dimensional numpy array, which is quite large in size. For each entry of the array, I need to produce a linearly spaced sub-array upto that entry value. Here is what I have as an example.

import numpy as np
a = np.array([2, 3])
b = np.array([np.linspace(0, i, 4) for i in a])

In this case there is linear space of size 4. The last statement in the above code involves a for loop which is rather slow if a is very large. Is there a trick to implement this in numpy itself?

Answer 1

You can phrase this as an outer product:

In [37]: a = np.arange(100000)

In [38]: %timeit np.array([np.linspace(0, i, 4) for i in a])
1 loop, best of 3: 1.3 s per loop

In [39]: %timeit np.outer(a, np.linspace(0, 1, 4))
1000 loops, best of 3: 1.44 ms per loop

The idea is to a take a unit linspace and then scale it separately by each element of a.

As you can see, this gives ~1000x speed up for n=100000.

For completeness, I'll mention that this code has slightly different roundoff properties than your original version (likely not an issue in practical applications):

In [52]: np.max(np.abs(np.array([np.linspace(0, i, 4) for i in a]) -
    ...:               np.outer(a, np.linspace(0, 1, 4))))
Out[52]: 1.4551915228366852e-11

P. S. An alternative way to express the idea is by using element-wise multiplication with broadcasting (based on a suggestion by @Scott Gigante):

In [55]: %timeit a[:, np.newaxis] * np.linspace(0, 1, 4)
1000 loops, best of 3: 1.48 ms per loop

P. P. S. See the comments below for further ideas on making this faster.