Custom Theano Op to do numerical integration

Question

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.

Baluster

In perform method, you should do only numerical computation. However, in grad method, you should build symbolic graph. – Angilaangina

Thanks for having a look at this. I presume you're referring to the theano.function and as_tensor_variable lines in perform()? The reason these are there is because integrate expects a python function as its first argument and returns an array whose 0th value is the value of the integral. Because the Op's signature expects a Theano variable in output I have to convert this, but I'm not sure I'm doing it right! – Baluster

It's not clear what you are trying to do. The integral in your example is just a number, and putting a number in theano isn't exactly difficult. :-) Do you want to code somthing like this: $f(x) = \int g(x, y)dy$? If so, you need to put the actual computation for a specific x in perform, and return ops that specify the derivatives in grad. – Gotten

@Gotten - good question, my thinking on this has actually been quite foggy. What I originally wanted to do was really a parametrised functional, i.e. I[f; t] = \int^{t}_{0} f(x) dx but I was trying to create a simple example by fixing the upper limit t. Ideally, I would be able to take a compute graph and integrate the function defined by the graph with respect to one of its variables. I think this might be achievable with SymPy, will give it a go and see... – Baluster

Answer 1

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)

Praxis

Yes, this looks excellent, thanks. Passing the function and the integration variable as parameters is the way to do this. I'll test it out on my likelihood integral later - I think we're looking at very similar problems. We can probably speed things up by trying a symbolic integration step when the node is initialised and falling back to numerics if that fails. – Baluster

Your code looks great! I’m currently working on an optimization problem that requires a numerical integration step and I would like to use Theano to solve it. In my case, the to-be-integrated input function is a matrix though. Do you think that your code can be extended to deal with the matrix case? Thanks! – Borodino

@Ludwig, is the function still a scalar function or are we talking about a matrix valued function? If it's still a scalar valued function, I believe this implementation should work. Otherwise, it should be possible provided the to-be-integrated variable is a scalar. Could you provide an example of the objective function? – Praxis

My objective function is something like $\int_0^T exp(At) dt$ where A is an n x n matrix, exp is the matrix exponential, and t is a scalar. So it is a matrix valued function where the to-be-integrated variable is a scalar. – Borodino

Upon further research (and correct me if I'm wrong), theano can't compute the gradient of a non-scalar function. Not sure if there's an easy work-around for your use case. – Praxis

@DanielWilson, Actually it seems to me that Theano can compute the gradient of some non-scalar functions, such as the matrix exponential. For further details: deeplearning.net/software/theano/library/tensor/slinalg.html – Borodino

Thank you for sharing. Did you apply this function to a pymc3 model? I would like to learn to do something similar if you have it available somewhere. – Coxcombry

Answer 2

SymPy is proving harder than anticipated, but in the meantime in case anyone's finding this useful, I'll also point out how to modify this Op to allow for changing the final timepoint without creating a new Op. This can be useful if you have a point process, or if you have uncertainty in your time measurements.

class integrateOut2(theano.Op):
    def __init__(self, f, int_var, *args,**kwargs):
        super(integrateOut2,self).__init__()
        self.f = f
        self.int_var = int_var

    def make_node(self, *inputs):
        tmax = inputs[0]
        self.fvars=list(inputs[1:])

        return theano.Apply(self, [tmax]+self.fvars, [T.dscalar().type()])

    def perform(self, node, inputs, output_storage):
        # Everything else is an argument to the quad function
        tmax = inputs[0]
        args = tuple(inputs[1:])

        # create a function to evaluate the integral
        f = theano.function([self.int_var]+self.fvars, self.f)

        # actually compute the integral
        output_storage[0][0] = quad(f, 0., tmax, args=args)[0]

    def grad(self, inputs, grads):
        tmax = inputs[0]
        param_grads = T.grad(self.f, self.fvars)

        ## Recall fundamental theorem of calculus
        ## d/dt \int^{t}_{0}f(x)dx = f(t)
        ## So sub in t_max to the graph
        FTC_grad = theano.clone(self.f, {self.int_var: tmax})

        grad_list = [FTC_grad*grads[0]] + \
                    [integrateOut2(grad_fn, self.int_var)(*inputs)*grads[0] \
                     for grad_fn in param_grads]

        return grad_list

Baluster

Thank you very much for sharing. How is your progress on the symbolic solution? – Coxcombry

Answer 3

I always use the following code where I generate B = 10000 samples of n = 30 observations from a normal distribution with µ = 1 and σ 2 = 2.25. For each sample, the parameters µ and σ are estimated and stored in a matrix. I hope this can help you.

loglik <- function(p,z){
sum(dnorm(z,mean=p[1],sd=p[2],log=TRUE))
}
set.seed(45)
n <- 30
x <- rnorm(n,mean=1,sd=1.5)
optim(c(mu=0,sd=1),loglik,control=list(fnscale=-1),z=x)

B <- 10000
bootstrap.results <- matrix(NA,nrow=B,ncol=3)
colnames(bootstrap.results) <- c("mu","sigma","convergence")
for (b in 1:B){
sample.b <- rnorm(n,mean=1,sd=1.5)
m.b <- optim(c(mu=0,sd=1),loglik,control=list(fnscale=-1),z=sample.b)
bootstrap.results[b,] <- c(m.b$par,m.b$convergence)
}

One can also obtain the ML estimate of λ and use the bootstrap to estimate the bias and the standard error of the estimate. First calculate the MLE of λ Then, we estimate the bias and the standard error of λˆ by a nonparametric bootstrap.

B <- 9999
lambda.B <- rep(NA,B)
n <- length(w.time)
for (b in 1:B){
b.sample <- sample(1:n,n,replace=TRUE)
lambda.B[b] <- 1/mean(w.time[b.sample])
}
bias <- mean(lambda.B-m$estimate)
sd(lambda.B)

In the second part we calculate a 95% confidence interval for the mean time between failures.

n <- length(w.time)
m <- mean(w.time)
se <- sd(w.time)/sqrt(n)
interval.1 <- m + se * qnorm(c(0.025,0.975))
interval.1

But we can also use the the assumption that the data are from an exponential distribution. In that case we have varX¯ = 1/(nλ^2) = θ^{2}/n which can be estimated by X¯^{2}/n.

sd.m <- sqrt(m^2/n)
interval.2 <- m + sd.m * qnorm(c(0.025,0.975))
interval.2

We can also estimate the standard error of ˆθ by means of a boostrap procedure. We use the nonparametric bootstrap, that is, we sample from the original sample with replacement.

B <- 9999
m.star <- rep(NA,B)
for (b in 1:B){
m.star[b] <- mean(sample(w.time,replace=TRUE))
}
sd.m.star <- sd(m.star)
interval.3 <- m + sd.m.star * qnorm(c(0.025,0.975))
interval.3
An interval not based on the assumption of normality of ˆθ is obtained by the percentile method:

interval.4 <- quantile(m.star, probs=c(0.025,0.975))
interval.4

Deathblow

welcome to stackoverflow ! please read how to write good answer? please provide the context to your answer so that what OP has done mistakenly and what have you corrected or may be suggestions etc.. – Harumscarum