Skip to main content

mesh_sieve/data/
section_layout.rs

1//! Closure-driven DOF layout helpers for sections and multi-sections.
2
3use crate::data::constrained_section::{ConstraintSet, DofConstraint};
4use crate::data::multi_section::MultiSection;
5use crate::data::section::Section;
6use crate::data::storage::Storage;
7use crate::mesh_error::MeshSieveError;
8use crate::topology::periodic::PeriodicMap;
9use crate::topology::point::PointId;
10use std::collections::{BTreeSet, HashMap};
11
12/// Contiguous DOF layout for a point set.
13#[derive(Clone, Debug, Default)]
14pub struct DofLayout {
15    offsets: Vec<Option<u64>>,
16    dof_lengths: Vec<Option<usize>>,
17    total_dofs: u64,
18}
19
20impl DofLayout {
21    /// Return the offset for a point in the layout.
22    pub fn offset(&self, point: PointId) -> Result<u64, MeshSieveError> {
23        let idx = point_index(point)?;
24        self.offsets
25            .get(idx)
26            .and_then(|val| *val)
27            .ok_or(MeshSieveError::PointNotInAtlas(point))
28    }
29
30    /// Return the DOF count for a point in the layout.
31    pub fn dof_len(&self, point: PointId) -> Result<usize, MeshSieveError> {
32        let idx = point_index(point)?;
33        self.dof_lengths
34            .get(idx)
35            .and_then(|val| *val)
36            .ok_or(MeshSieveError::PointNotInAtlas(point))
37    }
38
39    /// Return the global index for a local point/DOF pair.
40    pub fn global_index(&self, point: PointId, dof: usize) -> Result<u64, MeshSieveError> {
41        let len = self.dof_len(point)?;
42        if dof >= len {
43            return Err(MeshSieveError::ConstraintIndexOutOfBounds {
44                point,
45                index: dof,
46                len,
47            });
48        }
49        Ok(self.offset(point)? + dof as u64)
50    }
51
52    /// Total number of DOFs in the layout.
53    pub fn total_dofs(&self) -> u64 {
54        self.total_dofs
55    }
56}
57
58/// Compute a constraint-aware DOF count for a point slice.
59pub fn constrained_dof_len<V>(
60    point: PointId,
61    base_len: usize,
62    constraints: Option<&[DofConstraint<V>]>,
63) -> Result<usize, MeshSieveError> {
64    let Some(constraints) = constraints else {
65        return Ok(base_len);
66    };
67    let mut seen = BTreeSet::new();
68    for constraint in constraints {
69        if constraint.index >= base_len {
70            return Err(MeshSieveError::ConstraintIndexOutOfBounds {
71                point,
72                index: constraint.index,
73                len: base_len,
74            });
75        }
76        seen.insert(constraint.index);
77    }
78    Ok(base_len.saturating_sub(seen.len()))
79}
80
81/// Build a DOF layout using closures for DOF length and representative points.
82pub fn build_layout_with<I, F, R>(
83    points: I,
84    dof_len: F,
85    representative: R,
86) -> Result<DofLayout, MeshSieveError>
87where
88    I: IntoIterator<Item = PointId>,
89    F: Fn(PointId) -> Result<usize, MeshSieveError>,
90    R: Fn(PointId) -> PointId,
91{
92    let points_vec: Vec<PointId> = points.into_iter().collect();
93    let points_set: BTreeSet<PointId> = points_vec.iter().copied().collect();
94    let max_id = points_vec.iter().map(|p| p.get()).max().unwrap_or(0) as usize;
95    let mut layout = DofLayout {
96        offsets: vec![None; max_id],
97        dof_lengths: vec![None; max_id],
98        total_dofs: 0,
99    };
100
101    let mut rep_lengths: HashMap<PointId, usize> = HashMap::new();
102    for point in &points_vec {
103        let rep = representative(*point);
104        if !points_set.contains(&rep) {
105            return Err(MeshSieveError::InvalidGeometry(format!(
106                "representative {rep:?} missing from layout points",
107            )));
108        }
109        let len = dof_len(*point)?;
110        if let Some(existing) = rep_lengths.insert(rep, len)
111            && existing != len
112        {
113            return Err(MeshSieveError::InvalidGeometry(format!(
114                "periodic layout mismatch for {point:?}: expected {existing}, got {len}",
115            )));
116        }
117    }
118
119    let mut rep_offsets: HashMap<PointId, u64> = HashMap::new();
120    let mut total = 0u64;
121    for point in &points_vec {
122        let rep = representative(*point);
123        if rep_offsets.contains_key(&rep) {
124            continue;
125        }
126        let len = *rep_lengths
127            .get(&rep)
128            .ok_or_else(|| MeshSieveError::PointNotInAtlas(rep))? as u64;
129        rep_offsets.insert(rep, total);
130        total = total.saturating_add(len);
131    }
132
133    layout.total_dofs = total;
134    for point in &points_vec {
135        let idx = point_index(*point)?;
136        if idx >= layout.offsets.len() {
137            layout.offsets.resize(idx + 1, None);
138            layout.dof_lengths.resize(idx + 1, None);
139        }
140        let rep = representative(*point);
141        let offset = *rep_offsets
142            .get(&rep)
143            .ok_or_else(|| MeshSieveError::PointNotInAtlas(rep))?;
144        let len = *rep_lengths
145            .get(&rep)
146            .ok_or_else(|| MeshSieveError::PointNotInAtlas(rep))?;
147        layout.offsets[idx] = Some(offset);
148        layout.dof_lengths[idx] = Some(len);
149    }
150
151    Ok(layout)
152}
153
154/// Build a layout for a section using optional constraints and periodic identification.
155pub fn layout_for_section_with_constraints_and_periodic<V, S, C>(
156    section: &Section<V, S>,
157    constraints: &C,
158    periodic: Option<&PeriodicMap>,
159) -> Result<DofLayout, MeshSieveError>
160where
161    S: Storage<V>,
162    C: ConstraintSet<V>,
163{
164    let dof_len = |point: PointId| {
165        let (_, len) = section
166            .atlas()
167            .get(point)
168            .ok_or(MeshSieveError::PointNotInAtlas(point))?;
169        constrained_dof_len(point, len, constraints.constraints_for(point))
170    };
171    let rep = |point: PointId| {
172        periodic
173            .and_then(|map| map.master_of(point))
174            .unwrap_or(point)
175    };
176    build_layout_with(section.atlas().points(), dof_len, rep)
177}
178
179/// Build a layout for a multi-section using field constraints and optional periodic mapping.
180pub fn layout_for_multi_section_with_periodic<V, S>(
181    section: &MultiSection<V, S>,
182    periodic: Option<&PeriodicMap>,
183) -> Result<DofLayout, MeshSieveError>
184where
185    S: Storage<V>,
186{
187    let dof_len = |point: PointId| multi_section_dof_len_with_constraints(section, point);
188    let rep = |point: PointId| {
189        periodic
190            .and_then(|map| map.master_of(point))
191            .unwrap_or(point)
192    };
193    build_layout_with(section.atlas().points(), dof_len, rep)
194}
195
196/// Compute the constraint-aware DOF length for a multi-section point.
197pub fn multi_section_dof_len_with_constraints<V, S>(
198    section: &MultiSection<V, S>,
199    point: PointId,
200) -> Result<usize, MeshSieveError>
201where
202    S: Storage<V>,
203{
204    if section.atlas().get(point).is_none() {
205        return Err(MeshSieveError::PointNotInAtlas(point));
206    }
207    let mut total = 0usize;
208    for field in section.fields() {
209        let len = field
210            .section()
211            .atlas()
212            .get(point)
213            .map(|(_, len)| len)
214            .unwrap_or(0);
215        let constraints = field.constraints().get(&point).map(|c| c.as_slice());
216        total = total.saturating_add(constrained_dof_len(point, len, constraints)?);
217    }
218    Ok(total)
219}
220
221/// Allocate a zero-initialized local vector for a section.
222pub fn local_vector_for_section<V, S>(section: &Section<V, S>) -> Vec<V>
223where
224    V: Clone + Default,
225    S: Storage<V>,
226{
227    vec![V::default(); section.atlas().total_len()]
228}
229
230/// Allocate a zero-initialized local vector for a layout.
231pub fn local_vector_for_layout<V>(layout: &DofLayout) -> Vec<V>
232where
233    V: Clone + Default,
234{
235    vec![V::default(); layout.total_dofs as usize]
236}
237
238fn point_index(point: PointId) -> Result<usize, MeshSieveError> {
239    point
240        .get()
241        .checked_sub(1)
242        .ok_or(MeshSieveError::InvalidPointId)
243        .map(|idx| idx as usize)
244}