tip/src/dolfinx_access/function.cc

#include <array>
#include <cmath>
#include <span>

#include <basix/finite-element.h>
#include <dolfinx/fem/Function.h>
#include <nanobind/nanobind.h>

#include "dolfinx_access/function.h"
// #include "mpi_proto.h"
#include "dolfinx/geometry/BoundingBoxTree.h"
#include "pointsource_pde/src/dolfinx_access.rs.h"
#include "rust/cxx.h"

using namespace dolfinx::fem;
using namespace dolfinx::geometry;
using namespace basix::element;
using namespace dolfinx::graph;
using namespace dolfinx::la;

namespace dolfinx_access {
    extern void drop_Function_f64(Function_f64* f) {
        delete f;
    }

    FunctionInfo info_Function_f64(Function_f64 const* f) {
        auto fp = f->function_space();
        auto el = fp->element()->basix_element();
        return FunctionInfo{
            .domain_dim = (uint32_t)fp->mesh()->geometry().dim(),
            .codomain_dim = (uint32_t)fp->element()->value_size(),
            .order = (uint32_t)fp->element()->basix_element().degree(),
            .triangular_mesh = el.cell_type() == basix::cell::type::triangle,
        };
    }

    void check_compat_Function_f64(Function_f64 const* f, Function_f64 const* g) {
        if (f->function_space() != g->function_space()) {
            throw "Incompatible function spaces";
        }
    }

    std::map<std::weak_ptr<const dolfinx::mesh::Mesh<double>>,
             std::shared_ptr<BoundingBoxTree<double>>,
             std::owner_less<std::weak_ptr<const dolfinx::mesh::Mesh<double>>>>
        bb_tree_cache;

    /// Returns the cell containg x
    int32_t cell_Mesh_f64(std::shared_ptr<const dolfinx::mesh::Mesh<double>> mesh,
                          const std::array<double, 3>& x) {

        std::weak_ptr<const dolfinx::mesh::Mesh<double>> mesh_weak = mesh;
        std::shared_ptr<BoundingBoxTree<double>> bb_tree = bb_tree_cache[mesh_weak];
        if (bb_tree == nullptr) {
            auto topo = mesh->topology();
            // auto dim = mesh->geometry().dim();
            auto dim = topo->dim();

            // This fails as it inorrectly calls mesh.topology_mutable(), so need to
            // copy-paste the implementation below, with the appropriate fix.
            //*bb_tree = new BoundingBoxTree(*mesh, dim);
            // And this also doesn't work due to the privacy of `range`.
            //*bb_tree = new BoundingBoxTree(*mesh, dim, BoundingBoxTree::range(mesh->topology(),
            // dim
            auto index_map = topo->index_map(dim);
            int local_cells = index_map->size_local();
            // We don't really do MPI; pointsource_algs is not designed for it.
            assert(local_cells == index_map->size_global());
            std::vector<std::int32_t> r(local_cells);
            std::iota(r.begin(), r.end(), 0);
            bb_tree =
                std::make_shared<BoundingBoxTree<double>>(BoundingBoxTree(*mesh, dim, 0.0, r));
            bb_tree_cache[mesh_weak] = bb_tree;
        }
        // Find colliding bounding boxes. This is an AdjancencyTree.
        //
        // C++ and Dolfinx are just vomit-inducing 🤮. We need some weird const double span here,
        // but then just standard x in compute_first_collding_cell below.
        auto colliding_bbs_adj = compute_collisions(*bb_tree.get(), std::span<const double>(x));
        // Since we compute the collisions for just one node, the `array` of all links
        // of the adjancency list, will corresponding exactly the ccells whose bounding boxes
        // collide with the point.
        auto colliding_bbs = colliding_bbs_adj.array();
        // Within the colliding bounding boxes, find the first cell that actually collides.
        int32_t local_cell = compute_first_colliding_cell(*mesh, colliding_bbs, x, 1e-9);
        // local_cell may be -1 if it was not found by the local process, so combine
        // results with MPI_MAX.
        int32_t global_cell = local_cell;
        MPI_Allreduce(&local_cell, &global_cell, 1, MPI_INT32_T, MPI_MAX, mesh->comm());

        return global_cell;
    }

    int32_t cell_FunctionSpace_f64(FunctionSpace<double> const* v,
                                   const std::array<double, 3>& x) {
        return cell_Mesh_f64(v->mesh(), x);
    }

    int32_t cell_Function_f64(Function_f64 const* f, const std::array<double, 3>& x) {
        return cell_Mesh_f64(f->function_space()->mesh(), x);
    }

    std::array<double, 3 * 3> cell_coords_Function_f64_triangle(Function_f64 const* f,
                                                                int32_t cell) {
        auto geom = f->function_space()->mesh()->geometry();
        auto gx = geom.x();
        auto geom_dofmap = geom.dofmap();
        assert(geom_dofmap.extent(1) == 3);

        auto j0 = geom_dofmap(cell, 0);
        auto j1 = geom_dofmap(cell, 1);
        auto j2 = geom_dofmap(cell, 2);

        // Construct list of coordinates. Due to f->eval, we construct this as a single list
        // with all the Fenics weirdness of hard-coded 3D.
        return {gx[3 * j0], gx[3 * j0 + 1], 0, /* */
                gx[3 * j1], gx[3 * j1 + 1], 0, /* */
                gx[3 * j2], gx[3 * j2 + 1], 0};
    }

    // We assume `f` to be real-valued.
    double eval_Function_f64_cell(Function_f64 const* f, const std::array<double, 3>& x,
                                  int32_t cell) {
        if (cell < 0) {
            return 0.0;
        } else {
            std::array<double, 1> res;
            f->eval(x, {1, 3}, {{cell}}, std::span(res), {1, 1});
            return res[0];
        }
    }

    double eval_Function_f64(Function_f64 const* f, const std::array<double, 3>& x) {
        return eval_Function_f64_cell(f, x, cell_Function_f64(f, x));
    }

    std::array<double, 6> eval_Function_f64_cell_6(Function_f64 const* f,
                                                   const std::array<double, 3 * 6>& x,
                                                   int32_t cell) {
        std::array<double, 6> res;
        auto i = cell;
        f->eval(x, {6, 3}, {{i, i, i, i, i, i}}, std::span(res), {6, 1});
        return res;
    }

    rust::Slice<const double> data_Function_f64(Function_f64 const* f) {
        const std::vector<double>& v = f->x()->array();
        return rust::Slice<const double>(v.data(), v.size());
    }

    rust::Slice<double> data_mut_Function_f64(Function_f64* f) {
        std::vector<double>& v = f->x()->array();
        return rust::Slice<double>(v.data(), v.size());
    }

    Function_f64* similar_Function_f64(Function_f64 const* f) {
        return new Function_f64(f->function_space());
    }

    Function_f64* clone_Function_f64(Function_f64 const* f) {
        // return new Function_f64(f->function_space(),
        // std::make_shared<Vector<double>>(*f->x()));
        auto f_copy = new Function_f64(f->function_space());
        auto fx = f->x();
        auto x_new = f_copy->x();
        const std::vector<double>& a_old = fx->array();
        std::vector<double>& a_new = x_new->array();
        std::copy(a_old.begin(), a_old.end(), a_new.begin());
        x_new->scatter_fwd();
        return f_copy;
    }

    void scatter_fwd_Function_f64(Function_f64* f) {
        f->x()->scatter_fwd();
    }

    // PyObject* py_similar_Function_f64(Function_f64 const* f) {
    //     return nanobind::cast(Function_f64(f->function_space())).release().ptr();
    // }

} // namespace dolfinx_access