--- /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(()) +}