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;
namespace cell = basix::cell;
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() != 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(mesh::CellType::triangle, 0) != 3) {
            throw "A triangle should have three vertices";
        }

        // Check that we are dealing with a scalar function
        if (fp->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 != 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 (may be more than nodes, e.g., intermediate points)
        auto gx = geom.x();
        // auto num_points = gx.size() / 3;
        //  Map from cells to DOF indices
        auto geom_dofmap = geom.dofmap();
        assert(geom_dofmap.extent(0) == (size_t)n_cells);
        assert(geom_dofmap.extent(1) == 3);
        auto fp_dofmap = fp->dofmap();

        // auto dof_coords = fp->tabulate_dof_coordinates();

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

        // printf("%zd %zd\n", v->array().size(), gx.size());

        for (auto i = 0; i < n_cells; i++) {
            // Cell corners
            auto j0 = geom_dofmap(i, 0);
            auto j1 = geom_dofmap(i, 1);
            auto j2 = geom_dofmap(i, 2);

            // Construct list of coordinates. Due to f->eval, we construct this is a single list
            // with all the Fenics weirdness of hard-coded 3D.
            std::array<double, 9> x = {gx[3 * j0], gx[3 * j0 + 1], 0,
                                       gx[3 * j1], gx[3 * j1 + 1], 0,
                                       gx[3 * j2], gx[3 * j2 + 1], 0};

            // Find the solution on the simplex.
            auto newres = minmax_dolfinx_p2_cell(f, x, i, max);

            // printf("%f [%f, %f] [%f %f %f %f %f %f] <%f, %f> <%f, %f> <‸%f, %f>\n", newres.v,
            //        newres.x[0], newres.x[1], values[0], values[1], values[2], values[3],
            //        values[4], values[5], x[0], x[1], x[3], x[4], x[6], x[7]);
            if (newres.v < res.v) {
                res = newres;
            }

            // // Full DOF data
            // auto d = fp_dofmap->cell_dofs(i);
            // // TODO: array() gives just local part of vector
            // assert(d.size() == 6);
            // for (auto k = 0; k < d.size(); k++) {
            //     auto j = d[k];
            //     // printf("%d %d\n", 3 * j + 1, nodes.size());
            //     assert(2 * j + 1 <= nodes.size());
            //     auto x = {nodes[3 * j], nodes[3 * j + 1]};
            // }
        }

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

    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_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();