Fast 2d array numerical integration in Python

Question

(EDITED)

I need to perform numerical integration of:

formula

Where Pn - Legendre polynomial, λ – wavelength, ρ and z – matrices NxN size, J – zero kind Bessel function.

I wrote the following code:

import scipy.integrate as integrate
import datetime
from scipy.special import legendre, jn

image_size = 100
image_half_size = 50
scale = 10
cos_beta1 = 0.9999999
cos_beta2 = 0.01745
l_maxOrder = 10
wavelength = 660
k = 2 * np.pi / wavelength

x_center = 8 / scale
y_center = 14 / scale
x, y = np.meshgrid(np.arange(-image_half_size, image_half_size + 1) / scale,
                         np.arange(-image_half_size, image_half_size + 1) / scale,
                         sparse=False,
                         indexing='ij')

ro_p = np.sqrt((x-x_center)**2 + (y-y_center)**2 + 1e-4**2)

z = random.randint(-5, 5)
z_p = np.full((image_size + 1, image_size + 1), z)


def pi_plus_tau(n):
    def frst_deriv(arg):
        return  (n + 1) * (legendre(n+1)(arg) - arg*legendre(n)(arg)) / (arg**2-1)

    def sec_deriv(arg):
        return (n + 1) * (legendre(n)(arg)*((n-2)*(arg**4) + (3-n)*(arg**2)-1) / (arg**2-1) - legendre(n+1)(arg)*(5+2*n)*arg + legendre(n+2)(arg)*(n+2)) / (arg**2-1) ** 2

    def integrand(arg):
        coef = (-1 * jn(0, k * ro_p * (1-arg** 2) ** 0.5) * np.exp(1j * k * z_p * (1-arg)) * (arg**0.5))
        return (frst_deriv(arg) * (1 - arg) + (1 - arg ** 2) * sec_deriv(arg)) * coef

    return integrand


def intergration_plus(n, angle1, angle2):
    return (integrate.quad_vec(pi_plus_tau(int(n)), angle1, angle2, workers=1))[0]


for n in range(1, l_maxOrder):
    print(datetime.datetime.now())
    intergration_plus(n, cos_beta1, cos_beta2)
    print(datetime.datetime.now())

This works well, but for N=100, the calculation takes approx. 10 seconds and I have to perform a series of calculations for different n-s. And do it many many times. So 10 seconds is too long.

This math expression is part of a bigger expression. When I run code listed above in this question - it runs twice as fast as when I run it inside the whole program.

Could You advise me - how to do fast numerical integration with 2d arrays in Python? Some packages, to use cython, any hints.

Thank You

In general, you should always try to vectorize your code, i.e., make sure it works for input arguments of arbitrary shpae. quadpy does that, and it's perhaps a good reference. (Disclaimer: I wrote it.) — Nico Schlömer
– Nico Schlömer, Commented Jul 14, 2020 at 13:38
The jn/jv evaluation dominates the whole runtime. A parallelized cython version of your integrand looks to be the way to go. It would be very useful if you provide a runable example (this also includes some random inputs of the correct shape and datatype). — max9111
– max9111, Commented Jul 14, 2020 at 14:26

Raccoon 987 · Accepted Answer · 2020-10-21 14:55:47Z

1

The best that I could find to solve my problem is the quadpy library - https://pypi.org/project/quadpy/. Group your data into an N-dimensional array e.g. (batch_size, Image_width, Image_height, Legendre order) and use np.multiply.outer. Of all the options I tried, this provides the fastest calculation.

answered Oct 21, 2020 at 14:55

Raccoon 987

577 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Fast 2d array numerical integration in Python

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related