FEniCS + dolfin-adjoint in the Mosaic benchmark suite

Hi FEniCS folks,

We just released Mosaic (github.com/pasteurlabs/mosaic), a benchmark for differentiable physics solvers. Each solver runs in a container behind a uniform apply / vjp interface, so one harness compares gradient accuracy and cost across JAX, PyTorch, Julia, and C++ adjoint codes on shared tasks.

FEniCS (2019.1.0, with dolfin-adjoint) is one of the benchmarked solvers on the structural-mesh and thermal-mesh domains (results if you’re interested).

We did our best to do it justice, but we don’t know your code as well as you do and have surely made suboptimal calls. For example, we compute the gradient with dolfin-adjoint (ReducedFunctional.derivative()), but we call solve() with no explicit solver_parameters, so we’re just using DOLFIN’s defaults for the linear solve. If there’s a solver, preconditioner, or tolerance you’d pick for these elasticity and heat problems, or if we’re not even taping the right functional, we’d love to know.

The setup is self-contained under mosaic/tesseracts/structural-mesh/fenics-structural/ and mosaic/tesseracts/thermal-mesh/fenics-heat/. If something looks odd, corrections or a PR are very welcome, here or on the Tesseract Forum.

Mostly we wanted you to know your work is represented, and to get it right.

Cheers,
Dion on behalf of the Mosaic authors

@dionhaefner as I’ve been working on setting up DOLFINx-adjoint for a while, I had a go at rewriting your example from legacy DOLFIN-adjoint in DOLFINx adjoint:

I opted in on a different optimizer.
Under the hood, DOLFINx-adjoint caches alot more work than the old code, which is reflected in the run-time of the solvers.
On my laptop (one process, Intel® Core™ i7-7700HQ) I get the runtime of:

\[corner_load\] one tape recompute (Jhat(rho)): 4.7843e-02 s
\[corner_load\] one adjoint solve (Jhat.derivative()): 6.0980e-02 s

\[full_face_load\] one tape recompute (Jhat(rho)): 5.1036e-02 s
\[full_face_load\] one adjoint solve (Jhat.derivative()): 6.2825e-02 s

Compared to the timings for forward and adjoint in your code:
| FEniCS | 0.105 s @ N=16 | 0.358 s @ N=16 |
Let me know if VJP timed something different than computing the derivative.
|