Replacing a for loop related to a parameter inside function of two variables by numpy sum

user_9530017 0
Viewed 36

I am trying to speed up a function that can be minimally represented by:

import numpy as np

def simple_function(x, y, a1, a2, a3):
    return a1 + a2*x**2/(1 + a3*y**2)


def to_optimnize(x, y, a1, a2, a3, N):
    Sigma = 0
    for i in range(len(N)):
        yn = N[i]*y
        Sigma = Sigma + N[i]*simple_function(x, yn, a1, a2, a3)
    return Sigma


x = np.linspace(0, 10, 1000)
y = np.linspace(0, 4, 200)

N = np.random.random((180,))
a1, a2, a3 = 1, 2, 3


X, Y = np.meshgrid(x, y)
test = simple_function(X, Y, a1, a2, a3)
result = to_optimnize(X, Y, a1, a2, a3, N)

The for loop is mostly a cumulative sum, though I can not figure out how to use numpy here while keeping the 'vectorized' behavior allowing to call it using the grid resulting from np.meshgrid

vectorization numpy optimization function
1 Answers

This mostly seems to be a matter of getting the shapes for broadcasting correct. To understand what's going on here, note that given a 1-d array a, a[None, :] creates a 2-d array with a first dimension of length 1. a[:, None] creates a 2-d array with a second dimension of length 1.

def to_optimize_new(x, y, a1, a2, a3, n):
    n = n[None, None, :]
    y = y[:, :, None]
    x = x[:, :, None]
    yn = n * y
    series = n * simple_function(x, yn, a1, a2, a3)
    return series.sum(axis=2)

This gives correct results when tested with np.allclose:

>>> np.allclose(to_optimize_new(X, Y, a1, a2, a3, N), 
...             to_optimize(X, Y, a1, a2, a3, N))
True

This approach could be memory-intensive, since the results of all the operations are stored and summed at the very end. But for this example it works well.

By the way, if you have no reason for using meshgrid other than to enable broadcasting, then you can use the same reshaping trick to avoid the meshgrid call, like so:

def to_optimize_nomesh(x, y, a1, a2, a3, n):
    n = n[None, None, :]
    x = x[None, :, None]
    y = y[:, None, None]
    yn = n * y
    series = n * simple_function(x, yn, a1, a2, a3)
    return series.sum(axis=2)

If you want a fully generic function, just remove the reshaping commands entirely. This will work with scalar inputs.

def to_optimize_generic(x, y, a1, a2, a3, n):
    yn = n * y
    series = n * simple_function(x, yn, a1, a2, a3)
    return series

But if you want to use vector inputs, you have to give them the right shape outside the function and perform the sum after returning:

>>> np.allclose(
...     to_optimize_generic(x[None, :, None], 
...                         y[:, None, None], 
...                         a1, a2, a3, 
...                         N[None, None, :]).sum(axis=2), 
...     to_optimize(X, Y, a1, a2, a3, N))
True

You can tell which axis to sum based on which axis in N has the :. 

>>> np.allclose(
...     to_optimize_generic(x[None, None, :], 
...                         y[None, :, None], 
...                         a1, a2, a3, 
...                         N[:, None, None]).sum(axis=0), 
...     to_optimize(X, Y, a1, a2, a3, N))
True
user_5770880
Related