clean up and moving iter and threads to params

src/graph.rs

@@ -1,3 +1,4 @@
+use rand::Rng;
 use std::sync::{Arc, RwLock};
 pub struct Node {
     pub x: f32,
@@ -10,3 +11,26 @@ pub struct Edge {
 
 pub type EdgeMatrix = Arc<RwLock<Vec<Vec<Edge>>>>;
 pub type NodeVector = Arc<Vec<RwLock<Node>>>;
+
+pub fn new_edge_matrix(size: usize) -> EdgeMatrix {
+    let mut matrix = Vec::with_capacity(size);
+    for _ in 0..size {
+        let mut row = Vec::with_capacity(size);
+        for _ in 0..size {
+            row.push(Edge { weight: 0.0 });
+        }
+        matrix.push(row);
+    }
+    Arc::new(RwLock::new(matrix))
+}
+
+pub fn new_node_vector(size: usize) -> NodeVector {
+    let mut nodes = Vec::with_capacity(size);
+    for _ in 0..size {
+        nodes.push(RwLock::new(Node {
+            x: rand::thread_rng().gen_range(-0.5..0.5),
+            y: rand::thread_rng().gen_range(-0.5..0.5),
+        }));
+    }
+    Arc::new(nodes)
+}

src/layout.rs

@@ -1,94 +0,0 @@
-
-use rand::Rng;
-use std::sync::{Arc, RwLock};
-use std::thread;
-use crate::graph::{EdgeMatrix, NodeVector, Node, Edge};
-use crate::my_model;
-use crate::utils;
-
-
-fn nodes_list(size: usize) -> NodeVector {
-    let mut nodes = Vec::new();
-    for _ in 0..size {
-        let node = RwLock::new(Node {
-            x: rand::thread_rng().gen_range(-0.5..0.5),
-            y: rand::thread_rng().gen_range(-0.5..0.5),
-        });
-        nodes.push(node);
-    }
-    Arc::new(nodes)
-}
-
-fn connection_matrix(size: usize) -> EdgeMatrix {
-    let mut matrix = Vec::with_capacity(size);
-    for _ in 0..size {
-        let mut row = Vec::with_capacity(size);
-        for _ in 0..size {
-            row.push(Edge { weight: 0.0 });
-        }
-        matrix.push(row);
-    }
-    Arc::new(RwLock::new(matrix))
-}
-
-fn edge_matrix_from_edge_list(number_of_nodes: usize, edge_list: Vec<(u32, u32)>) -> EdgeMatrix {
-    let matrix_ptr = connection_matrix(number_of_nodes as usize);
-    {
-        let mut matrix = matrix_ptr.write().unwrap();
-        for (node_a, node_b) in edge_list {
-            matrix[node_a as usize][node_b as usize].weight = 1.0;
-            matrix[node_b as usize][node_a as usize].weight = 1.0;
-        }
-    }
-    matrix_ptr
-}
-
-pub fn layout(number_of_nodes: usize, edge_list: Vec<(u32, u32)>) -> Vec<(f32, f32)> {
-    const ITER: usize = 5000;
-    const THREADS: usize = 8;
-
-    // let edges = connection_matrix(size);
-    let edges = edge_matrix_from_edge_list(number_of_nodes, edge_list);
-    let mut nodes = nodes_list(number_of_nodes);
-    let mut nodes_next = nodes_list(number_of_nodes);
-
-    // let model = Arc::new(RwLock::new(spring_model::InitialModel::new(edges, number_of_nodes)));
-    let model = Arc::new(RwLock::new(my_model::MyModel::new(edges, number_of_nodes, ITER)));
-
-    let chunks = utils::gen_chunks(number_of_nodes, THREADS);
-    for epoch in 0..ITER {
-        model.write().unwrap().prepare(&nodes);
-        let mut handles = vec![];
-        for i in 0..THREADS {
-            let nodes = nodes.clone();
-            let nodes_next = nodes_next.clone();
-            let chunk = chunks[i].clone();
-            let model = model.clone();
-            let handle = thread::spawn(move || {
-                for n in chunk {
-                    let update = model.read().unwrap().step(&nodes,n);
-                    let mut result = nodes_next[n].write().unwrap();
-                    result.x = update.x;
-                    result.y = update.y;
-                }
-            });
-            handles.push(handle);
-        }
-
-        for handle in handles {
-            handle.join().unwrap();
-        }
-    
-        // swap nodes and nodes_next
-        let tmp = nodes.clone();
-        nodes = nodes_next.clone();
-        nodes_next = tmp.clone();
-    }
-
-    let mut result = vec![];
-    for node in nodes.iter() {
-        let node = node.read().unwrap();
-        result.push((node.x, node.y));
-    }
-    result
-}

src/lib.rs

@@ -1,16 +1,16 @@
-mod layout;
+mod runner;
 mod utils;
 mod spring_model;
-mod my_model;
 mod graph;
 
 use pyo3::prelude::*;
 
 /// Formats the sum of two numbers as string.
-#[pyfunction]
+#[pyfunction(number_of_nodes, edges, "*", iter=500, threads=0)]
-fn layout_from_edge_list(number_of_nodes: usize, edges: Vec<(u32, u32)>) -> PyResult<Vec<(f32, f32)>> {
+fn layout_from_edge_list(number_of_nodes: usize, edges: Vec<(u32, u32)>, iter: usize, threads: usize) -> PyResult<Vec<(f32, f32)>> {
+    let r = runner::Runner::new(iter, threads);
     Ok(
-        layout::layout(number_of_nodes, edges)
+        r.layout(number_of_nodes, edges)
     )
 }
 

src/my_model.rs

@@ -1,97 +0,0 @@
-use crate::graph::{EdgeMatrix, NodeVector, Node};
-
-pub struct MyModel {
-    edges: EdgeMatrix,
-    ranks: Vec<f32>,
-    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;
-        let mut ranks = vec![0.0; size];
-        {
-            let edges = edges.read().unwrap();
-            for i in 0..size {
-                ranks[i] = edges[i].iter().map(|e| e.weight).sum();
-            }
-        }
-
-        MyModel{
-            edges,
-            ranks,
-            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,
-        }        
-    }
-}

src/runner.rs

@@ -0,0 +1,80 @@
+
+use std::sync::{Arc, RwLock};
+use std::thread;
+use crate::graph::{EdgeMatrix, new_node_vector, new_edge_matrix};
+use crate::spring_model;
+use crate::utils;
+
+pub struct Runner {
+    iterations: usize,
+    threads: usize,
+}
+
+impl Runner {
+    pub fn new(iterations: usize, threads: usize) -> Self {
+        Self {
+            iterations,
+            threads,
+        }
+    }
+
+    pub fn layout(self: &Self, number_of_nodes: usize, edge_list: Vec<(u32, u32)>) -> Vec<(f32, f32)> {
+        // let edges = connection_matrix(size);
+        let edges = edge_matrix_from_edge_list(number_of_nodes, edge_list);
+        let mut nodes = new_node_vector(number_of_nodes);
+        let mut nodes_next = new_node_vector(number_of_nodes);
+
+        // let model = Arc::new(RwLock::new(spring_model::InitialModel::new(edges, number_of_nodes)));
+        let model = Arc::new(RwLock::new(spring_model::MyModel::new(edges, number_of_nodes, self.iterations)));
+
+        let chunks = utils::gen_chunks(number_of_nodes, self.threads);
+        for _epoch in 0..self.iterations {
+            model.write().unwrap().prepare(&nodes);
+            let mut handles = vec![];
+            for i in 0..self.threads {
+                let nodes = nodes.clone();
+                let nodes_next = nodes_next.clone();
+                let chunk = chunks[i].clone();
+                let model = model.clone();
+                let handle = thread::spawn(move || {
+                    for n in chunk {
+                        let update = model.read().unwrap().step(&nodes,n);
+                        let mut result = nodes_next[n].write().unwrap();
+                        result.x = update.x;
+                        result.y = update.y;
+                    }
+                });
+                handles.push(handle);
+            }
+
+            for handle in handles {
+                handle.join().unwrap();
+            }
+        
+            // swap nodes and nodes_next
+            let tmp = nodes.clone();
+            nodes = nodes_next.clone();
+            nodes_next = tmp.clone();
+        }
+
+        let mut result = vec![];
+        for node in nodes.iter() {
+            let node = node.read().unwrap();
+            result.push((node.x, node.y));
+        }
+        result
+    }
+}
+
+fn edge_matrix_from_edge_list(number_of_nodes: usize, edge_list: Vec<(u32, u32)>) -> EdgeMatrix {
+    let matrix_ptr = new_edge_matrix(number_of_nodes as usize);
+    {
+        let mut matrix = matrix_ptr.write().unwrap();
+        for (node_a, node_b) in edge_list {
+            matrix[node_a as usize][node_b as usize].weight = 1.0;
+            matrix[node_b as usize][node_a as usize].weight = 1.0;
+        }
+    }
+    matrix_ptr
+}
+

src/spring_model.rs

@@ -1,29 +1,30 @@
 use crate::graph::{EdgeMatrix, NodeVector, Node};
 
-
+pub struct MyModel {
-const C_REP: f32 = 0.1;
-const C_SPRING: f32 = 0.1;
-
-pub struct InitialModel {
     edges: EdgeMatrix,
     size: usize,
     opt_dist: f32,
-    t: f32,
+    c: f32,
-
+    dc: f32,
 }
 
-impl InitialModel {
+impl MyModel {
 
-    pub fn new(edges: EdgeMatrix, size: usize) -> InitialModel {
+    pub fn new(edges: EdgeMatrix, size: usize, iterations: usize) -> MyModel {
-        let opt_dist = 1.0 / (size as f32).sqrt();
+        let opt_dist = 1.0;
-        InitialModel{ edges, size, opt_dist, t: 0.1 }
+        let c = 0.1;
+
+        MyModel{
+            edges,
+            size,
+            opt_dist,
+            c: c,
+            dc: c / ((iterations + 1) as f32)
+        }
     }
 
-    pub fn prepare(& mut self, nodes: &NodeVector) {
+    pub fn prepare(& mut self, _nodes: &NodeVector) {
-        self.t = 0.1 * f32::max(
+        self.c -= self.dc;
-            nodes.iter().map(|n| n.read().unwrap().x.abs()).reduce(|a, b| a.max(b)).unwrap(),
-            nodes.iter().map(|n| n.read().unwrap().y.abs()).reduce(|a, b| a.max(b)).unwrap()
-        );
     }
 
     pub fn step(&self, nodes: &NodeVector, i_node: usize) -> Node {
@@ -31,8 +32,12 @@ impl InitialModel {
         let node = nodes[i_node].read().unwrap();
         let edges = self.edges.read().unwrap();
 
-        let mut node_x = node.x;
+        let node_x = node.x;
-        let mut node_y = node.y;
+        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;
@@ -47,30 +52,37 @@ impl InitialModel {
 
             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().max(0.01);
+            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 = (C_REP / (dist).powi(2)).min(self.t);
+                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;
 
-                node_x -= f_rep_x;
+                sum_x -= f_rep_x;
-                node_y -= f_rep_y;
+                sum_y -= f_rep_y;
             } else {
-                let f_spring = (C_SPRING * (dist / self.opt_dist).log(2.0)).min(self.t);
+                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;
-                node_x += f_spring_x;
+                sum_x += f_spring_x;
-                node_y += f_spring_y;
+                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,
+            x: node_x + sum_x,
-            y: node_y,
+            y: node_y + sum_y,
         }        
     }
 }

src/utils.rs

@@ -16,6 +16,7 @@ pub fn gen_chunks(n: usize, chunks: usize) -> Vec<Range<usize>> {
 }
 
 
+#[cfg(test)]
 mod test {
     use super::*;