dyna3/app-proto/src/assembly.rs

use nalgebra::{DMatrix, DVector, DVectorView, Vector3};
use rustc_hash::FxHashMap;
use slab::Slab;
use std::{collections::BTreeSet, sync::atomic::{AtomicU64, Ordering}};
use sycamore::prelude::*;
use web_sys::{console, wasm_bindgen::JsValue}; /* DEBUG */

use crate::engine::{realize_gram, ConfigSubspace, PartialMatrix, Q};

// the types of the keys we use to access an assembly's elements and constraints
pub type ElementKey = usize;
pub type ConstraintKey = 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);

#[derive(Clone, PartialEq)]
pub struct Element {
    pub id: String,
    pub label: String,
    pub color: ElementColor,
    pub representation: Signal<DVector<f64>>,
    pub constraints: Signal<BTreeSet<ConstraintKey>>,

    // 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
    column_index: usize
}

impl Element {
    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),
            constraints: create_signal(BTreeSet::default()),
            serial: serial,
            column_index: 0
        }
    }
    
    // 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
        }
    }
}

#[derive(Clone)]
pub struct Constraint {
    pub subjects: (ElementKey, ElementKey),
    pub lorentz_prod: Signal<f64>,
    pub lorentz_prod_text: Signal<String>,
    pub lorentz_prod_valid: Signal<bool>,
    pub active: Signal<bool>
}

// a complete, view-independent description of an assembly
#[derive(Clone)]
pub struct Assembly {
    // elements and constraints
    pub elements: Signal<Slab<Element>>,
    pub constraints: Signal<Slab<Constraint>>,
    
    // solution variety tangent space
    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()),
            constraints: create_signal(Slab::new()),
            tangent: create_signal(ConfigSubspace::zero()),
            elements_by_id: create_signal(FxHashMap::default())
        }
    }
    
    // --- inserting elements and constraints ---
    
    // insert an element 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_element_unchecked(&self, elt: Element) {
        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));
        
        // realize to update the tangent space
        /* KLUDGE */
        // since the newly inserted element is unconstrained, we should be able
        // to update the tangent space without recomputing the Hessian and its
        // eigendecomposition
        self.realize();
    }
    
    pub fn try_insert_element(&self, elt: Element) -> bool {
        let can_insert = self.elements_by_id.with_untracked(
            |elts_by_id| !elts_by_id.contains_key(&elt.id)
        );
        if can_insert {
            self.insert_element_unchecked(elt);
        }
        can_insert
    }
    
    pub fn insert_new_element(&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 new element
        self.insert_element_unchecked(
            Element::new(
                id,
                format!("Sphere {}", id_num),
                [0.75_f32, 0.75_f32, 0.75_f32],
                DVector::<f64>::from_column_slice(&[0.0, 0.0, 0.0, 0.5, -0.5])
            )
        );
    }
    
    pub fn insert_constraint(&self, constraint: Constraint) {
        let subjects = constraint.subjects;
        let key = self.constraints.update(|csts| csts.insert(constraint));
        let subject_constraints = self.elements.with(
            |elts| (elts[subjects.0].constraints, elts[subjects.1].constraints)
        );
        subject_constraints.0.update(|csts| csts.insert(key));
        subject_constraints.1.update(|csts| csts.insert(key));
    }
    
    // --- realization ---
    
    pub fn realize(&self) {
        // index the elements
        self.elements.update_silent(|elts| {
            for (index, (_, elt)) in elts.into_iter().enumerate() {
                elt.column_index = index;
            }
        });
        
        // set up the Gram matrix and the initial configuration matrix
        let (gram, guess) = self.elements.with_untracked(|elts| {
            // set up the off-diagonal part of the Gram matrix
            let mut gram_to_be = PartialMatrix::new();
            self.constraints.with_untracked(|csts| {
                for (_, cst) in csts {
                    if cst.active.get_untracked() && cst.lorentz_prod_valid.get_untracked() {
                        let subjects = cst.subjects;
                        let row = elts[subjects.0].column_index;
                        let col = elts[subjects.1].column_index;
                        gram_to_be.push_sym(row, col, cst.lorentz_prod.get_untracked());
                    }
                }
            });
            
            // set up the initial configuration matrix and the diagonal of the
            // Gram matrix
            let mut guess_to_be = DMatrix::<f64>::zeros(5, elts.len());
            for (_, elt) in elts {
                let index = elt.column_index;
                gram_to_be.push_sym(index, index, 1.0);
                guess_to_be.set_column(index, &elt.representation.get_clone_untracked());
            }
            
            (gram_to_be, guess_to_be)
        });
        
        /* DEBUG */
        // log the Gram matrix
        console::log_1(&JsValue::from("Gram matrix:"));
        gram.log_to_console();
        
        /* DEBUG */
        // log the initial configuration matrix
        console::log_1(&JsValue::from("Old configuration:"));
        for j in 0..guess.nrows() {
            let mut row_str = String::new();
            for k in 0..guess.ncols() {
                row_str.push_str(format!(" {:>8.3}", 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(
            &gram, guess, &[],
            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))
                );
            }
            
            // save the tangent space
            self.tangent.set_silent(tangent);
        }
    }
    
    // --- deformation ---
    
    pub fn deform(&self, element_motions: Vec<(ElementKey, DVectorView<f64>)>) {
        /* KLUDGE */
        // when the tangent space is zero, we currently need to avoid calling
        // its `proj` method, because it will panic rather than returning zero.
        // in the future, we'll want a more intentionally designed system for
        // handling this case
        if self.tangent.with(|tan| tan.dim() <= 0) {
            console::log_1(&JsValue::from("The assembly is rigid"));
            return;
        }
        
        const ELEMENT_DIM: usize = 5;
        let assembly_dim = self.elements.with(|elts| elts.len());
        let mut motion_proj = DMatrix::zeros(ELEMENT_DIM, assembly_dim);
        
        // project the element motions onto the tangent space of the solution
        // variety, and sum them to get a deformation of the whole assembly
        for (elt_key, elt_motion) in element_motions {
            let column_index = self.elements.with(|elts| elts[elt_key].column_index);
            motion_proj += self.tangent.with(|tan| tan.proj(&elt_motion, column_index));
        }
        
        // step each element along the mass shell geodesic that matches its
        // velocity in the deformation found above
        /* 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| {
                let rep_next = &*rep + motion_proj.column(elt.column_index);
                let normalizer = rep_next.dot(&(&*Q * &rep_next));
                rep.set_column(0, &(rep_next / normalizer));
            });
        }
        
        // bring the configuration back onto the solution variety
        self.realize();
    }
}