CODEWITHSUNDEEP
pythonarraysnumpy
user15770670
user15770670

Reputation: 105

numpy: summing arrays together in a shifted way

Suppose i have an input of n arrays with the shape k*k. Now i want to sum them together but also move the k*k array by some positiv integer (step) after each operation.

Its hard to actually explain it so heres a litle drawing.

enter image description here

at every part where the k*k arrays intersect they are just summed up.

my current way to implement it looks like this:

import numpy as np

# some input with shape (n ,k ,k)
x = np.arange(16).reshape(4 ,2 ,2)
# im intializing the output to zeros in the beginning and
# add the k*k windows to it
out = np.zeros((2 ,5))

# loops over the (n ,k ,k) input and adds the values
# to the output and shifts by 1 to the right
for i in range(4):
   # i could also be multipied by any integer to move 
   # the window by more then one 
   out[: ,i+0:i+2] += x[i]

so my question is if its possible to get rid of the loop or if you could use some builtin numpy function to accelerate this or more general a way to make it faster when the size of the input grows.

thanks for reading and have a nice day :)

Upvotes: 2

Views: 562

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114320

There are a couple of ways I can think of. The most brute-force is to create an (n, k, k + steps * (n - 1)) array, then just sum the first axis:

steps = ...
n, k, _ = x.shape

buf = np.zeros_like(x, shape=(n, k, k + steps * (n - 1)))
buf[np.arange(n)[:, None, None], np.arange(k)[:, None], np.arange(k) + step * np.arange(n)[:, None, None]] = x

result = buf.sum(0)

I suspect a simple loop will be faster, since you are wasting a lot of space on indexing here.

Here's a quick timing test:

def index_emplace(x, steps):
    k, m, n = x.shape
    buf = np.zeros_like(x, shape=(k, m, n + steps * (k - 1)))
    r = np.arange(max(x.shape))
    a = r[:k][:, None, None]
    b = r[:m][:, None]
    c = r[:n]
    buf[a, b, c + step * a] = x
    return buf.sum(0)

def loop_emplace(x, steps):
    k, m, n = x.shape
    result = np.zeros_like(x, shape=(m, n + steps * (k - 1)))
    for i in range(k):
        result[:, step * i:step * i + m] += x[i]
    return result

x = np.random.randint(10, shape=(100, 100, 100))
steps = 37

%timeit index_emplace(x, steps)
107 ms ± 970 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit loop_emplace(x, steps)
2.26 ms ± 14.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

(index_emplace(x, steps) == loop_emplace(x, steps)).all()
True

As you can see, for large arrays, the the loop can be as much as 50x faster.

Upvotes: 2

arnino
arnino

Reputation: 560

You could use numpy.concatenate to do it in two steps

The first step concatenates and adds all even arrays in your case the red (0) and blue (2) to the output. The second step concatenates and adds all odd arrays in your case the yellow (1) and green (3) to the output

out = np.zeros((2 ,5))

out[:, :-1] += np.concatenate(x[::2], axis=1)
out[:, 1:] += np.concatenate(x[1::2], axis=1)

Upvotes: 1

Related Questions