diff --git a/src/lib.rs b/src/lib.rs
index e6622dd..70b963a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,8 +1,11 @@
 mod graph;
+mod model;
 mod runner;
 mod spring_model;
 mod utils;
 
+use std::sync::{Arc, RwLock};
+
 use pyo3::prelude::*;
 
 /// Formats the sum of two numbers as string.
@@ -13,8 +16,9 @@ fn layout_from_edge_list(
     iter: usize,
     threads: usize,
 ) -> PyResult<Vec<(f32, f32)>> {
+    let model = Arc::new(RwLock::new(spring_model::SimpleSpringModel::new(1.0)));
     let r = runner::Runner::new(iter, threads);
-    Ok(r.layout(number_of_nodes, edges))
+    Ok(r.layout(number_of_nodes, edges, model))
 }
 
 /// A Python module implemented in Rust.
diff --git a/src/model.rs b/src/model.rs
new file mode 100644
index 0000000..d351a98
--- /dev/null
+++ b/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;
+}
diff --git a/src/runner.rs b/src/runner.rs
index 3bd4062..96358ae 100644
--- a/src/runner.rs
+++ b/src/runner.rs
@@ -1,5 +1,5 @@
 use crate::graph::{new_edge_matrix, new_node_vector, EdgeMatrix};
-use crate::spring_model;
+use crate::model::ForceModel;
 use crate::utils;
 use std::sync::{Arc, RwLock};
 use std::thread;
@@ -30,22 +30,21 @@ impl Runner {
         }
     }
 
-    pub fn layout(
+    pub fn layout<T: 'static + ForceModel + Sync + Send>(
         self: &Self,
         number_of_nodes: usize,
         edge_list: Vec<(u32, u32)>,
+        model: Arc<RwLock<T>>,
     ) -> 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,
-        )));
+        model
+            .write()
+            .unwrap()
+            .init(edges, number_of_nodes, self.iterations);
 
         let chunks = utils::gen_chunks(number_of_nodes, self.threads);
         for _epoch in 0..self.iterations {
diff --git a/src/spring_model.rs b/src/spring_model.rs
index 94168e1..12b2d84 100644
--- a/src/spring_model.rs
+++ b/src/spring_model.rs
@@ -1,32 +1,37 @@
-use crate::graph::{EdgeMatrix, Node, NodeVector};
-
-pub struct MyModel {
+use crate::graph::{new_edge_matrix, EdgeMatrix, Node, NodeVector};
+use crate::model::ForceModel;
+pub struct SimpleSpringModel {
     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),
+impl SimpleSpringModel {
+    pub fn new(c: f32) -> Self {
+        SimpleSpringModel {
+            edges: new_edge_matrix(0),
+            size: 0,
+            c,
+            dc: 0.0,
         }
     }
+}
 
-    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;