ndgrid/geometry/
geometry_map.rs

1//! Geometry map
2use crate::{traits::GeometryMap as GeometryMapTrait, types::Scalar};
3use ndelement::{reference_cell, traits::MappedFiniteElement, types::ReferenceCellType};
4use rlst::{Array, DynArray, RlstScalar, ValueArrayImpl};
5
6/// Single element geometry
7pub struct GeometryMap<'a, T, B2D, C2D> {
8    geometry_points: &'a Array<B2D, 2>,
9    entities: &'a Array<C2D, 2>,
10    tdim: usize,
11    gdim: usize,
12    table: DynArray<T, 4>,
13}
14
15fn norm<T: RlstScalar>(vector: &[T]) -> T {
16    vector.iter().map(|&i| i * i).sum::<T>().sqrt()
17}
18
19fn cross<T: RlstScalar>(mat: &[T], result: &mut [T]) {
20    match mat.len() {
21        0 => {}
22        2 => {
23            debug_assert!(result.len() == 2);
24            unsafe {
25                *result.get_unchecked_mut(0) = *mat.get_unchecked(1);
26                *result.get_unchecked_mut(1) = -*mat.get_unchecked(0);
27            }
28        }
29        6 => {
30            debug_assert!(result.len() == 3);
31            unsafe {
32                *result.get_unchecked_mut(0) = *mat.get_unchecked(1) * *mat.get_unchecked(5)
33                    - *mat.get_unchecked(2) * *mat.get_unchecked(4);
34                *result.get_unchecked_mut(1) = *mat.get_unchecked(2) * *mat.get_unchecked(3)
35                    - *mat.get_unchecked(0) * *mat.get_unchecked(5);
36                *result.get_unchecked_mut(2) = *mat.get_unchecked(0) * *mat.get_unchecked(4)
37                    - *mat.get_unchecked(1) * *mat.get_unchecked(3);
38            }
39        }
40        _ => {
41            unimplemented!();
42        }
43    }
44}
45
46impl<'a, T: Scalar, B2D: ValueArrayImpl<T, 2>, C2D: ValueArrayImpl<usize, 2>>
47    GeometryMap<'a, T, B2D, C2D>
48{
49    /// Create new
50    pub fn new<A2D: ValueArrayImpl<T, 2>>(
51        element: &impl MappedFiniteElement<CellType = ReferenceCellType, T = T>,
52        points: &Array<A2D, 2>,
53        geometry_points: &'a Array<B2D, 2>,
54        entities: &'a Array<C2D, 2>,
55    ) -> Self {
56        let tdim = reference_cell::dim(element.cell_type());
57        debug_assert!(points.shape()[0] == tdim);
58        let gdim = geometry_points.shape()[0];
59        let npoints = points.shape()[1];
60
61        let mut table = DynArray::<T, 4>::from_shape(element.tabulate_array_shape(1, npoints));
62        element.tabulate(points, 1, &mut table);
63
64        Self {
65            geometry_points,
66            entities,
67            tdim,
68            gdim,
69            table,
70        }
71    }
72}
73
74impl<T: Scalar, B2D: ValueArrayImpl<T, 2>, C2D: ValueArrayImpl<usize, 2>> GeometryMapTrait
75    for GeometryMap<'_, T, B2D, C2D>
76{
77    type T = T;
78
79    fn entity_topology_dimension(&self) -> usize {
80        self.tdim
81    }
82    fn geometry_dimension(&self) -> usize {
83        self.gdim
84    }
85    fn point_count(&self) -> usize {
86        self.table.shape()[1]
87    }
88    fn physical_points(&self, entity_index: usize, points: &mut [T]) {
89        let npts = self.table.shape()[1];
90        debug_assert!(points.len() == self.gdim * npts);
91
92        points.fill(T::default());
93        for i in 0..self.entities.shape()[0] {
94            let v = unsafe { self.entities.get_value_unchecked([i, entity_index]) };
95            for point_index in 0..npts {
96                let t = unsafe { *self.table.get_unchecked([0, point_index, i, 0]) };
97                for gd in 0..self.gdim {
98                    unsafe {
99                        *points.get_unchecked_mut(gd + self.gdim * point_index) +=
100                            self.geometry_points.get_value_unchecked([gd, v]) * t
101                    };
102                }
103            }
104        }
105    }
106    fn jacobians(&self, entity_index: usize, jacobians: &mut [T]) {
107        let npts = self.table.shape()[1];
108        debug_assert!(jacobians.len() == self.gdim * self.tdim * npts);
109
110        jacobians.fill(T::zero());
111        for i in 0..self.entities.shape()[0] {
112            let v = unsafe { self.entities.get_value_unchecked([i, entity_index]) };
113            for point_index in 0..npts {
114                for td in 0..self.tdim {
115                    let t = unsafe { self.table.get_value_unchecked([1 + td, point_index, i, 0]) };
116                    for gd in 0..self.gdim {
117                        unsafe {
118                            *jacobians.get_unchecked_mut(
119                                gd + self.gdim * td + self.gdim * self.tdim * point_index,
120                            ) += self.geometry_points.get_value_unchecked([gd, v]) * t
121                        };
122                    }
123                }
124            }
125        }
126    }
127
128    fn jacobians_dets_normals(
129        &self,
130        entity_index: usize,
131        jacobians: &mut [Self::T],
132        jdets: &mut [Self::T],
133        normals: &mut [Self::T],
134    ) {
135        if self.tdim + 1 != self.gdim {
136            panic!("Can only compute normal for entities where tdim + 1 == gdim");
137        }
138        let npts = self.table.shape()[1];
139        debug_assert!(jacobians.len() == self.gdim * self.tdim * npts);
140        debug_assert!(jdets.len() == npts);
141        debug_assert!(normals.len() == self.gdim * npts);
142
143        self.jacobians(entity_index, jacobians);
144
145        for point_index in 0..npts {
146            let j = unsafe {
147                &jacobians.get_unchecked(
148                    self.gdim * self.tdim * point_index..self.gdim * self.tdim * (point_index + 1),
149                )
150            };
151            let n = unsafe {
152                &mut normals
153                    .get_unchecked_mut(self.gdim * point_index..self.gdim * (point_index + 1))
154            };
155            let jd = unsafe { jdets.get_unchecked_mut(point_index) };
156            cross(j, n);
157            *jd = norm(n);
158            for n_i in n.iter_mut() {
159                *n_i /= *jd;
160            }
161        }
162    }
163}