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
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}{r_>^{\ell+1}} P_\ell(0) P_\ell(\cos\theta)$', y=1.03)
plt.xlabel(r'$x/c$', fontsize=20)
plt.ylabel(r'$z/c$', fontsize=20)


Here is the result: Equipotentials

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

  2. Use np.hypot instead of np.sqrt(x**2 + z**2).

  3. 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.

  • \$\begingroup\$ Thanks! it is faster. Now I can use a finer grid and truncate the sum at higher values of n :) \$\endgroup\$
    – vagoberto
    Jun 14 '16 at 19:39

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.