IvanYashchuk / fenics-pymc3

Differentiable interface to FEniCS for PyMC3
MIT License
20 stars 3 forks source link
adjoint fenics pymc3

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.

fenics-pymc3 · Build FEniCS Build Firedrake codecov DOI

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:

Example

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.02

Installation

First 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

Reporting bugs

If you found a bug, create an issue.

Asking questions and general discussion

If you have a question or anything else, create a new discussion. Using issues is also fine!

Contributing

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.