vectorized function for numpy arrays of different shape

Question

I would like to use vectorized operations for efficient code, where I do operations on arrays of different size such that slicing the arrays is necessary. However, I would like to adjust the function such that I can do calculations with flat numpy arrays as well with such having n rows. I have used simple numbers in this example, but will use random number generation in my project. Currently I have resolved it using an if statement to differentiate between slicing depending on shape. Is there a more efficient way without if statements?

import numpy as np

def func(a, b, c, d):
    if len(a.shape) == 1:
        cf = -c
        cf[:len(b)] += a - b
        cf[-len(d)] -= d
    else:
        cf = -c
        cf[:, :b.shape[1]] += a - b
        cf[:, -d.shape[1]:] -= d
    return cf


a = np.array([1, 1, 1, 1, 1, 1, 1])
b = np.array([2, 2, 2, 2, 2, 2, 2])
c = np.array([3, 3, 3, 3, 3, 3, 3, 3])
d = np.array([4])
func(a, b, c, d)

n = 10
a = np.zeros((n, 7))+1
b = np.zeros((n, 7))+2
c = np.zeros((n, 8))+3
d = np.zeros((n, 1))+4
func(a, b, c, d)

It might be that I just did not find a universal method of slicing here?!

Kinght 金 · Accepted Answer · 2017-12-13 12:18:58Z

5

This code works.

def func(a, b, c, d):
    cf = -c
    cf[..., :b.shape[-1]] += a - b
    cf[..., -d.shape[-1]:] -= d
    return cf

Testing:

## Test 1
a = np.array([1, 1, 1, 1, 1, 1, 1])
b = np.array([2, 2, 2, 2, 2, 2, 2])
c = np.array([3, 3, 3, 3, 3, 3, 3, 3])
d = np.array([4])
func(a, b, c, d)
"""
array([-4, -4, -4, -4, -4, -4, -4, -7])
"""

## Test 2
n = 2
a = np.zeros((n, 7))+1
b = np.zeros((n, 7))+2
c = np.zeros((n, 8))+3
d = np.zeros((n, 1))+4
func(a, b, c, d)
"""
array([[-4., -4., -4., -4., -4., -4., -4., -7.],
       [-4., -4., -4., -4., -4., -4., -4., -7.]])
"""

answered Dec 13, 2017 at 12:18

Kinght 金

18.4k5 gold badges63 silver badges77 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

Peterhack Over a year ago

Works great, thanks! Can you give a short explanation what the ... are for?

Kinght 金 Over a year ago

x[...,-1] : Ellipsis(...) expand to the number of : objects needed to make a selection tuple of the same length as x.ndim.

Peterhack Over a year ago

Thank you! Is the function written in the most efficient way?

Kinght 金 Over a year ago

I don't know what is the most efficient. You can improve it as your hope.

Collectives™ on Stack Overflow

vectorized function for numpy arrays of different shape

1 Answer 1

4 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

4 Comments

Your Answer

Sign up or log in

Post as a guest

Related