src/dolfinx_access.rs

changeset 1
a4137aedcb3a
child 3
c3a4f4bb87f7
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/dolfinx_access.rs	Thu Feb 26 09:32:12 2026 -0500
@@ -0,0 +1,209 @@
+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;
+use std::mem::MaybeUninit;
+
+// These are 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;
+
+#[allow(dead_code)]
+#[cxx::bridge(namespace = "dolfinx_access")]
+pub mod ffi {
+    #[derive(Copy, Clone, Debug)]
+    struct CoordValuePair {
+        x: [f64; 2],
+        v: f64,
+    }
+
+    #[derive(Copy, Clone, Debug)]
+    struct FunctionInfo {
+        domain_dim: u32,
+        codomain_dim: u32,
+        order: u32,
+        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");
+
+        type Function_f64;
+        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];
+
+        pub unsafe fn info_Function_f64(f: *const Function_f64) -> FunctionInfo;
+
+        pub unsafe fn min_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>;
+        pub unsafe fn max_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>;
+        pub unsafe fn minmax_Function_f64_p2(
+            f: *const Function_f64,
+            max: bool,
+        ) -> Result<CoordValuePair>;
+
+        pub unsafe fn check_Function_f64(o: *const PyObject) -> bool;
+        pub unsafe fn cast_Function_f64(o: *const PyObject) -> Result<*const Function_f64>;
+        pub unsafe fn cast_mut_Function_f64(o: *mut PyObject) -> Result<*mut Function_f64>;
+        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 function
+        pub unsafe fn similar_Function_f64(f: *const Function_f64) -> *mut Function_f64;
+
+        /// Clone the function
+        pub unsafe fn clone_Function_f64(f: *const Function_f64) -> *mut Function_f64;
+
+        pub unsafe fn wrap_Function_f64(f: *mut Function_f64) -> *mut PyObject;
+
+        pub unsafe fn drop_Function_f64(f: *mut Function_f64);
+    }
+
+    extern "Rust" {
+        unsafe fn minmax_dolfinx_p2_cell(
+            f: *const Function_f64,
+            simplex_coords: &[f64; 9],
+            cell: i32,
+            max: bool,
+        ) -> CoordValuePair;
+    }
+}
+
+use ffi::CoordValuePair;
+
+/// Convert an array of 2D oordinates obtained from Fenics, which hard-codes weird 3D coordinates.
+#[inline]
+fn coords_from_fenics(&[x1, x2, _, y1, y2, _, z1, z2, _]: &[f64; 9]) -> [Loc<2, f64>; 3] {
+    [Loc([x1, x2]), Loc([y1, y2]), Loc([z1, z2])]
+}
+// fn coords_from_fenics<const N: usize>(coords: &[[f64; 3]; N]) -> [Loc<F, 2>; N] {
+//     coords.map(|[x1, x2, _]| Loc([x1, x2]))
+// }
+
+/// Convert an array of 2D oordinates to Fenics, which hard-codes weird 3D coordinates.
+#[inline]
+fn coords_to_fenics([Loc([x1, x2]), Loc([y1, y2]), Loc([z1, z2])]: [Loc<2, f64>; 3]) -> [f64; 9] {
+    [x1, x2, 0.0, y1, y2, 0.0, z1, z2, 0.0]
+}
+// fn coords_to_fenics<const N: usize>(coords: [Loc<F, 2>; N]) -> [[f64; 3]; N] {
+//     coords.map(|Loc([x1, x2])| [x1, x2, 0.0])
+// }
+
+/// 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`.
+unsafe fn model_dolfinx_p2_cell(
+    f: *const ffi::Function_f64,
+    simplex_coords: &[f64; 9],
+    cell: i32,
+    max: bool,
+) -> (P2LocalModel<f64, 2, 3>, Simplex<f64, 2, 3>) {
+    // Form simplex
+    let simplex = Simplex(coords_from_fenics(simplex_coords));
+    // Form an array of all coordinates
+    let midpoints = simplex.midpoints();
+    let mut all_coords_uninit: [MaybeUninit<f64>; 6 * 3] = [const { MaybeUninit::uninit() }; 6 * 3];
+    all_coords_uninit
+        .iter_mut()
+        .zip(
+            simplex_coords
+                .iter()
+                .copied()
+                .chain(coords_to_fenics(midpoints).into_iter()),
+        )
+        .for_each(|(a, b)| {
+            a.write(b);
+        });
+    let all_coords = MaybeUninit::array_assume_init(all_coords_uninit);
+    // Evaluate the function, and invert the signs of maximising
+    let mut values = ffi::eval_Function_f64_cell_6(f, &all_coords, cell);
+    if max {
+        values.iter_mut().for_each(|v| *v = -*v);
+    }
+    // Form local model
+    let [ref n0, ref n1, ref n2] = simplex.0;
+    let [ref n01, ref n12, ref n20] = midpoints;
+    let nodes = [*n0, *n1, *n2, *n01, *n12, *n20];
+    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_cell(
+    f: *const ffi::Function_f64,
+    simplex_coords: &[f64; 9],
+    cell: i32,
+    max: bool,
+) -> ffi::CoordValuePair {
+    let (model, simplex) = model_dolfinx_p2_cell(f, simplex_coords, cell, 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(())
+}

mercurial