src/dolfinx_access/minmax_p2.cc

/// Minimisation code for second-order polynomial models from Fenics.
#include <array>
#include <cmath>

#include <basix/finite-element.h>
#include <dolfinx/fem/Function.h>
#include <dolfinx/mesh/cell_types.h>
#include <nanobind/nanobind.h>
#include <nanobind/stl/array.h> // needed for array conversions in nanobind bindings.

#include "dolfinx_access/minmax_p2.h"
#include "pointsource_pde/src/dolfinx_access.rs.h"

using namespace dolfinx::fem;
using namespace basix::element;
using namespace dolfinx_access;

namespace dolfinx_access {

    // Since Fenics documentation is poor, and it does not seem to provide an easy way to access
    // an internal P2 model (that would not involve vomotting your guts out 🤮 dealing with the
    // putrid horror that C++ has become (and its documentation written by lawyers and bureacrats,
    // which is even worse than that of Fenics), our idea here is to simply assume that f is a
    // second-order polynomial on each cell, and then evaluate it on that cell at 6 points, to
    // construct a new model based on our existing P2 model code that is also used for minimising
    // arbitrary sums.
    inline CoordValuePair minmax_Function_f64_p2(Function<double> const* f, bool max) {
        auto v = f->x();
        // Mesh
        auto fp = f->function_space();
        auto mesh = fp->mesh();
        auto geom = mesh->geometry();
        auto topo = mesh->topology();

        // Check that we're working with P2 elements
        //
        // Fenics defaults to GLL warped midpoint placement; see
        // https://docs.fenicsproject.org/dolfinx/v0.7.0.post0/python/demos/demo_lagrange_variants.html
        //  This is not a problem with our current implementation, which works with an
        //  second-order
        // polynomials due to duplication of effort (see above, why), but if we would directly
        // take the cell quadrature matrix from Fenics, we should restrict el.lagrange_variant() =
        // lagrange_variant::equispaced.
        auto el = fp->element()->basix_element();
        // printf("%d %d %d %d\n", el.degree(), el.cell_type(), el.lagrange_variant(),
        // el.family());
        if (el.degree() != 2 || el.cell_type() != basix::cell::type::triangle ||
            /*el.lagrange_variant() != lagrange_variant::equispaced ||*/ el.family() !=
                family::P) {
            throw "Only equispaced Lagrange second-order polynomial elements are supported";
        }
        if (cell_num_entities(dolfinx::mesh::CellType::triangle, 0) != 3) {
            throw "A triangle should have three vertices";
        }

        // Check that we are dealing with a scalar function
        if (fp->element()->value_shape().size() != 0) {
            throw "Only scalar functions are supported, obviously.";
        }

        // Check that we're in two dimensions
        if (topo->dim() != 2 || geom.dim() != 2) {
            throw "Only two-dimensional meshes are supported";
        }

        // Check that there are no other types of eleemnts
        auto entity_types = topo->entity_types(2);
        for (auto& t : entity_types) {
            if (t != dolfinx::mesh::CellType::triangle) {
                throw "Only triangular meshes are supported";
            }
        }

        auto adj = topo->connectivity(2, 0);
        auto n_cells = adj->num_nodes();

        // assert(n_cells == n_triangles);

        // printf("cell count? %d\n", n_cells);
        // for (auto i = 0; i < n_cells; i++) {
        //     // assert(adj->num_links(i) == 3); // Should not need this
        //     auto d = fp->dofmap()->cell_dofs(i);
        //     printf("%zd\n", d.size());
        // }

        // DOF coordinates and dofmap for the function space itself (not the mesh geometry).
        // For a scalar P2 space this gives us, per cell, the 6 nodal points (3 corners + 3 edge
        // midpoints) directly in basix's local ordering, so we no longer need to reconstruct
        // edge midpoints via f->eval.
        auto fp_dofmap = fp->dofmap();
        // transpose=false gives a flat vector laid out as (num_dofs, 3) row-major, i.e.
        // dof_coords[3*d + component].
        auto dof_coords = fp->tabulate_dof_coordinates(false);
        auto v_array = v->array(); // local coefficient vector

        CoordValuePair res{{{0.0, 0.0}}, INFINITY};

        for (auto i = 0; i < n_cells; i++) {
            auto dofs = fp_dofmap->cell_dofs(i);
            assert(dofs.size() == 6);

            // Coordinates of all 6 P2 nodes for this cell, in Fenics' hard-coded 3D layout.
            std::array<double, 18> x{};
            // Nodal values of f at those same 6 points.
            std::array<double, 6> vals{};

            for (std::size_t k = 0; k < 6; k++) {
                auto d = dofs[k];
                vals[k] = v_array[d];
                x[3 * k] = dof_coords[3 * d];
                x[3 * k + 1] = dof_coords[3 * d + 1];
                x[3 * k + 2] = 0.0;
            }

            // Find the solution on the simplex.
            auto newres = minmax_dolfinx_p2(x, vals, max);

            if (newres.v < res.v) {
                res = newres;
            }
        }

        // Negate result if maximising.
        if (max) {
            res.v = -res.v;
        }
        return res;
    }

    void cell_dofs_Function_f64_p2(Function<double> const* f, int32_t cell,
                                   std::array<double, 18>& x, std::array<double, 6>& vals) {
        auto fp = f->function_space();
        auto fp_dofmap = fp->dofmap();
        auto dofs = fp_dofmap->cell_dofs(cell);
        assert(dofs.size() == 6);
        auto dof_coords = fp->tabulate_dof_coordinates(false);
        auto v = f->x();
        auto v_array = v->array(); // local coefficient vector

        for (std::size_t k = 0; k < 6; k++) {
            auto d = dofs[k];
            vals[k] = v_array[d];
            x[3 * k] = dof_coords[3 * d];
            x[3 * k + 1] = dof_coords[3 * d + 1];
            x[3 * k + 2] = 0.0;
        }
    }

    CoordValuePair min_Function_f64_p2(Function<double> const* f) {
        return minmax_Function_f64_p2(f, false);
    }

    CoordValuePair max_Function_f64_p2(Function<double> const* f) {
        return minmax_Function_f64_p2(f, true);
    }

    // Create a Python module as well
    NB_MODULE(_dolfinx_access, m) {
        m.def("min_Function_f64_p2", &min_Function_f64_p2);
        m.def("max_Function_f64_p2", &max_Function_f64_p2);
        m.def("minmax_Function_f64_p2", &minmax_Function_f64_p2);
        m.def("eval_Function_f64", &eval_Function_f64);
        m.def("cell_Function_f64", &cell_Function_f64);
        m.def("cell_dofs_Function_f64_p2", &cell_dofs_Function_f64_p2);
        // m.def("cell_Mesh_f64", &cell_Mesh_f64);
        m.def("cell_FunctionSpace_f64", &cell_FunctionSpace_f64);
        nanobind::class_<CoordValuePair>(m, "CoordValuePair")
            .def_rw("x", &CoordValuePair::x)
            .def_rw("v", &CoordValuePair::v);
    }

} // namespace dolfinx_access

// Forward declaration for manual registration
extern "C" PyObject* PyInit__dolfinx_access();