Calculate electric field of a charged annulus This script calculates the shape of the electric field on a 2D grid perpendicular to a uniformly charged annulus using SciPy's dblquad (tutorial, documentation).

I've chosen a coarse grid and specified a very large relative error epsrel here so that this example will run in about 10 seconds, but I'll be using it on a much finer grid and with much lower specified error.

I'd like to know if there are ways this can be significantly sped up without loss of accuracy, which should be in the 1E-7 neighborhood.

I'll be avoiding pathological points manually, as I've done here by avoiding the z=0 plane. I've left out constants and units to reduce clutter.



def Excalc(r, th):
    x, y, z = r*np.cos(th), r*np.sin(th), 0.0
    return (x0-x) * ((x0-x)**2 + (y0-y)**2 + (z0-z)**2)**-1.5

def Eycalc(r, th):
    x, y, z = r*np.cos(th), r*np.sin(th), 0.0
    return (y0-y) * ((x0-x)**2 + (y0-y)**2 + (z0-z)**2)**-1.5

def Ezcalc(r, th):
    x, y, z = r*np.cos(th), r*np.sin(th), 0.0
    return (z0-z) * ((x0-x)**2 + (y0-y)**2 + (z0-z)**2)**-1.5


import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import dblquad

twopi = 2.*np.pi

# annulus of uniform, unit charge density
rmin, rmax   = 0.8, 1.2
thmin, thmax = 0,   twopi

x = np.arange(-2,    2.1, 0.1)
z = np.arange(-1.55, 1.6, 0.1) # avoid z=0 plane for integration.
X, Z = np.meshgrid(x, z)
s = X.shape

eps = 1E-4
y0  = 0

fields, errors = [], []
for x0, z0 in zip(X.flatten(), Z.flatten()):
    # use of lambda in lieu of function calls https://stackoverflow.com/a/49441680/3904031
    Ex, Exerr  = dblquad(Excalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
    Ey, Eyerr  = dblquad(Eycalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
    Ez, Ezerr  = dblquad(Ezcalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
    fields.append([Ex, Ey, Ez])
    errors.append([Exerr, Eyerr, Ezerr])

Evec = np.stack([np.array(thing).reshape(s) for thing in zip(*fields)])
Emag = np.sqrt((Evec**2).sum(axis=0))

Evecerr = np.stack([np.array(thing).reshape(s) for thing in zip(*errors)])
Emagerr = np.sqrt((Evecerr**2).sum(axis=0))

if True:
    names = 'Ex', 'Ey', 'Ez', 'Emag'
    plt.figure()
    for i, (E, name) in enumerate(zip([e for e in Evec] + [Emag], names)):
        smalls = np.abs(E) &lt; 1E-08
        E[smalls] = 0.     # fudge to keep Ey from showing 1E-14 fields
        plt.subplot(2, 2, i+1)
        plt.imshow(E, origin='lower')
        plt.colorbar()
        plt.title(name, fontsize=16)
    plt.show()

if True:
    names = 'Ex rel error', 'Ey rel error', 'Ez rel error', 'Emag rel error'
    plt.figure()
    for i, (err, name) in enumerate(zip([er for er in Evecerr] + [Emagerr], names)):
        plt.subplot(2, 2, i+1)
        plt.imshow(err/E, origin='lower')
        plt.colorbar()
        plt.title(name, fontsize=16)
    plt.show()
I think the main time consuming factor here is calculating the dbquad, that said there are some smaller improvements you can make

split into functions

lookup of local variables in a function is faster than a global lookup, so putthing each part in a function can speed up this process already. If you do this, beware that your E_calc functions use global state (x0, y0 and z0), so they should be defined in the function scope, so they can use the local variables

numpy

You can use numpy a bit more to assemble the final answers

Instead of  keeping 2 arrays fields and errors, you can  keep this in 1 array, and index the errors and values.

def calc2():
    twopi = 2.*np.pi

    # annulus of uniform, unit charge density
    rmin, rmax   = 0.8, 1.2
    thmin, thmax = 0,   twopi

    x = np.arange(-2,    2.1, 0.1)
    z = np.arange(-1.55, 1.6, 0.1) # avoid z=0 plane for integration.
    X, Z = np.meshgrid(x, z)
    s = X.shape

    eps = 1E-4
    y0  = 0
    def Excalc(r, th):
        x, y, z = r*np.cos(th), r*np.sin(th), 0.0
        return (x0-x) * ((x0-x)**2 + (y0-y)**2 + (z0-z)**2)**-1.5

    def Eycalc(r, th):
        x, y, z = r*np.cos(th), r*np.sin(th), 0.0
        return (y0-y) * ((x0-x)**2 + (y0-y)**2 + (z0-z)**2)**-1.5

    def Ezcalc(r, th):
        x, y, z = r*np.cos(th), r*np.sin(th), 0.0
        return (z0-z) * ((x0-x)**2 + (y0-y)**2 + (z0-z)**2)**-1.5
    values = []
    for x0, z0 in zip(X.flatten(), Z.flatten()):
        # use of lambda in lieu of function calls https://stackoverflow.com/a/49441680/3904031
        Ex, Exerr  = dblquad(Excalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
        Ey, Eyerr  = dblquad(Eycalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
        Ez, Ezerr  = dblquad(Ezcalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
        values.append([Ex, Ey, Ez, Exerr, Eyerr, Ezerr])
#         errors.append([Exerr, Eyerr, Ezerr])
    values = np.array(values).T.reshape(6, *X.shape)

    e_vec = values[:3]
    e_mag = np.linalg.norm(e_vec, axis=0)

    e_vec_err = values[3:]
    e_mag_err = np.linalg.norm(e_vec_err, axis=0)

    return Evec, Emag, Evecerr, Emagerr


This resulted in a very slight reduction in time, so will not save you heaps

lru_cache

refactoring the calc methods to reuse the most of the result actually slowed things down, this is how I tackled that

from functools import partial, lru_cache

@lru_cache(None)
def calc(r, th, coords_0):
    coords_0 = np.array(coords_0)
    diff = coords_0 - np.array((r*np.cos(th), r*np.sin(th), 0.0))
    return diff * ((diff)**2).sum()**-1.5

def calc3():
    rmin, rmax = 0.8, 1.2
    thmin, thmax = 0, 2.*np.pi

#     s = X.shape

    eps = 1E-4
    y0  = 0
    x = np.arange(-2,    2.1, 0.2)
    z = np.arange(-1.55, 1.6, 0.2) # avoid z=0 plane for integration.
    X, Z = np.meshgrid(x, z)

    calc_axis = lambda r, th, coords_0, axis: calc(r, th, coords_0)[axis]

    values = []
    for x0, z0 in zip(X.flatten(), Z.flatten()):
        coords_0 = (x0, y0, z0)
        Excalc, Eycalc, Ezcalc = (partial(calc_axis, coords_0=coords_0, axis=i) for i in range(3))
        # use of lambda in lieu of function calls https://stackoverflow.com/a/49441680/3904031
        Ex, Exerr  = dblquad(Excalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
        Ey, Eyerr  = dblquad(Eycalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
        Ez, Ezerr  = dblquad(Ezcalc, thmin, thmax, lambda t: rmin, lambda t: rmax, epsrel=eps)
        values.append((Ex, Ey, Ez, Exerr, Eyerr, Ezerr))
    print(calc.cache_info())
    values = np.array(values).T.reshape(6, *X.shape)

    e_vec = values[:3]
    e_mag = np.linalg.norm(e_vec, axis=0)

    e_vec_err = values[3:]
    e_mag_err = np.linalg.norm(e_vec_err, axis=0)
    return e_vec, e_mag, e_vec_err, e_mag_err, values


Apparently, about halve of the calls to calc were cached, but moving the calculation to this method resulted in a net slowdown.

numba

What did result in a speed increase was numba

just adding @jit before def calc(r, th, coords_0): made the calculation 3 times as fast