Note: This was a nice demo but PyMC changes its tensor backend quite frequently (first Theano->Aesara, then Aesara->PyTensor) hence this repository is now archived and read-only.
This package enables use of FEniCS or Firedrake for solving differentiable variational problems in PyMC3.
Automatic adjoint solvers for FEniCS programs are generated with dolfin-adjoint/pyadjoint. These solvers make it possible to use Theano's (PyMC3 backend) reverse mode automatic differentiation with FEniCS/Firedrake.
Current limitations:
- Differentiation wrt Dirichlet boundary conditions and mesh coordinates is not implemented yet.
Here is the demonstration of fitting coefficients of a variant of the Poisson's PDE using PyMC3's NUTS sampler.
import numpy as np
import fenics
fenics.set_log_level(fenics.LogLevel.ERROR)
import fenics_adjoint as fa
import ufl
from fenics_pymc3 import create_fem_theano_op
from fenics_pymc3 import to_numpy
# Create mesh for the unit square domain
n = 10
mesh = fa.UnitSquareMesh(n, n)
# Define discrete function spaces and functions
V = fenics.FunctionSpace(mesh, "CG", 1)
W = fenics.FunctionSpace(mesh, "DG", 0)
def solve_fenics(kappa0, kappa1):
# This function inside should be traceable by fenics_adjoint
f = fa.Expression(
"10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2
)
u = fa.Function(V)
bcs = [fa.DirichletBC(V, fa.Constant(0.0), "on_boundary")]
inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
v = fenics.TestFunction(V)
F = inner(kappa0*grad(u), grad(v)) * dx - kappa1 * f * v * dx
fa.solve(F == 0, u, bcs=bcs)
return u
# Let's generate artificial data
true_kappa0 = fa.Constant(1.25)
true_kappa1 = fa.Constant(0.55)
true_solution = solve_fenics(true_kappa0, true_kappa1)
true_solution_numpy = to_numpy(true_solution)
# Perturb the solution with noise
noise_level = 0.05
noise = np.random.normal(scale=noise_level * np.linalg.norm(true_solution_numpy), size=true_solution_numpy.size)
noisy_solution = true_solution_numpy + noise
# Define FEniCS template representation of Theano/NumPy input
templates = (fa.Constant(0.0), fa.Constant(0.0))
# Now let's create Theano wrapper of `solve_fenics` function
theano_fem_solver = create_fem_theano_op(templates)(solve_fenics)
# `theano_fem_solver` can now be used inside PyMC3's model
import pymc3 as pm
import theano.tensor as tt
with pm.Model() as fit_poisson:
sigma = pm.InverseGamma("sigma", alpha=3.0, beta=0.5)
kappa0 = pm.TruncatedNormal(
"kappa0", mu=1.0, sigma=0.5, lower=1e-5, upper=2.0, shape=(1,)
)
kappa1 = pm.TruncatedNormal(
"kappa1", mu=0.7, sigma=0.5, lower=1e-5, upper=2.0, shape=(1,)
)
predicted_solution = pm.Deterministic("pred_sol", theano_fem_solver(kappa0, kappa1))
d = pm.Normal("d", mu=predicted_solution, sd=sigma, observed=noisy_solution)
with fit_poisson:
trace = pm.sample(500, chains=4, cores=4)
pm.summary(trace)
# mean sd hdi_3% hdi_97% ... ess_sd ess_bulk ess_tail r_hat
# sigma 0.015 0.001 0.013 0.017 ... 689.0 715.0 723.0 1.00
# kappa0[0] 1.247 0.377 0.586 1.926 ... 334.0 331.0 462.0 1.02
# kappa1[0] 0.586 0.179 0.267 0.900 ... 352.0 352.0 582.0 1.02First install FEniCS or Firedrake. Then install pyadjoint with:
python -m pip install git+https://github.com/dolfin-adjoint/pyadjoint.git@master
Then install fecr with:
python -m pip install git+https://github.com/IvanYashchuk/fecr@master
Then install PyMC3 with:
python -m pip install pymc3
After that install fenics-pymc3 with:
python -m pip install git+https://github.com/IvanYashchuk/fenics-pymc3.git@master
If you found a bug, create an issue.
If you have a question or anything else, create a new discussion. Using issues is also fine!
Pull requests are welcome from everyone.
Fork, then clone the repository:
git clone https://github.com/IvanYashchuk/fenics-pymc3.git
Make your change. Add tests for your change. Make the tests pass:
pytest tests/
Check the formatting with black and flake8. Push to your fork and submit a pull request.