dyna3/app-proto/src/assembly.rs
Aaron Fenyes 620a6be918
All checks were successful
/ test (pull_request) Successful in 2m27s
Make regulator naming more consistent
Although the inversive distance regulator mostly acts as a general
Lorentz product regulator, calling it `InversiveDistanceRegulator`
explains how we intend to use it and demonstrates our naming convention.
2025-04-17 14:40:28 -07:00

643 lines
No EOL
23 KiB
Rust

use nalgebra::{DMatrix, DVector, DVectorView, Vector3};
use rustc_hash::FxHashMap;
use slab::Slab;
use std::{collections::BTreeSet, rc::Rc, sync::atomic::{AtomicU64, Ordering}};
use sycamore::prelude::*;
use web_sys::{console, wasm_bindgen::JsValue}; /* DEBUG */
use crate::{
engine::{
Q,
change_half_curvature,
local_unif_to_std,
realize_gram,
sphere,
ConfigSubspace,
ConstraintProblem
},
outline::OutlineItem,
specified::SpecifiedValue
};
// the types of the keys we use to access an assembly's elements and regulators
pub type ElementKey = usize;
pub type RegulatorKey = usize;
pub type ElementColor = [f32; 3];
/* KLUDGE */
// we should reconsider this design when we build a system for switching between
// assemblies. at that point, we might want to switch to hierarchical keys,
// where each each element has a key that identifies it within its assembly and
// each assembly has a key that identifies it within the sesssion
static NEXT_ELEMENT_SERIAL: AtomicU64 = AtomicU64::new(0);
pub trait ProblemPoser {
fn pose(&self, problem: &mut ConstraintProblem, elts: &Slab<Element>);
}
#[derive(Clone, PartialEq)]
pub struct Element {
pub id: String,
pub label: String,
pub color: ElementColor,
pub representation: Signal<DVector<f64>>,
// the regulators this element is subject to. the assembly that owns the
// element is responsible for keeping this set up to date
pub regulators: Signal<BTreeSet<RegulatorKey>>,
// a serial number, assigned by `Element::new`, that uniquely identifies
// each element
pub serial: u64,
// the configuration matrix column index that was assigned to this element
// last time the assembly was realized, or `None` if the element has never
// been through a realization
column_index: Option<usize>
}
impl Element {
const CURVATURE_COMPONENT: usize = 3;
pub fn new(
id: String,
label: String,
color: ElementColor,
representation: DVector<f64>
) -> Element {
// take the next serial number, panicking if that was the last number we
// had left. the technique we use to panic on overflow is taken from
// _Rust Atomics and Locks_, by Mara Bos
//
// https://marabos.nl/atomics/atomics.html#example-handle-overflow
//
let serial = NEXT_ELEMENT_SERIAL.fetch_update(
Ordering::SeqCst, Ordering::SeqCst,
|serial| serial.checked_add(1)
).expect("Out of serial numbers for elements");
Element {
id: id,
label: label,
color: color,
representation: create_signal(representation),
regulators: create_signal(BTreeSet::default()),
serial: serial,
column_index: None
}
}
// the smallest positive depth, represented as a multiple of `dir`, where
// the line generated by `dir` hits the element (which is assumed to be a
// sphere). returns `None` if the line misses the sphere. this function
// should be kept synchronized with `sphere_cast` in `inversive.frag`, which
// does essentially the same thing on the GPU side
pub fn cast(&self, dir: Vector3<f64>, assembly_to_world: &DMatrix<f64>) -> Option<f64> {
// if `a/b` is less than this threshold, we approximate
// `a*u^2 + b*u + c` by the linear function `b*u + c`
const DEG_THRESHOLD: f64 = 1e-9;
let rep = self.representation.with_untracked(|rep| assembly_to_world * rep);
let a = -rep[3] * dir.norm_squared();
let b = rep.rows_range(..3).dot(&dir);
let c = -rep[4];
let adjust = 4.0*a*c/(b*b);
if adjust < 1.0 {
// as long as `b` is non-zero, the linear approximation of
//
// a*u^2 + b*u + c
//
// at `u = 0` will reach zero at a finite depth `u_lin`. the root of
// the quadratic adjacent to `u_lin` is stored in `lin_root`. if
// both roots have the same sign, `lin_root` will be the one closer
// to `u = 0`
let square_rect_ratio = 1.0 + (1.0 - adjust).sqrt();
let lin_root = -(2.0*c)/b / square_rect_ratio;
if a.abs() > DEG_THRESHOLD * b.abs() {
if lin_root > 0.0 {
Some(lin_root)
} else {
let other_root = -b/(2.*a) * square_rect_ratio;
(other_root > 0.0).then_some(other_root)
}
} else {
(lin_root > 0.0).then_some(lin_root)
}
} else {
// the line through `dir` misses the sphere completely
None
}
}
}
impl ProblemPoser for Element {
fn pose(&self, problem: &mut ConstraintProblem, _elts: &Slab<Element>) {
if let Some(index) = self.column_index {
problem.gram.push_sym(index, index, 1.0);
problem.guess.set_column(index, &self.representation.get_clone_untracked());
} else {
panic!("Tried to write problem data from an unindexed element: \"{}\"", self.id);
}
}
}
pub trait Regulator: ProblemPoser + OutlineItem {
fn subjects(&self) -> Vec<ElementKey>;
fn measurement(&self) -> ReadSignal<f64>;
fn set_point(&self) -> Signal<SpecifiedValue>;
fn activate(&self, _assembly: &Assembly) {}
}
pub struct InversiveDistanceRegulator {
pub subjects: [ElementKey; 2],
pub measurement: ReadSignal<f64>,
pub set_point: Signal<SpecifiedValue>
}
impl InversiveDistanceRegulator {
pub fn new(subjects: [ElementKey; 2], assembly: &Assembly) -> InversiveDistanceRegulator {
let measurement = assembly.elements.map(
move |elts| {
let representations = subjects.map(|subj| elts[subj].representation);
representations[0].with(|rep_0|
representations[1].with(|rep_1|
rep_0.dot(&(&*Q * rep_1))
)
)
}
);
let set_point = create_signal(SpecifiedValue::from_empty_spec());
InversiveDistanceRegulator {
subjects: subjects,
measurement: measurement,
set_point: set_point
}
}
}
impl Regulator for InversiveDistanceRegulator {
fn subjects(&self) -> Vec<ElementKey> {
self.subjects.into()
}
fn measurement(&self) -> ReadSignal<f64> {
self.measurement
}
fn set_point(&self) -> Signal<SpecifiedValue> {
self.set_point
}
}
impl ProblemPoser for InversiveDistanceRegulator {
fn pose(&self, problem: &mut ConstraintProblem, elts: &Slab<Element>) {
self.set_point.with_untracked(|set_pt| {
if let Some(val) = set_pt.value {
let subject_column_indices = self.subjects.map(
|subj| elts[subj].column_index
);
if let [Some(row), Some(col)] = subject_column_indices {
problem.gram.push_sym(row, col, val);
} else {
panic!("Tried to write problem data from a regulator with an unindexed subject");
}
}
});
}
}
pub struct HalfCurvatureRegulator {
pub subject: ElementKey,
pub measurement: ReadSignal<f64>,
pub set_point: Signal<SpecifiedValue>
}
impl HalfCurvatureRegulator {
pub fn new(subject: ElementKey, assembly: &Assembly) -> HalfCurvatureRegulator {
let measurement = assembly.elements.map(
move |elts| elts[subject].representation.with(
|rep| rep[Element::CURVATURE_COMPONENT]
)
);
let set_point = create_signal(SpecifiedValue::from_empty_spec());
HalfCurvatureRegulator {
subject: subject,
measurement: measurement,
set_point: set_point
}
}
}
impl Regulator for HalfCurvatureRegulator {
fn subjects(&self) -> Vec<ElementKey> {
vec![self.subject]
}
fn measurement(&self) -> ReadSignal<f64> {
self.measurement
}
fn set_point(&self) -> Signal<SpecifiedValue> {
self.set_point
}
fn activate(&self, assembly: &Assembly) {
if let Some(half_curv) = self.set_point.with_untracked(|set_pt| set_pt.value) {
let representation = assembly.elements.with_untracked(
|elts| elts[self.subject].representation
);
representation.update(
|rep| change_half_curvature(rep, half_curv)
);
}
}
}
impl ProblemPoser for HalfCurvatureRegulator {
fn pose(&self, problem: &mut ConstraintProblem, elts: &Slab<Element>) {
self.set_point.with_untracked(|set_pt| {
if let Some(val) = set_pt.value {
if let Some(col) = elts[self.subject].column_index {
problem.frozen.push(Element::CURVATURE_COMPONENT, col, val);
} else {
panic!("Tried to write problem data from a regulator with an unindexed subject");
}
}
});
}
}
// the velocity is expressed in uniform coordinates
pub struct ElementMotion<'a> {
pub key: ElementKey,
pub velocity: DVectorView<'a, f64>
}
type AssemblyMotion<'a> = Vec<ElementMotion<'a>>;
// a complete, view-independent description of an assembly
#[derive(Clone)]
pub struct Assembly {
// elements and regulators
pub elements: Signal<Slab<Element>>,
pub regulators: Signal<Slab<Rc<dyn Regulator>>>,
// solution variety tangent space. the basis vectors are stored in
// configuration matrix format, ordered according to the elements' column
// indices. when you realize the assembly, every element that's present
// during realization gets a column index and is reflected in the tangent
// space. since the methods in this module never assign column indices
// without later realizing the assembly, we get the following invariant:
//
// (1) if an element has a column index, its tangent motions can be found
// in that column of the tangent space basis matrices
//
pub tangent: Signal<ConfigSubspace>,
// indexing
pub elements_by_id: Signal<FxHashMap<String, ElementKey>>
}
impl Assembly {
pub fn new() -> Assembly {
Assembly {
elements: create_signal(Slab::new()),
regulators: create_signal(Slab::new()),
tangent: create_signal(ConfigSubspace::zero(0)),
elements_by_id: create_signal(FxHashMap::default())
}
}
// --- inserting elements and regulators ---
// insert a sphere into the assembly without checking whether we already
// have an element with the same identifier. any element that does have the
// same identifier will get kicked out of the `elements_by_id` index
fn insert_sphere_unchecked(&self, elt: Element) -> ElementKey {
// insert the sphere
let id = elt.id.clone();
let key = self.elements.update(|elts| elts.insert(elt));
self.elements_by_id.update(|elts_by_id| elts_by_id.insert(id, key));
// regulate the sphere's curvature
self.insert_regulator(HalfCurvatureRegulator::new(key, &self));
key
}
pub fn try_insert_sphere(&self, elt: Element) -> Option<ElementKey> {
let can_insert = self.elements_by_id.with_untracked(
|elts_by_id| !elts_by_id.contains_key(&elt.id)
);
if can_insert {
Some(self.insert_sphere_unchecked(elt))
} else {
None
}
}
pub fn insert_new_sphere(&self) {
// find the next unused identifier in the default sequence
let mut id_num = 1;
let mut id = format!("sphere{}", id_num);
while self.elements_by_id.with_untracked(
|elts_by_id| elts_by_id.contains_key(&id)
) {
id_num += 1;
id = format!("sphere{}", id_num);
}
// create and insert a sphere
let _ = self.insert_sphere_unchecked(
Element::new(
id,
format!("Sphere {}", id_num),
[0.75_f32, 0.75_f32, 0.75_f32],
sphere(0.0, 0.0, 0.0, 1.0)
)
);
}
pub fn insert_regulator<T: Regulator + 'static>(&self, regulator: T) {
// add the regulator to the assembly's regulator list
let regulator_rc = Rc::new(regulator);
let key = self.regulators.update(
|regs| regs.insert(regulator_rc.clone())
);
// add the regulator to each subject's regulator list
let subjects = regulator_rc.subjects();
let subject_regulators: Vec<_> = self.elements.with_untracked(
|elts| subjects.into_iter().map(
|subj| elts[subj].regulators
).collect()
);
for regulators in subject_regulators {
regulators.update(|regs| regs.insert(key));
}
// update the realization when the regulator becomes a constraint, or is
// edited while acting as a constraint
let self_for_effect = self.clone();
create_effect(move || {
console::log_1(&JsValue::from(
format!("Updated regulator with subjects {:?}", regulator_rc.subjects())
));
if regulator_rc.set_point().with(|set_pt| set_pt.is_present()) {
regulator_rc.activate(&self_for_effect);
self_for_effect.realize();
}
});
/* DEBUG */
// print an updated list of regulators
console::log_1(&JsValue::from("Regulators:"));
self.regulators.with_untracked(|regs| {
for (_, reg) in regs.into_iter() {
console::log_1(&JsValue::from(format!(
" {:?}: {}",
reg.subjects(),
reg.set_point().with_untracked(
|set_pt| {
let spec = &set_pt.spec;
if spec.is_empty() {
"__".to_string()
} else {
spec.clone()
}
}
)
)));
}
});
}
// --- realization ---
pub fn realize(&self) {
// index the elements
self.elements.update_silent(|elts| {
for (index, (_, elt)) in elts.into_iter().enumerate() {
elt.column_index = Some(index);
}
});
// set up the constraint problem
let problem = self.elements.with_untracked(|elts| {
let mut problem = ConstraintProblem::new(elts.len());
for (_, elt) in elts {
elt.pose(&mut problem, elts);
}
self.regulators.with_untracked(|regs| {
for (_, reg) in regs {
reg.pose(&mut problem, elts);
}
});
problem
});
/* DEBUG */
// log the Gram matrix
console::log_1(&JsValue::from("Gram matrix:"));
problem.gram.log_to_console();
/* DEBUG */
// log the initial configuration matrix
console::log_1(&JsValue::from("Old configuration:"));
for j in 0..problem.guess.nrows() {
let mut row_str = String::new();
for k in 0..problem.guess.ncols() {
row_str.push_str(format!(" {:>8.3}", problem.guess[(j, k)]).as_str());
}
console::log_1(&JsValue::from(row_str));
}
// look for a configuration with the given Gram matrix
let (config, tangent, success, history) = realize_gram(
&problem, 1.0e-12, 0.5, 0.9, 1.1, 200, 110
);
/* DEBUG */
// report the outcome of the search
console::log_1(&JsValue::from(
if success {
"Target accuracy achieved!"
} else {
"Failed to reach target accuracy"
}
));
console::log_2(&JsValue::from("Steps:"), &JsValue::from(history.scaled_loss.len() - 1));
console::log_2(&JsValue::from("Loss:"), &JsValue::from(*history.scaled_loss.last().unwrap()));
console::log_2(&JsValue::from("Tangent dimension:"), &JsValue::from(tangent.dim()));
if success {
// read out the solution
for (_, elt) in self.elements.get_clone_untracked() {
elt.representation.update(
|rep| rep.set_column(0, &config.column(elt.column_index.unwrap()))
);
}
// save the tangent space
self.tangent.set_silent(tangent);
}
}
// --- deformation ---
// project the given motion to the tangent space of the solution variety and
// move the assembly along it. the implementation is based on invariant (1)
// from above and the following additional invariant:
//
// (2) if an element is affected by a constraint, it has a column index
//
// we have this invariant because the assembly gets realized each time you
// add a constraint
pub fn deform(&self, motion: AssemblyMotion) {
/* KLUDGE */
// when the tangent space is zero, deformation won't do anything, but
// the attempt to deform should be registered in the UI. this console
// message will do for now
if self.tangent.with(|tan| tan.dim() <= 0 && tan.assembly_dim() > 0) {
console::log_1(&JsValue::from("The assembly is rigid"));
}
// give a column index to each moving element that doesn't have one yet.
// this temporarily breaks invariant (1), but the invariant will be
// restored when we realize the assembly at the end of the deformation.
// in the process, we find out how many matrix columns we'll need to
// hold the deformation
let realized_dim = self.tangent.with(|tan| tan.assembly_dim());
let motion_dim = self.elements.update_silent(|elts| {
let mut next_column_index = realized_dim;
for elt_motion in motion.iter() {
let moving_elt = &mut elts[elt_motion.key];
if moving_elt.column_index.is_none() {
moving_elt.column_index = Some(next_column_index);
next_column_index += 1;
}
}
next_column_index
});
// project the element motions onto the tangent space of the solution
// variety and sum them to get a deformation of the whole assembly. the
// matrix `motion_proj` that holds the deformation has extra columns for
// any moving elements that aren't reflected in the saved tangent space
const ELEMENT_DIM: usize = 5;
let mut motion_proj = DMatrix::zeros(ELEMENT_DIM, motion_dim);
for elt_motion in motion {
// we can unwrap the column index because we know that every moving
// element has one at this point
let column_index = self.elements.with_untracked(
|elts| elts[elt_motion.key].column_index.unwrap()
);
if column_index < realized_dim {
// this element had a column index when we started, so by
// invariant (1), it's reflected in the tangent space
let mut target_columns = motion_proj.columns_mut(0, realized_dim);
target_columns += self.tangent.with(
|tan| tan.proj(&elt_motion.velocity, column_index)
);
} else {
// this element didn't have a column index when we started, so
// by invariant (2), it's unconstrained
let mut target_column = motion_proj.column_mut(column_index);
let unif_to_std = self.elements.with_untracked(
|elts| {
elts[elt_motion.key].representation.with_untracked(
|rep| local_unif_to_std(rep.as_view())
)
}
);
target_column += unif_to_std * elt_motion.velocity;
}
}
// step the assembly along the deformation. this changes the elements'
// normalizations, so we restore those afterward
/* KLUDGE */
// since our test assemblies only include spheres, we assume that every
// element is on the 1 mass shell
for (_, elt) in self.elements.get_clone_untracked() {
elt.representation.update_silent(|rep| {
match elt.column_index {
Some(column_index) => {
// step the assembly along the deformation
*rep += motion_proj.column(column_index);
// restore normalization by contracting toward the last
// coordinate axis
let q_sp = rep.fixed_rows::<3>(0).norm_squared();
let half_q_lt = -2.0 * rep[3] * rep[4];
let half_q_lt_sq = half_q_lt * half_q_lt;
let scaling = half_q_lt + (q_sp + half_q_lt_sq).sqrt();
rep.fixed_rows_mut::<4>(0).scale_mut(1.0 / scaling);
},
None => {
console::log_1(&JsValue::from(
format!("No velocity to unpack for fresh element \"{}\"", elt.id)
))
}
};
});
}
// bring the configuration back onto the solution variety. this also
// gets the elements' column indices and the saved tangent space back in
// sync
self.realize();
}
}
#[cfg(test)]
mod tests {
use crate::engine;
use super::*;
#[test]
#[should_panic(expected = "Tried to write problem data from an unindexed element: \"sphere\"")]
fn unindexed_element_test() {
let _ = create_root(|| {
Element::new(
"sphere".to_string(),
"Sphere".to_string(),
[1.0_f32, 1.0_f32, 1.0_f32],
engine::sphere(0.0, 0.0, 0.0, 1.0)
).pose(&mut ConstraintProblem::new(1), &Slab::new());
});
}
#[test]
#[should_panic(expected = "Tried to write problem data from a regulator with an unindexed subject")]
fn unindexed_subject_test() {
let _ = create_root(|| {
let mut elts = Slab::new();
let subjects = [0, 1].map(|k| {
elts.insert(
Element::new(
format!("sphere{k}"),
format!("Sphere {k}"),
[1.0_f32, 1.0_f32, 1.0_f32],
engine::sphere(0.0, 0.0, 0.0, 1.0)
)
)
});
elts[subjects[0]].column_index = Some(0);
InversiveDistanceRegulator {
subjects: subjects,
measurement: create_memo(|| 0.0),
set_point: create_signal(SpecifiedValue::try_from("0.0".to_string()).unwrap())
}.pose(&mut ConstraintProblem::new(2), &elts);
});
}
}