SUPG in dolfinx: how to replace C++ expressions?

Hi all,
we’re updating some old codes from fenics 2019 to dolfinx and would appreciate some pointers on how to implement functionalities for which in the past we used CompiledExpressions defined by C++ strings via the pybind11 interface. We would use these Expressions in UFL forms of the variational formulations which is not possible in dolfinx.

Here is a MWE using an expression for SUPG for illustrating the issue:

from dolfin import (CompiledExpression, Constant, FiniteElement, Function,
                    TestFunction, TrialFunction, UnitSquareMesh, FunctionSpace,
                    VectorFunctionSpace, assemble, compile_cpp_code, dx, grad,
                    inner, interpolate, as_backend_type)

tau_cpp = """
#include <pybind11/pybind11.h>
namespace py = pybind11;

#include <dolfin/function/GenericFunction.h>
#include <dolfin/function/Expression.h>
#include <dolfin/common/Array.h>
#include <dolfin/mesh/Cell.h>
#include <dolfin/mesh/Mesh.h>
#include <dolfin/function/FunctionSpace.h>

class tau : public dolfin::Expression
{
public:
    std::shared_ptr<dolfin::GenericFunction> viscosity, density;
    std::shared_ptr<dolfin::GenericFunction> u;

    tau() : dolfin::Expression() { }

    void eval(dolfin::Array<double>& values, const dolfin::Array<double>& x,
               const ufc::cell& c) const
    {
        const std::shared_ptr<const dolfin::Mesh> mesh = u->function_space()->mesh();
        const dolfin::Cell cell(*mesh, c.index);
        double h = cell.h();
        dolfin::Array<double> mu(viscosity->value_size());
        dolfin::Array<double> rho(density->value_size());
        viscosity->eval(mu, x, c);
        density->eval(rho, x, c);

        dolfin::Array<double> w(u->value_size());
        u->eval(w, x, c);
        double u_norm = 0.0;
        for (uint i = 0; i < w.size(); ++i)
            u_norm += w[i]*w[i];
        u_norm = sqrt(u_norm);

        // Peclet number
        double Pe = 0.5*u_norm*h*rho[0]/mu[0];
        // "doubly asymptotic" formula
        double xi = (Pe < 3.0) ? Pe/3. : 1.0;
        values[0] = (u_norm > 0) ? 0.5*h/u_norm*xi : 0;
    }
};

PYBIND11_MODULE(SIGNATURE, m)
{
    py::class_<tau, std::shared_ptr<tau>, dolfin::Expression>
    (m, "tau")
    .def(py::init<>())
    .def_readwrite("u", &tau::u)
    .def_readwrite("density", &tau::density)
    .def_readwrite("viscosity", &tau::viscosity);
}
"""

mesh = UnitSquareMesh(4, 4)
V = VectorFunctionSpace(mesh, "CG", 1)
u = TrialFunction(V)
v = TestFunction(V)
mu = Constant(0.1)
rho = Constant(1.0)
u_convection = interpolate(Constant((2, 2)), V)

tau = CompiledExpression(
    compile_cpp_code(tau_cpp).tau(),
    element=FiniteElement("DG", mesh.ufl_cell(), 0),
    domain=mesh,
)
tau.viscosity = mu.cpp_object()
tau.density = rho.cpp_object()
tau.u = u_convection.cpp_object()

a = tau * inner(grad(u) * u_convection, v) * dx
A = assemble(a)

Another rather different example we need to translate to dolfinx is a tensor expression of element metrics, in order to have element length scales in direction of the flow (this time not depending on a Function):

cppcode = """
#include <pybind11/pybind11.h>
namespace py = pybind11;

#include <Eigen/Dense>
#include <dolfin/function/Expression.h>
#include <dolfin/mesh/Mesh.h>
#include <dolfin/mesh/Cell.h>

class Metric : public dolfin::Expression
{{
    public:
        std::shared_ptr<const dolfin::Mesh> mesh;

        Metric() : dolfin::Expression({D}, {D}) {{ }}

    void eval(Eigen::Ref<Eigen::VectorXd> values,
              Eigen::Ref<const Eigen::VectorXd> x, const ufc::cell& c)
            const override
    {{
        const dolfin::Cell cell(*mesh, c.index);
        std::vector<double> coord;
        cell.get_vertex_coordinates(coord);
        Eigen::Map<Eigen::Matrix<double, {D1}, {D}, Eigen::RowMajor> > \
            XC(coord.data());
        Eigen::Matrix<double, {D}, {D}, Eigen::RowMajor> F, Finv;
        for (uint i=0; i < {D}; ++i)
            F.col(i) = XC.row(i+1) - XC.row(0);
        Finv = F.inverse();
        F = Finv.transpose() * Finv;
        values = Eigen::Map<Eigen::VectorXd>(F.data(), {DD});
    }}
}};

PYBIND11_MODULE(SIGNATURE, m)
{{
    py::class_<Metric, std::shared_ptr<Metric>, dolfin::Expression>
    (m, "Metric")
    .def(py::init<>())
    .def_readwrite("mesh", &Metric::mesh);
}}
"""
G = CompiledExpression(
    compile_cpp_code(
        cppcode.format(D=dim, D1=dim + 1, DD=dim**2)
    ).Metric(),
    element=TensorElement("DG", mesh.ufl_cell(), 0),
)

We would use G in a UFL form in inner(u, G*u).

I realize that this is quite specific, but would appreciate any hints on how to achieve this in dolfinx. What is the recommended way to handle such problems?

Thanks!

I would suggest using a dolfinx.fem.Expression, that can be interpolated into any function space, as for instance shown here: https://jsdokken.com/dolfinx-tutorial/chapter1/membrane_code.html?highlight=expression#interpolation-of-a-ufl-expression

You can use ufl.conditional for the conditionals in your output.

Thank you, dokken. I’m still uncertain how to replicate the second example, the mesh metric, which uses the coordinates of all vertices of a cell. I have this old python version, too, doing the same (but slowly) as the 2nd C++ expression in my first post (I think adapted from some post in the old forum):

        class Metric(Expression):
            def __init__(self, mesh, **kwargs):
                self.mesh = mesh
                self.dim = mesh.topology().dim()

            def eval_cell(self, values, x, cell):
                x = Cell(self.mesh, cell.index).get_vertex_coordinates()
                if self.dim == 2:
                    x0, x1, x2 = x.reshape(3, 2)
                    F = np.c_[x1 - x0, x2 - x0]
                elif self.dim == 3:
                    x0, x1, x2, x3 = x.reshape(4, 3)
                    F = np.c_[x1 - x0, x2 - x0, x3 - x0]
                Finv = np.linalg.inv(F)
                values[:] = (Finv.T.dot(Finv)).flatten()

            def value_shape(self):
                return (self.dim, self.dim)

        G = Metric(self.mesh, degree=0)
        Th = TensorFunctionSpace(self.mesh, "DG", 0)
        Gh = interpolate(G, Th)

Is there a way in dolfinx to do that?
Thank you!

Hi,
I think the ufl Jacobian function should answer your needs.
https://fenics.readthedocs.io/projects/ufl/en/latest/api-doc/ufl.html#ufl.classes.Jacobian

Oh, fantastic! I didn’t know this existed at all. Thank you :slight_smile:
(And it’s not even new…)