mesh_sieve/data/
multi_section.rs

1//! MultiSection: group multiple field sections with shared point offsets.
2
3use crate::data::atlas::Atlas;
4use crate::data::constrained_section::{
5    ConstrainedSection, DofConstraint, LabelConstraintSpec, apply_constraints_to_section,
6};
7use crate::data::section::Section;
8use crate::data::storage::Storage;
9use crate::mesh_error::MeshSieveError;
10use crate::topology::labels::LabelSet;
11use crate::topology::point::PointId;
12use std::collections::{BTreeMap, HashSet};
13
14/// A named field section with field-specific constraints.
15#[derive(Clone, Debug)]
16pub struct FieldSection<V, S: Storage<V>> {
17    name: String,
18    section: Section<V, S>,
19    constraints: BTreeMap<PointId, Vec<DofConstraint<V>>>,
20}
21
22/// Build a constrained section from explicit per-point dofs and label constraints.
23pub fn constrained_section_from_label_specs<V, S>(
24    point_dofs: &[(PointId, usize)],
25    labels: &LabelSet,
26    specs: &[LabelConstraintSpec],
27) -> Result<ConstrainedSection<V, S>, MeshSieveError>
28where
29    V: Clone + Default,
30    S: Storage<V> + Clone,
31{
32    let mut atlas = Atlas::default();
33    for (p, dof) in point_dofs {
34        atlas.try_insert(*p, *dof)?;
35    }
36    let section = Section::<V, S>::new(atlas);
37    let mut constrained = ConstrainedSection::new(section);
38    for spec in specs {
39        for p in labels.stratum_points(&spec.label, spec.value) {
40            let len = constrained
41                .section()
42                .atlas()
43                .get(p)
44                .map(|(_, l)| l)
45                .unwrap_or(0);
46            if len == 0 {
47                continue;
48            }
49            for &c in &spec.components {
50                if c < len {
51                    constrained.insert_constraint(p, c, V::default())?;
52                }
53            }
54        }
55    }
56    Ok(constrained)
57}
58
59impl<V, S> FieldSection<V, S>
60where
61    S: Storage<V>,
62{
63    /// Create a new field section with a name and data section.
64    pub fn new(name: impl Into<String>, section: Section<V, S>) -> Self {
65        Self {
66            name: name.into(),
67            section,
68            constraints: BTreeMap::new(),
69        }
70    }
71
72    /// Field name.
73    pub fn name(&self) -> &str {
74        &self.name
75    }
76
77    /// Access the underlying section.
78    pub fn section(&self) -> &Section<V, S> {
79        &self.section
80    }
81
82    /// Mutable access to the underlying section.
83    pub fn section_mut(&mut self) -> &mut Section<V, S> {
84        &mut self.section
85    }
86
87    /// Access the constraint map for this field.
88    pub fn constraints(&self) -> &BTreeMap<PointId, Vec<DofConstraint<V>>> {
89        &self.constraints
90    }
91
92    /// Mutable access to the constraint map for this field.
93    pub fn constraints_mut(&mut self) -> &mut BTreeMap<PointId, Vec<DofConstraint<V>>> {
94        &mut self.constraints
95    }
96
97    /// Insert or update a single constraint for a point.
98    pub fn insert_constraint(
99        &mut self,
100        point: PointId,
101        index: usize,
102        value: V,
103    ) -> Result<(), MeshSieveError> {
104        let len = self.section.try_restrict(point)?.len();
105        if index >= len {
106            return Err(MeshSieveError::ConstraintIndexOutOfBounds { point, index, len });
107        }
108        let entry = self.constraints.entry(point).or_default();
109        if let Some(existing) = entry.iter_mut().find(|c| c.index == index) {
110            existing.value = value;
111        } else {
112            entry.push(DofConstraint { index, value });
113        }
114        Ok(())
115    }
116
117    /// Remove all constraints for a point.
118    pub fn clear_constraints_for_point(&mut self, point: PointId) {
119        self.constraints.remove(&point);
120    }
121
122    /// Remove all constraints.
123    pub fn clear_constraints(&mut self) {
124        self.constraints.clear();
125    }
126
127    /// Apply all stored constraints to the underlying section.
128    pub fn apply_constraints(&mut self) -> Result<(), MeshSieveError>
129    where
130        V: Clone,
131    {
132        apply_constraints_to_section(&mut self.section, &self.constraints)
133    }
134}
135
136/// Group multiple field sections with consistent per-point offsets.
137#[derive(Clone, Debug)]
138pub struct MultiSection<V, S: Storage<V>> {
139    atlas: Atlas,
140    fields: Vec<FieldSection<V, S>>,
141}
142
143impl<V, S> MultiSection<V, S>
144where
145    S: Storage<V>,
146{
147    /// Construct a multi-section from field sections.
148    ///
149    /// All fields must share the same point set. Per-point total DOFs are the
150    /// sum of field DOFs, in field order.
151    pub fn new(fields: Vec<FieldSection<V, S>>) -> Result<Self, MeshSieveError> {
152        let atlas = Self::build_atlas(&fields)?;
153        Ok(Self { atlas, fields })
154    }
155
156    /// Borrow the combined atlas.
157    pub fn atlas(&self) -> &Atlas {
158        &self.atlas
159    }
160
161    /// Number of fields.
162    pub fn field_count(&self) -> usize {
163        self.fields.len()
164    }
165
166    /// Borrow all fields.
167    pub fn fields(&self) -> &[FieldSection<V, S>] {
168        &self.fields
169    }
170
171    /// Mutable access to all fields.
172    pub fn fields_mut(&mut self) -> &mut [FieldSection<V, S>] {
173        &mut self.fields
174    }
175
176    /// Borrow a field by index.
177    pub fn field(&self, index: usize) -> Option<&FieldSection<V, S>> {
178        self.fields.get(index)
179    }
180
181    /// Mutable access to a field by index.
182    pub fn field_mut(&mut self, index: usize) -> Option<&mut FieldSection<V, S>> {
183        self.fields.get_mut(index)
184    }
185
186    /// Borrow a field by name.
187    pub fn field_by_name(&self, name: &str) -> Option<&FieldSection<V, S>> {
188        self.fields.iter().find(|field| field.name == name)
189    }
190
191    /// Mutable access to a field by name.
192    pub fn field_by_name_mut(&mut self, name: &str) -> Option<&mut FieldSection<V, S>> {
193        self.fields.iter_mut().find(|field| field.name == name)
194    }
195
196    /// Total DOF count at point `p` across all fields.
197    pub fn dof(&self, p: PointId) -> Result<usize, MeshSieveError> {
198        let (_, len) = self
199            .atlas
200            .get(p)
201            .ok_or(MeshSieveError::PointNotInAtlas(p))?;
202        Ok(len)
203    }
204
205    /// Offset for point `p` in the combined layout.
206    pub fn offset(&self, p: PointId) -> Result<usize, MeshSieveError> {
207        let (offset, _) = self
208            .atlas
209            .get(p)
210            .ok_or(MeshSieveError::PointNotInAtlas(p))?;
211        Ok(offset)
212    }
213
214    /// DOF count for a specific field at point `p`.
215    pub fn field_dof(&self, p: PointId, field: usize) -> Result<usize, MeshSieveError> {
216        let section = self
217            .fields
218            .get(field)
219            .ok_or(MeshSieveError::SectionAccess {
220                point: p,
221                source: Box::new(std::io::Error::new(
222                    std::io::ErrorKind::InvalidInput,
223                    "field index out of bounds",
224                )),
225            })?
226            .section();
227        if self.atlas.get(p).is_none() {
228            return Err(MeshSieveError::PointNotInAtlas(p));
229        }
230        Ok(section.atlas().get(p).map(|(_, len)| len).unwrap_or(0))
231    }
232
233    /// Field offset for point `p` in the combined layout.
234    ///
235    /// This mirrors PetscSection field offsets: `offset(p) + sum_{i < field} dof_i(p)`.
236    pub fn field_offset(&self, p: PointId, field: usize) -> Result<usize, MeshSieveError> {
237        let base = self.offset(p)?;
238        let mut offset = base;
239        for idx in 0..field {
240            offset += self.field_dof(p, idx)?;
241        }
242        Ok(offset)
243    }
244
245    /// (offset, dof) pair for a field at point `p`.
246    pub fn field_span(&self, p: PointId, field: usize) -> Result<(usize, usize), MeshSieveError> {
247        Ok((self.field_offset(p, field)?, self.field_dof(p, field)?))
248    }
249
250    /// DOF count for a field at point `p`, by field name.
251    pub fn field_dof_by_name(&self, p: PointId, name: &str) -> Result<usize, MeshSieveError> {
252        let field = self
253            .fields
254            .iter()
255            .position(|field| field.name == name)
256            .ok_or(MeshSieveError::SectionAccess {
257                point: p,
258                source: Box::new(std::io::Error::new(
259                    std::io::ErrorKind::InvalidInput,
260                    "field name not found",
261                )),
262            })?;
263        self.field_dof(p, field)
264    }
265
266    /// Field offset for point `p`, by field name.
267    pub fn field_offset_by_name(&self, p: PointId, name: &str) -> Result<usize, MeshSieveError> {
268        let field = self
269            .fields
270            .iter()
271            .position(|field| field.name == name)
272            .ok_or(MeshSieveError::SectionAccess {
273                point: p,
274                source: Box::new(std::io::Error::new(
275                    std::io::ErrorKind::InvalidInput,
276                    "field name not found",
277                )),
278            })?;
279        self.field_offset(p, field)
280    }
281
282    /// Apply constraints for all fields.
283    pub fn apply_constraints(&mut self) -> Result<(), MeshSieveError>
284    where
285        V: Clone,
286    {
287        for field in &mut self.fields {
288            field.apply_constraints()?;
289        }
290        Ok(())
291    }
292
293    fn build_atlas(fields: &[FieldSection<V, S>]) -> Result<Atlas, MeshSieveError> {
294        if fields.is_empty() {
295            return Ok(Atlas::default());
296        }
297        let mut atlas = Atlas::default();
298        let mut points = Vec::new();
299        let mut seen = HashSet::new();
300        for field in fields {
301            for point in field.section().atlas().points() {
302                if seen.insert(point) {
303                    points.push(point);
304                }
305            }
306        }
307
308        for point in points {
309            let mut total = 0usize;
310            for field in fields {
311                if let Some((_, len)) = field.section().atlas().get(point) {
312                    total = total.saturating_add(len);
313                }
314            }
315            atlas.try_insert(point, total)?;
316        }
317        Ok(atlas)
318    }
319}
320
321impl<V, S> MultiSection<V, S>
322where
323    V: Clone,
324    S: Storage<V>,
325{
326    /// Gather all fields into a flat buffer in combined atlas order.
327    pub fn gather_in_order(&self) -> Result<Vec<V>, MeshSieveError> {
328        let mut out = Vec::with_capacity(self.atlas.total_len());
329        for point in self.atlas.points() {
330            for field in &self.fields {
331                if field.section().atlas().get(point).is_some() {
332                    let slice = field.section().try_restrict(point)?;
333                    out.extend_from_slice(slice);
334                }
335            }
336        }
337        if out.len() != self.atlas.total_len() {
338            return Err(MeshSieveError::ScatterLengthMismatch {
339                expected: self.atlas.total_len(),
340                found: out.len(),
341            });
342        }
343        Ok(out)
344    }
345
346    /// Scatter values from a flat buffer in combined atlas order.
347    pub fn try_scatter_in_order(&mut self, buf: &[V]) -> Result<(), MeshSieveError>
348    where
349        V: Clone + Default,
350    {
351        let total = self.atlas.total_len();
352        if buf.len() != total {
353            return Err(MeshSieveError::ScatterLengthMismatch {
354                expected: total,
355                found: buf.len(),
356            });
357        }
358        let points: Vec<PointId> = self.atlas.points().collect();
359        let mut cursor = 0usize;
360        for point in points {
361            for field in &mut self.fields {
362                if field.section().atlas().get(point).is_some() {
363                    let len = {
364                        let slice = field.section().try_restrict(point)?;
365                        slice.len()
366                    };
367                    let end =
368                        cursor
369                            .checked_add(len)
370                            .ok_or(MeshSieveError::ScatterChunkMismatch {
371                                offset: cursor,
372                                len,
373                            })?;
374                    let chunk =
375                        buf.get(cursor..end)
376                            .ok_or(MeshSieveError::ScatterChunkMismatch {
377                                offset: cursor,
378                                len,
379                            })?;
380                    field.section_mut().try_set(point, chunk)?;
381                    cursor = end;
382                }
383            }
384        }
385        if cursor != total {
386            return Err(MeshSieveError::ScatterLengthMismatch {
387                expected: total,
388                found: cursor,
389            });
390        }
391        Ok(())
392    }
393}
394
395#[cfg(test)]
396mod tests {
397    use super::{FieldSection, MultiSection, constrained_section_from_label_specs};
398    use crate::data::atlas::Atlas;
399    use crate::data::section::Section;
400    use crate::data::storage::VecStorage;
401    use crate::topology::point::PointId;
402
403    #[test]
404    fn field_offsets_match_stacked_layout() {
405        let mut atlas_a = Atlas::default();
406        let p1 = PointId::new(1).unwrap();
407        let p2 = PointId::new(2).unwrap();
408        atlas_a.try_insert(p1, 2).unwrap();
409        atlas_a.try_insert(p2, 1).unwrap();
410        let section_a = Section::<f64, VecStorage<f64>>::new(atlas_a);
411
412        let mut atlas_b = Atlas::default();
413        atlas_b.try_insert(p1, 1).unwrap();
414        atlas_b.try_insert(p2, 3).unwrap();
415        let section_b = Section::<f64, VecStorage<f64>>::new(atlas_b);
416
417        let fields = vec![
418            FieldSection::new("a", section_a),
419            FieldSection::new("b", section_b),
420        ];
421        let multi = MultiSection::new(fields).unwrap();
422
423        assert_eq!(multi.offset(p1).unwrap(), 0);
424        assert_eq!(multi.offset(p2).unwrap(), 3);
425        assert_eq!(multi.field_offset(p1, 0).unwrap(), 0);
426        assert_eq!(multi.field_offset(p1, 1).unwrap(), 2);
427        assert_eq!(multi.field_offset(p2, 0).unwrap(), 3);
428        assert_eq!(multi.field_offset(p2, 1).unwrap(), 4);
429        assert_eq!(multi.field_dof(p2, 1).unwrap(), 3);
430    }
431
432    #[test]
433    fn mixed_fields_have_independent_dofs() {
434        let p1 = PointId::new(1).unwrap();
435        let p2 = PointId::new(2).unwrap();
436        let p3 = PointId::new(3).unwrap();
437
438        let mut velocity_atlas = Atlas::default();
439        velocity_atlas.try_insert(p1, 3).unwrap();
440        velocity_atlas.try_insert(p2, 2).unwrap();
441        let velocity = Section::<f64, VecStorage<f64>>::new(velocity_atlas);
442
443        let mut pressure_atlas = Atlas::default();
444        pressure_atlas.try_insert(p2, 1).unwrap();
445        pressure_atlas.try_insert(p3, 1).unwrap();
446        let pressure = Section::<f64, VecStorage<f64>>::new(pressure_atlas);
447
448        let fields = vec![
449            FieldSection::new("velocity", velocity),
450            FieldSection::new("pressure", pressure),
451        ];
452        let multi = MultiSection::new(fields).unwrap();
453
454        assert_eq!(multi.offset(p1).unwrap(), 0);
455        assert_eq!(multi.offset(p2).unwrap(), 3);
456        assert_eq!(multi.offset(p3).unwrap(), 6);
457
458        assert_eq!(multi.field_dof(p1, 0).unwrap(), 3);
459        assert_eq!(multi.field_dof(p1, 1).unwrap(), 0);
460        assert_eq!(multi.field_offset(p1, 1).unwrap(), 3);
461
462        assert_eq!(multi.field_dof(p2, 0).unwrap(), 2);
463        assert_eq!(multi.field_dof(p2, 1).unwrap(), 1);
464        assert_eq!(multi.field_offset(p2, 1).unwrap(), 5);
465
466        assert_eq!(multi.field_dof(p3, 0).unwrap(), 0);
467        assert_eq!(multi.field_dof_by_name(p3, "pressure").unwrap(), 1);
468        assert_eq!(multi.field_offset_by_name(p3, "pressure").unwrap(), 6);
469    }
470
471    #[test]
472    fn constrained_layout_from_labels() {
473        use crate::data::constrained_section::LabelConstraintSpec;
474        use crate::topology::labels::LabelSet;
475
476        let p1 = PointId::new(1).unwrap();
477        let p2 = PointId::new(2).unwrap();
478        let mut labels = LabelSet::new();
479        labels.set_label(p2, "boundary", 1);
480        let cs = constrained_section_from_label_specs::<f64, VecStorage<f64>>(
481            &[(p1, 3), (p2, 3)],
482            &labels,
483            &[LabelConstraintSpec::new("boundary", 1, vec![0, 2])],
484        )
485        .unwrap();
486        assert_eq!(cs.section().atlas().total_len(), 6);
487        let c = cs.constraints().get(&p2).unwrap();
488        assert_eq!(c.len(), 2);
489        assert_eq!(c[0].index, 0);
490        assert_eq!(c[1].index, 2);
491    }
492}
mesh_sieve/data/multi_section.rs

mesh_sieve/data/
multi_section.rs
