mesh_sieve/data/
section_layout.rs1use 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#[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 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 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 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 pub fn total_dofs(&self) -> u64 {
54 self.total_dofs
55 }
56}
57
58pub 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
81pub 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
154pub 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
179pub 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
196pub 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
221pub 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
230pub 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}