Scipy: Speeding up calculation of a 2D complex integral

Question

I want to repeatedly calculate a two-dimensional complex integral using dblquad from scipy.integrate. As the number of evaluations will be quite high I would like to increase the evaluation speed of my code.

Dblquad does not seem to be able to handle complex integrands. Thus, I have split the complex integrand into a real and an imaginary part:

def integrand_real(x, y):
    R1=sqrt(x**2 + (y-y0)**2 + z**2)
    R2=sqrt(x**2 + y**2 + zxp**2)
    return real(exp(1j*k*(R1-R2)) * (-1j*z/lam/R2/R1**2) * (1+1j/k/R1))

def integrand_imag(x,y):
    R1=sqrt(x**2 + (y-y0)**2 + z**2)
    R2=sqrt(x**2 + y**2 + zxp**2)
    return imag(exp(1j*k*(R1-R2)) * (-1j*z/lam/R2/R1**2) * (1+1j/k/R1))

y0, z, zxp, k, and lam are variables defind in advance. To evaluate the integral over the area of a circle with radius ra I use the following commands:

from __future__ import division
from scipy.integrate import dblquad
from pylab import *

def ymax(x):
    return sqrt(ra**2-x**2)

lam = 0.000532
zxp = 5.
z = 4.94
k = 2*pi/lam
ra = 1.0

res_real = dblquad(integrand_real, -ra, ra, lambda x: -ymax(x), lambda x: ymax(x))
res_imag = dblquad(integrand_imag, -ra, ra, lambda x: -ymax(x), lambda x: ymax(x))
res = res_real[0]+ 1j*res_imag[0]

According to the profiler the two integrands are evaluated about 35000 times. The total calculation takes about one second, which is too long for the application I have in mind.

I am a beginner to scientific computing with Python and Scipy and would be happy about comments that point out ways of improving the evaluation speed. Are there ways of rewriting the commands in the integrand_real and integrand_complex functions that could lead to siginficant speed improvements?

Would it make sense to compile those functions using tools like Cython? If yes: Which tool would best fit this application?

Recapitulation

Your function is even in x. Simply changing the integration limits to (0, ra) cuts computation time by more than half. – Roselinerosella

Excellent comment Jaime! I just followed and am now at 50% of the original calculation time. Thanks! – Recapitulation

Answer 1

You can gain a factor of about 10 in speed by using Cython, see below:

In [87]: %timeit cythonmodule.doit(lam=lam, y0=y0, zxp=zxp, z=z, k=k, ra=ra)
1 loops, best of 3: 501 ms per loop
In [85]: %timeit doit()
1 loops, best of 3: 4.97 s per loop

This is probably not enough, and the bad news is that this is probably quite close (maybe factor of 2 at most) to everything-in-C/Fortran speed --- if using the same algorithm for adaptive integration. (scipy.integrate.quad itself is already in Fortran.)

To get further, you'd need to consider different ways to do the integration. This requires some thinking --- can't offer much from the top of my head now.

Alternatively, you can reduce the tolerance up to which the integral is evaluated.

# Do in Python
#
# >>> import pyximport; pyximport.install(reload_support=True)
# >>> import cythonmodule

cimport numpy as np
cimport cython

cdef extern from "complex.h":
    double complex csqrt(double complex z) nogil
    double complex cexp(double complex z) nogil
    double creal(double complex z) nogil
    double cimag(double complex z) nogil

from libc.math cimport sqrt

from scipy.integrate import dblquad

cdef class Params:
    cdef public double lam, y0, k, zxp, z, ra

    def __init__(self, lam, y0, k, zxp, z, ra):
        self.lam = lam
        self.y0 = y0
        self.k = k
        self.zxp = zxp
        self.z = z
        self.ra = ra

@cython.cdivision(True)
def integrand_real(double x, double y, Params p):
    R1 = sqrt(x**2 + (y-p.y0)**2 + p.z**2)
    R2 = sqrt(x**2 + y**2 + p.zxp**2)
    return creal(cexp(1j*p.k*(R1-R2)) * (-1j*p.z/p.lam/R2/R1**2) * (1+1j/p.k/R1))

@cython.cdivision(True)
def integrand_imag(double x, double y, Params p):
    R1 = sqrt(x**2 + (y-p.y0)**2 + p.z**2)
    R2 = sqrt(x**2 + y**2 + p.zxp**2)
    return cimag(cexp(1j*p.k*(R1-R2)) * (-1j*p.z/p.lam/R2/R1**2) * (1+1j/p.k/R1))

def ymax(double x, Params p):
    return sqrt(p.ra**2 + x**2)

def doit(lam, y0, k, zxp, z, ra):
    p = Params(lam=lam, y0=y0, k=k, zxp=zxp, z=z, ra=ra)
    rr, err = dblquad(integrand_real, -ra, ra, lambda x: -ymax(x, p), lambda x: ymax(x, p), args=(p,))
    ri, err = dblquad(integrand_imag, -ra, ra, lambda x: -ymax(x, p), lambda x: ymax(x, p), args=(p,))
    return rr + 1j*ri

Mezuzah

Wow, that is the kind of answer I was looking for. I definitely did not expect someone to rewrite my code in the rigorous way you did it. Thank you very much! I will test your code first thing tomorrow morning. – Recapitulation

One question: The last lines of your code are truncated. Would args=(p) be correct? – Recapitulation

It took me some time to get the compiler running but it was worth the effort. The code now runs ten times faster. Thanks again! – Recapitulation

Answer 2

Have you considered multiprocessing (multithreading)? It seems that you don't have a need to do a final integration (over the whole set) so simple parallel processing might be the answer. Even if you did have to integrate, you can wait for running threads to finish computation before doing the final integration. That is, you can block the main thread until all workers have completed.

http://docs.python.org/2/library/multiprocessing.html

Republicanize

Thanks for this advice. I want to evaluate the integral at various points in space and all those evaluations work independently of each other. Thus, I am pretty optimistic that I will be able to implement multiprocessing in the next version of my code. – Recapitulation

Answer 3

quadpy (a project of mine) supports many integration schemes for functions over disks. It supports complex-valued functions and is fully vectorized. For example with Peirce's scheme of order 83:

from numpy import sqrt, pi, exp
import quadpy

lam = 0.000532
zxp = 5.0
z = 4.94
k = 2 * pi / lam
ra = 1.0
y0 = 0.0


def f(X):
    x, y = X
    R1 = sqrt(x ** 2 + (y - y0) ** 2 + z ** 2)
    R2 = sqrt(x ** 2 + y ** 2 + zxp ** 2)
    return exp(1j * k * (R1 - R2)) * (-1j * z / lam / R2 / R1 ** 2) * (1 + 1j / k / R1)


scheme = quadpy.disk.peirce_1957(20)
val = scheme.integrate(f, [0.0, 0.0], ra)

print(val)

(18.57485726096671+9.619636385589759j)

Licastro