Skip to main content

mesh_sieve/data/
bundle.rs

1//! Bundle: Combines mesh topology, DOF storage, and data transfer rules.
2//!
3//! A `Bundle` ties together:
4//! 1. A **vertical stack** of mesh points → DOF points (`stack`),  
5//! 2. A **field section** storing per-point data (`section`),  
6//! 3. A **delta** strategy (`delta`) for refining/assembling data.
7//!
8//! This abstraction supports push (refine) and pull (assemble) of data
9//! across mesh hierarchy levels, as described in Knepley & Karpeev (2009).
10
11use 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
22/// Reducer combining multiple cap slices into an accumulator slice.
23///
24/// Implementors supply zero-initialization, per-slice accumulation, and an
25/// optional finalize step (e.g. for averaging). Callers are expected to ensure
26/// all slices share the same length before invoking [`SliceReducer::accumulate`].
27pub trait SliceReducer<V>: Sync {
28    /// Create a zero-initialized accumulator of length `len`.
29    fn make_zero(&self, len: usize) -> Vec<V>;
30
31    /// Accumulate `src` into `acc` element-wise.
32    ///
33    /// # Panics
34    /// Implementations may assume `acc.len() == src.len()`. Callers are
35    /// responsible for validating slice lengths before invoking this method.
36    fn accumulate(&self, acc: &mut [V], src: &[V])
37    -> Result<(), crate::mesh_error::MeshSieveError>;
38
39    /// Optional finalize step once all sources have been accumulated.
40    fn finalize(
41        &self,
42        _acc: &mut [V],
43        _count: usize,
44    ) -> Result<(), crate::mesh_error::MeshSieveError> {
45        Ok(())
46    }
47}
48
49/// Element-wise averaging reducer used by [`Bundle::assemble`].
50///
51/// # Preconditions
52/// Callers must ensure input slices have identical lengths before invoking
53/// [`SliceReducer::accumulate`]; see [`Bundle::assemble_with`].
54#[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
105/// `Bundle<V, S, D>` packages a mesh‐to‐DOF stack, a data section, and a `ValueDelta`-type.
106///
107/// - `V`: underlying data type stored at each DOF (e.g., `f64`, `i32`, …).
108/// - `S`: storage backend for the section (defaults to [`VecStorage`]).
109/// - `D`: overlap [`ValueDelta`](crate::overlap::delta::ValueDelta)<V> implementation guiding how
110///   values are reduced/merged across parts (defaults to [`CopyDelta`]).
111///   For per-slice permutation/orientation, see [`crate::data::refine::delta::SliceDelta`].
112///
113/// # Fields
114/// - `stack`: vertical arrows from base mesh points → cap (DOF) points,
115///      carrying a `Polarity` payload if needed.
116/// - `section`: contiguous storage of data `V` for each point in the atlas.
117/// - `delta`: rules for extracting (`restrict`) and merging (`fuse`) values.
118pub struct Bundle<V, S: Storage<V> = VecStorage<V>, D = CopyDelta> {
119    /// Vertical connectivity: base points → cap (DOF) points.
120    pub stack: InMemoryStack<PointId, PointId, crate::topology::arrow::Polarity>,
121    /// Field data storage, indexed by `PointId`.
122    pub section: Section<V, S>,
123    /// Delta strategy for refine/assemble operations.
124    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    /// **Refine**: push data *down* the stack (base → cap) using per-arrow orientation.
135    ///
136    /// For each base point in the sieve closure of `bases`, applies the
137    /// orientation delta from the base slice to each cap slice. Disjoint source
138    /// and destination slices are copied without allocation; overlapping slices
139    /// are temporarily buffered for safety.
140    ///
141    /// # Complexity
142    /// **O(Σ deg(base) · k)**, where `deg(base)` is the number of cap points per base
143    /// and `k` is the per-point slice length. One pass; no intermediate allocations.
144    ///
145    /// # Determinism
146    /// Deterministic **per cap point slice**, independent of traversal order, provided
147    /// the vertical mapping `base -> {caps}` has no duplicates. Polarity handling
148    /// is local to each write.
149    ///
150    /// # Errors
151    /// Propagates errors from [`Section::try_apply_delta_between_points`].
152    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            // Validate base slice exists even if no caps are present
158            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    /// Apply constraints to the underlying section.
168    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    /// **Refine** with constraints applied after the refinement step.
180    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    /// **Assemble**: pull data *up* the stack (cap → base) using element-wise averaging.
194    ///
195    /// For each base point, gathers slices from all cap points and replaces the base
196    /// slice with the element-wise average. Cap slices must match the base slice
197    /// length.
198    ///
199    /// # Complexity
200    /// **O(Σ deg(base) · k)**; one pass.
201    ///
202    /// # Determinism
203    /// Deterministic; each cap contributes at most once.
204    ///
205    /// # Errors
206    /// Returns an error if any cap slice length differs from the base slice, in
207    /// which case the [`MeshSieveError::SliceLengthMismatch`] reports the
208    /// offending cap `PointId`. Also propagates reducer-specific errors such as
209    /// primitive conversion failures.
210    ///
211    /// # Behavior
212    /// Validates slice lengths using the base slice as ground truth, collects all
213    /// cap slices once, and reduces them element-wise into a fresh accumulator
214    /// before writing the result back to the base slice.
215    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            // Stream the cap ids; avoid per-base allocation.
224            let mut caps_iter = self.stack.lift(b).map(|(cap, _)| cap);
225
226            // Empty? nothing to assemble for this base.
227            let first_cap = match caps_iter.next() {
228                Some(c) => c,
229                None => continue,
230            };
231
232            // Base slice defines the expected length
233            let base_len = self.section.try_restrict(b)?.len();
234
235            // Initialize accumulator using the first cap slice after length validation.
236            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    /// Assemble using the provided reducer, then apply constraints.
269    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    /// Backward-compatible assemble: element-wise average of cap slices.
285    ///
286    /// # Migration
287    /// Prefer [`Bundle::assemble_with`] for explicit reduction control.
288    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    /// Backward-compatible assemble with constraints using element-wise averaging.
303    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    /// Iterate over `(cap_point, &[V])` pairs for all DOFs attached to base `p`.
320    ///
321    /// # Errors
322    /// Returns an error for any cap point missing in the underlying Section.
323    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(); // cap DOF for 1
348        atlas.try_insert(PointId::new(102).unwrap(), 1).unwrap(); // cap DOF for 2
349        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        // Refine: push base values to cap
388        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        // Assemble: pull cap values back to base
406        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        // Should not panic, nothing to do
444        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        // Both slots should be copied
484        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        // Should get reversed [2,1]
517        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        // base receives average (5+7)/2
569        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}