CODEWITHSUNDEEP
pythonnumpyscipyarray-broadcasting
Walter
Walter

Reputation: 45414

How to use scipy.integrate.fixed_quad for computing many integrals at once?

Given a function func(x,y,z), I want to provide a function

def integral_over_z(func,x,y,zmin=0,zmax=1,n=16):
    lambda_func = z,x,y: ???
    return scipy.integrate.fixed_quad(lambda_func,a=zmin,b=zmax,args=(x,y),n=n)

that computes its integral over z for user provided (x,y) inputs using scipy.integrate.fixed_quad. The input (x,y) can be each be a single float or an array of floats (when both are arrays, their shapes are identical).

scipy.integrate.fixed_quad supports integrating vector-valued functions. To this end, the function func must return a corresponding array of higher dimension: "If integrating a vector-valued function, the returned array must have shape (..., len(x))" (from the docs).

My question therefore is how to generate the corresponding output array of the lambda_func (which may be implemented using a special-purpose class).

EDIT: to help understand my question, here is an implementation that works, but is not vectorized over z (and hence doesn't use scipy.integrate.fixed_quad).

def integral_over_z(func,x,y,zmin,zmax,n=16):
    z,w = scipy.special.roots_legendre(n)
    dz = 0.5*(zmax-zmin)
    z = zmin + (np.real(z)+1) * dz
    w = np.real(w) * dz
    result = w[0] * func(x,y,z[0])
    for i in range(1,len(z)):
        result += w[i] * func(x,y,z[i])
    return result

The problem is: how to vectorize it, such that it works for any valid input (x and/or y floats or arrays).

ANOTHER EDIT: For the implementation via scipy.integrate.fixed_quad, the integrand function must take a 1D array of z of shape (nz). The inputs x and y must broadcast together, when the broadcasted shape of them could be anything, say (n0,n1,..,nk) Then the return from func must have shape (n0,n1,..,nk,nz) -- how to I generated that?

Upvotes: 3

Views: 415

Answers (2)

flawr
flawr

Reputation: 11628

It seems as a vector valued function the vector values must be in the 0th dimension, and the integration arguments (in your case z) must come last (that what they mean with (..., len(x)), their x is your z), I think this comes from the broadcasting rules. Following example worked fine for me - the key here is that x and y must have the right shape for the broadcasting to work

import numpy as np
import scipy.integrate
def integral_over_z(func,x,y,n=16):
    lambda_func = lambda z, x, y: func(x[..., None],y[..., None],z)  # the last dimension of (x,y) needs to be size 1, but you can have as many leading dimensions as you want
    return scipy.integrate.fixed_quad(lambda_func,a=0,b=1,args=(x,y),n=n)
func = lambda x,y,z: 1 + 0*x + 0*y + 0*z  # make sure that the output has the right (broadcast) shape
x = np.zeros((5,))
y =  np.arange(5)
print(integral_over_z(func, x, y, 2))

Upvotes: 2

Walter
Walter

Reputation: 45414

After the (incomplete) answer by flawr and reading about numpy broadcasting, I found a solution. I'd be happy to learn whether this can still be improved and/or if this is really correct, i.e. works for any valid input (it does for my tests sofar).

The important point is to adapt the shapes of x and y such that

  • func(x,y,z) works just fine, i.e. x, y, and z are jointly broadcastable;

  • after summing the output of func over the last (z) dimension, the result has the joint broadcasted shape of x and y.

Here is my solution:

def integral_over_z(func,x,y,zmin=0,zmax=1,n=16):
    xe = x
    ye = y
    if type(xe) is np.ndarray or type(ye) is np.ndarray:
        xe,ye = np.broadcast_arrays(x,y)   # replace x,y by their joint broadcast
        xe = np.expand_dims(xe, xe.ndim)   # expand by an extra dimension for z
        ye = np.expand_dims(ye, ye.ndim)   # expand by an extra dimension for z
    return scipy.integrate.fixed_quad(lambda z : func(xe,ye,z), a=zmin, b=zmax, n=n)

Upvotes: 0

Related Questions