diff --git a/fa2/fa2util.pxd b/fa2/fa2util.pxd index fe95db3..04b1aff 100644 --- a/fa2/fa2util.pxd +++ b/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) \ No newline at end of file +cpdef dict adjustSpeedAndApplyForces(list nodes, double speed, double speedEfficiency, double jitterTolerance, bint anticollision=*) \ No newline at end of file diff --git a/fa2/fa2util.py b/fa2/fa2util.py index 81fb3fd..be00a2d 100644 --- a/fa2/fa2util.py +++ b/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) diff --git a/fa2/forceatlas2.py b/fa2/forceatlas2.py index 29f16db..1d7eace 100644 --- a/fa2/forceatlas2.py +++ b/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))