Run a vectorised function on an array in Python with NumPy

Question

I have a function in python that is meant to operate on a scalar input, but multiplies matrices in the process. The exact code is shown below:

def f(t, n):
    T = np.pi
    a_0 = 0.5

    n = np.arange(1, n + 1)
    # Calculate the Fourier series of f(t)
    a_n = np.sin(n*T) / (n * np.pi)
    b_n = (1 - np.cos(n * T)) / (n * np.pi)
    res = a_0 + np.sum(a_n * np.cos(n*t)) + np.sum(b_n * np.sin(n*t))
    return res

Now, I want this to operate on a vector of inputs t, and for the implementation to stay vectorised (not to use for loops). I can see that making a matrix of dimensions len(t) x n where the initial vector n is just stacked vertically len(t) times, and then performing elementwise multiplication with t would be a solution, but what would be the proper way to implement this function?

what's T? it's not defined. also, b_n is defined but not used — Paul Panzer
– Paul Panzer, Commented Feb 28, 2017 at 18:24
Was too hasty with writing the function down. T is just a constant. — Bruno KM
– Bruno KM, Commented Feb 28, 2017 at 18:29

Divakar · Accepted Answer · 2017-02-28 19:10:30Z

1

Here's one vectorized approach that accepts a vector of inputs as t making use of broadcasting and sum-reduction for the final step with matrix-multiplication using np.dot -

def f_vectorized(t, n): # where t is an array
    t2D = t[:,None]    

    T = np.pi
    a_0 = 0.5

    n = np.arange(1, n + 1)
    a_n = np.sin(n*T) / (n * np.pi)
    b_n = (1 - np.cos(n * T)) / (n * np.pi)

    nt2D = n*t2D
    return a_0 + np.cos(nt2D).dot(a_n) + np.sin(nt2D).dot(b_n)

Sample run -

In [142]: t
Out[142]: array([8, 1, 8, 0, 2, 7, 8, 8])

In [143]: n = 5

In [144]: f_vectorized(t,n)
Out[144]: 
array([ 1.03254608,  0.94354963,  1.03254608,  0.5       ,  0.95031599,
        1.04127659,  1.03254608,  1.03254608])

edited Feb 28, 2017 at 19:10

answered Feb 28, 2017 at 18:48

Divakar

222k19 gold badges273 silver badges374 bronze badges

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

Comments

Paul Panzer · Accepted Answer · 2017-02-28 18:41:02Z

1

Here is a formulaic "vectorisation". Note that only a handful of changes were necessary. First line and last but one.

First line: the asanyarray allows to accept array-like inputs, i.e. scalars, arrays, nested lists etc. and treat them all the same. The indexing adds one axis at the very end. That is the space for the Fourier coefficients. Conveniently, these will be automatically broadcast since they occupy the last dimension and missing axes are inserted on the left. This is why the code works almost unchanged.

Only the summations in the end have to be restricted to the Fourier axis, which is what the ..., axis=-1) kwargs do.

def f(t, n):
    t = np.asanyarray(t)[..., None]
    T = np.pi
    a_0 = 0.5 
    n = np.arange(1, n + 1)
    # Calculate the Fourier series of f(t)
    a_n = np.sin(n*T) / (n * np.pi)
    b_n = (1 - np.cos(n * T)) / (n * np.pi)
    res = a_0 + np.sum(a_n * np.cos(n*t), axis=-1) + np.sum(b_n * np.sin(n*t), axis=-1)
    return res

edited Feb 28, 2017 at 18:41

answered Feb 28, 2017 at 18:35

Paul Panzer

53.3k3 gold badges59 silver badges103 bronze badges

2 Comments

Bruno KM Over a year ago

What does the "..." ellipsis operator do?

Paul Panzer Over a year ago

It basically means :, :, :, : where the number of repeats equals the number of dimensions not indexed which in this case is all of them. Since we do not know the shape of t beforehand this is convenient. The number is also allowed to be zero, so this also works for scalar t

Collectives™ on Stack Overflow

Run a vectorised function on an array in Python with NumPy

2 Answers 2

Comments

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related