src/dolfinx_access.rs

/*!
 Utility functions accessing C++ side Dolfinx.
*/
use alg_tools::error::DynResult;
use alg_tools::fe_model::{
    base::RealLocalModel,
    p2_local_model::{P2LocalModel, Simplex},
};
use alg_tools::loc::Loc;
use anyhow::bail;
use pyo3::ffi::PyImport_AppendInittab;
use std::ffi::CString;

/// Import helper required for the linking to the sys crates to actually happen.
#[allow(unused_imports)]
mod dummy_import {
    use dolfinx_sys;
    use nanobind_sys;
}

mod function;
pub use function::DolfinxPyFunction_f64;

/// C++ side routines
#[allow(dead_code)]
#[cxx::bridge(namespace = "dolfinx_access")]
pub mod ffi {
    /// Structure for coordinte-value pairs
    #[derive(Copy, Clone, Debug)]
    struct CoordValuePair {
        /// 2D coordinate
        x: [f64; 2],
        /// Value
        v: f64,
    }

    /// Structure for information about a dolfinx `Function`
    #[derive(Copy, Clone, Debug)]
    struct FunctionInfo {
        /// Domain dimension
        domain_dim: u32,
        /// Codomain dimension
        codomain_dim: u32,
        /// Element order
        order: u32,
        /// Whether the mesh is triangular
        triangular_mesh: bool,
    }

    unsafe extern "C++" {
        include!("pointsource_pde/include/dolfinx_access/nanobind_helpers.h");
        include!("pointsource_pde/include/dolfinx_access/function.h");
        include!("pointsource_pde/include/dolfinx_access/minmax_p2.h");

        /// C++ side Dolfinx `Function<double>`
        type Function_f64;
        /// C-side python object.
        type PyObject;

        /// Find the cell containing `x` (in 1D, 2D, or 3D, following Fenics weirdness
        /// of hard-coding vectors to 3D).
        pub unsafe fn cell_Function_f64(f: *const Function_f64, x: &[f64; 3]) -> i32;

        /// Evaluate `f` at `x`.
        pub unsafe fn eval_Function_f64(f: *const Function_f64, x: &[f64; 3]) -> f64;

        /// Evaluates `f` at `x` (in 1D, 2D, or 3D, following Fenics weirdness
        /// of hard-coding vectors to 3D), assuming `x` to belong to `cell`.
        pub unsafe fn eval_Function_f64_cell(
            f: *const Function_f64,
            x: &[f64; 3],
            cell: i32,
        ) -> f64;

        /// Evaluate the function at 6 coordinates known to belong to given cell
        pub unsafe fn eval_Function_f64_cell_6(
            f: *const Function_f64,
            all_coords: &[f64; 18], // 6*3
            cell: i32,
        ) -> [f64; 6];

        /// Get information about a dolfinx `Function_f64`.
        pub unsafe fn info_Function_f64(f: *const Function_f64) -> FunctionInfo;

        /// Returns the minimum value of the Dolfinx `Function`,
        /// if supported (P2 elements on triangular mesh)
        pub unsafe fn min_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>;

        /// Returns the maximum value of the Dolfinx `Function`,
        /// if supported (P2 elements on triangular mesh)
        pub unsafe fn max_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>;

        /// Returns the minimum or maximum value of the Dolfinx `Function`,
        /// if supported (P2 elements on triangular mesh)
        pub unsafe fn minmax_Function_f64_p2(
            f: *const Function_f64,
            max: bool,
        ) -> Result<CoordValuePair>;

        pub unsafe fn cell_dofs_Function_f64_p2(f: *const Function_f64, cell:i32, coords: &mut [f64; 18], vals: &mut [f64; 6]);

        /// Checks if a Dolfinx `Function` is supported (P2 elements on triangular mesh)
        pub unsafe fn check_Function_f64(o: *const PyObject) -> bool;

        /// Casts a Python object into a Dolfinx `Function` wrapper.
        pub unsafe fn cast_Function_f64(o: *const PyObject) -> Result<*const Function_f64>;

        /// Casts a Python object into a mutable Dolfinx `Function` wrapper.
        pub unsafe fn cast_mut_Function_f64(o: *mut PyObject) -> Result<*mut Function_f64>;

        /// Gets the 3D coordinates of the nodes of a triangular cell
        pub unsafe fn cell_coords_Function_f64_triangle(
            f: *const Function_f64,
            cell: i32,
        ) -> [f64; 9];

        /// Returns the local weights
        pub unsafe fn data_Function_f64<'a>(f: *const Function_f64) -> &'a [f64];

        /// Returns the local mutable weights
        pub unsafe fn data_mut_Function_f64<'a>(f: *mut Function_f64) -> &'a mut [f64];

        /// Check the compatibility of function spaces of two functions.
        pub unsafe fn check_compat_Function_f64(
            f: *const Function_f64,
            g: *const Function_f64,
        ) -> Result<()>;

        /// Create a new similar Dolfinx `Function`
        pub unsafe fn similar_Function_f64(f: *const Function_f64) -> *mut Function_f64;

        /// Clone the Dolfinx `Function`
        pub unsafe fn clone_Function_f64(f: *const Function_f64) -> *mut Function_f64;

        /// Wraps a Dolfinx `Function` into a Python object
        pub unsafe fn wrap_Function_f64(f: *mut Function_f64) -> *mut PyObject;

        /// Drop helper
        pub unsafe fn drop_Function_f64(f: *mut Function_f64);

        /// Scatter helper. Not truly supported.
        pub unsafe fn scatter_fwd_Function_f64(f: *mut Function_f64);
    }

    extern "Rust" {
        /// Rust-side helper for minimising P2 elements. This bypasses poorly documented basix,
        /// and does things at a low level directly.
        /// `f` is the function, `simplex_coords` the coorindates of the triangle
        /// see [`cell_coords_Function_f64_triangle`], `cell` is the cell number,
        /// and `max` indicates whether to maximise or minimise.
        unsafe fn minmax_dolfinx_p2(
            simplex_coords: [f64; 18],
            vals: [f64; 6],
            max: bool,
        ) -> CoordValuePair;
    }
}

use ffi::CoordValuePair;

/// Given a dolphinx `Function` `f`, as well as 6 coordinates on a simplex in its mesh
/// (the corners and midpoints), construct a new internal P2 model on `cell`.
/// If `max` is true, the model is negated compared to `f`.
#[inline]
unsafe fn model_dolfinx_p2_cell(
    [p0x, p0y, _, p1x, p1y, _, p2x, p2y, _, p12x, p12y, _, p20x, p20y, _, p01x, p01y, _]: [f64; 18],
    [v0, v1, v2, v12, v20, v01]: [f64; 6],
    max: bool,
) -> (P2LocalModel<f64, 2, 3>, Simplex<f64, 2, 3>) {
    // Form simplex
    let coords = [Loc([p0x, p0y]), Loc([p1x, p1y]), Loc([p2x, p2y])];
    let simplex = Simplex(coords);
    let values = if max {
        [-v0, -v1, -v2, -v01, -v12, -v20]
    } else {
        [v0, v1, v2, v01, v12, v20]
    };
    // Form local model
    let nodes = [Loc([p0x, p0y]), Loc([p1x, p1y]), Loc([p2x, p2y]), Loc([p01x, p01y]), Loc([p12x, p12y]), Loc([p20x, p20y])];
    let model = P2LocalModel::<f64, 2, 3>::new(&nodes, &values);

    return (model, simplex);
}

/// This is a helper routine for `minmax_dolfinx_p2`. It reconstructs a P2
/// model on the cell by evaluating f at 6 points (the simplex corners given
/// in `simplex_doords`, as well as edge midpoints). Then it minimises this
/// model. The result should be accurate, given that `f` is assumed to have a
/// P2 model on the simplex.
pub unsafe fn minmax_dolfinx_p2(
    simplex_coords: [f64; 18],
    vals: [f64; 6],
    max: bool,
) -> ffi::CoordValuePair {
    let (model, simplex) = model_dolfinx_p2_cell(simplex_coords, vals, max);
    let (x, v) = model.minimise(&simplex);
    CoordValuePair { x: x.into(), v }
}

pub const NAME: &str = "dolfinx_access";

unsafe extern "C" {
    pub fn PyInit__dolfinx_access() -> *mut pyo3::ffi::PyObject;
}

/// Register the module in Python
pub fn register_python_ffi() -> DynResult<()> {
    let cname = CString::new(NAME).unwrap();
    // We need to use into_raw() to transfer ownership; otherwise this stops working.
    if unsafe { PyImport_AppendInittab(cname.into_raw(), Some(PyInit__dolfinx_access)) } != 0 {
        bail!("Failed to add embedded nanobind dolfinx_access module");
    }
    Ok(())
}