Theano Op grad() with FEM Solver

Questions
1

Is it possible to define a grad() routine in a Theano Op which requires an FEM solver tool? I’m trying to set up a model which is compatible with the NUTS sampling method but my gradient calculation in non-trivial. To calculate the gradient i need to solve a PDE using esys-escript (an FEM programming tool) which works with its own style of data objects. Currently i have the gradient routine set up so it accepts a numpy vector as the parameter. The parameter vector is then converted into an escript data object so the FEM solver can be executed. After some post-processing of the FEM results, the gradient routine returns a numpy vector of gradient values for each entry in the parameter vector.

Thanks in advance.

2

Yes it is possible - you can have a look at @aseyboldt’s WIP notebook doing exactly that: https://gist.github.com/aseyboldt/9e7c8dafaec8325d5ec1f8b912d45b1b

1 Like
3

Thanks junpenglao! Its taken me a while but i have the grad() routine up and running now although i formulated my Theano Op slightly differently to @aseyboldt’s WIP notebook - i explicitly calculate the Jacobian matrix. When i run my PyMC3 model with this Op, however, i get the following error message:

Auto-assigning NUTS sampler…
Initializing NUTS using jitter+adapt_diag…
Initializing NUTS failed. Falling back to elementwise auto-assignment.
Sequential sampling (2 chains in 1 job)
CompoundStep
NUTS: [sigma_log__]
Slice: [pr_vec]

The results produced by the model are good despite this error but i’m wondering if i can improve them. Once the sampling is complete (2 chains with 1000 samples each), i get the following messages:

The acceptance probability does not match the target. It is 0.7066969406166785, but should be close to 0.8. Try to increase the number of tuning steps.
The acceptance probability does not match the target. It is 0.7208052675795203, but should be close to 0.8. Try to increase the number of tuning steps.
The gelman-rubin statistic is larger than 1.4 for some parameters. The sampler did not converge.
The estimated number of effective samples is smaller than 200 for some parameters.

Do you know why the Initialising NUTS step fails and how i might be able to fix it? Below is the perform and grad functions from my Theano Op if it helps.

def perform(self, node, inputs, outputs):
pr_vec, = inputs
self.updateReducedParameter(pr_vec)
self.updateReducedStiffness()
self.updateReducedSolution()
yr = self.getValuesAtSensors(self.ur_field)
outputs[0][0] = yr # Models current estimate of the observed data

def grad(self, inputs, output_gradients):
self.updateGradVU()
jac = theano.shared(npy.transpose(self.jac()))
g, = output_gradients
print g.type.shape
return [tt.dot(jac, g[0])]

Additionally, below is my code for building and running the PyMC3 model:

with pm.Model() as GroundWaterModel:
pr_vec = pm.Normal(‘pr_vec’, mu=0, sd=1, shape=np, testval=npy.ones(np))
sigma = pm.HalfNormal(‘sigma’, sd=10)
yr = pm.Deterministic(‘yr’, reducedOp(pr_vec))
pm.Normal(‘y_rec’, mu=yr, sd=sigma, observed=yd) # recovered y values
trace = pm.sample()

Any assistance greatly appreciated.

Cheers,
Nick

4

This error is usually an indication that the gradient implementation is not correct. You should try validating the gradient - have a look at this doc http://docs.pymc.io/advanced_theano.html near the bottom there is a cell block showing how to test the gradient implementation using theano.tests.unittest_tools.verify_grad

1 Like