Hello everyone!
I have put online my work on integrating PyMC3 with FEniCS Project which is a library for solving partial differential equations (PDEs) using the finite element method.
The API is just one function create_fenics_theano_op which turns a normal Python function, which expects FEniCS inputs and outputs a solution to the problem, into a differentiable Theano Op that can be directly used in a PyMC3 model.
This obviously allows fitting parameters (or solving inverse problems) of PDE models using the NUTS sampler. See an example in README. I anticipate that one not so immediately obvious way of using this new bridge would be embedding Gaussian Random Fields based on Stochastic PDE approach in PyMC3 model (someone still needs to develop that package).
It looks lightweight because most of the logic is hidden in the dependencies numpy-fenics-adjoint and pyadjoint. It is possible to remove the pyadjoint dependency though, at the expense that every FEniCS elemental operation would need to be wrapped as Theano Op. But that’s less user friendly since currently there is a clear separation of FEniCS part and PyMC3 part.
I think this work can well exist in a separate package. To make it more discoverable I’d make a better case study than the example from README and contribute that to the PyMC3’s notebooks gallery.
Thanks for sharing, and congrats on this work @ivan, it looks super nice and useful
Pinging @mikkelbue and @gmingas, as they are working on adding the MLDA sampler to PyMC3 and some use-cases need FeniCS, so they could be interested in further integration
@AlexAndorra Thanks for connecting these two branches of work.
@ivan Amazing work, I think this is super useful and can help a lot of people!
I am part of the team that is working on a PR to add the MLDA sampler to PyMC3. Some of the main use cases for this sampler are inverse problems which require complex PDE solves when calculating the forward model. We are particularly interested in high dimensional, computationally intensive models. For those reasons, FEniCS is something we believe a lot of potential users will want to combine with MLDA and we have been using it in some of our examples.
We think we could use your library in the examples to make them cleaner and easier to read. We plan to try this when we have finished with the PR and the other features we are working on at the moment.
Moreover, one of the ways we want to evolve MLDA is to allow it to correct biased coarse models using information from the fine model. In order to do that, we need users to write their custom theano ops so that it sets a pymc3 model variable within the perform method, using info from the forward model. Do you think automating something like this could be a possible extension in your library for use within MLDA?
Finally, it would be interesting to combine MLDA with gradient-based samplers like NUTS by taking advantage of your library. Currently MLDA uses a base sampler in the bottom level of the chains hierarchy. This is a Metropolis method but it would be very cool to try a gradient based sampler as it is probably going to help a lot with difficult multidimensional posteriors. This is a longer term thing though as we are currently working on the basic features of the implementation. It would be good to connect and discuss more on this one.
I have quickly gone through the MLDA PR if I did get correctly the usage then some simple FEniCS code with MLDA sampler would look something like
from fenics import *
from fenics_adjoint import *
# import pymc3
coarse_mesh1 = UnitSquareMesh(5, 5)
coarse_mesh2 = UnitSquareMesh(10, 10)
fine_mesh = UnitSquareMesh(20, 20)
# Define appropriate FunctionSpace objects etc.
def coarse_solve1(...):
# write the solver here
def coarse_solve2(...):
# write the solver here
def fine_solve(...):
# write the solver here
# Wrap above solver functions with Theano Op
theano_coarse_solve1 = create_fenics_theano_op(...)
# ...
with Model() as coarse_model1:
y = pm.Deterministic("y", theano_coarse_solve1(...))
obs = Normal('obs', mu=y, sigma=1., observed=data)
with Model() as coarse_model2:
y = pm.Deterministic("y", theano_coarse_solve2(...))
obs = Normal('obs', mu=y, sigma=1., observed=data)
with Model() as model:
y = pm.Deterministic("y", theano_fine_solve(...))
obs = Normal('obs', mu=y, sigma=1., observed=data)
with model:
sample(10, MLDA(coarse_models=[coarse_model1, coarse_model2]))
So fenics-pymc3 can be combined now with multi-level sampler by defining needed coarse systems separately.
What kind of information has to be in this custom Op?
Currently create_fenics_theano_op is functional wrapper for creating FenicsOp. FenicsOp has perform method defined that itself is a wrapper around evaluate_primal function from numpy-fenics-adjoint package. What evaluate_primal does is transforms NumPy array input into FEniCS types calls the user-defined function with FEniCS inputs and transforms FEniCS output to NumPy array.
Of course, FenicsOp can be taken as a template for any further work that requires some additional computations/operations in perform, but does it have to be all in one place? It’s important to keep Ops as one unit of computation so that grad or L_op methods are simple.
If there are any specific questions about the usage of fenics-pymc3 interface, feel free to ask here or GitHub issues.
Thanks for the reply and sorry for the delay in replying.
This is indeed how it would look and we will work on an example when we finish with the current work we are doing.
Seeing your code, I tend to agree that the thing we want to do is a bit specific and it is understandable to want to keep your code general. I am thinking possibly and extension (e.g. inherited class) could be used in this case, offering the ability to define extra computations. Specifically, we want to be able to write into pm.Data variables in the pymc3 model from within perform, e.g. model.variable_name.set_value(...).
In any case, the part of the code that requires this functionality is still not ready and there is some work left to finish it so will come back to you when this happens. I will also let you know if we manage to build a nice example notebook using the current version of fenics-pymc3 with MLDA.
