Search code examples
pythontheanopymc3

Custom Theano Op to do numerical integration


I'm attempting to write a custom Theano Op which numerically integrates a function between two values. The Op is a custom likelihood for PyMC3 which involves the numerical evaluation of some integrals. I can't simply use the @as_op decorator as I need to use HMC to do the MCMC step. Any help would be much appreciated, as this question seems to have come up several times but has never been solved (e.g. https://stackoverflow.com/questions/36853015/using-theano-with-numerical-integration, Theano: implementing an integral function).

Clearly one solution would be to write a numerical integrator within Theano, but this seems like a waste of effort when very good integrators are already available, for example through scipy.integrate.

To keep this as a minimal example, let's just try and integrate a function between 0 and 1 inside an Op. The following integrates a Theano function outside of an Op, and produces correct results as far as my testing has gone.

import theano
import theano.tensor as tt
from scipy.integrate import quad

x = tt.dscalar('x')
y = x**4 # integrand
f = theano.function([x], y)

print f(0)
print f(1)

ans = integrate.quad(f, 0, 1)[0]

print ans

However, attempting to do integration within an Op appears much harder. My current best effort is:

import numpy as np
import theano
import theano.tensor as tt
from scipy import integrate

class IntOp(theano.Op):
    __props__ = ()

    def make_node(self, x):
        x = tt.as_tensor_variable(x)
        return theano.Apply(self, [x], [x.type()])

    def perform(self, node, inputs, output_storage):
        x = inputs[0]
        z = output_storage[0]

        f_to_int = theano.function([x], x)
        z[0] = tt.as_tensor_variable(integrate.quad(f_to_int, 0, 1)[0])

    def infer_shape(self, node, i0_shapes):
        return i0_shapes

    def grad(self, inputs, output_grads):
        ans = integrate.quad(output_grads[0], 0, 1)[0]
        return [ans]

intOp = IntOp()

x = tt.dmatrix('x')
y = intOp(x)

f = theano.function([x], y)

inp = np.asarray([[2, 4], [6, 8]], dtype=theano.config.floatX)
out = f(inp)

print inp
print out

Which gives the following error:

Traceback (most recent call last):
  File "stackoverflow.py", line 35, in <module>
    out = f(inp)
  File "/usr/local/lib/python2.7/dist-packages/theano/compile/function_module.py", line 871, in __call__
    storage_map=getattr(self.fn, 'storage_map', None))
  File "/usr/local/lib/python2.7/dist-packages/theano/gof/link.py", line 314, in raise_with_op
    reraise(exc_type, exc_value, exc_trace)
  File "/usr/local/lib/python2.7/dist-packages/theano/compile/function_module.py", line 859, in __call__
    outputs = self.fn()
  File "/usr/local/lib/python2.7/dist-packages/theano/gof/op.py", line 912, in rval
    r = p(n, [x[0] for x in i], o)
  File "stackoverflow.py", line 17, in perform
    f_to_int = theano.function([x], x)
  File "/usr/local/lib/python2.7/dist-packages/theano/compile/function.py", line 320, in function
    output_keys=output_keys)
  File "/usr/local/lib/python2.7/dist-packages/theano/compile/pfunc.py", line 390, in pfunc
    for p in params]
  File "/usr/local/lib/python2.7/dist-packages/theano/compile/pfunc.py", line 489, in _pfunc_param_to_in
    raise TypeError('Unknown parameter type: %s' % type(param))
TypeError: Unknown parameter type: <type 'numpy.ndarray'>
Apply node that caused the error: IntOp(x)
Toposort index: 0
Inputs types: [TensorType(float64, matrix)]
Inputs shapes: [(2, 2)]
Inputs strides: [(16, 8)]
Inputs values: [array([[ 2.,  4.],
       [ 6.,  8.]])]
Outputs clients: [['output']]

Backtrace when the node is created(use Theano flag traceback.limit=N to make it longer):
  File "stackoverflow.py", line 30, in <module>
    y = intOp(x)
  File "/usr/local/lib/python2.7/dist-packages/theano/gof/op.py", line 611, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "stackoverflow.py", line 11, in make_node
    return theano.Apply(self, [x], [x.type()])

HINT: Use the Theano flag 'exception_verbosity=high' for a debugprint and storage map footprint of this apply node.

I'm surprised by this, especially the TypeError, as I thought I had converted the output_storage variable into a tensor but it appears to believe here that it is still an ndarray.


Solution

  • I found your question because I'm trying to build a random variable in PyMC3 that represents a general point process (Hawkes, Cox, Poisson, etc) and the likelihood function has an integral. I really want to be able to use Hamiltonian Monte Carlo or NUTS samplers, so I needed that integral with respect to time to be differentiable.

    Starting off of your attempt, I made an integrateOut theano Op that seems to work correctly with the behavior I need. I've tested it out on a few different inputs (not on my stats model just yet, but it appears promising!). I'm a total theano n00b, so pardon any stupidity. I would greatly appreciate feedback if anyone has any. Not sure it's exactly what you're looking for, but here's my solution (example at the bottom and in the doc strings). *EDIT: simplified some remnants of screwing around with ways to do this.

    import theano
    import theano.tensor as T
    from scipy.integrate import quad
    
    class integrateOut(theano.Op):
        """
        Integrate out a variable from an expression, computing
        the definite integral w.r.t. the variable specified
        !!! Only implemented in this for scalars !!!
    
    
        Parameters
        ----------
        f : scalar
            input 'function' to integrate
        t : scalar
            the variable to integrate out
        t0: float
            lower integration limit
        tf: float
            upper integration limit
    
        Returns
        -------
        scalar
            a new scalar with the 't' integrated out
    
        Notes
        -----
    
        usage of this looks like:
        x = T.dscalar('x')
        y = T.dscalar('y')
        t = T.dscalar('t')
    
        z = (x**2 + y**2)*t
    
        # integrate z w.r.t. t as a function of (x,y)
        intZ = integrateOut(z,t,0.0,5.0)(x,y)
        gradIntZ = T.grad(intZ,[x,y])
    
        funcIntZ = theano.function([x,y],intZ)
        funcGradIntZ = theano.function([x,y],gradIntZ)
    
        """
        def __init__(self,f,t,t0,tf,*args,**kwargs):
            super(integrateOut,self).__init__()
            self.f = f
            self.t = t
            self.t0 = t0
            self.tf = tf
    
        def make_node(self,*inputs):
            self.fvars=list(inputs)
            # This will fail when taking the gradient... don't be concerned
            try:
                self.gradF = T.grad(self.f,self.fvars)
            except:
                self.gradF = None
            return theano.Apply(self,self.fvars,[T.dscalar().type()])
    
        def perform(self,node, inputs, output_storage):
            # Everything else is an argument to the quad function
            args = tuple(inputs)
            # create a function to evaluate the integral
            f = theano.function([self.t]+self.fvars,self.f)
            # actually compute the integral
            output_storage[0][0] = quad(f,self.t0,self.tf,args=args)[0]
    
        def grad(self,inputs,grads):
            return [integrateOut(g,self.t,self.t0,self.tf)(*inputs)*grads[0] \
                for g in self.gradF]
    
    x = T.dscalar('x')
    y = T.dscalar('y')
    t = T.dscalar('t')
    
    z = (x**2+y**2)*t
    
    intZ = integrateOut(z,t,0,1)(x,y)
    gradIntZ = T.grad(intZ,[x,y])
    funcIntZ = theano.function([x,y],intZ)
    funcGradIntZ = theano.function([x,y],gradIntZ)
    print funcIntZ(2,2)
    print funcGradIntZ(2,2)