Added first implementation of adjustSizes. Compiles but needs testing

fa2/fa2util.pxd

@@ -20,6 +20,7 @@ cdef class Node:
     cdef public double old_dx, old_dy
     cdef public double dx, dy
     cdef public double x, y
+    cdef public double size
 
 # This is not in the original java function, but it makes it easier to
 # deal with edges.
@@ -35,7 +36,7 @@ cdef class Edge:
                yDist = cython.double, 
                distance2 = cython.double, 
                factor = cython.double)
-cdef void linRepulsion(Node n1, Node n2, double coefficient=*)
+cdef void linRepulsion(Node n1, Node n2, double coefficient=*, bint anticollision=*)
 
 @cython.locals(xDist = cython.double,
                yDist = cython.double,
@@ -59,19 +60,19 @@ cdef void strongGravity(Node n, double g, double coefficient=*)
 @cython.locals(xDist = cython.double, 
                yDist = cython.double, 
                factor = cython.double)
-cpdef void linAttraction(Node n1, Node n2, double e, bint distributedAttraction, double coefficient=*)
+cpdef void linAttraction(Node n1, Node n2, double e, bint distributedAttraction, double coefficient=*,  bint anticollision=*)
 
 @cython.locals(i = cython.int,
                j = cython.int,
                n1 = Node,
                n2 = Node)
-cpdef void apply_repulsion(list nodes, double coefficient)
+cpdef void apply_repulsion(list nodes, double coefficient, bint anticollision=*)
 
 @cython.locals(n = Node)
 cpdef void apply_gravity(list nodes, double gravity, bint useStrongGravity=*)
 
 @cython.locals(edge = Edge)
-cpdef void apply_attraction(list nodes, list edges, bint distributedAttraction, double coefficient, double edgeWeightInfluence)
+cpdef void apply_attraction(list nodes, list edges, bint distributedAttraction, double coefficient, double edgeWeightInfluence, bint anticollision=*)
 
 cdef class Region:
     cdef public double mass
@@ -99,10 +100,10 @@ cdef class Region:
 
     @cython.locals(distance = cython.double,
                    subregion = Region)
-    cdef void applyForce(self, Node n, double theta, double coefficient=*)
+    cdef void applyForce(self, Node n, double theta, double coefficient=*, bint anticollision=*)
 
     @cython.locals(n = Node)
-    cpdef applyForceOnNodes(self, list nodes, double theta, double coefficient=*)
+    cpdef applyForceOnNodes(self, list nodes, double theta, double coefficient=*, bint anticollision=*)
 
 @cython.locals(totalSwinging = cython.double,
                totalEffectiveTraction = cython.double,
@@ -119,4 +120,4 @@ cdef class Region:
                maxRise = cython.double,
                factor = cython.double,
                values = dict)
-cpdef dict adjustSpeedAndApplyForces(list nodes, double speed, double speedEfficiency, double jitterTolerance)
+cpdef dict adjustSpeedAndApplyForces(list nodes, double speed, double speedEfficiency, double jitterTolerance, bint anticollision=*)

fa2/fa2util.py

@@ -23,6 +23,7 @@ def __init__(self):
         self.dy = 0.0
         self.x = 0.0
         self.y = 0.0
+        self.size = 0.
 
 
 # This is not in the original java code, but it makes it easier to deal with edges
@@ -38,18 +39,27 @@ def __init__(self):
 
 # Repulsion function.  `n1` and `n2` should be nodes.  This will
 # adjust the dx and dy values of `n1`  `n2`
-def linRepulsion(n1, n2, coefficient=0):
+def linRepulsion(n1, n2, coefficient=0, anticollision=False):
     xDist = n1.x - n2.x
     yDist = n1.y - n2.y
-    distance2 = xDist * xDist + yDist * yDist  # Distance squared
-
-    if distance2 > 0:
-        factor = coefficient * n1.mass * n2.mass / distance2
-        n1.dx += xDist * factor
-        n1.dy += yDist * factor
-        n2.dx -= xDist * factor
-        n2.dy -= yDist * factor
-
+
+    distance = sqrt(xDist * xDist + yDist * yDist) 
+
+    if anticollision:
+        distance -=  n1.size + n2.size
+
+    if distance > 0: # Clearly distance is always positive without collision detection
+        factor = coefficient * n1.mass * n2.mass / distance**2
+    elif distance < 0: # If the distance is smaller than the sum of radiuses then increase the repulsion
+        factor = 100 * coefficient * n1.mass * n2.mass
+
+    else: # If distance is 0 do nothing
+        return
+    # Apply the force
+    n1.dx += xDist * factor
+    n1.dy += yDist * factor
+    n2.dx -= xDist * factor
+    n2.dy -= yDist * factor
 
 # Repulsion function. 'n' is node and 'r' is region
 def linRepulsion_region(n, r, coefficient=0):
@@ -94,30 +104,36 @@ def strongGravity(n, g, coefficient=0):
 # Attraction function.  `n1` and `n2` should be nodes.  This will
 # adjust the dx and dy values of `n1` and `n2`.  It does
 # not return anything.
-def linAttraction(n1, n2, e, distributedAttraction, coefficient=0):
+def linAttraction(n1, n2, e, distributedAttraction, coefficient=0, anticollision=False):
     xDist = n1.x - n2.x
     yDist = n1.y - n2.y
-    if not distributedAttraction:
-        factor = -coefficient * e
-    else:
-        factor = -coefficient * e / n1.mass
-    n1.dx += xDist * factor
-    n1.dy += yDist * factor
-    n2.dx -= xDist * factor
-    n2.dy -= yDist * factor
-
+
+    distance = 1.
+    if anticollision:
+        # Check if the nodes are colliding
+        distance = sqrt(xDist * xDist + yDist * yDist) - n1.size - n2.size
+
+    if distance > 0:
+        if not distributedAttraction:
+            factor = -coefficient * e
+        else:
+            factor = -coefficient * e / n1.mass
+        n1.dx += xDist * factor
+        n1.dy += yDist * factor
+        n2.dx -= xDist * factor
+        n2.dy -= yDist * factor
 
 # The following functions iterate through the nodes or edges and apply
 # the forces directly to the node objects.  These iterations are here
 # instead of the main file because Python is slow with loops.
-def apply_repulsion(nodes, coefficient):
+def apply_repulsion(nodes, coefficient, anticollision=False):
     i = 0
     for n1 in nodes:
         j = i
         for n2 in nodes:
             if j == 0:
                 break
-            linRepulsion(n1, n2, coefficient)
+            linRepulsion(n1, n2, coefficient, anticollision)
             j -= 1
         i += 1
 
@@ -131,18 +147,18 @@ def apply_gravity(nodes, gravity, useStrongGravity=False):
             strongGravity(n, gravity)
 
 
-def apply_attraction(nodes, edges, distributedAttraction, coefficient, edgeWeightInfluence):
+def apply_attraction(nodes, edges, distributedAttraction, coefficient, edgeWeightInfluence, anticollision=False):
     # Optimization, since usually edgeWeightInfluence is 0 or 1, and pow is slow
     if edgeWeightInfluence == 0:
         for edge in edges:
-            linAttraction(nodes[edge.node1], nodes[edge.node2], 1, distributedAttraction, coefficient)
+            linAttraction(nodes[edge.node1], nodes[edge.node2], 1, distributedAttraction, coefficient, anticollision=anticollision)
     elif edgeWeightInfluence == 1:
         for edge in edges:
-            linAttraction(nodes[edge.node1], nodes[edge.node2], edge.weight, distributedAttraction, coefficient)
+            linAttraction(nodes[edge.node1], nodes[edge.node2], edge.weight, distributedAttraction, coefficient, anticollision=anticollision)
     else:
         for edge in edges:
             linAttraction(nodes[edge.node1], nodes[edge.node2], pow(edge.weight, edgeWeightInfluence),
-                          distributedAttraction, coefficient)
+                          distributedAttraction, coefficient, anticollision=anticollision)
 
 
 # For Barnes Hut Optimization
@@ -239,24 +255,24 @@ def buildSubRegions(self):
             for subregion in self.subregions:
                 subregion.buildSubRegions()
 
-    def applyForce(self, n, theta, coefficient=0):
+    def applyForce(self, n, theta, coefficient=0, anticollision=False):
         if len(self.nodes) < 2:
-            linRepulsion(n, self.nodes[0], coefficient)
+            linRepulsion(n, self.nodes[0], coefficient, anticollision=anticollision)
         else:
             distance = sqrt((n.x - self.massCenterX) ** 2 + (n.y - self.massCenterY) ** 2)
             if distance * theta > self.size:
                 linRepulsion_region(n, self, coefficient)
             else:
                 for subregion in self.subregions:
-                    subregion.applyForce(n, theta, coefficient)
+                    subregion.applyForce(n, theta, coefficient, anticollision=anticollision)
 
-    def applyForceOnNodes(self, nodes, theta, coefficient=0):
+    def applyForceOnNodes(self, nodes, theta, coefficient=0, anticollision=False): 
         for n in nodes:
-            self.applyForce(n, theta, coefficient)
+            self.applyForce(n, theta, coefficient, anticollision=anticollision)
 
 
 # Adjust speed and apply forces step
-def adjustSpeedAndApplyForces(nodes, speed, speedEfficiency, jitterTolerance):
+def adjustSpeedAndApplyForces(nodes, speed, speedEfficiency, jitterTolerance, anticollision=False):
     # Auto adjust speed.
     totalSwinging = 0.0  # How much irregular movement
     totalEffectiveTraction = 0.0  # How much useful movement
@@ -303,7 +319,14 @@ def adjustSpeedAndApplyForces(nodes, speed, speedEfficiency, jitterTolerance):
     # implemented.
     for n in nodes:
         swinging = n.mass * sqrt((n.old_dx - n.dx) * (n.old_dx - n.dx) + (n.old_dy - n.dy) * (n.old_dy - n.dy))
-        factor = speed / (1.0 + sqrt(speed * swinging))
+
+        if anticollision:
+            factor = 0.1 * speed / (1.0 + sqrt(speed * swinging))
+            df = sqrt(n.dx**2 + n.dy**2)
+            factor = min(factor * df, 10.) / df
+        else:
+            factor = speed / (1.0 + sqrt(speed * swinging))
+
         n.x = n.x + (n.dx * factor)
         n.y = n.y + (n.dy * factor)
 

fa2/forceatlas2.py

@@ -65,7 +65,7 @@ def __init__(self,
 
                  # Log
                  verbose=True):
-        assert linLogMode == adjustSizes == multiThreaded == False, "You selected a feature that has not been implemented yet..."
+        assert linLogMode == multiThreaded == False, "You selected a feature that has not been implemented yet..."
         self.outboundAttractionDistribution = outboundAttractionDistribution
         self.linLogMode = linLogMode
         self.adjustSizes = adjustSizes
@@ -80,7 +80,8 @@ def __init__(self,
 
     def init(self,
              G,  # a graph in 2D numpy ndarray format (or) scipy sparse matrix format
-             pos=None  # Array of initial positions
+             pos=None,  # Array of initial positions
+             sizes=None # Array of node sizes 
              ):
         isSparse = False
         if isinstance(G, numpy.ndarray):
@@ -96,6 +97,8 @@ def init(self,
             isSparse = True
         else:
             assert False, "G is not numpy ndarray or scipy sparse matrix"
+
+        assert isinstance(sizes, numpy.ndarray) or (sizes is None), "Invalid node sizes"
 
         # Put nodes into a data structure we can understand
         nodes = []
@@ -115,6 +118,9 @@ def init(self,
             else:
                 n.x = pos[i][0]
                 n.y = pos[i][1]
+
+            if sizes is not None:
+                n.size = sizes[i]
             nodes.append(n)
 
         # Put edges into a data structure we can understand
@@ -149,6 +155,7 @@ def init(self,
     def forceatlas2(self,
                     G,  # a graph in 2D numpy ndarray format (or) scipy sparse matrix format
                     pos=None,  # Array of initial positions
+                    sizes=None, # Array of node sizes (if self.adjustSizes=True)
                     iterations=100  # Number of times to iterate the main loop
                     ):
         # Initializing, initAlgo()
@@ -159,7 +166,7 @@ def forceatlas2(self,
         # algorithm runs to help ensure convergence.
         speed = 1.0
         speedEfficiency = 1.0
-        nodes, edges = self.init(G, pos)
+        nodes, edges = self.init(G, pos, sizes)
         outboundAttCompensation = 1.0
         if self.outboundAttractionDistribution:
             outboundAttCompensation = numpy.mean([n.mass for n in nodes])
@@ -196,9 +203,9 @@ def forceatlas2(self,
             repulsion_timer.start()
             # parallelization should be implemented here
             if self.barnesHutOptimize:
-                rootRegion.applyForceOnNodes(nodes, self.barnesHutTheta, self.scalingRatio)
+                rootRegion.applyForceOnNodes(nodes, self.barnesHutTheta, self.scalingRatio, anticollision=self.adjustSizes)
             else:
-                fa2util.apply_repulsion(nodes, self.scalingRatio)
+                fa2util.apply_repulsion(nodes, self.scalingRatio, anticollision=self.adjustSizes)
             repulsion_timer.stop()
 
             # Gravitational forces
@@ -209,12 +216,12 @@ def forceatlas2(self,
             # If other forms of attraction were implemented they would be selected here.
             attraction_timer.start()
             fa2util.apply_attraction(nodes, edges, self.outboundAttractionDistribution, outboundAttCompensation,
-                                     self.edgeWeightInfluence)
+                                     self.edgeWeightInfluence, anticollision=self.adjustSizes)
             attraction_timer.stop()
 
             # Adjust speeds and apply forces
             applyforces_timer.start()
-            values = fa2util.adjustSpeedAndApplyForces(nodes, speed, speedEfficiency, self.jitterTolerance)
+            values = fa2util.adjustSpeedAndApplyForces(nodes, speed, speedEfficiency, self.jitterTolerance, anticollision=self.adjustSizes)
             speed = values['speed']
             speedEfficiency = values['speedEfficiency']
             applyforces_timer.stop()
@@ -233,14 +240,24 @@ def forceatlas2(self,
     #
     # This function returns a NetworkX layout, which is really just a
     # dictionary of node positions (2D X-Y tuples) indexed by the node name.
-    def forceatlas2_networkx_layout(self, G, pos=None, iterations=100):
+    def forceatlas2_networkx_layout(self, G, pos=None, sizes=None, iterations=100):
         import networkx
         assert isinstance(G, networkx.classes.graph.Graph), "Not a networkx graph"
         assert isinstance(pos, dict) or (pos is None), "pos must be specified as a dictionary, as in networkx"
+        if self.adjustSizes:
+            isinstance(sizes, dict), "adjustSizes=True requires a sizes dict to be passed"
+
         M = networkx.to_scipy_sparse_matrix(G, dtype='f', format='lil')
+
+        if sizes is not None:
+            sizelist =  numpy.asarray([sizes[i] for i in G.nodes()])
+        else:
+            sizelist = None
+
         if pos is None:
-            l = self.forceatlas2(M, pos=None, iterations=iterations)
+            l = self.forceatlas2(M, pos=None, sizes=sizelist, iterations=iterations)
         else:
             poslist = numpy.asarray([pos[i] for i in G.nodes()])
-            l = self.forceatlas2(M, pos=poslist, iterations=iterations)
+            l = self.forceatlas2(M, pos=poslist, sizes=sizelist, iterations=iterations)
+
         return dict(zip(G.nodes(), l))