242 lines
8.6 KiB
Rust
242 lines
8.6 KiB
Rust
use super::InteractionGroups;
|
|
use crate::math::Real;
|
|
use na::DVector;
|
|
|
|
pub(crate) trait ConstraintTypes {
|
|
type OneBody;
|
|
type TwoBodies;
|
|
type GenericOneBody;
|
|
type GenericTwoBodies;
|
|
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
type SimdOneBody;
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
type SimdTwoBodies;
|
|
|
|
type BuilderOneBody;
|
|
type BuilderTwoBodies;
|
|
|
|
type GenericBuilderOneBody;
|
|
type GenericBuilderTwoBodies;
|
|
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
type SimdBuilderOneBody;
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
type SimdBuilderTwoBodies;
|
|
}
|
|
|
|
pub enum AnyConstraintMut<'a, Constraints: ConstraintTypes> {
|
|
OneBody(&'a mut Constraints::OneBody),
|
|
TwoBodies(&'a mut Constraints::TwoBodies),
|
|
GenericOneBody(&'a mut Constraints::GenericOneBody),
|
|
GenericTwoBodies(&'a mut Constraints::GenericTwoBodies),
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
SimdOneBody(&'a mut Constraints::SimdOneBody),
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
SimdTwoBodies(&'a mut Constraints::SimdTwoBodies),
|
|
}
|
|
|
|
pub(crate) struct SolverConstraintsSet<Constraints: ConstraintTypes> {
|
|
pub generic_jacobians: DVector<Real>,
|
|
pub two_body_interactions: Vec<usize>,
|
|
pub one_body_interactions: Vec<usize>,
|
|
pub generic_two_body_interactions: Vec<usize>,
|
|
pub generic_one_body_interactions: Vec<usize>,
|
|
pub interaction_groups: InteractionGroups,
|
|
pub one_body_interaction_groups: InteractionGroups,
|
|
|
|
pub velocity_constraints: Vec<Constraints::TwoBodies>,
|
|
pub generic_velocity_constraints: Vec<Constraints::GenericTwoBodies>,
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
pub simd_velocity_constraints: Vec<Constraints::SimdTwoBodies>,
|
|
pub velocity_one_body_constraints: Vec<Constraints::OneBody>,
|
|
pub generic_velocity_one_body_constraints: Vec<Constraints::GenericOneBody>,
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
pub simd_velocity_one_body_constraints: Vec<Constraints::SimdOneBody>,
|
|
|
|
pub velocity_constraints_builder: Vec<Constraints::BuilderTwoBodies>,
|
|
pub generic_velocity_constraints_builder: Vec<Constraints::GenericBuilderTwoBodies>,
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
pub simd_velocity_constraints_builder: Vec<Constraints::SimdBuilderTwoBodies>,
|
|
pub velocity_one_body_constraints_builder: Vec<Constraints::BuilderOneBody>,
|
|
pub generic_velocity_one_body_constraints_builder: Vec<Constraints::GenericBuilderOneBody>,
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
pub simd_velocity_one_body_constraints_builder: Vec<Constraints::SimdBuilderOneBody>,
|
|
}
|
|
|
|
impl<Constraints: ConstraintTypes> SolverConstraintsSet<Constraints> {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
generic_jacobians: DVector::zeros(0),
|
|
two_body_interactions: vec![],
|
|
one_body_interactions: vec![],
|
|
generic_two_body_interactions: vec![],
|
|
generic_one_body_interactions: vec![],
|
|
interaction_groups: InteractionGroups::new(),
|
|
one_body_interaction_groups: InteractionGroups::new(),
|
|
|
|
velocity_constraints: vec![],
|
|
generic_velocity_constraints: vec![],
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
simd_velocity_constraints: vec![],
|
|
velocity_one_body_constraints: vec![],
|
|
generic_velocity_one_body_constraints: vec![],
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
simd_velocity_one_body_constraints: vec![],
|
|
|
|
velocity_constraints_builder: vec![],
|
|
generic_velocity_constraints_builder: vec![],
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
simd_velocity_constraints_builder: vec![],
|
|
velocity_one_body_constraints_builder: vec![],
|
|
generic_velocity_one_body_constraints_builder: vec![],
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
simd_velocity_one_body_constraints_builder: vec![],
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)] // Useful for debuging.
|
|
pub fn print_counts(&self) {
|
|
println!("Solver constraints:");
|
|
println!(
|
|
"|__ two_body_interactions: {}",
|
|
self.two_body_interactions.len()
|
|
);
|
|
println!(
|
|
"|__ one_body_interactions: {}",
|
|
self.one_body_interactions.len()
|
|
);
|
|
println!(
|
|
"|__ generic_two_body_interactions: {}",
|
|
self.generic_two_body_interactions.len()
|
|
);
|
|
println!(
|
|
"|__ generic_one_body_interactions: {}",
|
|
self.generic_one_body_interactions.len()
|
|
);
|
|
|
|
println!(
|
|
"|__ velocity_constraints: {}",
|
|
self.velocity_constraints.len()
|
|
);
|
|
println!(
|
|
"|__ generic_velocity_constraints: {}",
|
|
self.generic_velocity_constraints.len()
|
|
);
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
println!(
|
|
"|__ simd_velocity_constraints: {}",
|
|
self.simd_velocity_constraints.len()
|
|
);
|
|
println!(
|
|
"|__ velocity_one_body_constraints: {}",
|
|
self.velocity_one_body_constraints.len()
|
|
);
|
|
println!(
|
|
"|__ generic_velocity_one_body_constraints: {}",
|
|
self.generic_velocity_one_body_constraints.len()
|
|
);
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
println!(
|
|
"|__ simd_velocity_one_body_constraints: {}",
|
|
self.simd_velocity_one_body_constraints.len()
|
|
);
|
|
|
|
println!(
|
|
"|__ velocity_constraints_builder: {}",
|
|
self.velocity_constraints_builder.len()
|
|
);
|
|
println!(
|
|
"|__ generic_velocity_constraints_builder: {}",
|
|
self.generic_velocity_constraints_builder.len()
|
|
);
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
println!(
|
|
"|__ simd_velocity_constraints_builder: {}",
|
|
self.simd_velocity_constraints_builder.len()
|
|
);
|
|
println!(
|
|
"|__ velocity_one_body_constraints_builder: {}",
|
|
self.velocity_one_body_constraints_builder.len()
|
|
);
|
|
println!(
|
|
"|__ generic_velocity_one_body_constraints_builder: {}",
|
|
self.generic_velocity_one_body_constraints_builder.len()
|
|
);
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
println!(
|
|
"|__ simd_velocity_one_body_constraints_builder: {}",
|
|
self.simd_velocity_one_body_constraints_builder.len()
|
|
);
|
|
}
|
|
|
|
pub fn clear_constraints(&mut self) {
|
|
self.generic_jacobians.fill(0.0);
|
|
self.velocity_constraints.clear();
|
|
self.velocity_one_body_constraints.clear();
|
|
self.generic_velocity_constraints.clear();
|
|
self.generic_velocity_one_body_constraints.clear();
|
|
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
{
|
|
self.simd_velocity_constraints.clear();
|
|
self.simd_velocity_one_body_constraints.clear();
|
|
}
|
|
}
|
|
|
|
pub fn clear_builders(&mut self) {
|
|
self.velocity_constraints_builder.clear();
|
|
self.velocity_one_body_constraints_builder.clear();
|
|
self.generic_velocity_constraints_builder.clear();
|
|
self.generic_velocity_one_body_constraints_builder.clear();
|
|
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
{
|
|
self.simd_velocity_constraints_builder.clear();
|
|
self.simd_velocity_one_body_constraints_builder.clear();
|
|
}
|
|
}
|
|
|
|
// Returns the generic jacobians and a mutable iterator through all the constraints.
|
|
pub fn iter_constraints_mut(
|
|
&mut self,
|
|
) -> (
|
|
&DVector<Real>,
|
|
impl Iterator<Item = AnyConstraintMut<Constraints>>,
|
|
) {
|
|
let jac = &self.generic_jacobians;
|
|
let a = self
|
|
.velocity_constraints
|
|
.iter_mut()
|
|
.map(AnyConstraintMut::TwoBodies);
|
|
let b = self
|
|
.generic_velocity_constraints
|
|
.iter_mut()
|
|
.map(AnyConstraintMut::GenericTwoBodies);
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
let c = self
|
|
.simd_velocity_constraints
|
|
.iter_mut()
|
|
.map(AnyConstraintMut::SimdTwoBodies);
|
|
let d = self
|
|
.velocity_one_body_constraints
|
|
.iter_mut()
|
|
.map(AnyConstraintMut::OneBody);
|
|
let e = self
|
|
.generic_velocity_one_body_constraints
|
|
.iter_mut()
|
|
.map(AnyConstraintMut::GenericOneBody);
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
let f = self
|
|
.simd_velocity_one_body_constraints
|
|
.iter_mut()
|
|
.map(AnyConstraintMut::SimdOneBody);
|
|
|
|
#[cfg(feature = "simd-is-enabled")]
|
|
return (jac, a.chain(b).chain(c).chain(d).chain(e).chain(f));
|
|
|
|
#[cfg(not(feature = "simd-is-enabled"))]
|
|
return (jac, a.chain(b).chain(d).chain(e));
|
|
}
|
|
}
|