Is there a way to do array broadcasting faster, or without using loops?

Question

I would like to optimize the following code. In real applications, the number of snapshots can be 10000, and for each of them, the lengths can be length_list = np.arange(0, 13, 0.01). So, it takes ages for it to calculate the properties of matter that I want.

import numpy as np

list_shot = list(np.arange(0, 100, 50)) 
list_lens = list(np.arange(10))
list_total = []

for shot in list_shot:
    print("-", shot)
    list_out = []

    for lens in list_lens:
        val = shot * lens
        list_out.append(val)
    print(list_out)

    list_total.append(list_out) 

Basically, the above code is just a simple representation of a code that I need running much faster than it's presently running.

Answer 1

What you're trying to do can actually be done very quickly, without any loops, and can even be ported to make valuable use of a GPU.

Here's a quick example that you can test yourself in octave-online:

shots=0:50:250
lens = 0:10
bsxfun(@times,shots',lens)

The outputs are as follows:

shots=[0    50   100   150   200   250]
lens =[0    1    2    3    4    5    6    7    8    9   10]
bsxfun(@times,shots',lens)
ans =

      0      0      0      0      0      0      0      0      0      0      0
      0     50    100    150    200    250    300    350    400    450    500
      0    100    200    300    400    500    600    700    800    900   1000
      0    150    300    450    600    750    900   1050   1200   1350   1500
      0    200    400    600    800   1000   1200   1400   1600   1800   2000
      0    250    500    750   1000   1250   1500   1750   2000   2250   2500

For shots up to 10,000 in steps of 50, the calculation takes 5.6e-5 seconds on a CPU:

tic;shots=0:50:10000;lens=0:10;bsxfun(@times,shots',lens);toc;
Elapsed time is 5.79357e-05 seconds.

With an extremely un-optimized FOR LOOP in octave-online, I repeated the calculation 100,000 times and the total CPU time was 1.97 seconds:

tic;for i=1:1e5;shots=0:50:10000;lens=0:10;bsxfun(@times,shots',lens);end;toc;
Elapsed time is 1.96675 seconds.

Believe it or not, this single function called bsxfun which stands for "binary singleton expansion function" was the basis for porting my code for calculating Feynman integrals for open quantum systems numerically from CPUs onto GPUs (though I used bsxfun for the CPU version too since this saves enormous time compared to using loops or other functions), which lead to huge speed-ups (more and more as the arrays got bigger). A paper about this: "Dattani (2013) FeynDyn: A MATLAB program for fast numerical Feynman integral calculations for open quantum system dynamics on GPUs" explains the story.

If you need to run this on a GPU you can ask a separate question about that and I'll answer that (it's been a while since I've done development for GPUs).

In Python

In Python the task accomplished by bsxfun is known as "broadcasting" and is implemented in numpy.broadcast so that when you multiply two arrays, you automatically get nearly exactly the result that you desire here:

shots = np.array([1.0, 2.0, 3.0])
lens = np.array([2.0, 2.0, 2.0])
shots * lens
array([ 2.,  4.,  6.])

The numpy.broadcast documentation will then tell you how to generalize to work with more elements in the way that bsxfun does in MATLAB/Octave.

Answer 2

Assuming that your sample code is really what you are looking for (which I doubt judging by the full code example you had posted originally), the answer by @vtan707 would work, but is very awkward because of its requirement for specifically shaped 2D arrays. Having to do nested calls of np.array on a single-element list of a function output kinda gives it away.

In fact, your example code is essentially a verbose version of the outer vector product where you multiply each element of one vector with each element of another one. You could simplify it to:

shot = np.arange(...)
lens = np.arange(...)

total = np.outer(shot, lens)

---

As per @NikeDattani's request in the comments, here are some ipython timings with their larger arrays. The first one is using np.outer, the second one reshapes the arrays into row and column vectors and lets numpy broadcast the result when using the * operator:

shots = np.arange(0,11) # 0 to 10 including
lens = np.arange(0,10001,50) # 0 to 10000 including

%timeit -n 100000 np.outer(shots, lens)
4.33 µs ± 17.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit -n 100000 shots[:,None] * lens[None,:]
3.38 µs ± 19.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

^{(I've used the slightly obscure indexing using None to create new dimensions here. For more info, see Dimensional Indexing Tools in the numpy docs.)}

Because numpy is heavily optimized behind the scenes, there is no big performance difference between np.outer and the reshaped example -- presumably because it falls back to the exact same operations anyway. The slight overhead we see might come from function calls.

Interesting side note: The processing times increase slightly for both strategies (to about 5 to 6 µs) if I reverse the two arrays.

(Thanks also to @wflynny for doing some more timings in their answer.)

Answer 3

You can use dot multiplication to achieve this:

import numpy as np

shot = np.array([np.arange(0, 100, 50)]).T
lens = np.array([np.arange(10)])

total = np.dot(shot,lens)
print(total)

This gives you the below results:

[[  0   0   0   0   0   0   0   0   0   0]
 [  0  50 100 150 200 250 300 350 400 450]]

Answer 4

Just adding some timings for @Antimon's answer. Using numpy.outer is definitely the way to go IMO.

def list_version(shots, lens):
    list_total = []
    for shot in list_shot:
        list_out = []

        for lens in list_lens:
            val = shot * lens
            list_out.append(val)

        list_total.append(list_out) 
    return list_total

def numpy_version1(shots, lens):
    _shots = np.array(shots)
    _lens = np.array(lens)
    return np.outer(_shots, _lens)

def numpy_version2(shots, lens):
    return np.outer(shots, lens)

Depending on the input type (meaning the original arrays are numpy arrays and don't need to be converted to lists), numpy.outer can give a substantial (300 fold) increase in performance:

arr_shot = np.arange(0, 10000)
arr_lens = np.linspace(0, 13, 0.01)
list_shot = list(arr_shot)
list_lens = list(arr_lens)

%timeit list_version(list_shot, list_lens)
21.5 s ± 24.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit numpy_version1(list_shot, list_lens)
73 ms ± 75.2 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit numpy_version2(arr_shot, arr_lens)
72.2 ms ± 46.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Method	Time (ms)
Original	21500
Native Numpy	72.2
Numpy convert from list	73

---

In response to @Nike's comment below, I timed the numpy version over several scales of input array sizes:

results = []
for m in np.logspace(1, 4, 10, base=10, dtype=int):
    for n in np.logspace(1, 4, 10, base=10, dtype=int):
        x = np.random.rand(m)
        y = np.random.rand(n)
        t = timeit.timeit("numpy_version2(x, y)", number=10, globals=globals())
        results.append((m, n, t))
data = pd.DataFrame(results, columns=["size_shots", "size_lens", "time_numpy"])
data["size_total"] = data.size_shots * data.size_lens

This should produce 2 measurements of 10 iterations at every combination of array sizes (whose log10s are equally spaced) from 10 to 10,000.

TLDR: OP's arrays are of the order of 10^5 and 10^3 elements which should take <10s (10^8 output elements takes <10s above).