1use crate::data::constrained_section::{ConstraintSet, apply_constraints_to_section};
12#[allow(unused_imports)]
13use crate::data::refine::delta::SliceDelta;
14use crate::data::section::Section;
15use crate::data::storage::{Storage, VecStorage};
16use crate::overlap::delta::CopyDelta;
17use crate::topology::point::PointId;
18use crate::topology::sieve::Sieve;
19use crate::topology::stack::{InMemoryStack, Stack};
20use core::marker::PhantomData;
21
22pub trait SliceReducer<V>: Sync {
28 fn make_zero(&self, len: usize) -> Vec<V>;
30
31 fn accumulate(&self, acc: &mut [V], src: &[V])
37 -> Result<(), crate::mesh_error::MeshSieveError>;
38
39 fn finalize(
41 &self,
42 _acc: &mut [V],
43 _count: usize,
44 ) -> Result<(), crate::mesh_error::MeshSieveError> {
45 Ok(())
46 }
47}
48
49#[derive(Copy, Clone, Debug, Default)]
55pub struct AverageReducer;
56
57impl<V> SliceReducer<V> for AverageReducer
58where
59 V: Clone
60 + Default
61 + num_traits::FromPrimitive
62 + core::ops::AddAssign
63 + core::ops::Div<Output = V>,
64{
65 fn make_zero(&self, len: usize) -> Vec<V> {
66 vec![V::default(); len]
67 }
68
69 fn accumulate(
70 &self,
71 acc: &mut [V],
72 src: &[V],
73 ) -> Result<(), crate::mesh_error::MeshSieveError> {
74 use crate::mesh_error::MeshSieveError;
75 if acc.len() != src.len() {
76 return Err(MeshSieveError::ReducerLengthMismatch {
77 expected: acc.len(),
78 found: src.len(),
79 });
80 }
81 for (dst, s) in acc.iter_mut().zip(src.iter()) {
82 *dst += s.clone();
83 }
84 Ok(())
85 }
86
87 fn finalize(
88 &self,
89 acc: &mut [V],
90 count: usize,
91 ) -> Result<(), crate::mesh_error::MeshSieveError> {
92 use crate::mesh_error::MeshSieveError;
93 if count == 0 {
94 return Ok(());
95 }
96 let denom: V = num_traits::FromPrimitive::from_usize(count)
97 .ok_or(MeshSieveError::SievedArrayPrimitiveConversionFailure(count))?;
98 for v in acc.iter_mut() {
99 *v = v.clone() / denom.clone();
100 }
101 Ok(())
102 }
103}
104
105pub struct Bundle<V, S: Storage<V> = VecStorage<V>, D = CopyDelta> {
119 pub stack: InMemoryStack<PointId, PointId, crate::topology::arrow::Polarity>,
121 pub section: Section<V, S>,
123 pub delta: D,
125 #[doc(hidden)]
126 pub _marker: PhantomData<V>,
127}
128
129impl<V, S: Storage<V>, D> Bundle<V, S, D>
130where
131 V: Clone + Default,
132 D: crate::overlap::delta::ValueDelta<V, Part = V>,
133{
134 pub fn refine(
153 &mut self,
154 bases: impl IntoIterator<Item = PointId>,
155 ) -> Result<(), crate::mesh_error::MeshSieveError> {
156 for b in self.stack.base().closure(bases) {
157 self.section.try_restrict(b)?;
159 for (cap, orientation) in self.stack.lift(b) {
160 self.section
161 .try_apply_delta_between_points(b, cap, &orientation)?;
162 }
163 }
164 Ok(())
165 }
166
167 pub fn apply_constraints<C>(
169 &mut self,
170 constraints: &C,
171 ) -> Result<(), crate::mesh_error::MeshSieveError>
172 where
173 V: Clone,
174 C: ConstraintSet<V>,
175 {
176 apply_constraints_to_section(&mut self.section, constraints)
177 }
178
179 pub fn refine_with_constraints<C>(
181 &mut self,
182 bases: impl IntoIterator<Item = PointId>,
183 constraints: &C,
184 ) -> Result<(), crate::mesh_error::MeshSieveError>
185 where
186 V: Clone,
187 C: ConstraintSet<V>,
188 {
189 self.refine(bases)?;
190 self.apply_constraints(constraints)
191 }
192
193 pub fn assemble_with<R: SliceReducer<V>>(
216 &mut self,
217 bases: impl IntoIterator<Item = PointId>,
218 reducer: &R,
219 ) -> Result<(), crate::mesh_error::MeshSieveError> {
220 use crate::mesh_error::MeshSieveError;
221
222 for b in self.stack.base().closure(bases) {
223 let mut caps_iter = self.stack.lift(b).map(|(cap, _)| cap);
225
226 let first_cap = match caps_iter.next() {
228 Some(c) => c,
229 None => continue,
230 };
231
232 let base_len = self.section.try_restrict(b)?.len();
234
235 let first_slice = self.section.try_restrict(first_cap)?;
237 if first_slice.len() != base_len {
238 return Err(MeshSieveError::SliceLengthMismatch {
239 point: first_cap,
240 expected: base_len,
241 found: first_slice.len(),
242 });
243 }
244
245 let mut acc = reducer.make_zero(base_len);
246 reducer.accumulate(&mut acc, first_slice)?;
247 let mut count = 1usize;
248
249 for cap in caps_iter {
250 let sl = self.section.try_restrict(cap)?;
251 if sl.len() != base_len {
252 return Err(MeshSieveError::SliceLengthMismatch {
253 point: cap,
254 expected: base_len,
255 found: sl.len(),
256 });
257 }
258 reducer.accumulate(&mut acc, sl)?;
259 count += 1;
260 }
261
262 reducer.finalize(&mut acc, count)?;
263 self.section.try_set(b, &acc)?;
264 }
265 Ok(())
266 }
267
268 pub fn assemble_with_constraints<R, C>(
270 &mut self,
271 bases: impl IntoIterator<Item = PointId>,
272 reducer: &R,
273 constraints: &C,
274 ) -> Result<(), crate::mesh_error::MeshSieveError>
275 where
276 V: Clone,
277 R: SliceReducer<V>,
278 C: ConstraintSet<V>,
279 {
280 self.assemble_with(bases, reducer)?;
281 self.apply_constraints(constraints)
282 }
283
284 pub fn assemble(
289 &mut self,
290 bases: impl IntoIterator<Item = PointId>,
291 ) -> Result<(), crate::mesh_error::MeshSieveError>
292 where
293 V: Clone
294 + Default
295 + num_traits::FromPrimitive
296 + std::ops::AddAssign
297 + std::ops::Div<Output = V>,
298 {
299 self.assemble_with(bases, &AverageReducer)
300 }
301
302 pub fn assemble_with_constraints_default<C>(
304 &mut self,
305 bases: impl IntoIterator<Item = PointId>,
306 constraints: &C,
307 ) -> Result<(), crate::mesh_error::MeshSieveError>
308 where
309 V: Clone
310 + Default
311 + num_traits::FromPrimitive
312 + std::ops::AddAssign
313 + std::ops::Div<Output = V>,
314 C: ConstraintSet<V>,
315 {
316 self.assemble_with_constraints(bases, &AverageReducer, constraints)
317 }
318
319 pub fn dofs<'a>(
324 &'a self,
325 p: PointId,
326 ) -> impl Iterator<Item = Result<(PointId, &'a [V]), crate::mesh_error::MeshSieveError>> + 'a
327 {
328 self.stack
329 .lift(p)
330 .map(move |(cap, _)| self.section.try_restrict(cap).map(|sl| (cap, sl)))
331 }
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337 use crate::data::atlas::Atlas;
338 use crate::data::storage::VecStorage;
339 use crate::overlap::delta::CopyDelta;
340 use crate::topology::arrow::Polarity;
341 use core::marker::PhantomData;
342 #[test]
343 fn bundle_basic_refine_and_assemble() {
344 let mut atlas = Atlas::default();
345 atlas.try_insert(PointId::new(1).unwrap(), 1).unwrap();
346 atlas.try_insert(PointId::new(2).unwrap(), 1).unwrap();
347 atlas.try_insert(PointId::new(101).unwrap(), 1).unwrap(); atlas.try_insert(PointId::new(102).unwrap(), 1).unwrap(); let mut section = Section::<i32, VecStorage<i32>>::new(atlas.clone());
350 section.try_set(PointId::new(1).unwrap(), &[10]).unwrap();
351 section.try_set(PointId::new(2).unwrap(), &[20]).unwrap();
352 let mut stack = InMemoryStack::<PointId, PointId, Polarity>::new();
353 stack
354 .base_mut()
355 .unwrap()
356 .add_arrow(PointId::new(1).unwrap(), PointId::new(1).unwrap(), ());
357 stack
358 .base_mut()
359 .unwrap()
360 .add_arrow(PointId::new(2).unwrap(), PointId::new(2).unwrap(), ());
361 stack.cap_mut().unwrap().add_arrow(
362 PointId::new(101).unwrap(),
363 PointId::new(101).unwrap(),
364 (),
365 );
366 stack.cap_mut().unwrap().add_arrow(
367 PointId::new(102).unwrap(),
368 PointId::new(102).unwrap(),
369 (),
370 );
371 let _ = stack.add_arrow(
372 PointId::new(1).unwrap(),
373 PointId::new(101).unwrap(),
374 Polarity::Forward,
375 );
376 let _ = stack.add_arrow(
377 PointId::new(2).unwrap(),
378 PointId::new(102).unwrap(),
379 Polarity::Forward,
380 );
381 let mut bundle = Bundle {
382 stack,
383 section,
384 delta: CopyDelta,
385 _marker: PhantomData,
386 };
387 bundle
389 .refine([PointId::new(1).unwrap(), PointId::new(2).unwrap()])
390 .unwrap();
391 assert_eq!(
392 bundle
393 .section
394 .try_restrict(PointId::new(101).unwrap())
395 .unwrap(),
396 &[10]
397 );
398 assert_eq!(
399 bundle
400 .section
401 .try_restrict(PointId::new(102).unwrap())
402 .unwrap(),
403 &[20]
404 );
405 bundle
407 .section
408 .try_set(PointId::new(101).unwrap(), &[30])
409 .unwrap();
410 bundle
411 .section
412 .try_set(PointId::new(102).unwrap(), &[40])
413 .unwrap();
414 bundle
415 .assemble([PointId::new(1).unwrap(), PointId::new(2).unwrap()])
416 .unwrap();
417 assert_eq!(
418 bundle
419 .section
420 .try_restrict(PointId::new(1).unwrap())
421 .unwrap(),
422 &[30]
423 );
424 assert_eq!(
425 bundle
426 .section
427 .try_restrict(PointId::new(2).unwrap())
428 .unwrap(),
429 &[40]
430 );
431 }
432 #[test]
433 fn empty_bundle_noop() {
434 let atlas = Atlas::default();
435 let section = Section::<i32, VecStorage<i32>>::new(atlas.clone());
436 let stack = InMemoryStack::<PointId, PointId, Polarity>::new();
437 let mut bundle = Bundle {
438 stack,
439 section,
440 delta: CopyDelta,
441 _marker: PhantomData,
442 };
443 bundle.refine(std::iter::empty::<PointId>()).unwrap();
445 bundle.assemble(std::iter::empty::<PointId>()).unwrap();
446 }
447
448 #[test]
449 fn multiple_dofs_only_first_moved() {
450 let mut atlas = Atlas::default();
451 atlas.try_insert(PointId::new(1).unwrap(), 2).unwrap();
452 atlas.try_insert(PointId::new(101).unwrap(), 2).unwrap();
453 let mut section = Section::<i32, VecStorage<i32>>::new(atlas.clone());
454 section
455 .try_set(PointId::new(1).unwrap(), &[10, 20])
456 .unwrap();
457 let mut stack = InMemoryStack::<PointId, PointId, Polarity>::new();
458 stack
459 .base_mut()
460 .unwrap()
461 .add_arrow(PointId::new(1).unwrap(), PointId::new(1).unwrap(), ());
462 stack.cap_mut().unwrap().add_arrow(
463 PointId::new(101).unwrap(),
464 PointId::new(101).unwrap(),
465 (),
466 );
467 let _ = stack.add_arrow(
468 PointId::new(1).unwrap(),
469 PointId::new(101).unwrap(),
470 Polarity::Forward,
471 );
472 let mut bundle = Bundle {
473 stack,
474 section,
475 delta: CopyDelta,
476 _marker: PhantomData,
477 };
478 bundle.refine([PointId::new(1).unwrap()]).unwrap();
479 let vals = bundle
480 .section
481 .try_restrict(PointId::new(101).unwrap())
482 .unwrap();
483 assert_eq!(vals, &[10, 20]);
485 }
486
487 #[test]
488 fn reverse_orientation_refine() {
489 let mut atlas = Atlas::default();
490 atlas.try_insert(PointId::new(1).unwrap(), 2).unwrap();
491 atlas.try_insert(PointId::new(101).unwrap(), 2).unwrap();
492 let mut section = Section::<i32, VecStorage<i32>>::new(atlas.clone());
493 section.try_set(PointId::new(1).unwrap(), &[1, 2]).unwrap();
494 let mut stack = InMemoryStack::<PointId, PointId, Polarity>::new();
495 stack
496 .base_mut()
497 .unwrap()
498 .add_arrow(PointId::new(1).unwrap(), PointId::new(1).unwrap(), ());
499 stack.cap_mut().unwrap().add_arrow(
500 PointId::new(101).unwrap(),
501 PointId::new(101).unwrap(),
502 (),
503 );
504 let _ = stack.add_arrow(
505 PointId::new(1).unwrap(),
506 PointId::new(101).unwrap(),
507 Polarity::Reverse,
508 );
509 let mut bundle = Bundle {
510 stack,
511 section,
512 delta: CopyDelta,
513 _marker: PhantomData,
514 };
515 bundle.refine([PointId::new(1).unwrap()]).unwrap();
516 assert_eq!(
518 bundle
519 .section
520 .try_restrict(PointId::new(101).unwrap())
521 .unwrap(),
522 &[2, 1]
523 );
524 }
525
526 #[test]
527 fn assemble_with_add_delta() {
528 use crate::overlap::delta::AddDelta;
529 let mut atlas = Atlas::default();
530 atlas.try_insert(PointId::new(1).unwrap(), 1).unwrap();
531 atlas.try_insert(PointId::new(101).unwrap(), 1).unwrap();
532 atlas.try_insert(PointId::new(102).unwrap(), 1).unwrap();
533 let mut section = Section::<i32, VecStorage<i32>>::new(atlas.clone());
534 section.try_set(PointId::new(101).unwrap(), &[5]).unwrap();
535 section.try_set(PointId::new(102).unwrap(), &[7]).unwrap();
536 let mut stack = InMemoryStack::<PointId, PointId, Polarity>::new();
537 stack
538 .base_mut()
539 .unwrap()
540 .add_arrow(PointId::new(1).unwrap(), PointId::new(1).unwrap(), ());
541 stack.cap_mut().unwrap().add_arrow(
542 PointId::new(101).unwrap(),
543 PointId::new(101).unwrap(),
544 (),
545 );
546 stack.cap_mut().unwrap().add_arrow(
547 PointId::new(102).unwrap(),
548 PointId::new(102).unwrap(),
549 (),
550 );
551 let _ = stack.add_arrow(
552 PointId::new(1).unwrap(),
553 PointId::new(101).unwrap(),
554 Polarity::Forward,
555 );
556 let _ = stack.add_arrow(
557 PointId::new(1).unwrap(),
558 PointId::new(102).unwrap(),
559 Polarity::Forward,
560 );
561 let mut bundle = Bundle {
562 stack,
563 section,
564 delta: AddDelta,
565 _marker: PhantomData,
566 };
567 bundle.assemble([PointId::new(1).unwrap()]).unwrap();
568 assert_eq!(
570 bundle
571 .section
572 .try_restrict(PointId::new(1).unwrap())
573 .unwrap(),
574 &[6]
575 );
576 }
577
578 #[test]
579 fn dofs_iterator() {
580 let mut atlas = Atlas::default();
581 atlas.try_insert(PointId::new(1).unwrap(), 1).unwrap();
582 atlas.try_insert(PointId::new(101).unwrap(), 1).unwrap();
583 atlas.try_insert(PointId::new(102).unwrap(), 1).unwrap();
584 let mut section = Section::<i32, VecStorage<i32>>::new(atlas.clone());
585 section.try_set(PointId::new(101).unwrap(), &[8]).unwrap();
586 section.try_set(PointId::new(102).unwrap(), &[9]).unwrap();
587 let mut stack = InMemoryStack::<PointId, PointId, Polarity>::new();
588 stack
589 .base_mut()
590 .unwrap()
591 .add_arrow(PointId::new(1).unwrap(), PointId::new(1).unwrap(), ());
592 stack.cap_mut().unwrap().add_arrow(
593 PointId::new(101).unwrap(),
594 PointId::new(101).unwrap(),
595 (),
596 );
597 stack.cap_mut().unwrap().add_arrow(
598 PointId::new(102).unwrap(),
599 PointId::new(102).unwrap(),
600 (),
601 );
602 let _ = stack.add_arrow(
603 PointId::new(1).unwrap(),
604 PointId::new(101).unwrap(),
605 Polarity::Forward,
606 );
607 let _ = stack.add_arrow(
608 PointId::new(1).unwrap(),
609 PointId::new(102).unwrap(),
610 Polarity::Forward,
611 );
612 let bundle = Bundle {
613 stack,
614 section,
615 delta: CopyDelta,
616 _marker: PhantomData,
617 };
618 let vec: Vec<_> = bundle.dofs(PointId::new(1).unwrap()).collect();
619 let mut vec = vec.into_iter().collect::<Result<Vec<_>, _>>().unwrap();
620 vec.sort_by_key(|(cap, _)| cap.get());
621 assert_eq!(
622 vec,
623 vec![
624 (PointId::new(101).unwrap(), &[8][..]),
625 (PointId::new(102).unwrap(), &[9][..]),
626 ]
627 );
628 }
629
630 #[test]
631 fn refine_unknown_base_errors() {
632 let atlas = Atlas::default();
633 let section = Section::<i32, VecStorage<i32>>::new(atlas.clone());
634 let stack = InMemoryStack::<PointId, PointId, Polarity>::new();
635 let mut bundle = Bundle {
636 stack,
637 section,
638 delta: CopyDelta,
639 _marker: PhantomData,
640 };
641 let err = bundle.refine([PointId::new(999).unwrap()]).unwrap_err();
642 assert!(
643 matches!(err, crate::mesh_error::MeshSieveError::PointNotInAtlas(pid) if pid.get() == 999)
644 );
645 }
646}