1use crate::traits::FunctionSpace;
3use itertools::izip;
4use ndelement::{
5 reference_cell,
6 traits::{ElementFamily, MappedFiniteElement},
7 types::ReferenceCellType,
8};
9use ndmesh::{
10 traits::{Entity, Mesh, Topology},
11 types::Ownership,
12};
13use rlst::RlstScalar;
14use std::collections::HashMap;
15use std::fmt::Debug;
16use std::hash::Hash;
17
18pub struct FunctionSpaceImpl<
20 'a,
21 T: RlstScalar,
22 TMesh: RlstScalar,
23 E: Debug + PartialEq + Eq + Clone + Copy + Hash,
24 G: Mesh<EntityDescriptor = E, T = TMesh>,
25 F: MappedFiniteElement<CellType = E, T = T>,
26> {
27 mesh: &'a G,
28 elements: Vec<F>,
29 entity_dofs: HashMap<E, Vec<Vec<usize>>>,
30 entity_closure_dofs: HashMap<E, Vec<Vec<usize>>>,
31 ndofs: usize,
32 entities_by_element: HashMap<E, Vec<usize>>,
33}
34
35impl<
36 'a,
37 T: RlstScalar,
38 TMesh: RlstScalar,
39 G: Mesh<EntityDescriptor = ReferenceCellType, T = TMesh>,
40 F: MappedFiniteElement<CellType = ReferenceCellType, T = T>,
41> FunctionSpaceImpl<'a, T, TMesh, ReferenceCellType, G, F>
42{
43 pub fn new<EF: ElementFamily<FiniteElement = F, CellType = ReferenceCellType>>(
45 mesh: &'a G,
46 family: &EF,
47 ) -> Self {
48 let elements = mesh
49 .entity_types(mesh.topology_dim())
50 .iter()
51 .map(|e| family.element(*e))
52 .collect::<Vec<_>>();
53 let mut entity_dofs = HashMap::new();
54 let mut entity_closure_dofs = HashMap::new();
55 let mut entities_by_element = HashMap::new();
56
57 for d in 0..=mesh.topology_dim() {
58 for e in mesh.entity_types(d) {
59 entity_dofs.insert(*e, vec![vec![]; mesh.entity_count(*e)]);
60 entity_closure_dofs.insert(*e, vec![vec![]; mesh.entity_count(*e)]);
61 }
62 }
63
64 let mut ndofs = 0;
65
66 for (cell_type, element) in izip!(mesh.entity_types(mesh.topology_dim()), &elements) {
67 let sub_entity_types = reference_cell::entity_types(*cell_type);
68 let sub_entity_indices = reference_cell::entity_indices(*cell_type);
69 let mut cell_indices = vec![];
70 for cell in mesh.entity_iter(*cell_type) {
71 cell_indices.push(cell.local_index());
72 let mut cell_dofs = vec![]; for (d, (types, indices)) in
74 izip!(&sub_entity_types, &sub_entity_indices).enumerate()
75 {
76 for (t, i) in izip!(types, indices) {
77 let reference_entity_dofs = element.entity_dofs(d, *i).unwrap();
78 if !reference_entity_dofs.is_empty() && !reference_entity_dofs.is_empty() {
79 let ed = &mut entity_dofs.get_mut(t).unwrap()
80 [cell.topology().sub_entity(*t, *i)];
81 while ed.len() < reference_entity_dofs.len() {
82 ed.push(ndofs);
83 ndofs += 1;
84 }
85 for dof in ed {
86 cell_dofs.push(*dof);
87 }
88 }
89 entity_closure_dofs.get_mut(t).unwrap()
90 [cell.topology().sub_entity(*t, *i)] = element
91 .entity_closure_dofs(d, *i)
92 .unwrap()
93 .iter()
94 .map(|e| cell_dofs[*e])
95 .collect::<Vec<_>>();
96 }
97 }
98 }
99 entities_by_element.insert(*cell_type, cell_indices);
100 }
101
102 Self {
103 mesh,
104 elements,
105 entity_dofs,
106 entity_closure_dofs,
107 ndofs,
108 entities_by_element,
109 }
110 }
111}
112
113impl<
114 'a,
115 T: RlstScalar,
116 TMesh: RlstScalar,
117 E: Debug + PartialEq + Eq + Clone + Copy + Hash,
118 G: Mesh<EntityDescriptor = E, T = TMesh>,
119 F: MappedFiniteElement<CellType = E, T = T>,
120> FunctionSpace for FunctionSpaceImpl<'a, T, TMesh, E, G, F>
121{
122 type T = T;
123 type TMesh = TMesh;
124 type EntityDescriptor = E;
125 type Mesh = G;
126 type FiniteElement = F;
127
128 fn mesh(&self) -> &G {
129 self.mesh
130 }
131
132 fn elements(&self) -> &[F] {
133 &self.elements
134 }
135
136 fn entities_by_element(&self, element_index: usize) -> Option<&[usize]> {
137 self.entities_by_element
138 .get(&self.elements[element_index].cell_type())
139 .map(|v| &**v)
140 }
141
142 fn entity_dofs(&self, entity_type: E, entity_number: usize) -> Option<&[usize]> {
143 if let Some(i) = self.entity_dofs.get(&entity_type) {
144 if let Some(j) = i.get(entity_number) {
145 Some(j)
146 } else {
147 None
148 }
149 } else {
150 None
151 }
152 }
153
154 fn entity_closure_dofs(&self, entity_type: E, entity_number: usize) -> Option<&[usize]> {
155 if let Some(i) = self.entity_closure_dofs.get(&entity_type) {
156 if let Some(j) = i.get(entity_number) {
157 Some(j)
158 } else {
159 None
160 }
161 } else {
162 None
163 }
164 }
165
166 fn process_size(&self) -> usize {
167 self.ndofs
168 }
169
170 fn process_owned_size(&self) -> usize {
171 self.ndofs
172 }
173
174 fn global_size(&self) -> usize {
175 self.ndofs
176 }
177
178 fn global_dof_index(&self, process_dof_index: usize) -> usize {
179 process_dof_index
180 }
181
182 fn ownership(&self, _process_dof_index: usize) -> Ownership {
183 Ownership::Owned
184 }
185}
186
187#[cfg(test)]
188mod test {
189 use super::*;
190 use ndelement::{
191 ciarlet::{LagrangeElementFamily, LagrangeVariant},
192 types::{Continuity, ReferenceCellType},
193 };
194 use ndmesh::shapes::unit_cube_boundary;
195
196 #[test]
197 fn test_dp0() {
198 let mesh = unit_cube_boundary::<f64>(2, 2, 2, ReferenceCellType::Triangle);
199 let family = LagrangeElementFamily::<f64>::new(
200 0,
201 Continuity::Discontinuous,
202 LagrangeVariant::Equispaced,
203 );
204
205 let space = FunctionSpaceImpl::new(&mesh, &family);
206
207 assert_eq!(
208 space.process_size(),
209 mesh.entity_count(ReferenceCellType::Triangle)
210 );
211
212 for cell in mesh.entity_iter(ReferenceCellType::Point) {
213 assert_eq!(
214 space
215 .entity_dofs(ReferenceCellType::Point, cell.local_index())
216 .unwrap()
217 .len(),
218 0
219 );
220 assert_eq!(
221 space
222 .entity_closure_dofs(ReferenceCellType::Point, cell.local_index())
223 .unwrap()
224 .len(),
225 0
226 );
227 }
228 for cell in mesh.entity_iter(ReferenceCellType::Interval) {
229 assert_eq!(
230 space
231 .entity_dofs(ReferenceCellType::Interval, cell.local_index())
232 .unwrap()
233 .len(),
234 0
235 );
236 assert_eq!(
237 space
238 .entity_closure_dofs(ReferenceCellType::Interval, cell.local_index())
239 .unwrap()
240 .len(),
241 0
242 );
243 }
244 for cell in mesh.entity_iter(ReferenceCellType::Triangle) {
245 assert_eq!(
246 space
247 .entity_dofs(ReferenceCellType::Triangle, cell.local_index())
248 .unwrap()
249 .len(),
250 1
251 );
252 assert_eq!(
253 space
254 .entity_closure_dofs(ReferenceCellType::Triangle, cell.local_index())
255 .unwrap()
256 .len(),
257 1
258 );
259 }
260 }
261
262 #[test]
263 fn test_p1() {
264 let mesh = unit_cube_boundary::<f64>(2, 2, 2, ReferenceCellType::Triangle);
265 let family =
266 LagrangeElementFamily::<f64>::new(1, Continuity::Standard, LagrangeVariant::Equispaced);
267
268 let space = FunctionSpaceImpl::new(&mesh, &family);
269
270 assert_eq!(
271 space.process_size(),
272 mesh.entity_count(ReferenceCellType::Point)
273 );
274
275 for cell in mesh.entity_iter(ReferenceCellType::Point) {
276 assert_eq!(
277 space
278 .entity_dofs(ReferenceCellType::Point, cell.local_index())
279 .unwrap()
280 .len(),
281 1
282 );
283 assert_eq!(
284 space
285 .entity_closure_dofs(ReferenceCellType::Point, cell.local_index())
286 .unwrap()
287 .len(),
288 1
289 );
290 }
291 for cell in mesh.entity_iter(ReferenceCellType::Interval) {
292 assert_eq!(
293 space
294 .entity_dofs(ReferenceCellType::Interval, cell.local_index())
295 .unwrap()
296 .len(),
297 0
298 );
299 assert_eq!(
300 space
301 .entity_closure_dofs(ReferenceCellType::Interval, cell.local_index())
302 .unwrap()
303 .len(),
304 2
305 );
306 }
307 for cell in mesh.entity_iter(ReferenceCellType::Triangle) {
308 assert_eq!(
309 space
310 .entity_dofs(ReferenceCellType::Triangle, cell.local_index())
311 .unwrap()
312 .len(),
313 0
314 );
315 assert_eq!(
316 space
317 .entity_closure_dofs(ReferenceCellType::Triangle, cell.local_index())
318 .unwrap()
319 .len(),
320 3
321 );
322 }
323 }
324
325 #[test]
326 fn test_dp1() {
327 let mesh = unit_cube_boundary::<f64>(2, 2, 2, ReferenceCellType::Triangle);
328 let family = LagrangeElementFamily::<f64>::new(
329 1,
330 Continuity::Discontinuous,
331 LagrangeVariant::Equispaced,
332 );
333
334 let space = FunctionSpaceImpl::new(&mesh, &family);
335
336 assert_eq!(
337 space.process_size(),
338 3 * mesh.entity_count(ReferenceCellType::Triangle)
339 );
340
341 for cell in mesh.entity_iter(ReferenceCellType::Point) {
342 assert_eq!(
343 space
344 .entity_dofs(ReferenceCellType::Point, cell.local_index())
345 .unwrap()
346 .len(),
347 0
348 );
349 assert_eq!(
350 space
351 .entity_closure_dofs(ReferenceCellType::Point, cell.local_index())
352 .unwrap()
353 .len(),
354 0
355 );
356 }
357 for cell in mesh.entity_iter(ReferenceCellType::Interval) {
358 assert_eq!(
359 space
360 .entity_dofs(ReferenceCellType::Interval, cell.local_index())
361 .unwrap()
362 .len(),
363 0
364 );
365 assert_eq!(
366 space
367 .entity_closure_dofs(ReferenceCellType::Interval, cell.local_index())
368 .unwrap()
369 .len(),
370 0
371 );
372 }
373 for cell in mesh.entity_iter(ReferenceCellType::Triangle) {
374 assert_eq!(
375 space
376 .entity_dofs(ReferenceCellType::Triangle, cell.local_index())
377 .unwrap()
378 .len(),
379 3
380 );
381 assert_eq!(
382 space
383 .entity_closure_dofs(ReferenceCellType::Triangle, cell.local_index())
384 .unwrap()
385 .len(),
386 3
387 );
388 }
389 }
390
391 #[test]
392 fn test_p2() {
393 let mesh = unit_cube_boundary::<f64>(2, 2, 2, ReferenceCellType::Triangle);
394 let family =
395 LagrangeElementFamily::<f64>::new(2, Continuity::Standard, LagrangeVariant::Equispaced);
396
397 let space = FunctionSpaceImpl::new(&mesh, &family);
398
399 assert_eq!(
400 space.process_size(),
401 mesh.entity_count(ReferenceCellType::Point)
402 + mesh.entity_count(ReferenceCellType::Interval)
403 );
404
405 for cell in mesh.entity_iter(ReferenceCellType::Point) {
406 assert_eq!(
407 space
408 .entity_dofs(ReferenceCellType::Point, cell.local_index())
409 .unwrap()
410 .len(),
411 1
412 );
413 assert_eq!(
414 space
415 .entity_closure_dofs(ReferenceCellType::Point, cell.local_index())
416 .unwrap()
417 .len(),
418 1
419 );
420 }
421 for cell in mesh.entity_iter(ReferenceCellType::Interval) {
422 assert_eq!(
423 space
424 .entity_dofs(ReferenceCellType::Interval, cell.local_index())
425 .unwrap()
426 .len(),
427 1
428 );
429 assert_eq!(
430 space
431 .entity_closure_dofs(ReferenceCellType::Interval, cell.local_index())
432 .unwrap()
433 .len(),
434 3
435 );
436 }
437 for cell in mesh.entity_iter(ReferenceCellType::Triangle) {
438 assert_eq!(
439 space
440 .entity_dofs(ReferenceCellType::Triangle, cell.local_index())
441 .unwrap()
442 .len(),
443 0
444 );
445 assert_eq!(
446 space
447 .entity_closure_dofs(ReferenceCellType::Triangle, cell.local_index())
448 .unwrap()
449 .len(),
450 6
451 );
452 }
453 }
454
455 #[test]
456 fn test_p3() {
457 let mesh = unit_cube_boundary::<f64>(2, 2, 2, ReferenceCellType::Triangle);
458 let family =
459 LagrangeElementFamily::<f64>::new(3, Continuity::Standard, LagrangeVariant::Equispaced);
460
461 let space = FunctionSpaceImpl::new(&mesh, &family);
462
463 assert_eq!(
464 space.process_size(),
465 mesh.entity_count(ReferenceCellType::Point)
466 + 2 * mesh.entity_count(ReferenceCellType::Interval)
467 + mesh.entity_count(ReferenceCellType::Triangle)
468 );
469
470 for cell in mesh.entity_iter(ReferenceCellType::Point) {
471 assert_eq!(
472 space
473 .entity_dofs(ReferenceCellType::Point, cell.local_index())
474 .unwrap()
475 .len(),
476 1
477 );
478 assert_eq!(
479 space
480 .entity_closure_dofs(ReferenceCellType::Point, cell.local_index())
481 .unwrap()
482 .len(),
483 1
484 );
485 }
486 for cell in mesh.entity_iter(ReferenceCellType::Interval) {
487 assert_eq!(
488 space
489 .entity_dofs(ReferenceCellType::Interval, cell.local_index())
490 .unwrap()
491 .len(),
492 2
493 );
494 assert_eq!(
495 space
496 .entity_closure_dofs(ReferenceCellType::Interval, cell.local_index())
497 .unwrap()
498 .len(),
499 4
500 );
501 }
502 for cell in mesh.entity_iter(ReferenceCellType::Triangle) {
503 assert_eq!(
504 space
505 .entity_dofs(ReferenceCellType::Triangle, cell.local_index())
506 .unwrap()
507 .len(),
508 1
509 );
510 assert_eq!(
511 space
512 .entity_closure_dofs(ReferenceCellType::Triangle, cell.local_index())
513 .unwrap()
514 .len(),
515 10
516 );
517 }
518 }
519
520 #[test]
521 fn test_p1_quad() {
522 let mesh = unit_cube_boundary::<f64>(2, 2, 2, ReferenceCellType::Quadrilateral);
523 let family =
524 LagrangeElementFamily::<f64>::new(1, Continuity::Standard, LagrangeVariant::Equispaced);
525
526 let space = FunctionSpaceImpl::new(&mesh, &family);
527
528 assert_eq!(
529 space.process_size(),
530 mesh.entity_count(ReferenceCellType::Point)
531 );
532
533 for cell in mesh.entity_iter(ReferenceCellType::Point) {
534 assert_eq!(
535 space
536 .entity_dofs(ReferenceCellType::Point, cell.local_index())
537 .unwrap()
538 .len(),
539 1
540 );
541 assert_eq!(
542 space
543 .entity_closure_dofs(ReferenceCellType::Point, cell.local_index())
544 .unwrap()
545 .len(),
546 1
547 );
548 }
549 for cell in mesh.entity_iter(ReferenceCellType::Interval) {
550 assert_eq!(
551 space
552 .entity_dofs(ReferenceCellType::Interval, cell.local_index())
553 .unwrap()
554 .len(),
555 0
556 );
557 assert_eq!(
558 space
559 .entity_closure_dofs(ReferenceCellType::Interval, cell.local_index())
560 .unwrap()
561 .len(),
562 2
563 );
564 }
565 for cell in mesh.entity_iter(ReferenceCellType::Quadrilateral) {
566 assert_eq!(
567 space
568 .entity_dofs(ReferenceCellType::Quadrilateral, cell.local_index())
569 .unwrap()
570 .len(),
571 0
572 );
573 assert_eq!(
574 space
575 .entity_closure_dofs(ReferenceCellType::Quadrilateral, cell.local_index())
576 .unwrap()
577 .len(),
578 4
579 );
580 }
581 }
582
583 #[test]
584 fn test_p2_quad() {
585 let mesh = unit_cube_boundary::<f64>(2, 2, 2, ReferenceCellType::Quadrilateral);
586 let family =
587 LagrangeElementFamily::<f64>::new(2, Continuity::Standard, LagrangeVariant::Equispaced);
588
589 let space = FunctionSpaceImpl::new(&mesh, &family);
590
591 assert_eq!(
592 space.process_size(),
593 mesh.entity_count(ReferenceCellType::Point)
594 + mesh.entity_count(ReferenceCellType::Interval)
595 + mesh.entity_count(ReferenceCellType::Quadrilateral)
596 );
597
598 for cell in mesh.entity_iter(ReferenceCellType::Point) {
599 assert_eq!(
600 space
601 .entity_dofs(ReferenceCellType::Point, cell.local_index())
602 .unwrap()
603 .len(),
604 1
605 );
606 assert_eq!(
607 space
608 .entity_closure_dofs(ReferenceCellType::Point, cell.local_index())
609 .unwrap()
610 .len(),
611 1
612 );
613 }
614 for cell in mesh.entity_iter(ReferenceCellType::Interval) {
615 assert_eq!(
616 space
617 .entity_dofs(ReferenceCellType::Interval, cell.local_index())
618 .unwrap()
619 .len(),
620 1
621 );
622 assert_eq!(
623 space
624 .entity_closure_dofs(ReferenceCellType::Interval, cell.local_index())
625 .unwrap()
626 .len(),
627 3
628 );
629 }
630 for cell in mesh.entity_iter(ReferenceCellType::Quadrilateral) {
631 assert_eq!(
632 space
633 .entity_dofs(ReferenceCellType::Quadrilateral, cell.local_index())
634 .unwrap()
635 .len(),
636 1
637 );
638 assert_eq!(
639 space
640 .entity_closure_dofs(ReferenceCellType::Quadrilateral, cell.local_index())
641 .unwrap()
642 .len(),
643 9
644 );
645 }
646 }
647}