Split tests for fenics and firedrake

tests/.coveragerc

@@ -0,0 +1,15 @@
+# .coveragerc to control coverage.py
+[run]
+branch = True
+source = fenics_pymc3
+
+[report]
+ignore_errors = True
+omit =
+    examples/*
+    notebooks/*
+    tests/*
+    setup.py
+
+[html]
+directory = coverage_html

tests/test_assemble.py → tests/fenics/test_assemble.py

@@ -7,14 +7,13 @@
 
 import theano
 
-theano.config.optimizer = "fast_compile"
-theano.config.compute_test_value = "ignore"
-
 from fenics_pymc3 import create_fenics_theano_op
 from fenics_pymc3 import FenicsVJPOp
 
 from fenics_numpy import evaluate_primal, evaluate_vjp
 
+theano.config.optimizer = "fast_compile"
+theano.config.compute_test_value = "ignore"
 
 mesh = fa.UnitSquareMesh(3, 2)
 V = fenics.FunctionSpace(mesh, "P", 1)

tests/test_pymc3.py → tests/fenics/test_pymc3.py

@@ -3,16 +3,18 @@
 import numpy as np
 
 import fenics
-
-fenics.set_log_level(fenics.LogLevel.ERROR)
 import fenics_adjoint as fa
 import ufl
 
 import fdm
 
-# from fenics_pymc3 import fem_eval, vjp_fem_eval_impl
 from fenics_pymc3 import create_fenics_theano_op
-from fenics_pymc3 import fenics_to_numpy, numpy_to_fenics
+from fenics_pymc3 import to_numpy
+
+import pymc3 as pm
+import theano.tensor as tt
+
+fenics.set_log_level(fenics.LogLevel.ERROR)
 
 n = 25
 mesh = fa.UnitSquareMesh(n, n)
@@ -51,7 +53,7 @@ def solve_fenics(kappa0, kappa1):
 # true_kappa1.interpolate(fa.Constant(0.55))
 
 true_solution = solve_fenics(true_kappa0, true_kappa1)
-true_solution_numpy = fenics_to_numpy(true_solution)
+true_solution_numpy = to_numpy(true_solution)
 
 # perturb state solution and create synthetic measurements
 noise_level = 0.05
@@ -61,9 +63,6 @@ def solve_fenics(kappa0, kappa1):
 
 theano_fem = create_fenics_theano_op(templates)(solve_fenics)
 
-import pymc3 as pm
-import theano.tensor as tt
-
 with pm.Model() as fit_diffusion:
     sigma = pm.InverseGamma("sigma", alpha=3.0, beta=0.5)
 
@@ -81,4 +80,4 @@ def solve_fenics(kappa0, kappa1):
 
 def test_run_sample():
     with fit_diffusion:
-        trace = pm.sample(5, tune=5, chains=1)
+        pm.sample(5, tune=5, chains=1)

tests/firedrake/test_assemble.py

@@ -0,0 +1,66 @@
+from pytest_check import check
+import numpy as np
+
+import firedrake
+import firedrake_adjoint
+import ufl
+
+import theano
+
+from fenics_pymc3 import create_fenics_theano_op
+from fenics_pymc3 import FenicsVJPOp
+
+from fecr import evaluate_primal, evaluate_pullback
+
+theano.config.optimizer = "fast_compile"
+theano.config.compute_test_value = "ignore"
+
+mesh = firedrake.UnitSquareMesh(3, 2)
+V = firedrake.FunctionSpace(mesh, "P", 1)
+
+
+def assemble_firedrake(u, kappa0, kappa1):
+
+    x = firedrake.SpatialCoordinate(mesh)
+    f = x[0]
+
+    inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
+    J_form = 0.5 * inner(kappa0 * grad(u), grad(u)) * dx - kappa1 * f * u * dx
+    J = firedrake.assemble(J_form)
+    return J
+
+
+templates = (firedrake.Function(V), firedrake.Constant(0.0), firedrake.Constant(0.0))
+inputs = (np.ones(V.dim()), np.ones(1) * 0.5, np.ones(1) * 0.6)
+
+
+def test_theano_primal():
+    theano.config.compute_test_value = "ignore"
+    hh = create_fenics_theano_op(templates)(assemble_firedrake)
+    x = theano.tensor.vector()
+    y = theano.tensor.vector()
+    z = theano.tensor.vector()
+    f = theano.function([x, y, z], hh(x, y, z))
+    theano_output = f(*inputs)
+    numpy_putput = evaluate_primal(assemble_firedrake, templates, *inputs)[0]
+    assert np.isclose(theano_output, numpy_putput)
+
+
+def test_theano_vjp():
+    theano.config.compute_test_value = "ignore"
+    numpy_output, fenics_output, fenics_inputs, tape = evaluate_primal(
+        assemble_firedrake, templates, *inputs
+    )
+    vjp_op = FenicsVJPOp(
+        assemble_firedrake, templates, fenics_output, tuple(fenics_inputs), tape
+    )
+    g = theano.tensor.vector()
+    f = theano.function([g], vjp_op(g))
+    theano_output = f(np.ones(1))
+
+    numpy_output = evaluate_pullback(
+        fenics_output, tuple(fenics_inputs), tape, np.ones(1)
+    )
+    for to, no in zip(theano_output, numpy_output):
+        with check:
+            assert np.allclose(to, no)

tests/firedrake/test_pymc3.py

@@ -0,0 +1,80 @@
+import pytest
+
+import numpy as np
+
+import firedrake
+import firedrake_adjoint
+import ufl
+
+import pymc3 as pm
+import theano.tensor as tt
+
+import fdm
+
+from fenics_pymc3 import create_fenics_theano_op
+from fenics_pymc3 import to_numpy
+
+n = 25
+mesh = firedrake.UnitSquareMesh(n, n)
+V = firedrake.FunctionSpace(mesh, "P", 1)
+DG = firedrake.FunctionSpace(mesh, "DG", 0)
+
+
+def solve_firedrake(kappa0, kappa1):
+
+    x = firedrake.SpatialCoordinate(mesh)
+    f = x[0]
+
+    u = firedrake.Function(V)
+    bcs = [firedrake.DirichletBC(V, firedrake.Constant(0.0), "on_boundary")]
+
+    inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
+    JJ = 0.5 * inner(kappa0 * grad(u), grad(u)) * dx - kappa1 * f * u * dx
+    v = firedrake.TestFunction(V)
+    F = firedrake.derivative(JJ, u, v)
+    firedrake.solve(F == 0, u, bcs=bcs)
+    return u
+
+
+templates = (firedrake.Constant(0.0), firedrake.Constant(0.0))
+# inputs = (np.ones(1) * 0.5, np.ones(1) * 0.6)
+
+# templates = (firedrake.Function(DG), firedrake.Function(DG))
+
+true_kappa0 = firedrake.Constant(1.25)
+true_kappa1 = firedrake.Constant(0.55)
+
+# true_kappa0 = firedrake.Function(DG)
+# true_kappa0.interpolate(firedrake.Constant(1.25))
+# true_kappa1 = firedrake.Function(DG)
+# true_kappa1.interpolate(firedrake.Constant(0.55))
+
+true_solution = solve_firedrake(true_kappa0, true_kappa1)
+true_solution_numpy = to_numpy(true_solution)
+
+# perturb state solution and create synthetic measurements
+noise_level = 0.05
+MAX = np.linalg.norm(true_solution_numpy)
+noise = np.random.normal(scale=noise_level * MAX, size=true_solution_numpy.size)
+noisy_solution = true_solution_numpy + noise
+
+theano_fem = create_fenics_theano_op(templates)(solve_firedrake)
+
+with pm.Model() as fit_diffusion:
+    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,)
+    )  # truncated(Normal(1.0, 0.5), 1e-5, 2)
+    kappa1 = pm.TruncatedNormal(
+        "kappa1", mu=0.7, sigma=0.5, lower=1e-5, upper=2.0, shape=(1,)
+    )  # truncated(Normal(0.7, 0.5), 1e-5, 2)
+
+    predicted_solution = pm.Deterministic("pred_sol", theano_fem(kappa0, kappa1))
+
+    d = pm.Normal("d", mu=predicted_solution, sd=sigma, observed=noisy_solution)
+
+
+def test_run_sample():
+    with fit_diffusion:
+        pm.sample(5, tune=5, chains=1)