1use crate::data::hanging_node_constraints::{
4 HangingNodeConstraints, apply_hanging_constraints_to_section,
5};
6use crate::data::section::{FallibleMap, Section};
7use crate::data::storage::Storage;
8use crate::mesh_error::MeshSieveError;
9use crate::topology::cache::InvalidateCache;
10use crate::topology::point::PointId;
11use std::collections::BTreeMap;
12
13#[derive(Clone, Debug, PartialEq, Eq)]
15pub struct LabelConstraintSpec {
16 pub label: String,
18 pub value: i32,
20 pub components: Vec<usize>,
22}
23
24impl LabelConstraintSpec {
25 pub fn new(label: impl Into<String>, value: i32, components: Vec<usize>) -> Self {
27 Self {
28 label: label.into(),
29 value,
30 components,
31 }
32 }
33}
34
35#[derive(Clone, Debug, PartialEq)]
37pub struct DofConstraint<V> {
38 pub index: usize,
40 pub value: V,
42}
43
44impl<V> DofConstraint<V> {
45 pub fn new(index: usize, value: V) -> Self {
47 Self { index, value }
48 }
49}
50
51pub trait ConstraintSet<V> {
53 fn points(&self) -> Box<dyn Iterator<Item = PointId> + '_>;
55
56 fn constraints_for(&self, point: PointId) -> Option<&[DofConstraint<V>]>;
58}
59
60#[derive(Clone, Debug)]
62pub struct ConstrainedSection<V, S: Storage<V>> {
63 section: Section<V, S>,
64 constraints: BTreeMap<PointId, Vec<DofConstraint<V>>>,
65 hanging_constraints: HangingNodeConstraints<V>,
66}
67
68impl<V, S> ConstrainedSection<V, S>
69where
70 S: Storage<V>,
71{
72 pub fn new(section: Section<V, S>) -> Self {
74 Self {
75 section,
76 constraints: BTreeMap::new(),
77 hanging_constraints: HangingNodeConstraints::default(),
78 }
79 }
80
81 pub fn section(&self) -> &Section<V, S> {
83 &self.section
84 }
85
86 pub fn section_mut(&mut self) -> &mut Section<V, S> {
88 &mut self.section
89 }
90
91 pub fn into_section(self) -> Section<V, S> {
93 self.section
94 }
95
96 pub fn constraints(&self) -> &BTreeMap<PointId, Vec<DofConstraint<V>>> {
98 &self.constraints
99 }
100
101 pub fn constraints_mut(&mut self) -> &mut BTreeMap<PointId, Vec<DofConstraint<V>>> {
103 &mut self.constraints
104 }
105
106 pub fn hanging_constraints(&self) -> &HangingNodeConstraints<V> {
108 &self.hanging_constraints
109 }
110
111 pub fn hanging_constraints_mut(&mut self) -> &mut HangingNodeConstraints<V> {
113 &mut self.hanging_constraints
114 }
115
116 pub fn insert_constraint(
118 &mut self,
119 point: PointId,
120 index: usize,
121 value: V,
122 ) -> Result<(), MeshSieveError> {
123 let len = self.section.try_restrict(point)?.len();
124 if index >= len {
125 return Err(MeshSieveError::ConstraintIndexOutOfBounds { point, index, len });
126 }
127 let entry = self.constraints.entry(point).or_default();
128 if let Some(existing) = entry.iter_mut().find(|c| c.index == index) {
129 existing.value = value;
130 } else {
131 entry.push(DofConstraint { index, value });
132 }
133 Ok(())
134 }
135
136 pub fn clear_constraints_for_point(&mut self, point: PointId) {
138 self.constraints.remove(&point);
139 }
140
141 pub fn clear_constraints(&mut self) {
143 self.constraints.clear();
144 }
145
146 pub fn remove_constraint(&mut self, point: PointId, index: usize) -> Option<DofConstraint<V>> {
148 let entry = self.constraints.get_mut(&point)?;
149 let pos = entry.iter().position(|c| c.index == index)?;
150 Some(entry.remove(pos))
151 }
152
153 pub fn apply_constraints_for_point(&mut self, point: PointId) -> Result<(), MeshSieveError>
155 where
156 V: Clone,
157 {
158 if let Some(constraints) = self.constraints.get(&point) {
159 apply_constraints_to_slice(&mut self.section, point, constraints)?;
160 self.section.invalidate_cache();
161 }
162 Ok(())
163 }
164
165 pub fn apply_constraints(&mut self) -> Result<(), MeshSieveError>
167 where
168 V: Clone,
169 {
170 apply_constraints_to_section(&mut self.section, &self.constraints)
171 }
172
173 pub fn apply_hanging_constraints(&mut self) -> Result<(), MeshSieveError>
175 where
176 V: Clone + Default + core::ops::AddAssign + core::ops::Mul<Output = V>,
177 {
178 apply_hanging_constraints_to_section(&mut self.section, &self.hanging_constraints)
179 }
180
181 pub fn apply_all_constraints(&mut self) -> Result<(), MeshSieveError>
183 where
184 V: Clone + Default + core::ops::AddAssign + core::ops::Mul<Output = V>,
185 {
186 self.apply_hanging_constraints()?;
187 self.apply_constraints()
188 }
189}
190
191impl<V, S> ConstraintSet<V> for ConstrainedSection<V, S>
192where
193 S: Storage<V>,
194{
195 fn points(&self) -> Box<dyn Iterator<Item = PointId> + '_> {
196 Box::new(self.constraints.keys().copied())
197 }
198
199 fn constraints_for(&self, point: PointId) -> Option<&[DofConstraint<V>]> {
200 self.constraints
201 .get(&point)
202 .map(|constraints| constraints.as_slice())
203 }
204}
205
206impl<V> ConstraintSet<V> for BTreeMap<PointId, Vec<DofConstraint<V>>> {
207 fn points(&self) -> Box<dyn Iterator<Item = PointId> + '_> {
208 Box::new(self.keys().copied())
209 }
210
211 fn constraints_for(&self, point: PointId) -> Option<&[DofConstraint<V>]> {
212 self.get(&point).map(|constraints| constraints.as_slice())
213 }
214}
215
216impl<V, S> FallibleMap<V> for ConstrainedSection<V, S>
217where
218 S: Storage<V>,
219{
220 #[inline]
221 fn try_get(&self, p: PointId) -> Result<&[V], MeshSieveError> {
222 self.section.try_restrict(p)
223 }
224
225 #[inline]
226 fn try_get_mut(&mut self, p: PointId) -> Result<&mut [V], MeshSieveError> {
227 self.section.try_restrict_mut(p)
228 }
229}
230
231impl<V, S> InvalidateCache for ConstrainedSection<V, S>
232where
233 S: Storage<V>,
234{
235 fn invalidate_cache(&mut self) {
236 self.section.invalidate_cache();
237 }
238}
239
240pub fn apply_constraints_to_section<V, S, C>(
242 section: &mut Section<V, S>,
243 constraints: &C,
244) -> Result<(), MeshSieveError>
245where
246 V: Clone,
247 S: Storage<V>,
248 C: ConstraintSet<V>,
249{
250 for point in constraints.points() {
251 if let Some(list) = constraints.constraints_for(point) {
252 apply_constraints_to_slice(section, point, list)?;
253 }
254 }
255 section.invalidate_cache();
256 Ok(())
257}
258
259fn apply_constraints_to_slice<V, S>(
260 section: &mut Section<V, S>,
261 point: PointId,
262 constraints: &[DofConstraint<V>],
263) -> Result<(), MeshSieveError>
264where
265 V: Clone,
266 S: Storage<V>,
267{
268 let slice = section.try_restrict_mut(point)?;
269 let len = slice.len();
270 for constraint in constraints {
271 if constraint.index >= len {
272 return Err(MeshSieveError::ConstraintIndexOutOfBounds {
273 point,
274 index: constraint.index,
275 len,
276 });
277 }
278 slice[constraint.index] = constraint.value.clone();
279 }
280 Ok(())
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286 use crate::data::atlas::Atlas;
287 use crate::data::storage::VecStorage;
288
289 #[test]
290 fn apply_constraints_sets_values() {
291 let mut atlas = Atlas::default();
292 let p = PointId::new(1).unwrap();
293 atlas.try_insert(p, 3).unwrap();
294 let mut section = Section::<i32, VecStorage<i32>>::new(atlas);
295 section.try_set(p, &[1, 2, 3]).unwrap();
296 let mut constrained = ConstrainedSection::new(section);
297 constrained.insert_constraint(p, 0, 10).unwrap();
298 constrained.insert_constraint(p, 2, 30).unwrap();
299 constrained.apply_constraints().unwrap();
300 assert_eq!(constrained.section().try_restrict(p).unwrap(), &[10, 2, 30]);
301 }
302}