Vectorizing a nested for-loop in python for index-dependent function

Question

I'm currently porting a C++ program to Python using Numpy arrays. I'm looking for a way to implement, if possible, the following loops in a more Pythonic way:

for (int j = start_y; j < end_y; j++)
    {
        for (int i = start_x; i < end_x; i++)
        {
            plasmaFreq[i][j] = plasmaFreq_0*(tanh((i - 50)/10) - tanh((i - (nx - 50))/10))/2.0;
        }
    }

Above, plasmaFreq_0 is a constant passed into the surrounding function, as is nx. Obviously it's easy to vectorize the loop bounds to operate on a particular region of a numpy array, but this leaves me with the issue of how to map the above index-dependent function across the array.

Your calculation of plasmaFreq[i][j] currently has no reference to j; is that correct? — Phil H
– Phil H, Commented Mar 17, 2014 at 16:39
That's correct for this loop. It's necessarily a two-dimensional array elsewhere in the program, but here, I am establishing a profile in the x-direction that is constant in the y-direction. — Scott
– Scott, Commented Mar 17, 2014 at 16:46

mgilson · Accepted Answer · 2014-03-17 16:35:25Z

3

You'll need an array i,

i = np.arange(start_x, end_x)
plasmaFreq[start_x:end_x, start_y: end_y] = plasmaFreq_0 *(np.tanh((i - 50)/10) - np.tanh((i - (nx - 50))/10))/2.0

I think that broadcasting should take it from there.

Note that your original code is quite inefficient¹... First, you're calculating the right hand side for each j, but it doesn't depend on j, so you only really need to calculate it once. Second, your inner loop is over the slow index so you won't be effectively using your cache. I would probably write it as:

for (int i = start_x; i < end_x; i++)
{
    rhs = plasmaFreq_0*(tanh((i - 50)/10) - tanh((i - (nx - 50))/10))/2.0;
    for (int j = start_y; j < end_y; j++)
    {
        plasmaFreq[i][j] = rhs;
    }
}

^{¹How inefficient depends on how well the compiler can do figuring out the loops. Someday maybe some compilers could generate the same code from yours and mine}

edited Mar 17, 2014 at 16:35

answered Mar 17, 2014 at 16:30

mgilson

312k70 gold badges656 silver badges722 bronze badges

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

2 Comments

Scott Over a year ago

Is this more efficient than directly implementing the loops?

mgilson Over a year ago

@Scott -- Usually, but when it comes to efficiency, usually it's best to forget what you think you know and just timeit. Even the pros are frequently surprised by the results.

Collectives™ on Stack Overflow

Vectorizing a nested for-loop in python for index-dependent function

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related