version bump

Reviewed-on: #1

.github/workflows/rust.yml

@@ -0,0 +1,39 @@
+name: Rust
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python 3.11
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.11
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m venv venv
+        source venv/bin/activate
+        pip install -r requirements.txt
+    - name: Build
+      run: |
+        source venv/bin/activate
+        maturin build --verbose
+    - name: Run tests
+      run: cargo test --verbose
+    - name: Run E2E tests
+      run: |
+        source venv/bin/activate
+        maturin develop --release
+        pytest

.gitignore

@@ -2,3 +2,4 @@
 result/
 
 venv/
+*/__pycache__/

Cargo.lock

@@ -33,7 +33,7 @@ dependencies = [
 
 [[package]]
 name = "graph_force"
-version = "0.1.0"
+version = "0.2.2"
 dependencies = [
  "pyo3",
  "rand",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "graph_force"
-version = "0.1.0"
+version = "0.2.2"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

DEVELOPMENT.md

@@ -0,0 +1,13 @@
+## Development
+
+### Setup
+```
+python3 -m venv venv
+source venv/bin/activate
+pip install -r requirements.txt
+```
+
+### Build
+```
+maturin develop
+```

README.md

@@ -2,24 +2,88 @@
 
 A python/rust library for embedding graphs in 2D space, using force-directed layouts.
 
-## Development
+## Installation
 
-### Setup
+```bash
-```
+pip install graph_force
-python3 -m venv venv
-source venv/bin/activate
-pip install -r requirements.txt
-```
-
-### Build
-```
-maturin develop
 ```
 
 ## Usage
+
+The first parameter defines the number of nodes in graph.
+The second parameter is an iterable of edges, where each edge is a tuple of two integers representing the nodes it connects. Node ids start at 0. 
+
+
 ```python
 import graph_force
 
 edges = [(0, 1), (1, 2), (2, 3), (3, 0)]
 pos = graph_force.layout_from_edge_list(4, edges)
 ```
+
+### Example with networkx
+
+This library does not have a function to consume a networkx graph directly, but it is easy to convert it to an edge list.
+
+```python
+import networkx as nx
+import graph_force
+
+G = nx.grid_2d_graph(10, 10)
+# we have to map the names to integers
+# as graph_force only supports integers as node ids at the moment
+edges = []
+mapping = {n: i for i, n in enumerate(G.nodes)}
+i = 0
+for edge in G.edges:
+    edges.append((mapping[edge[0]], mapping[edge[1]]))
+
+pos = graph_force.layout_from_edge_list(len(G.nodes), edges, iter=1000)
+nx.draw(G, {n: pos[i] for n, i in mapping.items()}, node_size=2, width=0.1)
+```
+
+### Example with edge file
+
+This methods can be used with large graphs, where the edge list does not fit into memory.
+
+Format of the file:
+- Little endian
+- 4 bytes: number of nodes(int)
+- 12 bytes: nodeA(int), nodeB(int), weight(float)
+
+```python
+import graph_force
+import struct
+
+with open("edges.bin", "rb") as f:
+    n = 10
+    f.write(struct.pack("i", n))
+    for x in range(n-1):
+        f.write(struct.pack("iif", x, x+1, 1))
+
+pos = graph_force.layout_from_edge_file("edges.bin", iter=50)
+```
+
+
+### Options
+
+`iter`, `threads` and `model`, `initial_pos` are optional parameters, supported by `layout_from_edge_list` and `layout_from_edge_file`.
+
+```python
+pos = graph_force.layout_from_edge_list(
+    number_of_nodes,
+    edges,
+    iter=500,  # number of iterations, default 500
+    threads=0,  # number of threads, default 0 (all available)
+    model="spring_model", # model to use, default "spring_model", other option is "networkx_model"
+    initial_pos=[(0.4, 0.7), (0.7, 0.2), ...], # initial positions, default None (random)
+)
+```
+#### Available models
+
+- `spring_model`: A simple spring model (my own implementation)
+- `networkx_model`: Reimplementation of the [spring model from networkx](https://networkx.org/documentation/stable/reference/generated/networkx.drawing.layout.spring_layout.html)
+
+## Contributing
+
+- [Development](DEVELOPMENT.md)

pyproject.toml

@@ -4,11 +4,21 @@ build-backend = "maturin"
 
 [project]
 name = "graph_force"
-requires-python = ">=3.7"
+description = "A library for embedding graphs in 2D space, using force-directed layouts."
+keywords = ["graph", "layout", "force-directed"]
+authors = [ { email = "niko@rerere.org", name = "Niko Abeler" } ]
+maintainers = [ { email = "niko@rerere.org", name = "Niko Abeler" } ]
+readme = "README.md"
+requires-python = ">=3.11"
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
+    "Topic :: Scientific/Engineering :: Visualization",
+    "License :: OSI Approved :: MIT License",    
 ]
 
 
+[project.urls]
+homepage = "https://github.com/H4kor/graph-force/"
+repository = "https://github.com/H4kor/graph-force/"

requirements.txt

@@ -1 +1,2 @@
 maturin==0.14.1
+pytest==7.2.0

src/graph.rs

@@ -34,3 +34,46 @@ pub fn new_node_vector(size: usize) -> NodeVector {
     }
     Arc::new(nodes)
 }
+
+pub fn add_edge(matrix: &mut EdgeMatrix, i: usize, j: usize) {
+    let mut edges = matrix.write().unwrap();
+    edges[i][j].weight = 1.0;
+    edges[j][i].weight = 1.0;
+}
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+
+    #[test]
+    fn test_new_edge_matrix() {
+        let matrix = new_edge_matrix(5);
+        let edges = matrix.read().unwrap();
+        assert_eq!(edges.len(), 5);
+        for row in edges.iter() {
+            assert_eq!(row.len(), 5);
+            for edge in row.iter() {
+                assert_eq!(edge.weight, 0.0);
+            }
+        }
+    }
+
+    #[test]
+    fn test_new_node_vector() {
+        let nodes = new_node_vector(5);
+        assert_eq!(nodes.len(), 5);
+    }
+
+    #[test]
+    fn test_add_edge() {
+        let mut matrix = new_edge_matrix(5);
+        add_edge(&mut matrix, 0, 1);
+        let edges = matrix.read().unwrap();
+        assert_eq!(edges[0][1].weight, 1.0);
+        assert_eq!(edges[1][0].weight, 1.0);
+
+        assert_eq!(edges[0][0].weight, 0.0);
+        assert_eq!(edges[1][1].weight, 0.0);
+    }
+}

src/lib.rs

@@ -1,25 +1,125 @@
 mod graph;
+mod model;
+mod networkx_model;
+mod reader;
 mod runner;
 mod spring_model;
 mod utils;
 
-use pyo3::prelude::*;
+use graph::NodeVector;
+use pyo3::exceptions;
+use pyo3::{prelude::*, types::PyIterator};
 
-/// Formats the sum of two numbers as string.
+fn pick_model(model: &str) -> Result<Box<dyn model::ForceModel + Send + Sync>, PyErr> {
-#[pyfunction(number_of_nodes, edges, "*", iter = 500, threads = 0)]
+    match model {
-fn layout_from_edge_list(
+        "spring_model" => Ok(Box::new(spring_model::SimpleSpringModel::new(1.0))),
-    number_of_nodes: usize,
+        "networkx_model" => Ok(Box::new(networkx_model::NetworkXModel::new())),
-    edges: Vec<(u32, u32)>,
+        _ => {
-    iter: usize,
+            return Err(PyErr::new::<exceptions::PyValueError, _>(
-    threads: usize,
+                "model must be either 'spring_model' or 'networkx_model'",
-) -> PyResult<Vec<(f32, f32)>> {
+            ))
-    let r = runner::Runner::new(iter, threads);
+        }
-    Ok(r.layout(number_of_nodes, edges))
+    }
+}
+
+fn initial_pos_to_node_vector(initial_pos: Option<Vec<(f32, f32)>>) -> Option<NodeVector> {
+    match initial_pos {
+        Some(pos) => {
+            let nodes = graph::new_node_vector(pos.len());
+            for (i, (x, y)) in pos.iter().enumerate() {
+                let mut node = nodes[i].write().unwrap();
+                node.x = *x;
+                node.y = *y;
+            }
+            Some(nodes)
+        }
+        None => None,
+    }
+}
+
+#[pyfunction(
+    file_path,
+    "*",
+    iter = 500,
+    threads = 0,
+    model = "\"spring_model\"",
+    initial_pos = "None"
+)]
+fn layout_from_edge_file(
+    file_path: &str,
+    iter: usize,
+    threads: usize,
+    model: &str,
+    initial_pos: Option<Vec<(f32, f32)>>,
+) -> PyResult<Vec<(f32, f32)>> {
+    let (size, matrix) = reader::read_graph(file_path);
+    let model = pick_model(model)?;
+
+    let r = runner::Runner::new(iter, threads);
+    Ok(r.layout(size, matrix, model, initial_pos_to_node_vector(initial_pos)))
+}
+
+#[pyfunction(
+    number_of_nodes,
+    edges,
+    "*",
+    iter = 500,
+    threads = 0,
+    model = "\"spring_model\"",
+    initial_pos = "None"
+)]
+fn layout_from_edge_list(
+    number_of_nodes: usize,
+    edges: &PyAny,
+    iter: usize,
+    threads: usize,
+    model: &str,
+    initial_pos: Option<Vec<(f32, f32)>>,
+) -> PyResult<Vec<(f32, f32)>> {
+    let model: Box<dyn model::ForceModel + Send + Sync> = pick_model(model)?;
+
+    let mut edge_matrix = graph::new_edge_matrix(number_of_nodes);
+    match edges.extract::<&PyIterator>() {
+        Ok(iter) => {
+            iter.iter()?
+                .map(|edge| edge.and_then(PyAny::extract::<(usize, usize)>))
+                .for_each(|edge| {
+                    if let Ok((u, v)) = edge {
+                        graph::add_edge(&mut edge_matrix, u, v);
+                    }
+                });
+            for edge in iter {
+                let edge = edge?;
+                let edge = edge.extract::<(usize, usize)>()?;
+                graph::add_edge(&mut edge_matrix, edge.0, edge.1);
+            }
+        }
+        Err(_) => match edges.extract::<Vec<(usize, usize)>>() {
+            Ok(edge) => {
+                for edge in edge {
+                    graph::add_edge(&mut edge_matrix, edge.0, edge.1);
+                }
+            }
+            Err(_) => {
+                return Err(PyErr::new::<exceptions::PyTypeError, _>(
+                    "Edges must be an iterable of (int, int)",
+                ));
+            }
+        },
+    }
+
+    let r = runner::Runner::new(iter, threads);
+    Ok(r.layout(
+        number_of_nodes,
+        edge_matrix,
+        model,
+        initial_pos_to_node_vector(initial_pos),
+    ))
 }
 
-/// A Python module implemented in Rust.
 #[pymodule]
 fn graph_force(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(layout_from_edge_list, m)?)?;
+    m.add_function(wrap_pyfunction!(layout_from_edge_file, m)?)?;
     Ok(())
 }

src/model.rs

@@ -0,0 +1,7 @@
+use crate::graph::{EdgeMatrix, Node, NodeVector};
+
+pub trait ForceModel {
+    fn init(&mut self, edges: EdgeMatrix, size: usize, iterations: usize);
+    fn prepare(&mut self, _nodes: &NodeVector);
+    fn step(&self, nodes: &NodeVector, i_node: usize) -> Node;
+}

src/networkx_model.rs

@@ -0,0 +1,133 @@
+use crate::graph::{new_edge_matrix, EdgeMatrix, Node, NodeVector};
+use crate::model::ForceModel;
+
+/**
+ * Translation of the NetworkX spring_layout function.
+ * https://networkx.org/documentation/stable/reference/generated/networkx.drawing.layout.spring_layout.html
+ */
+pub struct NetworkXModel {
+    edges: EdgeMatrix,
+    size: usize,
+    k: f32,  // optimal distance
+    t: f32,  // temperature
+    dt: f32, // temperature decrease
+}
+
+impl NetworkXModel {
+    pub fn new() -> Self {
+        NetworkXModel {
+            edges: new_edge_matrix(0),
+            size: 0,
+            k: 1.0,
+            t: 1.0,
+            dt: 0.0,
+        }
+    }
+}
+
+impl ForceModel for NetworkXModel {
+    fn init(&mut self, edges: EdgeMatrix, size: usize, iterations: usize) {
+        self.k = (1.0 / size as f32).sqrt();
+        self.t = 0.1;
+        self.dt = self.t / ((iterations + 1) as f32);
+        self.t += self.dt; // prepare is called before the first step
+        self.edges = edges;
+        self.size = size;
+    }
+
+    fn prepare(&mut self, _nodes: &NodeVector) {
+        self.t -= self.dt; // decrease temperature
+    }
+
+    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 displacement_x = 0.0;
+        let mut displacement_y = 0.0;
+
+        for o in 0..self.size {
+            if o == i_node {
+                continue;
+            }
+
+            let edge = edges[i_node][o].weight;
+
+            let o_x: f32;
+            let o_y: f32;
+            {
+                let other = nodes[o].read().unwrap();
+                o_x = other.x;
+                o_y = other.y;
+            }
+            // difference between node and other
+            let delta_x = node_x - o_x;
+            let delta_y = node_y - o_y;
+            // distance between node and other
+            let dist = (delta_x * delta_x + delta_y * delta_y).sqrt();
+            // enforce minimum distance of 0.01
+            let dist = if dist < 0.01 { 0.01 } else { dist };
+
+            // displacement "force"
+            displacement_x += delta_x * (self.k.powi(2) / dist.powi(2) - edge * dist / self.k);
+            displacement_y += delta_y * (self.k.powi(2) / dist.powi(2) - edge * dist / self.k);
+        }
+        // update positions
+        let length = (displacement_x * displacement_x + displacement_y * displacement_y).sqrt();
+        let length = if length < 0.01 { 0.01 } else { length };
+        let delta_pos_x = displacement_x * self.t / length;
+        let delta_pos_y = displacement_y * self.t / length;
+
+        Node {
+            x: node_x + delta_pos_x,
+            y: node_y + delta_pos_y,
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::graph::{new_edge_matrix, new_node_vector};
+
+    #[test]
+    fn test_simple_spring_model_attraction() {
+        let mut model = NetworkXModel::new();
+        let edges = new_edge_matrix(2);
+        edges.write().unwrap()[0][1].weight = 1.0;
+        edges.write().unwrap()[1][0].weight = 1.0;
+        model.init(edges, 2, 1);
+
+        let nodes = new_node_vector(2);
+        nodes[0].write().unwrap().x = 0.0;
+        nodes[0].write().unwrap().y = 0.0;
+        nodes[1].write().unwrap().x = 1.0;
+        nodes[1].write().unwrap().y = 1.0;
+
+        model.prepare(&nodes);
+        let node = model.step(&nodes, 0);
+        assert!(node.x > 0.0);
+        assert!(node.y > 0.0);
+    }
+
+    #[test]
+    fn test_simple_spring_model_repulsion() {
+        let mut model = NetworkXModel::new();
+        let edges = new_edge_matrix(2);
+        model.init(edges, 2, 1);
+
+        let nodes = new_node_vector(2);
+        nodes[0].write().unwrap().x = 0.0;
+        nodes[0].write().unwrap().y = 0.0;
+        nodes[1].write().unwrap().x = 1.0;
+        nodes[1].write().unwrap().y = 1.0;
+
+        model.prepare(&nodes);
+        let node = model.step(&nodes, 0);
+        assert!(node.x < 0.0);
+        assert!(node.y < 0.0);
+    }
+}

src/reader.rs

@@ -1,3 +1,7 @@
+use std::{fs::File, io::Read};
+
+use crate::graph::{add_edge, new_edge_matrix, EdgeMatrix};
+
 /**
  * Read Graph data from file
  * Format:
@@ -5,22 +9,20 @@
  *  - 4 bytes: number of nodes(int)
  *  - 12 bytes: nodeA(int), nodeB(int), weight(float)
  */
-fn read_graph(file_name: &str) -> (usize, EdgeMatrix) {
+pub fn read_graph(file_name: &str) -> (usize, EdgeMatrix) {
     let mut file = File::open(file_name).expect("file not found");
     let mut size_buffer = [0; 4];
     file.read_exact(&mut size_buffer).expect("buffer overflow");
     let size = u32::from_le_bytes(size_buffer) as usize;
-    let matrix_ptr = connection_matrix(size);
+    let mut matrix = new_edge_matrix(size);
     {
-        let mut matrix = matrix_ptr.write().unwrap();
         let mut buffer = [0; 12];
         while file.read_exact(&mut buffer).is_ok() {
             let node_a = u32::from_le_bytes(buffer[0..4].try_into().unwrap()) as usize;
             let node_b = u32::from_le_bytes(buffer[4..8].try_into().unwrap()) as usize;
-            let weight = f32::from_le_bytes(buffer[8..12].try_into().unwrap());
+            let _weight = f32::from_le_bytes(buffer[8..12].try_into().unwrap());
-            matrix[node_a][node_b].weight = weight;
+            add_edge(&mut matrix, node_a, node_b);
-            matrix[node_b][node_a].weight = weight;
         }
     }
-    (size, matrix_ptr)
+    (size, matrix)
 }

src/runner.rs

@@ -1,5 +1,5 @@
-use crate::graph::{new_edge_matrix, new_node_vector, EdgeMatrix};
+use crate::graph::{new_node_vector, EdgeMatrix, NodeVector};
-use crate::spring_model;
+use crate::model::ForceModel;
 use crate::utils;
 use std::sync::{Arc, RwLock};
 use std::thread;
@@ -33,19 +33,23 @@ impl Runner {
     pub fn layout(
         self: &Self,
         number_of_nodes: usize,
-        edge_list: Vec<(u32, u32)>,
+        edges: EdgeMatrix,
+        model: Box<dyn ForceModel + Send + Sync>,
+        initial_pos: Option<NodeVector>,
     ) -> Vec<(f32, f32)> {
         // let edges = connection_matrix(size);
-        let edges = edge_matrix_from_edge_list(number_of_nodes, edge_list);
+        let mut nodes = match initial_pos {
-        let mut nodes = new_node_vector(number_of_nodes);
+            Some(pos) => pos,
+            None => 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(model));
-        let model = Arc::new(RwLock::new(spring_model::MyModel::new(
+
-            edges,
+        model
-            number_of_nodes,
+            .write()
-            self.iterations,
+            .unwrap()
-        )));
+            .init(edges, number_of_nodes, self.iterations);
 
         let chunks = utils::gen_chunks(number_of_nodes, self.threads);
         for _epoch in 0..self.iterations {
@@ -86,14 +90,35 @@ impl Runner {
     }
 }
 
-fn edge_matrix_from_edge_list(number_of_nodes: usize, edge_list: Vec<(u32, u32)>) -> EdgeMatrix {
+#[cfg(test)]
-    let matrix_ptr = new_edge_matrix(number_of_nodes as usize);
+mod test {
-    {
+    use crate::graph::{self, Node, NodeVector};
-        let mut matrix = matrix_ptr.write().unwrap();
+
-        for (node_a, node_b) in edge_list {
+    use super::*;
-            matrix[node_a as usize][node_b as usize].weight = 1.0;
+
-            matrix[node_b as usize][node_a as usize].weight = 1.0;
+    struct MockModel {
+        counter: usize,
+    }
+
+    impl ForceModel for MockModel {
+        fn init(&mut self, _edges: EdgeMatrix, _size: usize, _iterations: usize) {}
+        fn prepare(&mut self, _nodes: &NodeVector) {
+            self.counter += 1;
+        }
+        fn step(&self, _nodes: &NodeVector, i_node: usize) -> Node {
+            Node {
+                x: i_node as f32,
+                y: self.counter as f32,
+            }
         }
     }
-    matrix_ptr
+
+    #[test]
+    fn test_layout() {
+        let model = Box::new(MockModel { counter: 0 });
+        let runner = Runner::new(10, 1);
+        let edges = graph::new_edge_matrix(3);
+        let result = runner.layout(3, edges, model, None);
+        assert_eq!(result, vec![(0.0, 10.0), (1.0, 10.0), (2.0, 10.0)]);
+    }
 }

src/spring_model.rs

@@ -1,32 +1,37 @@
-use crate::graph::{EdgeMatrix, Node, NodeVector};
+use crate::graph::{new_edge_matrix, EdgeMatrix, Node, NodeVector};
-
+use crate::model::ForceModel;
-pub struct MyModel {
+pub struct SimpleSpringModel {
     edges: EdgeMatrix,
     size: usize,
-    opt_dist: f32,
     c: f32,
     dc: f32,
 }
 
-impl MyModel {
+impl SimpleSpringModel {
-    pub fn new(edges: EdgeMatrix, size: usize, iterations: usize) -> MyModel {
+    pub fn new(c: f32) -> Self {
-        let opt_dist = 1.0;
+        SimpleSpringModel {
-        let c = 0.1;
+            edges: new_edge_matrix(0),
-
+            size: 0,
-        MyModel {
+            c,
-            edges,
+            dc: 0.0,
-            size,
-            opt_dist,
-            c: c,
-            dc: c / ((iterations + 1) as f32),
         }
     }
+}
 
-    pub fn prepare(&mut self, _nodes: &NodeVector) {
+impl ForceModel for SimpleSpringModel {
+    fn init(&mut self, edges: EdgeMatrix, size: usize, iterations: usize) {
+        let dc = self.c / ((iterations + 1) as f32);
+
+        self.edges = edges;
+        self.size = size;
+        self.dc = dc;
+    }
+
+    fn prepare(&mut self, _nodes: &NodeVector) {
         self.c -= self.dc;
     }
 
-    pub fn step(&self, nodes: &NodeVector, i_node: usize) -> Node {
+    fn step(&self, nodes: &NodeVector, i_node: usize) -> Node {
         let node = nodes[i_node].read().unwrap();
         let edges = self.edges.read().unwrap();
 
@@ -57,14 +62,14 @@ impl MyModel {
             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 = dist.powi(2).recip().min(1.0);
                 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 = 0.5 * (dist - 1.0);
                 let f_spring_x = f_spring * unit_x;
                 let f_spring_y = f_spring * unit_y;
                 sum_x += f_spring_x;
@@ -84,3 +89,47 @@ impl MyModel {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::graph::{new_edge_matrix, new_node_vector};
+
+    #[test]
+    fn test_simple_spring_model_attraction() {
+        let mut model = SimpleSpringModel::new(1.0);
+        let edges = new_edge_matrix(2);
+        edges.write().unwrap()[0][1].weight = 1.0;
+        edges.write().unwrap()[1][0].weight = 1.0;
+        model.init(edges, 2, 1);
+
+        let nodes = new_node_vector(2);
+        nodes[0].write().unwrap().x = 0.0;
+        nodes[0].write().unwrap().y = 0.0;
+        nodes[1].write().unwrap().x = 1.0;
+        nodes[1].write().unwrap().y = 1.0;
+
+        model.prepare(&nodes);
+        let node = model.step(&nodes, 0);
+        assert!(node.x > 0.0);
+        assert!(node.y > 0.0);
+    }
+
+    #[test]
+    fn test_simple_spring_model_repulsion() {
+        let mut model = SimpleSpringModel::new(1.0);
+        let edges = new_edge_matrix(2);
+        model.init(edges, 2, 1);
+
+        let nodes = new_node_vector(2);
+        nodes[0].write().unwrap().x = 0.0;
+        nodes[0].write().unwrap().y = 0.0;
+        nodes[1].write().unwrap().x = 1.0;
+        nodes[1].write().unwrap().y = 1.0;
+
+        model.prepare(&nodes);
+        let node = model.step(&nodes, 0);
+        assert!(node.x < 0.0);
+        assert!(node.y < 0.0);
+    }
+}

src/utils.rs

@@ -3,11 +3,18 @@ use std::ops::Range;
 pub fn gen_chunks(n: usize, chunks: usize) -> Vec<Range<usize>> {
     let mut borders = vec![];
     let chunk_size = n / chunks;
+    let mut remainder = n % chunks;
     let mut start = 0;
     for _ in 0..chunks {
         let end = start + chunk_size;
-        borders.push(start..end);
+        if remainder > 0 {
-        start = end;
+            remainder -= 1;
+            borders.push(start..end + 1);
+            start = end + 1;
+        } else {
+            borders.push(start..end);
+            start = end;
+        }
     }
     // Change the last chunk to include the remainder
     let size = borders.len();
@@ -22,7 +29,7 @@ mod test {
     #[test]
     fn test_gen_chunks() {
         let borders = gen_chunks(10, 3);
-        assert_eq!(borders, vec![(0..3), (3..6), (6..10)]);
+        assert_eq!(borders, vec![(0..4), (4..7), (7..10)]);
     }
 
     #[test]
@@ -36,4 +43,10 @@ mod test {
         let borders = gen_chunks(10, 1);
         assert_eq!(borders, vec![(0..10)]);
     }
+
+    #[test]
+    fn test_gen_chunks4() {
+        let borders = gen_chunks(11, 3);
+        assert_eq!(borders, vec![(0..4), (4..8), (8..11)]);
+    }
 }

tests/test_e2e.py

@@ -0,0 +1,50 @@
+import graph_force
+import struct
+
+def test_list_of_edges():
+    edges = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]
+    pos = graph_force.layout_from_edge_list(7, edges)
+    assert pos is not None
+    assert len(pos) == 7
+
+
+def test_iterator_of_edges():
+    pos = graph_force.layout_from_edge_list(
+        7,
+        ((0, i + 1) for i in range(6))
+    )
+    assert pos is not None
+    assert len(pos) == 7
+
+def test_tuple_of_edges():
+    pos = graph_force.layout_from_edge_list(
+        7,
+        ((0,1), (1,2), (2,3), (3,4), (4,5), (5,6))
+    )
+    assert pos is not None
+    assert len(pos) == 7
+
+def test_model_selection():
+    edges = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]
+    pos = graph_force.layout_from_edge_list(7, edges, model='networkx_model')
+    assert pos is not None
+    assert len(pos) == 7
+
+def test_from_file():
+    with open("/tmp/edges.bin", "wb") as f:
+        n = 10
+        f.write(struct.pack("i", n))
+        for x in range(n-1):
+            f.write(struct.pack("iif", x, x+1, 1))
+
+    pos = graph_force.layout_from_edge_file('/tmp/edges.bin')
+    assert pos is not None
+    assert len(pos) == 10
+
+def test_initial_pos():
+    edges = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]
+    initial = [(i, i) for i in range(7)]
+    pos = graph_force.layout_from_edge_list(7, edges, iter=0, initial_pos=initial)
+    assert pos is not None
+    assert len(pos) == 7
+    assert pos == initial