[WIP] optimization support
This commit is contained in:
parent
16cf295c0a
commit
858a9264fd
752
Cargo.lock
generated
752
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
13
Cargo.toml
13
Cargo.toml
@ -9,7 +9,7 @@ edition = "2021"
|
||||
argmin = { version = "0.8.1", features = [] }
|
||||
argmin-math = { version = "0.3.0", features = ["nalgebra_v0_32"] }
|
||||
bevy_ecs = "0.10.1"
|
||||
eframe = { version = "0.21.3" }
|
||||
eframe = { version = "0.21.3", default-features = false, features = ["default_fonts", "glow"] }
|
||||
indexmap = "1.8.1"
|
||||
levenberg-marquardt = "0.13.0"
|
||||
nalgebra = { version = "0.32.0", features = ["rand"] }
|
||||
@ -17,13 +17,18 @@ nalgebra-sparse = "0.9.0"
|
||||
nohash-hasher = "0.2.0"
|
||||
rand = "0.8.5"
|
||||
rand_distr = "0.4.3"
|
||||
tracing = { version = "0.1.37", optional = true }
|
||||
slotmap = "1.0.6"
|
||||
tracing = "0.1.37"
|
||||
tracing-chrome = { version = "0.7.1", optional = true }
|
||||
tracing-error = { version = "0.2.0", optional = true }
|
||||
tracing-subscriber = { version = "0.3.17", features = ["env-filter"], optional = true }
|
||||
tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
|
||||
|
||||
[dev-dependencies]
|
||||
criterion = "0.4.0"
|
||||
|
||||
[features]
|
||||
trace = ["tracing", "tracing-subscriber", "tracing-chrome", "tracing-error", "bevy_ecs/trace"]
|
||||
trace = ["tracing-chrome", "tracing-error", "bevy_ecs/trace"]
|
||||
|
||||
[[bench]]
|
||||
name = "my_benchmark"
|
||||
harness = false
|
||||
|
||||
5
src/lib.rs
Normal file
5
src/lib.rs
Normal file
@ -0,0 +1,5 @@
|
||||
#![feature(generic_const_exprs)]
|
||||
|
||||
pub mod geometry;
|
||||
pub mod optimization;
|
||||
// mod relations;
|
||||
48
src/main.rs
48
src/main.rs
@ -1,19 +1,41 @@
|
||||
#![feature(generic_const_exprs)]
|
||||
|
||||
use bevy_ecs::prelude::*;
|
||||
use eframe::egui;
|
||||
use rand::seq::SliceRandom;
|
||||
use rand_distr::Distribution;
|
||||
|
||||
use sketchrs::geometry;
|
||||
|
||||
mod geometry;
|
||||
pub mod optimization;
|
||||
mod relations;
|
||||
#[cfg(feature = "trace")]
|
||||
mod tracing;
|
||||
mod trace;
|
||||
mod ui;
|
||||
|
||||
fn init(mut commands: Commands) {
|
||||
let p1 = geometry::insert_point_at(&mut commands, (10., 30.));
|
||||
let p2 = geometry::insert_point_at(&mut commands, (-20., 15.));
|
||||
geometry::insert_point_at(&mut commands, (0., -10.));
|
||||
// let p1 = geometry::insert_point_at(&mut commands, (10., 30.));
|
||||
// let p2 = geometry::insert_point_at(&mut commands, (-20., 15.));
|
||||
// geometry::insert_point_at(&mut commands, (0., -10.));
|
||||
// commands.spawn(geometry::LineBundle::new(p1, p2));
|
||||
|
||||
// Insert 1000 random points
|
||||
let dist = rand_distr::Uniform::new(-100., 100.);
|
||||
let mut rng = rand::thread_rng();
|
||||
let mut points = Vec::new();
|
||||
for _ in 0..1000 {
|
||||
let p = geometry::insert_point_at(
|
||||
&mut commands,
|
||||
(dist.sample(&mut rng), dist.sample(&mut rng)),
|
||||
);
|
||||
points.push(p);
|
||||
}
|
||||
|
||||
// Insert 1000 random lines
|
||||
for _ in 0..1000 {
|
||||
let p1 = *points.choose(&mut rng).unwrap();
|
||||
let p2 = *points.choose(&mut rng).unwrap();
|
||||
commands.spawn(geometry::LineBundle::new(p1, p2));
|
||||
}
|
||||
}
|
||||
|
||||
struct MyApp {
|
||||
world: World,
|
||||
@ -31,9 +53,9 @@ impl MyApp {
|
||||
let mut world = World::default();
|
||||
|
||||
#[cfg(feature = "trace")]
|
||||
tracing::setup(&mut world);
|
||||
trace::setup(&mut world);
|
||||
|
||||
world.init_resource::<ui::Tool>();
|
||||
ui::init(&mut world);
|
||||
|
||||
let mut init_sched = Schedule::new();
|
||||
init_sched.add_system(init);
|
||||
@ -51,6 +73,9 @@ impl MyApp {
|
||||
|
||||
impl eframe::App for MyApp {
|
||||
fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
|
||||
let span = tracing::info_span!("update");
|
||||
let _enter = span.enter();
|
||||
|
||||
self.world.insert_resource(ui::ContextRes(ctx.clone()));
|
||||
|
||||
self.update_schedule.run(&mut self.world);
|
||||
@ -60,6 +85,11 @@ impl eframe::App for MyApp {
|
||||
}
|
||||
|
||||
fn main() {
|
||||
tracing_subscriber::fmt()
|
||||
.with_max_level(tracing::Level::INFO)
|
||||
.with_target(false)
|
||||
.init();
|
||||
|
||||
let options = eframe::NativeOptions::default();
|
||||
eframe::run_native("sketchrs", options, Box::new(|_cc| Box::<MyApp>::default())).unwrap();
|
||||
}
|
||||
|
||||
@ -1,4 +1,9 @@
|
||||
use std::{borrow::Borrow, fmt::Debug};
|
||||
use std::{borrow::Borrow, fmt::Debug, ops::Index};
|
||||
|
||||
mod params;
|
||||
mod residuals;
|
||||
|
||||
use params::ReadParams;
|
||||
|
||||
use argmin_math;
|
||||
use levenberg_marquardt::{self, TerminationReason};
|
||||
@ -7,11 +12,154 @@ use nalgebra_sparse::CooMatrix;
|
||||
|
||||
use crate::geometry::Scalar;
|
||||
|
||||
pub type Idx = u32;
|
||||
|
||||
use self::{params::{Parameters, Param}, residuals::WriteResiduals};
|
||||
|
||||
type SVector<const D: usize> = nalgebra::SVector<Scalar, D>;
|
||||
type SMatrix<const R: usize, const C: usize> = nalgebra::SMatrix<Scalar, R, C>;
|
||||
type DVector = nalgebra::DVector<Scalar>;
|
||||
type DMatrix = nalgebra::DMatrix<Scalar>;
|
||||
|
||||
|
||||
trait CostFunction {
|
||||
type Params;
|
||||
// const NUM_RESIDUALS: usize;
|
||||
type Residuals;
|
||||
type Jacobian;
|
||||
|
||||
fn evaluate(
|
||||
&self,
|
||||
params: &Self::Params,
|
||||
residual: Option<&mut Self::Residuals>,
|
||||
jacobian: Option<&mut Self::Jacobian>,
|
||||
);
|
||||
|
||||
fn residual(&self, x: &Self::Params) -> Self::Residuals
|
||||
where
|
||||
Self::Residuals: Default,
|
||||
{
|
||||
let mut residual = Self::Residuals::default();
|
||||
self.evaluate(x, Some(&mut residual), None);
|
||||
residual
|
||||
}
|
||||
|
||||
fn jacobian(&self, x: &Self::Params) -> Self::Jacobian
|
||||
where
|
||||
Self::Jacobian: Default,
|
||||
{
|
||||
let mut jacobian = Self::Jacobian::default();
|
||||
self.evaluate(x, None, Some(&mut jacobian));
|
||||
jacobian
|
||||
}
|
||||
|
||||
fn with_indices<PI, RI, JI>(
|
||||
self,
|
||||
param_indices: PI,
|
||||
residual_indices: RI,
|
||||
jacobian_indices: JI,
|
||||
) -> IndexedCostFunctionImpl<Self, PI, RI, JI>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
IndexedCostFunctionImpl {
|
||||
residual: self,
|
||||
param_indices,
|
||||
residual_indices,
|
||||
jacobian_indices,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
trait IndexedCostFunction {
|
||||
fn evaluate(
|
||||
&self,
|
||||
params: &Parameters,
|
||||
residuals: Option<&mut [Scalar]>,
|
||||
jacobian: Option<&mut [Scalar]>,
|
||||
);
|
||||
}
|
||||
|
||||
struct IndexedCostFunctionImpl<R, PI, RI, JI> {
|
||||
residual: R,
|
||||
param_indices: PI,
|
||||
residual_indices: RI,
|
||||
jacobian_indices: JI,
|
||||
}
|
||||
|
||||
impl<R, PI, RI, JI> IndexedCostFunction for IndexedCostFunctionImpl<R, PI, RI, JI>
|
||||
where
|
||||
R: CostFunction,
|
||||
R::Params: ReadParams<PI>,
|
||||
R::Residuals: WriteResiduals<RI> + Default,
|
||||
R::Jacobian: WriteResiduals<JI> + Default,
|
||||
{
|
||||
fn evaluate(
|
||||
&self,
|
||||
params: &Parameters,
|
||||
residuals: Option<&mut [Scalar]>,
|
||||
jacobian: Option<&mut [Scalar]>,
|
||||
) {
|
||||
// let p = R::Params::read(params, self.param_indices);
|
||||
|
||||
// let mut res = residuals.and_then(|_| Some(R::Residuals::default()));
|
||||
// let mut jac = jacobian.and_then(|_| Some(R::Jacobian::default()));
|
||||
// self.residual.evaluate(&p, res.as_mut(), jac.as_mut());
|
||||
|
||||
// if let (Some(res), Some(residuals)) = (res, residuals) {
|
||||
// res.write(self.residual_indices, residuals);
|
||||
// }
|
||||
// if let (Some(jac), Some(jacobian)) = (jac, jacobian) {
|
||||
// jac.write(self.jacobian_indices, jacobian);
|
||||
// }
|
||||
}
|
||||
}
|
||||
|
||||
slotmap::new_key_type!{ struct ParamKey; }
|
||||
|
||||
struct ParamIndex(Idx);
|
||||
|
||||
#[derive(Default)]
|
||||
struct CostFunctions {
|
||||
functions: Vec<Box<dyn IndexedCostFunction>>,
|
||||
param_indices: slotmap::SlotMap<ParamKey, Idx>,
|
||||
num_residuals: Idx,
|
||||
}
|
||||
|
||||
impl CostFunctions {
|
||||
fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
// fn add_param<P: Param>(&mut self) -> ParamIndex {
|
||||
// // let index = self.params_size;
|
||||
// // self.params_size += P::SIZE;
|
||||
// ParamIndex(index)
|
||||
// }
|
||||
|
||||
// fn add<R, PI, RI, JI>(&mut self, residual: R, param_indices: PI, residual_indices: RI, jacobian_indices: JI)
|
||||
// where
|
||||
// R: CostFunction,
|
||||
// R::Params: ReadParams<PI>,
|
||||
// R::Residuals: WriteResiduals<RI> + Default,
|
||||
// R::Jacobian: WriteResiduals<JI> + Default,
|
||||
// {
|
||||
// self.functions.push(Box::new(
|
||||
// residual
|
||||
// .with_indices(param_indices, residual_indices, jacobian_indices),
|
||||
// ));
|
||||
// self.param_indices.push(param_indices.into());
|
||||
// self.residual_indices.push(residual_indices.into());
|
||||
// self.jacobian_indices.push(jacobian_indices.into());
|
||||
// }
|
||||
|
||||
fn evaluate(&self, params: &Parameters, residuals: Option<&mut [Scalar]>, jacobian: Option<&mut [Scalar]>) {
|
||||
// for (i, f) in self.functions.iter().enumerate() {
|
||||
// f.evaluate(params, residuals, jacobian);
|
||||
// }
|
||||
}
|
||||
}
|
||||
|
||||
trait SResBlock<const NIN: usize, const NOUT: usize = 1>: Clone + Debug {
|
||||
fn evaluate(
|
||||
&self,
|
||||
@ -39,20 +187,20 @@ trait SResBlock<const NIN: usize, const NOUT: usize = 1>: Clone + Debug {
|
||||
}
|
||||
|
||||
// Probably should move to own trait
|
||||
fn into_indexed(
|
||||
self,
|
||||
in_indices: [Index; NIN],
|
||||
out_indices: [Index; NOUT],
|
||||
) -> Box<dyn IndexedResBlock>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Box::new(IndexedResBlockImpl {
|
||||
residual: self,
|
||||
in_indices,
|
||||
out_indices,
|
||||
})
|
||||
}
|
||||
// fn into_indexed(
|
||||
// self,
|
||||
// in_indices: [Index; NIN],
|
||||
// out_indices: [Index; NOUT],
|
||||
// ) -> Box<dyn IndexedResBlock>
|
||||
// where
|
||||
// Self: Sized + 'static,
|
||||
// {
|
||||
// Box::new(IndexedResBlockImpl {
|
||||
// residual: self,
|
||||
// in_indices,
|
||||
// out_indices,
|
||||
// })
|
||||
// }
|
||||
}
|
||||
|
||||
fn square(x: Scalar) -> Scalar {
|
||||
@ -84,6 +232,36 @@ impl SResBlock<2> for Equal1D {
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AxisDistance(Scalar);
|
||||
|
||||
/*
|
||||
impl ResBlock for AxisDistance {
|
||||
fn parameter_sizes(&self) -> &[usize] {
|
||||
&[2]
|
||||
}
|
||||
|
||||
fn n_residuals(&self) -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
fn evaluate(
|
||||
&self,
|
||||
x: &[&[Scalar]],
|
||||
residual: Option<&mut [Scalar]>,
|
||||
jacobian: Option<&mut [&mut [Scalar]]>,
|
||||
) {
|
||||
let [x1, x2] = x[0] else { panic!() };
|
||||
let dx = x2 - x1;
|
||||
let r = self.0;
|
||||
if let Some(residual) = residual {
|
||||
residual[0] = square(dx) - square(r);
|
||||
}
|
||||
if let Some(jacobian) = jacobian {
|
||||
jacobian[0][0] = -2. * dx;
|
||||
jacobian[0][1] = 2. * dx;
|
||||
}
|
||||
}
|
||||
}
|
||||
*/
|
||||
|
||||
impl SResBlock<2> for AxisDistance {
|
||||
fn evaluate(
|
||||
&self,
|
||||
@ -107,6 +285,34 @@ impl SResBlock<2> for AxisDistance {
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PointPointDistance(Scalar);
|
||||
|
||||
impl CostFunction for PointPointDistance {
|
||||
type Params = (SVector<2>, SVector<2>);
|
||||
type Residuals = Scalar;
|
||||
type Jacobian = (SMatrix<2, 1>, SMatrix<2, 1>);
|
||||
|
||||
fn evaluate(
|
||||
&self,
|
||||
(p1, p2): &Self::Params,
|
||||
residual: Option<&mut Self::Residuals>,
|
||||
jacobian: Option<&mut Self::Jacobian>,
|
||||
) {
|
||||
let d = p2 - p1;
|
||||
if let Some(residual) = residual {
|
||||
*residual = square(d.x) + square(d.y) - square(self.0);
|
||||
}
|
||||
if let Some(jacobian) = jacobian {
|
||||
jacobian.0.copy_from(&(2. * d));
|
||||
jacobian.1.copy_from(&(-2. * d));
|
||||
// jacobian.row_mut(0).copy_from(&(2. * d).transpose());
|
||||
// jacobian.row_mut(1).copy_from(&(-2. * d).transpose());
|
||||
// jacobian[0] = -2. * d.x;
|
||||
// jacobian[1] = -2. * d.y;
|
||||
// jacobian[2] = 2. * d.x;
|
||||
// jacobian[3] = 2. * d.y;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SResBlock<4> for PointPointDistance {
|
||||
fn evaluate(
|
||||
&self,
|
||||
@ -116,9 +322,8 @@ impl SResBlock<4> for PointPointDistance {
|
||||
) {
|
||||
let [x1, y1, x2, y2] = *x.as_ref();
|
||||
let (dx, dy) = (x2 - x1, y2 - y1);
|
||||
let r = self.0;
|
||||
if let Some(residual) = residual {
|
||||
residual[0] = square(dx) + square(dy) - square(r);
|
||||
residual[0] = square(dx) + square(dy) - square(self.0);
|
||||
}
|
||||
if let Some(jacobian) = jacobian {
|
||||
jacobian[0] = -2. * dx;
|
||||
@ -153,6 +358,33 @@ impl SResBlock<2> for PointDistance {
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Collinear;
|
||||
|
||||
impl CostFunction for Collinear {
|
||||
type Params = (SVector<2>, SVector<2>, SVector<2>);
|
||||
type Residuals = Scalar;
|
||||
type Jacobian = (SVector<2>, SVector<2>, SVector<2>);
|
||||
|
||||
fn evaluate(
|
||||
&self,
|
||||
(p1, p2, p3): &Self::Params,
|
||||
residual: Option<&mut Self::Residuals>,
|
||||
jacobian: Option<&mut Self::Jacobian>,
|
||||
) {
|
||||
let d21 = p2 - p1;
|
||||
let d32 = p3 - p2;
|
||||
if let Some(residual) = residual {
|
||||
*residual = d21.y * d32.x - d32.y * d21.x;
|
||||
}
|
||||
if let Some(jacobian) = jacobian {
|
||||
jacobian.0.x = d32.y;
|
||||
jacobian.0.y = -d32.x;
|
||||
jacobian.1.x = p1.y - p3.y;
|
||||
jacobian.1.y = p3.x - p1.x;
|
||||
jacobian.2.x = d21.y;
|
||||
jacobian.2.y = -d21.x;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SResBlock<6> for Collinear {
|
||||
fn evaluate(
|
||||
&self,
|
||||
@ -171,15 +403,14 @@ impl SResBlock<6> for Collinear {
|
||||
if let Some(jacobian) = jacobian {
|
||||
jacobian[0] = dy32;
|
||||
jacobian[1] = -dx32;
|
||||
jacobian[2] = -dy32;
|
||||
jacobian[3] = -dx21;
|
||||
jacobian[2] = y1 - y3;
|
||||
jacobian[3] = x3 - x1;
|
||||
jacobian[4] = dy21;
|
||||
jacobian[5] = -dx21;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type Index = u32;
|
||||
|
||||
trait ScatterResBlock: Debug {
|
||||
fn shape(&self) -> (usize, usize);
|
||||
@ -194,9 +425,9 @@ trait ScatterResBlock: Debug {
|
||||
#[derive(Clone, Debug)]
|
||||
struct ScatterResBlockImpl<const NIN: usize, const NOUT: usize, R> {
|
||||
residual: R,
|
||||
in_indices: [Index; NIN],
|
||||
residual_indices: [Index; NOUT],
|
||||
jacobian_indices: [[Index; NIN]; NOUT],
|
||||
in_indices: [Idx; NIN],
|
||||
residual_indices: [Idx; NOUT],
|
||||
jacobian_indices: [[Idx; NIN]; NOUT],
|
||||
}
|
||||
|
||||
impl<const NIN: usize, const NOUT: usize, R> ScatterResBlockImpl<NIN, NOUT, R> {
|
||||
@ -278,6 +509,7 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
trait IndexedResBlock {
|
||||
fn shape(&self) -> (usize, usize);
|
||||
fn in_indices(&self) -> &[Index];
|
||||
@ -323,20 +555,22 @@ where
|
||||
})
|
||||
}
|
||||
}
|
||||
*/
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Residuals {
|
||||
struct ResidualBlocks {
|
||||
res_blocks: Vec<Box<dyn ScatterResBlock>>,
|
||||
shape: (usize, usize),
|
||||
}
|
||||
|
||||
impl Residuals {
|
||||
fn new(res_blocks: Vec<Box<dyn IndexedResBlock>>, shape: (usize, usize)) -> Self {
|
||||
let res_blocks = res_blocks
|
||||
.into_iter()
|
||||
.map(|r| r.into_scatter(shape))
|
||||
.collect();
|
||||
Self { res_blocks, shape }
|
||||
impl ResidualBlocks {
|
||||
fn new(res_blocks: Vec<Box<dyn IndexedCostFunction>>, shape: (usize, usize)) -> Self {
|
||||
// let res_blocks = res_blocks
|
||||
// .into_iter()
|
||||
// .map(|r| r.into_scatter(shape))
|
||||
// .collect();
|
||||
// Self { res_blocks, shape }
|
||||
todo!()
|
||||
}
|
||||
|
||||
fn residuals_into<X: AsRef<[Scalar]>>(&self, x: X, residuals: &mut [Scalar]) {
|
||||
@ -379,7 +613,7 @@ impl Residuals {
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Problem {
|
||||
residuals: Residuals,
|
||||
residuals: ResidualBlocks,
|
||||
x: DVector,
|
||||
}
|
||||
|
||||
@ -412,18 +646,18 @@ pub fn test() {
|
||||
|
||||
let nvar = 4;
|
||||
let nres = 3;
|
||||
let residuals = Residuals::new(
|
||||
vec![
|
||||
PointDistance(SVector::<2>::new(1., 0.), 0.75).into_indexed([0, 1], [0]),
|
||||
PointPointDistance(1.0).into_indexed([0, 1, 2, 3], [2]),
|
||||
PointDistance(SVector::<2>::new(0., 0.), 1.25).into_indexed([2, 3], [1]),
|
||||
],
|
||||
(nres, nvar),
|
||||
);
|
||||
// let residuals = CostFunction::new(
|
||||
// vec![
|
||||
// PointDistance(SVector::<2>::new(1., 0.), 0.75).into_indexed([0, 1], [0]),
|
||||
// PointPointDistance(1.0).into_indexed([0, 1, 2, 3], [2]),
|
||||
// PointDistance(SVector::<2>::new(0., 0.), 1.25).into_indexed([2, 3], [1]),
|
||||
// ],
|
||||
// (nres, nvar),
|
||||
// );
|
||||
|
||||
let dist = rand_distr::Normal::new(0., 1.0).unwrap();
|
||||
let prob = Problem {
|
||||
residuals,
|
||||
residuals: todo!(),
|
||||
x: DVector::from_distribution_generic(
|
||||
Dyn(nvar),
|
||||
Const::<1>,
|
||||
|
||||
173
src/optimization/params.rs
Normal file
173
src/optimization/params.rs
Normal file
@ -0,0 +1,173 @@
|
||||
use nalgebra::{SMatrix, SVector};
|
||||
|
||||
use crate::geometry::Scalar;
|
||||
|
||||
use super::Idx;
|
||||
|
||||
pub struct Parameters {
|
||||
values: Vec<Scalar>,
|
||||
}
|
||||
|
||||
pub trait Param {
|
||||
const SIZE: usize;
|
||||
type Jacobian<const N: usize>;
|
||||
|
||||
fn from_array(array: [Scalar; Self::SIZE]) -> Self;
|
||||
|
||||
fn from_slice(slice: &[Scalar]) -> Self
|
||||
where
|
||||
[Scalar; Self::SIZE]: Sized,
|
||||
Self: Sized,
|
||||
{
|
||||
let array =
|
||||
TryInto::<[Scalar; Self::SIZE]>::try_into(slice).expect("Slice has wrong length");
|
||||
Self::from_array(array)
|
||||
}
|
||||
|
||||
fn copy_to_array(&self, array: &mut [Scalar; Self::SIZE]);
|
||||
|
||||
fn copy_to_slice(&self, slice: &mut [Scalar])
|
||||
where
|
||||
[Scalar; Self::SIZE]: Sized,
|
||||
{
|
||||
let array =
|
||||
TryInto::<&mut [Scalar; Self::SIZE]>::try_into(slice).expect("Slice has wrong length");
|
||||
self.copy_to_array(array)
|
||||
}
|
||||
}
|
||||
|
||||
trait WithJacobian<const N: usize> {
|
||||
type Jacobian<const M: usize>;
|
||||
}
|
||||
|
||||
trait Jacobian {}
|
||||
|
||||
impl Param for Scalar {
|
||||
const SIZE: usize = 1;
|
||||
|
||||
type Jacobian<const N: usize> = SVector<Scalar, N>;
|
||||
|
||||
fn from_slice(slice: &[Scalar]) -> Self {
|
||||
let [x] = slice else {
|
||||
panic!("Slice has wrong length");
|
||||
};
|
||||
*x
|
||||
}
|
||||
|
||||
fn copy_to_slice(&self, slice: &mut [Scalar]) {
|
||||
let [x] = slice else {
|
||||
panic!("Slice has wrong length");
|
||||
};
|
||||
*x = *self;
|
||||
}
|
||||
|
||||
fn copy_to_array(&self, array: &mut [Scalar; Self::SIZE]) {
|
||||
array[0] = *self;
|
||||
}
|
||||
|
||||
fn from_array(array: [Scalar; Self::SIZE]) -> Self {
|
||||
array[0]
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ReadParams<I> {
|
||||
fn read(params: &Parameters, index: I) -> Self;
|
||||
}
|
||||
|
||||
// impl ReadParams<Idx> for Scalar {
|
||||
// fn read(params: &Parameters, index: Idx) -> Self {
|
||||
// params.values[index as usize]
|
||||
// }
|
||||
// }
|
||||
|
||||
impl<const M: usize> Param for [Scalar; M] {
|
||||
const SIZE: usize = M;
|
||||
type Jacobian<const N: usize> = SMatrix<Scalar, M, N>;
|
||||
|
||||
fn from_array(array: [Scalar; Self::SIZE]) -> Self {
|
||||
unsafe { std::mem::transmute_copy(&array) }
|
||||
}
|
||||
|
||||
fn copy_to_array(&self, slice: &mut [Scalar; Self::SIZE]) {
|
||||
slice.copy_from_slice(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<const M: usize> Param for SMatrix<Scalar, M, 1> {
|
||||
const SIZE: usize = M;
|
||||
type Jacobian<const N: usize> = SMatrix<Scalar, M, N>;
|
||||
|
||||
fn from_slice(slice: &[Scalar]) -> Self {
|
||||
Self::from_column_slice(slice)
|
||||
}
|
||||
|
||||
fn from_array(array: [Scalar; Self::SIZE]) -> Self {
|
||||
Self::from_slice(&array)
|
||||
}
|
||||
|
||||
fn copy_to_slice(&self, slice: &mut [Scalar]) {
|
||||
slice.copy_from_slice(self.as_slice())
|
||||
}
|
||||
|
||||
fn copy_to_array(&self, array: &mut [Scalar; Self::SIZE]) {
|
||||
array.copy_from_slice(self.as_slice())
|
||||
}
|
||||
}
|
||||
|
||||
impl<P: Param> ReadParams<Idx> for P
|
||||
where
|
||||
[Scalar; P::SIZE]: Sized,
|
||||
{
|
||||
fn read(params: &Parameters, index: Idx) -> Self {
|
||||
Self::from_slice(¶ms.values[index as usize..][..Self::SIZE])
|
||||
}
|
||||
}
|
||||
|
||||
// impl<const R: usize, const C: usize> ReadParams<Idx> for SMatrix<Scalar, R, C> {
|
||||
// fn read(params: &Parameters, index: Idx) -> Self {
|
||||
// SMatrix::from_column_slice(¶ms.values[index as usize..][..R * C])
|
||||
// }
|
||||
// }
|
||||
|
||||
// TODO: Implement for tuples of arbitrary length using a macro
|
||||
impl<T1, T2, I1, I2> ReadParams<(I1, I2)> for (T1, T2)
|
||||
where
|
||||
T1: ReadParams<I1>,
|
||||
T2: ReadParams<I2>,
|
||||
{
|
||||
fn read(params: &Parameters, index: (I1, I2)) -> Self {
|
||||
(T1::read(params, index.0), T2::read(params, index.1))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T1, T2, T3, I1, I2, I3> ReadParams<(I1, I2, I3)> for (T1, T2, T3)
|
||||
where
|
||||
T1: ReadParams<I1>,
|
||||
T2: ReadParams<I2>,
|
||||
T3: ReadParams<I3>,
|
||||
{
|
||||
fn read(params: &Parameters, index: (I1, I2, I3)) -> Self {
|
||||
(
|
||||
T1::read(params, index.0),
|
||||
T2::read(params, index.1),
|
||||
T3::read(params, index.2),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T1, T2, T3, T4, I1, I2, I3, I4> ReadParams<(I1, I2, I3, I4)> for (T1, T2, T3, T4)
|
||||
where
|
||||
T1: ReadParams<I1>,
|
||||
T2: ReadParams<I2>,
|
||||
T3: ReadParams<I3>,
|
||||
T4: ReadParams<I4>,
|
||||
{
|
||||
fn read(params: &Parameters, index: (I1, I2, I3, I4)) -> Self {
|
||||
(
|
||||
T1::read(params, index.0),
|
||||
T2::read(params, index.1),
|
||||
T3::read(params, index.2),
|
||||
T4::read(params, index.3),
|
||||
)
|
||||
}
|
||||
}
|
||||
64
src/optimization/residuals.rs
Normal file
64
src/optimization/residuals.rs
Normal file
@ -0,0 +1,64 @@
|
||||
use nalgebra::SMatrix;
|
||||
|
||||
use crate::geometry::Scalar;
|
||||
|
||||
pub trait WriteResiduals<Index> {
|
||||
fn write(&self, index: Index, residuals: &mut [Scalar]);
|
||||
}
|
||||
|
||||
impl<I> WriteResiduals<I> for Scalar
|
||||
where
|
||||
usize: From<I>,
|
||||
{
|
||||
fn write(&self, index: I, residuals: &mut [Scalar]) {
|
||||
residuals[usize::from(index)] = *self;
|
||||
}
|
||||
}
|
||||
|
||||
impl<I, const R: usize, const C: usize> WriteResiduals<I> for SMatrix<Scalar, R, C>
|
||||
where
|
||||
usize: From<I>,
|
||||
{
|
||||
fn write(&self, index: I, residuals: &mut [Scalar]) {
|
||||
residuals[usize::from(index)..][..R * C].copy_from_slice(self.as_slice());
|
||||
}
|
||||
}
|
||||
|
||||
impl<T1, T2, I1, I2> WriteResiduals<(I1, I2)> for (T1, T2)
|
||||
where
|
||||
T1: WriteResiduals<I1>,
|
||||
T2: WriteResiduals<I2>,
|
||||
{
|
||||
fn write(&self, index: (I1, I2), residuals: &mut [Scalar]) {
|
||||
self.0.write(index.0, residuals);
|
||||
self.1.write(index.1, residuals);
|
||||
}
|
||||
}
|
||||
|
||||
impl<T1, T2, T3, I1, I2, I3> WriteResiduals<(I1, I2, I3)> for (T1, T2, T3)
|
||||
where
|
||||
T1: WriteResiduals<I1>,
|
||||
T2: WriteResiduals<I2>,
|
||||
T3: WriteResiduals<I3>,
|
||||
{
|
||||
fn write(&self, index: (I1, I2, I3), residuals: &mut [Scalar]) {
|
||||
self.0.write(index.0, residuals);
|
||||
self.1.write(index.1, residuals);
|
||||
self.2.write(index.2, residuals);
|
||||
}
|
||||
}
|
||||
|
||||
impl<T1, T2, T3, T4, I1, I2, I3, I4> WriteResiduals<(I1, I2, I3, I4)> for (T1, T2, T3, T4)
|
||||
where
|
||||
T1: WriteResiduals<I1>,
|
||||
T2: WriteResiduals<I2>,
|
||||
T3: WriteResiduals<I3>,
|
||||
T4: WriteResiduals<I4>,
|
||||
{
|
||||
fn write(&self, index: (I1, I2, I3, I4), residuals: &mut [Scalar]) {
|
||||
self.0.write(index.0, residuals);
|
||||
self.1.write(index.1, residuals);
|
||||
self.2.write(index.2, residuals);
|
||||
self.3.write(index.3, residuals);
|
||||
}
|
||||
}
|
||||
165
src/ui.rs
165
src/ui.rs
@ -1,11 +1,12 @@
|
||||
use std::ops::Deref;
|
||||
use std::{cmp::Ordering, ops::Deref};
|
||||
|
||||
use bevy_ecs::{prelude::*, query::WorldQuery, system::ReadOnlySystem};
|
||||
use eframe::{
|
||||
egui::{self, CursorIcon, Painter, Response, Sense},
|
||||
egui::{self, plot::Plot, CursorIcon, Painter, Response, ScrollArea, Sense},
|
||||
emath::RectTransform,
|
||||
epaint::{Color32, Hsva, Pos2, Rect, Stroke, Vec2},
|
||||
epaint::{Color32, Hsva, Pos2, Rect, Rounding, Stroke, Vec2},
|
||||
};
|
||||
use tracing::info;
|
||||
|
||||
use crate::geometry::{self};
|
||||
|
||||
@ -25,12 +26,46 @@ impl Default for Tool {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Resource, Clone)]
|
||||
struct Viewport {
|
||||
scale: f32,
|
||||
center: Pos2,
|
||||
}
|
||||
|
||||
impl Default for Viewport {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
scale: 1.0,
|
||||
center: Pos2::ZERO,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn update_viewport(mut viewport: ResMut<Viewport>, ui: Res<UiRes>) {
|
||||
// if ui.response.hovered() && ui.response.ctx.input(|i| i.modifiers.command_only()) {
|
||||
if ui.response.hovered() {
|
||||
let delta = ui.response.ctx.input(|i| i.scroll_delta);
|
||||
if delta != Vec2::ZERO {
|
||||
// viewport.scale *= 1.0 + delta.y * 0.05;
|
||||
viewport.scale *= (1.0f32 - 0.005).powf(-delta.y);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Resource)]
|
||||
struct UiRes {
|
||||
response: Response,
|
||||
painter: Painter,
|
||||
to_screen: ToScreen,
|
||||
}
|
||||
|
||||
impl UiRes {
|
||||
fn to_screen(&self) -> ToScreen {
|
||||
self.to_screen
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct ToScreen(RectTransform);
|
||||
|
||||
impl ToScreen {
|
||||
@ -64,15 +99,36 @@ fn color_for_var_status(status: geometry::VarStatus) -> Hsva {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default, Resource)]
|
||||
struct SelectionState {
|
||||
start_point: Option<Pos2>,
|
||||
end_point: Option<Pos2>,
|
||||
}
|
||||
|
||||
impl SelectionState {
|
||||
fn selection_rect(&self) -> Option<Rect> {
|
||||
if let (Some(start), Some(end)) = (self.start_point, self.end_point) {
|
||||
Some(Rect::from_two_pos(start, end))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const POINT_RADIUS: f32 = 3.0;
|
||||
|
||||
fn update_hover_point(
|
||||
ui: Res<UiRes>,
|
||||
points: Query<(geometry::PointId, &geometry::ComputedPointPos)>,
|
||||
sel: Res<SelectionState>,
|
||||
mut commands: Commands,
|
||||
) {
|
||||
points.for_each(|(id, pos)| {
|
||||
let hovered = if let Some(hover_pos) = ui.response.hover_pos() {
|
||||
let hovered = if let Some(sel_rect) = sel.selection_rect() {
|
||||
let center = ui.to_screen.transform_pos(pos);
|
||||
|
||||
sel_rect.contains(center)
|
||||
} else if let Some(hover_pos) = ui.response.hover_pos() {
|
||||
let center = ui.to_screen.transform_pos(pos);
|
||||
|
||||
(hover_pos - center).length() < (POINT_RADIUS * 3.)
|
||||
@ -91,15 +147,18 @@ fn update_hover_point(
|
||||
fn update_hover_line(
|
||||
ui: Res<UiRes>,
|
||||
lines: Query<(Entity, &geometry::ComputedLinePos)>,
|
||||
sel: Res<SelectionState>,
|
||||
mut commands: Commands,
|
||||
) {
|
||||
lines.for_each(|(id, pos)| {
|
||||
let hovered = if let Some(hover_pos) = ui.response.hover_pos() {
|
||||
let points = [
|
||||
ui.to_screen.transform_pos(&pos.start),
|
||||
ui.to_screen.transform_pos(&pos.end),
|
||||
];
|
||||
|
||||
let hovered = if let Some(sel_rect) = sel.selection_rect() {
|
||||
// TODO: check if line intersects rect
|
||||
sel_rect.contains(points[0]) || sel_rect.contains(points[1])
|
||||
} else if let Some(hover_pos) = ui.response.hover_pos() {
|
||||
let b = points[1] - points[0];
|
||||
let a = hover_pos - points[0];
|
||||
let p = a.dot(b) / b.dot(b);
|
||||
@ -121,27 +180,44 @@ fn select_tool(
|
||||
ui: Res<UiRes>,
|
||||
hovered: Query<Entity, With<Hovered>>,
|
||||
selected: Query<Entity, With<Selected>>,
|
||||
mut sel: ResMut<SelectionState>,
|
||||
mut commands: Commands,
|
||||
) {
|
||||
ui.response.ctx.output_mut(|output| {
|
||||
output.cursor_icon = if !hovered.is_empty() {
|
||||
output.cursor_icon = if sel.selection_rect().is_some() {
|
||||
CursorIcon::Grabbing
|
||||
} else if !hovered.is_empty() {
|
||||
CursorIcon::PointingHand
|
||||
} else {
|
||||
CursorIcon::Default
|
||||
}
|
||||
});
|
||||
|
||||
if ui.response.clicked() {
|
||||
if ui.response.drag_released() {
|
||||
if !ui.response.ctx.input(|input| input.modifiers.shift) {
|
||||
selected.for_each(|selected| {
|
||||
commands.entity(selected).remove::<Selected>();
|
||||
});
|
||||
}
|
||||
// TODO: choose which to select
|
||||
if let Some(hovered) = hovered.iter().next() {
|
||||
commands.entity(hovered).insert(Selected);
|
||||
for entity in hovered.iter() {
|
||||
commands.entity(entity).insert(Selected);
|
||||
// Only select one if not dragging
|
||||
// if sel.selection_rect().is_none() {
|
||||
// break;
|
||||
// }
|
||||
}
|
||||
}
|
||||
if ui.response.dragged() {
|
||||
// TODO: choose which to select
|
||||
if sel.start_point.is_none() {
|
||||
sel.start_point = ui.response.hover_pos();
|
||||
} else if ui.response.hover_pos() != sel.start_point {
|
||||
sel.end_point = ui.response.hover_pos();
|
||||
}
|
||||
} else {
|
||||
sel.start_point = None;
|
||||
sel.end_point = None;
|
||||
}
|
||||
|
||||
if ui
|
||||
.response
|
||||
@ -162,6 +238,17 @@ fn select_tool(
|
||||
}
|
||||
}
|
||||
|
||||
fn paint_selection(ui: Res<UiRes>, sel: Res<SelectionState>) {
|
||||
if let Some(rect) = sel.selection_rect() {
|
||||
ui.painter.rect(
|
||||
rect,
|
||||
Rounding::none(),
|
||||
Hsva::new(0.0, 0.0, 0.5, 0.5),
|
||||
Stroke::new(1.0, Color32::WHITE),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct DragDelta(Vec2);
|
||||
|
||||
@ -199,7 +286,7 @@ fn move_tool(
|
||||
selected.for_each(|selected| {
|
||||
commands.entity(selected.0).remove::<Selected>();
|
||||
});
|
||||
drag_delta.0 = Vec2::ZERO;
|
||||
(*drag_delta).0 = Vec2::ZERO;
|
||||
None
|
||||
} else {
|
||||
selected.iter().next()
|
||||
@ -208,7 +295,7 @@ fn move_tool(
|
||||
if let Some((_, point)) = selected {
|
||||
if response.drag_started() {
|
||||
let drag_point_pos = point_pos.get(point).unwrap();
|
||||
drag_delta.0 = hover_pos - to_screen.transform_pos(&drag_point_pos);
|
||||
(*drag_delta).0 = hover_pos - to_screen.transform_pos(&drag_point_pos);
|
||||
}
|
||||
let move_to = to_screen.inverse_transform_to_point(hover_pos - drag_delta.0);
|
||||
point_pos.set(point, move_to).unwrap();
|
||||
@ -330,8 +417,11 @@ fn paint_lines(
|
||||
|
||||
let mut color = color_for_var_status(pos.status);
|
||||
color.v -= 0.6;
|
||||
if hovered.contains(id) || selected.contains(id) {
|
||||
if selected.contains(id) {
|
||||
color.s -= 0.8;
|
||||
color.v += 0.2;
|
||||
} else if hovered.contains(id) {
|
||||
color.s -= 0.4;
|
||||
}
|
||||
|
||||
let stroke = Stroke::new(2.0, color);
|
||||
@ -350,9 +440,11 @@ fn paint_points(
|
||||
let center = ui.to_screen.transform_pos(pos);
|
||||
|
||||
let color = color_for_var_status(pos.status);
|
||||
let stroke = if selected.contains(id) || hovered.contains(id) {
|
||||
// color.s -= 0.8;
|
||||
let stroke = if selected.contains(id) {
|
||||
Stroke::new(1.0, Color32::WHITE)
|
||||
} else if hovered.contains(id) {
|
||||
// color.s -= 0.8;
|
||||
Stroke::new(1.0, Color32::WHITE.linear_multiply(0.1))
|
||||
} else {
|
||||
Stroke::default()
|
||||
};
|
||||
@ -414,6 +506,7 @@ pub enum ShowEntitiesStage {
|
||||
PostTools,
|
||||
ToolsFlush,
|
||||
Paint,
|
||||
PostPaint,
|
||||
}
|
||||
|
||||
impl Default for ShowEntitiesSchedule {
|
||||
@ -430,11 +523,14 @@ impl Default for ShowEntitiesSchedule {
|
||||
ShowEntitiesStage::ToolsFlush,
|
||||
ShowEntitiesStage::PostTools,
|
||||
ShowEntitiesStage::Paint,
|
||||
ShowEntitiesStage::PostPaint,
|
||||
)
|
||||
.chain(),
|
||||
)
|
||||
.add_system(update_viewport.in_base_set(ShowEntitiesStage::Update))
|
||||
.add_systems(
|
||||
(geometry::update_point_pos, geometry::update_line_pos)
|
||||
.chain()
|
||||
.in_base_set(ShowEntitiesStage::Update),
|
||||
)
|
||||
.add_system(apply_system_buffers.in_base_set(ShowEntitiesStage::UpdateFlush))
|
||||
@ -456,6 +552,7 @@ impl Default for ShowEntitiesSchedule {
|
||||
.add_system(geometry::remove_dangling_lines.in_base_set(ShowEntitiesStage::PostTools))
|
||||
.add_systems(
|
||||
(geometry::update_point_pos, geometry::update_line_pos)
|
||||
.chain()
|
||||
.in_base_set(ShowEntitiesStage::PostTools),
|
||||
)
|
||||
.add_system(apply_system_buffers.in_base_set(ShowEntitiesStage::ToolsFlush))
|
||||
@ -463,7 +560,8 @@ impl Default for ShowEntitiesSchedule {
|
||||
(paint_lines, paint_points)
|
||||
.chain()
|
||||
.in_base_set(ShowEntitiesStage::Paint),
|
||||
);
|
||||
)
|
||||
.add_system((paint_selection).in_base_set(ShowEntitiesStage::PostPaint));
|
||||
Self(schedule)
|
||||
}
|
||||
}
|
||||
@ -474,13 +572,18 @@ fn central_panel(world: &mut World, mut schedule: Local<ShowEntitiesSchedule>) {
|
||||
.frame(egui::Frame::none())
|
||||
.show(&ctx, |ui| {
|
||||
let sense = match *world.resource::<Tool>() {
|
||||
Tool::Move => Sense::drag(),
|
||||
Tool::Select | Tool::AddPoint | Tool::AddLine | Tool::AddRelation => Sense::click(),
|
||||
Tool::Select | Tool::Move => Sense::drag(),
|
||||
Tool::AddPoint | Tool::AddLine | Tool::AddRelation => Sense::click(),
|
||||
};
|
||||
info!("render");
|
||||
|
||||
let (response, painter) = ui.allocate_painter(ui.available_size(), sense);
|
||||
let viewport = world.get_resource::<Viewport>().unwrap().clone();
|
||||
let to_screen = ToScreen(RectTransform::from_to(
|
||||
Rect::from_center_size(Pos2::ZERO, response.rect.size() / 2.0),
|
||||
Rect::from_center_size(
|
||||
viewport.center,
|
||||
response.rect.size() * (viewport.scale / 2.0),
|
||||
),
|
||||
response.rect,
|
||||
));
|
||||
world.insert_resource(UiRes {
|
||||
@ -496,10 +599,10 @@ fn central_panel(world: &mut World, mut schedule: Local<ShowEntitiesSchedule>) {
|
||||
}
|
||||
|
||||
#[derive(WorldQuery)]
|
||||
struct SelectableEntity<'a> {
|
||||
struct SelectableEntity {
|
||||
id: Entity,
|
||||
point: Option<&'a geometry::Point>,
|
||||
line: Option<&'a geometry::Line>,
|
||||
point: Option<&'static geometry::Point>,
|
||||
line: Option<&'static geometry::Line>,
|
||||
}
|
||||
|
||||
fn side_panel_ui(
|
||||
@ -508,10 +611,14 @@ fn side_panel_ui(
|
||||
tool: Res<Tool>,
|
||||
) {
|
||||
let tool = *tool;
|
||||
ui.vertical(|ui| match tool {
|
||||
ui.heading("Selection:");
|
||||
ScrollArea::vertical().show(ui, |ui| {
|
||||
match tool {
|
||||
Tool::Select => {
|
||||
let mut count = 0;
|
||||
selected.for_each(|sel| {
|
||||
let mut selection: Vec<_> = selected.iter().collect();
|
||||
selection.sort_by(|a, b| a.id.index().cmp(&b.id.index()));
|
||||
selection.iter().for_each(|sel| {
|
||||
count += 1;
|
||||
if sel.point.is_some() {
|
||||
ui.label(format!("Selected point {}", sel.id.index()));
|
||||
@ -546,7 +653,9 @@ fn side_panel_ui(
|
||||
Tool::AddRelation => {
|
||||
ui.label("Click to add a relation");
|
||||
}
|
||||
};
|
||||
});
|
||||
ui.separator();
|
||||
}
|
||||
|
||||
fn side_panel(
|
||||
@ -569,6 +678,12 @@ fn bottom_panel(ctx: Res<ContextRes>) {
|
||||
});
|
||||
}
|
||||
|
||||
pub fn init(world: &mut World) {
|
||||
world.init_resource::<Tool>();
|
||||
world.init_resource::<SelectionState>();
|
||||
world.init_resource::<Viewport>();
|
||||
}
|
||||
|
||||
pub fn add_to_schedule(schedule: &mut Schedule) {
|
||||
schedule.add_systems((prepare, toolbar, side_panel, bottom_panel, central_panel).chain());
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user