Mercurial > repos > pointsource_pde / file comparison

comparison: src/dolfinx_access.rs

src/dolfinx_access.rs

changeset 1
a4137aedcb3a
child 3
c3a4f4bb87f7
equal deleted inserted replaced
0:7ec1cfe19a24 1:a4137aedcb3a
1 use alg_tools::error::DynResult;
2 use alg_tools::fe_model::{
3 base::RealLocalModel,
4 p2_local_model::{P2LocalModel, Simplex},
5 };
6 use alg_tools::loc::Loc;
7 use anyhow::bail;
8 use pyo3::ffi::PyImport_AppendInittab;
9 use std::ffi::CString;
10 use std::mem::MaybeUninit;
11
12 // These are required for the linking to the sys crates to actually happen.
13 #[allow(unused_imports)]
14 mod dummy_import {
15 use dolfinx_sys;
16 use nanobind_sys;
17 }
18
19 mod function;
20 pub use function::DolfinxPyFunction_f64;
21
22 #[allow(dead_code)]
23 #[cxx::bridge(namespace = "dolfinx_access")]
24 pub mod ffi {
25 #[derive(Copy, Clone, Debug)]
26 struct CoordValuePair {
27 x: [f64; 2],
28 v: f64,
29 }
30
31 #[derive(Copy, Clone, Debug)]
32 struct FunctionInfo {
33 domain_dim: u32,
34 codomain_dim: u32,
35 order: u32,
36 triangular_mesh: bool,
37 }
38
39 unsafe extern "C++" {
40 include!("pointsource_pde/include/dolfinx_access/nanobind_helpers.h");
41 include!("pointsource_pde/include/dolfinx_access/function.h");
42 include!("pointsource_pde/include/dolfinx_access/minmax_p2.h");
43
44 type Function_f64;
45 type PyObject;
46
47 /// Find the cell containing `x` (in 1D, 2D, or 3D, following Fenics weirdness
48 /// of hard-coding vectors to 3D).
49 pub unsafe fn cell_Function_f64(f: *const Function_f64, x: &[f64; 3]) -> i32;
50
51 /// Evaluate `f` at `x`.
52 pub unsafe fn eval_Function_f64(f: *const Function_f64, x: &[f64; 3]) -> f64;
53
54 /// Evaluates `f` at `x` (in 1D, 2D, or 3D, following Fenics weirdness
55 /// of hard-coding vectors to 3D), assuming `x` to belong to `cell`.
56 pub unsafe fn eval_Function_f64_cell(
57 f: *const Function_f64,
58 x: &[f64; 3],
59 cell: i32,
60 ) -> f64;
61
62 /// Evaluate the function at 6 coordinates known to belong to given cell
63 pub unsafe fn eval_Function_f64_cell_6(
64 f: *const Function_f64,
65 all_coords: &[f64; 18], // 6*3
66 cell: i32,
67 ) -> [f64; 6];
68
69 pub unsafe fn info_Function_f64(f: *const Function_f64) -> FunctionInfo;
70
71 pub unsafe fn min_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>;
72 pub unsafe fn max_Function_f64_p2(f: *const Function_f64) -> Result<CoordValuePair>;
73 pub unsafe fn minmax_Function_f64_p2(
74 f: *const Function_f64,
75 max: bool,
76 ) -> Result<CoordValuePair>;
77
78 pub unsafe fn check_Function_f64(o: *const PyObject) -> bool;
79 pub unsafe fn cast_Function_f64(o: *const PyObject) -> Result<*const Function_f64>;
80 pub unsafe fn cast_mut_Function_f64(o: *mut PyObject) -> Result<*mut Function_f64>;
81 pub unsafe fn cell_coords_Function_f64_triangle(
82 f: *const Function_f64,
83 cell: i32,
84 ) -> [f64; 9];
85
86 /// Returns the local weights
87 pub unsafe fn data_Function_f64<'a>(f: *const Function_f64) -> &'a [f64];
88
89 /// Returns the local mutable weights
90 pub unsafe fn data_mut_Function_f64<'a>(f: *mut Function_f64) -> &'a mut [f64];
91
92 /// Check the compatibility of function spaces of two functions.
93 pub unsafe fn check_compat_Function_f64(
94 f: *const Function_f64,
95 g: *const Function_f64,
96 ) -> Result<()>;
97
98 /// Create a new similar function
99 pub unsafe fn similar_Function_f64(f: *const Function_f64) -> *mut Function_f64;
100
101 /// Clone the function
102 pub unsafe fn clone_Function_f64(f: *const Function_f64) -> *mut Function_f64;
103
104 pub unsafe fn wrap_Function_f64(f: *mut Function_f64) -> *mut PyObject;
105
106 pub unsafe fn drop_Function_f64(f: *mut Function_f64);
107 }
108
109 extern "Rust" {
110 unsafe fn minmax_dolfinx_p2_cell(
111 f: *const Function_f64,
112 simplex_coords: &[f64; 9],
113 cell: i32,
114 max: bool,
115 ) -> CoordValuePair;
116 }
117 }
118
119 use ffi::CoordValuePair;
120
121 /// Convert an array of 2D oordinates obtained from Fenics, which hard-codes weird 3D coordinates.
122 #[inline]
123 fn coords_from_fenics(&[x1, x2, _, y1, y2, _, z1, z2, _]: &[f64; 9]) -> [Loc<2, f64>; 3] {
124 [Loc([x1, x2]), Loc([y1, y2]), Loc([z1, z2])]
125 }
126 // fn coords_from_fenics<const N: usize>(coords: &[[f64; 3]; N]) -> [Loc<F, 2>; N] {
127 // coords.map(|[x1, x2, _]| Loc([x1, x2]))
128 // }
129
130 /// Convert an array of 2D oordinates to Fenics, which hard-codes weird 3D coordinates.
131 #[inline]
132 fn coords_to_fenics([Loc([x1, x2]), Loc([y1, y2]), Loc([z1, z2])]: [Loc<2, f64>; 3]) -> [f64; 9] {
133 [x1, x2, 0.0, y1, y2, 0.0, z1, z2, 0.0]
134 }
135 // fn coords_to_fenics<const N: usize>(coords: [Loc<F, 2>; N]) -> [[f64; 3]; N] {
136 // coords.map(|Loc([x1, x2])| [x1, x2, 0.0])
137 // }
138
139 /// Given a dolphinx `Function` `f`, as well as 6 coordinates on a simplex in its mesh
140 /// (the corners and midpoints), construct a new internal P2 model on `cell`.
141 /// If `max` is true, the model is negated compared to `f`.
142 unsafe fn model_dolfinx_p2_cell(
143 f: *const ffi::Function_f64,
144 simplex_coords: &[f64; 9],
145 cell: i32,
146 max: bool,
147 ) -> (P2LocalModel<f64, 2, 3>, Simplex<f64, 2, 3>) {
148 // Form simplex
149 let simplex = Simplex(coords_from_fenics(simplex_coords));
150 // Form an array of all coordinates
151 let midpoints = simplex.midpoints();
152 let mut all_coords_uninit: [MaybeUninit<f64>; 6 * 3] = [const { MaybeUninit::uninit() }; 6 * 3];
153 all_coords_uninit
154 .iter_mut()
155 .zip(
156 simplex_coords
157 .iter()
158 .copied()
159 .chain(coords_to_fenics(midpoints).into_iter()),
160 )
161 .for_each(|(a, b)| {
162 a.write(b);
163 });
164 let all_coords = MaybeUninit::array_assume_init(all_coords_uninit);
165 // Evaluate the function, and invert the signs of maximising
166 let mut values = ffi::eval_Function_f64_cell_6(f, &all_coords, cell);
167 if max {
168 values.iter_mut().for_each(|v| *v = -*v);
169 }
170 // Form local model
171 let [ref n0, ref n1, ref n2] = simplex.0;
172 let [ref n01, ref n12, ref n20] = midpoints;
173 let nodes = [*n0, *n1, *n2, *n01, *n12, *n20];
174 let model = P2LocalModel::<f64, 2, 3>::new(&nodes, &values);
175
176 return (model, simplex);
177 }
178
179 /// This is a helper routine for `minmax_dolfinx_p2`. It reconstructs a P2
180 /// model on the cell by evaluating f at 6 points (the simplex corners given
181 /// in `simplex_doords`, as well as edge midpoints). Then it minimises this
182 /// model. The result should be accurate, given that `f` is assumed to have a
183 /// P2 model on the simplex.
184 pub unsafe fn minmax_dolfinx_p2_cell(
185 f: *const ffi::Function_f64,
186 simplex_coords: &[f64; 9],
187 cell: i32,
188 max: bool,
189 ) -> ffi::CoordValuePair {
190 let (model, simplex) = model_dolfinx_p2_cell(f, simplex_coords, cell, max);
191 let (x, v) = model.minimise(&simplex);
192 CoordValuePair { x: x.into(), v }
193 }
194
195 pub const NAME: &str = "dolfinx_access";
196
197 unsafe extern "C" {
198 pub fn PyInit__dolfinx_access() -> *mut pyo3::ffi::PyObject;
199 }
200
201 /// Register the module in Python
202 pub fn register_python_ffi() -> DynResult<()> {
203 let cname = CString::new(NAME).unwrap();
204 // We need to use into_raw() to transfer ownership; otherwise this stops working.
205 if unsafe { PyImport_AppendInittab(cname.into_raw(), Some(PyInit__dolfinx_access)) } != 0 {
206 bail!("Failed to add embedded nanobind dolfinx_access module");
207 }
208 Ok(())
209 }

Mercurial Repository: pointsource_pde

mercurial