graph-force/src/spring_model.rs

87 lines
2.2 KiB
Rust

use crate::graph::{EdgeMatrix, Node, NodeVector};
pub struct MyModel {
edges: EdgeMatrix,
size: usize,
opt_dist: f32,
c: f32,
dc: f32,
}
impl MyModel {
pub fn new(edges: EdgeMatrix, size: usize, iterations: usize) -> MyModel {
let opt_dist = 1.0;
let c = 0.1;
MyModel {
edges,
size,
opt_dist,
c: c,
dc: c / ((iterations + 1) as f32),
}
}
pub fn prepare(&mut self, _nodes: &NodeVector) {
self.c -= self.dc;
}
pub fn step(&self, nodes: &NodeVector, i_node: usize) -> Node {
let node = nodes[i_node].read().unwrap();
let edges = self.edges.read().unwrap();
let node_x = node.x;
let node_y = node.y;
let mut sum_x = 0.0;
let mut sum_y = 0.0;
for o in 0..self.size {
if o == i_node {
continue;
}
let o_x: f32;
let o_y: f32;
{
let other = nodes[o].read().unwrap();
o_x = other.x;
o_y = other.y;
}
let d_x = o_x - node_x;
let d_y = o_y - node_y;
let dist = (d_x * d_x + d_y * d_y).sqrt();
let unit_x = d_x / dist;
let unit_y = d_y / dist;
let edge = edges[i_node][o].weight;
if edge == 0.0 {
let f_rep = dist.powi(2).recip().min(self.opt_dist);
let f_rep_x = f_rep * unit_x;
let f_rep_y = f_rep * unit_y;
sum_x -= f_rep_x;
sum_y -= f_rep_y;
} else {
let f_spring = 0.5 * (dist - self.opt_dist);
let f_spring_x = f_spring * unit_x;
let f_spring_y = f_spring * unit_y;
sum_x += f_spring_x;
sum_y += f_spring_y;
}
}
// limit the movement
// TODO: find a good upper bound
let sum_l = (sum_x * sum_x + sum_y * sum_y).sqrt().max(1e-6).recip() * self.c;
let sum_x = sum_x * sum_l;
let sum_y = sum_y * sum_l;
Node {
x: node_x + sum_x,
y: node_y + sum_y,
}
}
}