Evaluating a series of Legendre polynomials

The following function represents the electrostatic potential, in spherical coordinates, due to a ring of charge \$q=1\$ and radius \$R=1\$, placed in the plane \$x\$-\$y\$:

$$\phi(r,\theta) = \sum_{n=0}^\infty \frac{r_\min^n}{r_\max^{n+1}} P_n(\cos\theta) P_n(0) $$

where \$P_n\$ are the Legendre Polynomials of degree \$n\$, \$r_\min = \min(r,1)\$, and \$r_\max = \max(r,1)\$.

I wrote a code in python to plot this function, with the sum truncated to \$0\leq n \leq 35\$, in a cartesian grid of 50x50 points. But this takes several minutes to complete. The same code in other languages gives almost instantaneous results in mi computer.

This is the code I want to improve:

import matplotlib.pyplot as plt         # plotting routines
import scipy.special as sp              # legendre polynomials
import numpy as np
from timeit import default_timer as timer

# Potential from a ring of radius c=1 and charge q=1
#    phi(r,th) = \sum_{n=0}^\infty (rmin^n/rmax^{n+1}) P_n(cos(th)) P_n(0.0)
#
# where rmin=min(r,1) and rmax=max(r,1).
# The potential is normalized with respect to q/c, where q is the charge of the ring.                                 
# The radii is normalized with respect to c.

def phi(x, z):
    r = np.sqrt(x**2 + z**2)
    costh = z/r

    rmin = min(r,1)
    rmax = max(r,1)

    dr = rmin/rmax

    suma = 0.0
    for n in range(0,35):
        suma += dr**n * sp.legendre(n)(costh) * sp.legendre(n)(0.0)

    return suma/rmax

phivec = np.vectorize(phi)

# grid for plotting
X, Y = np.meshgrid( np.linspace(-2,2,50),
                    np.linspace(-2,2,50) )

# evaluating phi(X,Y)
start = timer()

Z = phivec(X,Y)

end = timer()
print "time = ", end-start

# plotting phi
plt.figure()
CS = plt.contour(X, Y, Z)
plt.clabel(CS, inline=1, fontsize=10)
plt.title(r'$\psi(r,\theta) = \sum_{\ell=0}^\infty \frac{r_^{\ell+1}} P_\ell(0) P_\ell(\cos\theta)

Here is the result:, y=1.03)
plt.xlabel(r'$x/c

Here is the result:, fontsize=20)
plt.ylabel(r'$z/c

Here is the result:, fontsize=20)

plt.show()

Here is the result:

-
Compute legendre(n) once per loop iteration instead of twice.

-
Use np.hypot instead of np.sqrt(x2 + z2).

-
Don't use phivec, just pass the arrays X and Y directly to phi, which is already almost vectorized. The only change that's needed is to use np.minimum and np.maximum instead of min and max respectively.

Revised code:

def phi(x, z, nlimit=35):
    r = np.hypot(x, z)
    costh = z / r
    rmin = np.minimum(r, 1)
    rmax = np.maximum(r, 1)
    dr = rmin / rmax
    suma = 0.0
    for n in range(nlimit):
        Ln = sp.legendre(n)
        suma += dr**n * Ln(costh) * Ln(0.0)
    return suma / rmax

X, Y = np.meshgrid(*[np.linspace(-2, 2, 50)] * 2)
Z = phi(X, Y)

I find that this takes about 0.05 seconds, which ought to be fast enough for you.

StackExchange Code Review Q#131984, answer score: 5