Skip to main content

kcl_ezpz/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::collections::HashSet;
4
5pub use crate::analysis::FreedomAnalysis;
6use crate::analysis::{Analysis, NoAnalysis, SolveOutcomeAnalysis};
7pub use crate::constraint_request::ConstraintRequest;
8pub use crate::constraints::Constraint;
9use crate::constraints::ConstraintEntry;
10pub use crate::error::*;
11pub use crate::solver::Config;
12// Only public for now so that I can benchmark it.
13// TODO: Replace this with an end-to-end benchmark,
14// or find a different way to structure modules.
15pub use crate::id::{Id, IdGenerator};
16use crate::solver::Model;
17pub use solve_outcome::{FailureOutcome, SolveOutcome, SolveOutcomeFreedomAnalysis};
18pub use warnings::{Warning, WarningContent};
19
20mod analysis;
21mod constraint_request;
22/// Each kind of constraint we support.
23mod constraints;
24/// Geometric data (lines, points, etc).
25pub mod datatypes;
26mod error;
27/// IDs of various entities, points, scalars etc.
28mod id;
29/// Residual field visualization (optional).
30#[cfg(feature = "residual-viz")]
31pub mod residual_viz;
32mod solve_outcome;
33/// Numeric solver using sparse matrices.
34mod solver;
35/// Unit tests
36#[cfg(test)]
37mod tests;
38/// Parser for textual representation of these problems.
39pub mod textual;
40mod vector;
41mod warnings;
42
43const EPSILON: f64 = 1e-4;
44
45/// Given some initial guesses, constrain them.
46/// Returns the same variables in the same order, but constrained.
47/// ```
48/// use kcl_ezpz::{Config, solve, Constraint, ConstraintRequest, datatypes::inputs::DatumPoint, IdGenerator};
49///
50/// // Define the geometry.
51/// let mut ids = IdGenerator::default();
52/// let p = DatumPoint::new(&mut ids);
53/// let q = DatumPoint::new(&mut ids);
54///
55/// // Define constraints on the geometry.
56/// let requests = [
57///     // Fix P to the origin
58///     ConstraintRequest::highest_priority(Constraint::Fixed(p.id_x(), 0.0)),
59///     ConstraintRequest::highest_priority(Constraint::Fixed(p.id_y(), 0.0)),
60///     // P and Q should be 4 units apart.
61///     ConstraintRequest::highest_priority(Constraint::Distance(p, q, 4.0)),
62/// ];
63///
64/// // Provide some initial guesses to the solver for their locations.
65/// let initial_guesses = vec![
66///     (p.id_x(), 0.0),
67///     (p.id_y(), -0.02),
68///     (q.id_x(), 4.39),
69///     (q.id_y(), 4.38),
70/// ];
71///
72/// // Run the solver!
73/// let _outcome = solve(
74///     &requests,
75///     initial_guesses,
76///     // You can customize the config, but for this example, we'll just use the default.
77///     Config::default(),
78/// );
79/// ```
80pub fn solve(
81    reqs: &[ConstraintRequest],
82    initial_guesses: Vec<(Id, f64)>,
83    config: Config,
84) -> Result<SolveOutcome, FailureOutcome> {
85    let out = solve_with_priority_inner::<NoAnalysis>(reqs, initial_guesses, config)?;
86    Ok(out.outcome)
87}
88
89/// Just like [`solve`] except it also does some expensive analysis steps
90/// at the end. This lets it calculate helpful data for the user, like degrees of freedom.
91/// Should not be called on every iteration of a system when you change the initial values!
92/// Just call this when you change the constraint structure.
93/// ```
94/// use kcl_ezpz::{Config, solve_analysis, Constraint, ConstraintRequest, datatypes::inputs::DatumPoint, IdGenerator};
95///
96/// // Define the geometry.
97/// let mut ids = IdGenerator::default();
98/// let p = DatumPoint::new(&mut ids);
99/// let q = DatumPoint::new(&mut ids);
100///
101/// // Define constraints on the geometry.
102/// let requests = [
103///     // Fix P to the origin
104///     ConstraintRequest::highest_priority(Constraint::Fixed(p.id_x(), 0.0)),
105///     ConstraintRequest::highest_priority(Constraint::Fixed(p.id_y(), 0.0)),
106///     // P and Q should be 4 units apart.
107///     ConstraintRequest::highest_priority(Constraint::Distance(p, q, 4.0)),
108/// ];
109///
110/// // Provide some initial guesses to the solver for their locations.
111/// let initial_guesses = vec![
112///     (p.id_x(), 0.0),
113///     (p.id_y(), -0.02),
114///     (q.id_x(), 4.39),
115///     (q.id_y(), 4.38),
116/// ];
117///
118/// // Run the solver!
119/// let solver_res = solve_analysis(
120///     &requests,
121///     initial_guesses,
122///     Config::default(),
123/// );
124/// let analysis = solver_res.unwrap().analysis;
125/// let underconstrained_vars = analysis.underconstrained();
126/// // P is fully constrained, because it's completely fixed to the origin.
127/// assert!(!underconstrained_vars.contains(&p.id_x()));
128/// assert!(!underconstrained_vars.contains(&p.id_y()));
129/// // Q is underconstrained. It has to be 4 units away from Q, but there's many
130/// // possible positions for Q. It could be at (4, 0), (-4, 0), etc etc.
131/// assert!(underconstrained_vars.contains(&q.id_x()));
132/// assert!(underconstrained_vars.contains(&q.id_y()));
133/// ```
134pub fn solve_analysis(
135    reqs: &[ConstraintRequest],
136    initial_guesses: Vec<(Id, f64)>,
137    config: Config,
138) -> Result<SolveOutcomeFreedomAnalysis, FailureOutcome> {
139    let out = solve_with_priority_inner::<FreedomAnalysis>(reqs, initial_guesses, config)?;
140    Ok(SolveOutcomeFreedomAnalysis {
141        analysis: out.analysis,
142        outcome: out.outcome,
143    })
144}
145
146/// Given some initial guesses, constrain them.
147/// Returns the same variables in the same order, but constrained.
148pub(crate) fn solve_with_priority_inner<A: Analysis>(
149    reqs: &[ConstraintRequest],
150    initial_guesses: Vec<(Id, f64)>,
151    config: Config,
152) -> Result<SolveOutcomeAnalysis<A>, FailureOutcome> {
153    // When there's no constraints, return early.
154    // Use the initial guesses as the final values.
155    if reqs.is_empty() {
156        return Ok(SolveOutcomeAnalysis {
157            analysis: A::no_constraints(),
158            outcome: SolveOutcome {
159                unsatisfied: Vec::new(),
160                final_values: initial_guesses
161                    .into_iter()
162                    .map(|(_id, guess)| guess)
163                    .collect(),
164                iterations: 0,
165                warnings: Vec::new(),
166                priority_solved: 0,
167            },
168        });
169    }
170
171    let reqs: Vec<_> = reqs
172        .iter()
173        .enumerate()
174        .map(|(id, c)| ConstraintEntry {
175            constraint: c.constraint(),
176            priority: c.priority(),
177            id,
178        })
179        .collect();
180
181    // Find all the priority levels, and put them into order from highest to lowest priority.
182    let priorities: HashSet<_> = reqs.iter().map(|c| c.priority).collect();
183    let mut priorities: Vec<_> = priorities.into_iter().collect();
184    let lowest_priority = priorities.iter().min().copied().unwrap_or(0);
185    priorities.sort();
186
187    // Handle the case with 0 constraints.
188    // (this gets used below, if the per-constraint loop never returns).
189    let mut res = None;
190    let total_constraints = reqs.len();
191
192    // Try solving, starting with only the highest priority constraints,
193    // adding more and more until we eventually either finish all constraints,
194    // or cannot find a solution that satisfies all of them.
195    let mut constraint_subset: Vec<ConstraintEntry<'_>> = Vec::with_capacity(total_constraints);
196
197    for curr_max_priority in priorities {
198        constraint_subset.clear();
199        for req in &reqs {
200            if req.priority <= curr_max_priority {
201                constraint_subset.push(req.to_owned()); // Notice: this clones.
202            }
203        }
204        let solve_res = solve_inner(
205            constraint_subset.as_slice(),
206            initial_guesses.clone(),
207            config,
208        );
209
210        match solve_res {
211            Ok(outcome) => {
212                // If there were unsatisfied constraints, then there's no point trying to add more lower-priority constraints,
213                // just return now.
214                if outcome.outcome.is_unsatisfied() {
215                    return Ok(res.unwrap_or(outcome));
216                }
217                // Otherwise, continue the loop again, adding higher-priority constraints.
218                res = Some(outcome);
219            }
220            // If this constraint couldn't be solved,
221            Err(e) => {
222                // then return a previous solved system with fewer (higher-priority) constraints,
223                // or if there was no such previous system, then this was the first run,
224                // and we should just return the error.
225                return res.ok_or(e);
226            }
227        }
228    }
229    // The unwrap default value is used when
230    // there were 0 constraints.
231    Ok(res.unwrap_or(SolveOutcomeAnalysis {
232        analysis: A::no_constraints(),
233        outcome: SolveOutcome {
234            unsatisfied: Vec::new(),
235            final_values: initial_guesses
236                .into_iter()
237                .map(|(_id, guess)| guess)
238                .collect(),
239            iterations: 0,
240            warnings: Vec::new(),
241            priority_solved: lowest_priority,
242        },
243    }))
244}
245
246fn solve_inner<A: Analysis>(
247    constraints: &[ConstraintEntry<'_>],
248    initial_guesses: Vec<(Id, f64)>,
249    config: Config,
250) -> Result<SolveOutcomeAnalysis<A>, FailureOutcome> {
251    let num_vars = initial_guesses.len();
252    let num_eqs = constraints
253        .iter()
254        .map(|c| c.constraint.residual_dim())
255        .sum();
256    let (all_variables, mut values): (Vec<Id>, Vec<f64>) = initial_guesses.into_iter().unzip();
257    let mut warnings = warnings::lint(constraints);
258    let initial_values = values.clone();
259
260    let mut model = match Model::new(constraints, all_variables, initial_values, config) {
261        Ok(o) => o,
262        Err(error) => {
263            return Err(FailureOutcome {
264                error,
265                warnings,
266                num_vars,
267                num_eqs,
268            });
269        }
270    };
271
272    let mut unsatisfied: Vec<usize> = Vec::new();
273    let outcome = model.solve_gauss_newton(&mut values, config);
274    warnings.extend(model.warnings.lock().unwrap().drain(..));
275    let success = match outcome {
276        Ok(o) => o,
277        Err(error) => {
278            return Err(FailureOutcome {
279                error,
280                warnings,
281                num_vars,
282                num_eqs,
283            });
284        }
285    };
286    let cs: Vec<_> = constraints.iter().map(|c| c.constraint).collect();
287    let layout = solver::Layout::new(&Vec::new(), cs.as_slice(), config);
288    for constraint in constraints {
289        let mut residual0 = 0.0;
290        let mut residual1 = 0.0;
291        let mut residual2 = 0.0;
292        let mut degenerate = false;
293        constraint.constraint.residual(
294            &layout,
295            &values,
296            &mut residual0,
297            &mut residual1,
298            &mut residual2,
299            &mut degenerate,
300        );
301        let satisfied = is_satisfied(
302            constraint.constraint.residual_dim(),
303            [residual0, residual1, residual2],
304        );
305        if !satisfied {
306            unsatisfied.push(constraint.id);
307        }
308    }
309    let analysis = match A::analyze(model) {
310        Ok(o) => o,
311        Err(error) => {
312            return Err(FailureOutcome {
313                error,
314                warnings,
315                num_vars,
316                num_eqs,
317            });
318        }
319    };
320
321    let lowest_priority = constraints
322        .iter()
323        .map(|c| c.priority)
324        .max()
325        .unwrap_or_default();
326    Ok(SolveOutcomeAnalysis {
327        outcome: SolveOutcome {
328            priority_solved: lowest_priority,
329            unsatisfied,
330            final_values: values,
331            iterations: success.iterations,
332            warnings,
333        },
334        analysis,
335    })
336}
337
338fn is_satisfied(residual_dim: usize, residuals: [f64; 3]) -> bool {
339    let sat0 = residuals[0].abs() < EPSILON;
340    let sat1 = residuals[1].abs() < EPSILON;
341    let sat2 = residuals[2].abs() < EPSILON;
342    match residual_dim {
343        1 => sat0,
344        2 => sat0 && sat1,
345        3 => sat0 && sat1 && sat2,
346        other => unreachable!(
347            "Unsupported number of residuals {other}, the `residual` method must be modified."
348        ),
349    }
350}
351
352#[cfg(test)]
353mod basic_tests {
354    use super::*;
355
356    #[test]
357    fn test_is_satisfied_0() {
358        let actual = is_satisfied(1, [1e-8, 44.0, 44.0]);
359        let expected = true;
360        assert_eq!(actual, expected);
361    }
362
363    #[test]
364    fn test_is_satisfied_1() {
365        let actual = is_satisfied(2, [1e-8, 1e-8, 44.0]);
366        let expected = true;
367        assert_eq!(actual, expected);
368    }
369
370    #[test]
371    fn test_is_satisfied_2() {
372        let actual = is_satisfied(3, [1e-8, 1e-8, 1e-8]);
373        let expected = true;
374        assert_eq!(actual, expected);
375    }
376
377    #[test]
378    fn test_is_unsatisfied_0() {
379        let actual = is_satisfied(1, [44.0, 44.0, 44.0]);
380        let expected = false;
381        assert_eq!(actual, expected);
382    }
383
384    #[test]
385    fn test_is_unsatisfied_1() {
386        let actual = is_satisfied(2, [1e-8, 44.0, 44.0]);
387        let expected = false;
388        assert_eq!(actual, expected);
389    }
390
391    #[test]
392    fn test_is_unsatisfied_2() {
393        let actual = is_satisfied(3, [44.0, 1e-8, 1e-8]);
394        let expected = false;
395        assert_eq!(actual, expected);
396    }
397}