#[cfg(feature = "simd")]
use super::simd::{calculate_leaf_force_simd, SimdPositions};
use crate::graph::Node;
use nalgebra::Vector2;
#[derive(Debug, Clone, Copy)]
pub struct Rect {
pub x: f64,
pub y: f64,
pub width: f64,
pub height: f64,
}
impl Rect {
pub fn new(x: f64, y: f64, width: f64, height: f64) -> Self {
Self {
x,
y,
width,
height,
}
}
pub fn contains(&self, x: f64, y: f64) -> bool {
x >= self.x && x < self.x + self.width && y >= self.y && y < self.y + self.height
}
pub fn center(&self) -> (f64, f64) {
(self.x + self.width / 2.0, self.y + self.height / 2.0)
}
pub fn size(&self) -> f64 {
self.width.max(self.height)
}
}
#[derive(Debug, Clone)]
pub struct NodeRef {
pub id: String,
pub position: Vector2<f64>,
pub mass: f64,
}
impl NodeRef {
pub fn from_node(node: &Node) -> Self {
Self {
id: node.id.clone(),
position: Vector2::new(node.visual.x, node.visual.y),
mass: node.visual.weight.max(1.0), }
}
}
#[derive(Debug)]
pub struct QuadTree {
bounds: Rect,
nodes: Vec<NodeRef>,
children: Option<Box<[QuadTree; 4]>>,
center_of_mass: Vector2<f64>,
total_mass: f64,
capacity: usize,
}
impl QuadTree {
pub fn new(bounds: Rect, capacity: usize) -> Self {
Self {
bounds,
nodes: Vec::new(),
children: None,
center_of_mass: Vector2::zeros(),
total_mass: 0.0,
capacity,
}
}
pub fn build(nodes: &[&Node]) -> Self {
if nodes.is_empty() {
return Self::new(Rect::new(0.0, 0.0, 1000.0, 1000.0), 1);
}
let mut min_x = f64::INFINITY;
let mut min_y = f64::INFINITY;
let mut max_x = f64::NEG_INFINITY;
let mut max_y = f64::NEG_INFINITY;
for node in nodes {
min_x = min_x.min(node.visual.x);
min_y = min_y.min(node.visual.y);
max_x = max_x.max(node.visual.x);
max_y = max_y.max(node.visual.y);
}
let padding = 0.2;
let width = (max_x - min_x) * (1.0 + padding);
let height = (max_y - min_y) * (1.0 + padding);
let x = min_x - width * padding / 2.0;
let y = min_y - height * padding / 2.0;
let size = width.max(height);
let bounds = Rect::new(x, y, size, size);
let mut tree = Self::new(bounds, 1);
for node in nodes {
tree.insert(NodeRef::from_node(node));
}
tree
}
pub fn insert(&mut self, node_ref: NodeRef) {
let new_total_mass = self.total_mass + node_ref.mass;
if new_total_mass > 0.0 {
self.center_of_mass = (self.center_of_mass * self.total_mass
+ node_ref.position * node_ref.mass)
/ new_total_mass;
}
self.total_mass = new_total_mass;
let quadrant = self.get_quadrant(node_ref.position.x, node_ref.position.y);
if let Some(children) = self.children.as_mut() {
children[quadrant].insert(node_ref);
return;
}
self.nodes.push(node_ref);
if self.nodes.len() > self.capacity {
self.subdivide();
}
}
fn subdivide(&mut self) {
let half_width = self.bounds.width / 2.0;
let half_height = self.bounds.height / 2.0;
let children = Box::new([
QuadTree::new(
Rect::new(self.bounds.x, self.bounds.y, half_width, half_height),
self.capacity,
),
QuadTree::new(
Rect::new(
self.bounds.x + half_width,
self.bounds.y,
half_width,
half_height,
),
self.capacity,
),
QuadTree::new(
Rect::new(
self.bounds.x,
self.bounds.y + half_height,
half_width,
half_height,
),
self.capacity,
),
QuadTree::new(
Rect::new(
self.bounds.x + half_width,
self.bounds.y + half_height,
half_width,
half_height,
),
self.capacity,
),
]);
self.children = Some(children);
let nodes = std::mem::take(&mut self.nodes);
for node in nodes {
let quadrant = self.get_quadrant(node.position.x, node.position.y);
self.children.as_mut().unwrap()[quadrant].insert(node);
}
}
fn get_quadrant(&self, x: f64, y: f64) -> usize {
let (center_x, center_y) = self.bounds.center();
match (x >= center_x, y >= center_y) {
(false, false) => 0, (true, false) => 1, (false, true) => 2, (true, true) => 3, }
}
pub fn calculate_force(
&self,
node_pos: Vector2<f64>,
node_id: &str,
theta: f64,
strength: f64,
) -> Vector2<f64> {
if self.total_mass == 0.0 {
return Vector2::zeros();
}
if self.children.is_none() {
if self.nodes.len() == 1 && self.nodes[0].id == node_id {
return Vector2::zeros();
}
return self.calculate_exact_force(node_pos, node_id, strength);
}
let delta = self.center_of_mass - node_pos;
let distance = delta.norm();
if distance < 0.0001 {
return Vector2::zeros();
}
let s = self.bounds.size();
if s / distance < theta {
let force_magnitude = strength * self.total_mass / (distance * distance);
return delta.normalize() * force_magnitude;
}
if let Some(ref children) = self.children {
children
.iter()
.map(|child| child.calculate_force(node_pos, node_id, theta, strength))
.fold(Vector2::zeros(), |acc, f| acc + f)
} else {
Vector2::zeros()
}
}
#[allow(unsafe_code)]
fn calculate_exact_force(
&self,
node_pos: Vector2<f64>,
node_id: &str,
strength: f64,
) -> Vector2<f64> {
let sources: Vec<&NodeRef> = self
.nodes
.iter()
.filter(|node| node.id != node_id)
.collect();
if sources.is_empty() {
return Vector2::zeros();
}
#[cfg(feature = "simd")]
{
if sources.len() >= 4 {
let positions: Vec<(f64, f64, f64)> = sources
.iter()
.map(|n| (n.position.x, n.position.y, n.mass))
.collect();
let simd_sources = SimdPositions::from_positions_weighted(&positions);
#[cfg(target_arch = "wasm32")]
{
let (fx, fy) = unsafe {
calculate_leaf_force_simd(
node_pos.x as f32,
node_pos.y as f32,
&simd_sources,
strength as f32,
)
};
return Vector2::new(fx as f64, fy as f64);
}
#[cfg(not(target_arch = "wasm32"))]
{
let (fx, fy) = calculate_leaf_force_simd(
node_pos.x as f32,
node_pos.y as f32,
&simd_sources,
strength as f32,
);
return Vector2::new(fx as f64, fy as f64);
}
}
}
sources
.iter()
.map(|node| {
let delta = node_pos - node.position;
let distance_sq = delta.norm_squared();
if distance_sq < 0.0001 {
let offset_x = ((node_pos.x + node.position.x) * 7.0).sin() * 0.01;
let offset_y = ((node_pos.y + node.position.y) * 11.0).cos() * 0.01;
Vector2::new(offset_x, offset_y)
} else {
let force_magnitude = strength * node.mass / distance_sq;
delta.normalize() * force_magnitude
}
})
.fold(Vector2::zeros(), |acc, f| acc + f)
}
pub fn count_nodes(&self) -> usize {
if let Some(ref children) = self.children {
children.iter().map(|c| c.count_nodes()).sum()
} else {
self.nodes.len()
}
}
pub fn max_depth(&self) -> usize {
if let Some(ref children) = self.children {
1 + children.iter().map(|c| c.max_depth()).max().unwrap_or(0)
} else {
0
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::graph::{Node, NodeType, SemanticAttributes, VisualAttributes};
fn create_test_node(id: &str, x: f64, y: f64) -> Node {
Node {
id: id.to_string(),
label: id.to_string(),
node_type: NodeType::Class,
visual: VisualAttributes {
x,
y,
vx: 0.0,
vy: 0.0,
fixed: false,
weight: 1.0,
color: None,
visible: true,
},
semantic: SemanticAttributes::default(),
}
}
#[test]
fn test_rect_contains() {
let rect = Rect::new(0.0, 0.0, 100.0, 100.0);
assert!(rect.contains(50.0, 50.0));
assert!(rect.contains(0.0, 0.0));
assert!(!rect.contains(100.0, 100.0));
assert!(!rect.contains(-10.0, 50.0));
}
#[test]
fn test_rect_center() {
let rect = Rect::new(0.0, 0.0, 100.0, 100.0);
assert_eq!(rect.center(), (50.0, 50.0));
let rect2 = Rect::new(-50.0, -50.0, 100.0, 100.0);
assert_eq!(rect2.center(), (0.0, 0.0));
}
#[test]
fn test_quadtree_creation() {
let bounds = Rect::new(0.0, 0.0, 1000.0, 1000.0);
let tree = QuadTree::new(bounds, 1);
assert_eq!(tree.total_mass, 0.0);
assert_eq!(tree.nodes.len(), 0);
assert!(tree.children.is_none());
}
#[test]
fn test_quadtree_insert_single() {
let bounds = Rect::new(0.0, 0.0, 1000.0, 1000.0);
let mut tree = QuadTree::new(bounds, 1);
let node = create_test_node("node1", 500.0, 500.0);
tree.insert(NodeRef::from_node(&node));
assert_eq!(tree.total_mass, 1.0);
assert_eq!(tree.nodes.len(), 1);
}
#[test]
fn test_quadtree_subdivision() {
let bounds = Rect::new(0.0, 0.0, 1000.0, 1000.0);
let mut tree = QuadTree::new(bounds, 1);
let node1 = create_test_node("node1", 250.0, 250.0);
let node2 = create_test_node("node2", 750.0, 750.0);
tree.insert(NodeRef::from_node(&node1));
assert!(tree.children.is_none());
tree.insert(NodeRef::from_node(&node2));
assert!(tree.children.is_some()); assert_eq!(tree.nodes.len(), 0); }
#[test]
fn test_quadtree_build() {
let nodes = [
create_test_node("n1", 0.0, 0.0),
create_test_node("n2", 100.0, 100.0),
create_test_node("n3", 200.0, 200.0),
create_test_node("n4", 50.0, 150.0),
];
let node_refs: Vec<&Node> = nodes.iter().collect();
let tree = QuadTree::build(&node_refs);
assert_eq!(tree.count_nodes(), 4);
assert_eq!(tree.total_mass, 4.0);
}
#[test]
fn test_quadtree_center_of_mass() {
let bounds = Rect::new(0.0, 0.0, 1000.0, 1000.0);
let mut tree = QuadTree::new(bounds, 10);
let node1 = create_test_node("node1", 0.0, 0.0);
let node2 = create_test_node("node2", 100.0, 0.0);
tree.insert(NodeRef::from_node(&node1));
tree.insert(NodeRef::from_node(&node2));
assert!((tree.center_of_mass.x - 50.0).abs() < 0.001);
assert!((tree.center_of_mass.y - 0.0).abs() < 0.001);
}
#[test]
fn test_barnes_hut_force_approximation() {
let nodes = [
create_test_node("n1", 0.0, 0.0),
create_test_node("n2", 1000.0, 1000.0),
create_test_node("n3", 1010.0, 1010.0),
];
let node_refs: Vec<&Node> = nodes.iter().collect();
let tree = QuadTree::build(&node_refs);
let force = tree.calculate_force(
Vector2::new(0.0, 0.0),
"n1",
0.9, -1000.0, );
assert!(force.norm() > 0.0);
}
#[test]
fn test_avoid_self_force() {
let nodes = [
create_test_node("n1", 0.0, 0.0),
create_test_node("n2", 100.0, 100.0),
];
let node_refs: Vec<&Node> = nodes.iter().collect();
let tree = QuadTree::build(&node_refs);
let force = tree.calculate_force(Vector2::new(0.0, 0.0), "n1", 0.9, -1000.0);
assert!(force.norm() > 0.0);
}
#[test]
fn test_empty_tree_force() {
let tree = QuadTree::build(&[]);
let force = tree.calculate_force(Vector2::new(0.0, 0.0), "n1", 0.9, -1000.0);
assert_eq!(force, Vector2::zeros());
}
#[test]
fn test_large_graph_performance() {
let mut nodes = Vec::new();
for i in 0..100 {
let x = (i as f64 * 37.0) % 1000.0;
let y = (i as f64 * 73.0) % 1000.0;
nodes.push(create_test_node(&format!("n{}", i), x, y));
}
let node_refs: Vec<&Node> = nodes.iter().collect();
let tree = QuadTree::build(&node_refs);
assert_eq!(tree.count_nodes(), 100);
let depth = tree.max_depth();
assert!(depth > 0 && depth < 20);
}
}