Question

I have to create a 2D numpy array from the values x, y from a function return to plot using contourf from matplotlib, and so far I’m using a “C” like structure, that it seems to be very inefficient in Python:

    dim_x = np.linspace(self.min_x, self.max_x, self.step)
    dim_y = np.linspace(self.min_y, self.max_y, self.step)
    X, Y = np.meshgrid(dim_x, dim_y)

    len_x = len(dim_x)
    len_y = len(dim_y)


    a = np.zeros([len_x, len_y], dtype=complex)

    for i, y in enumerate(dim_y):
        for j, x in enumerate(dim_x):
            a[i][j] = aux_functions.final_potential(complex(x, y), element_list)

cs = plt.contourf(X, Y, (a.real), 100)

How can this be done in a more pythonic way?

Thanks!

Answer 1

It would be ideal if you could rewrite final_potential as a vectorized function. A simple and, perhaps, too-obvious example:

>>> dim_x = np.linspace(0, 2, 5)
>>> dim_y = np.linspace(0, 2, 5)
>>> X * Y
array([[ 0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
       [ 0.  ,  0.25,  0.5 ,  0.75,  1.  ],
       [ 0.  ,  0.5 ,  1.  ,  1.5 ,  2.  ],
       [ 0.  ,  0.75,  1.5 ,  2.25,  3.  ],
       [ 0.  ,  1.  ,  2.  ,  3.  ,  4.  ]])

But if you really can’t do that, you could also vectorize:

>>> np.vectorize(lambda x, y: x * y + 2)(X, Y)
array([[ 2.  ,  2.  ,  2.  ,  2.  ,  2.  ],
       [ 2.  ,  2.25,  2.5 ,  2.75,  3.  ],
       [ 2.  ,  2.5 ,  3.  ,  3.5 ,  4.  ],
       [ 2.  ,  2.75,  3.5 ,  4.25,  5.  ],
       [ 2.  ,  3.  ,  4.  ,  5.  ,  6.  ]])

In your case, it would probably look something like this:

def wrapper(x, y): 
    return aux_functions.final_potential(complex(x, y), element_list)

a = np.vectorize(wrapper)(X, Y)

This will probably be a bit faster than nested for loops, although the overhead of a python function call will negate a lot of numpy’s efficiency. In tests I’ve done in the past, using vectorize provided a modest 5x speedup. (That’s compared to a 100x or 1000x speedup for pure numpy operations, as in the X * Y example.)