← Back to branch summary

~mwshinn/forceatlas2-python/trunk

« back to all changes in this revision

Viewing changes to fa2util.py

  • Committer: Max Shinn
  • Date: 2016-04-20 14:32:02 UTC
  • Revision ID: trombonechamp@gmail.com-20160420143202-bu6qscmc5nyp4p6k
Packaged into a module, speed improvements

Show diffs side-by-side

added added

removed removed

Lines of Context:
fa2util.py
 
1
# This file allows separating the most CPU intensive routines from the
 
2
# main code.  This allows them to be optimized with Cython.  If you
 
3
# don't have Cython, this will run normally.  However, if you use
 
4
# Cython, you'll get speed boosts from 10-100x automatically.
 
5
#
 
6
# The only catch is that IF YOU MODIFY THIS FILE, YOU MUST ALSO MODIFY
 
7
# fa2util.pxd TO REFLECT ANY CHANGES IN FUNCTION DEFINITIONS!
 
8
#
 
9
# Copyright 2016 Max Shinn <mws41@cam.ac.uk>
 
10
#
 
11
# Available under the GPLv3
 
12
 
 
13
from math import sqrt
 
14
 
 
15
# This will substitute for the nLayout object
 
16
class Node:
 
17
    def __init__(self):
 
18
        self.mass = 0
 
19
        self.old_dx = 0
 
20
        self.old_dy = 0
 
21
        self.dx = 0
 
22
        self.dy = 0
 
23
        self.x = 0
 
24
        self.y = 0
 
25
 
 
26
# This is not in the original java code, but it makes it easier to
 
27
# deal with edges.
 
28
class Edge:
 
29
    def __init__(self):
 
30
        self.node1 = -1
 
31
        self.node2 = -1
 
32
        self.weight = 0
 
33
 
 
34
 
 
35
 
 
36
# Here are some functions from ForceFactor.java
 
37
# =============================================
 
38
 
 
39
# Repulsion function.  `n1` and `n2` should be nodes.  This will
 
40
# adjust the dx and dy values of `n1` (and optionally `n2`).  It does
 
41
# not return anything.
 
42
def linRepulsion(n1, n2, coefficient=0):
 
43
    xDist = n1.x - n2.x
 
44
    yDist = n1.y - n2.y
 
45
    distance2 = xDist*xDist + yDist*yDist # Distance squared
 
46
 
 
47
    if distance2 > 0:
 
48
        factor = coefficient * n1.mass * n2.mass / distance2
 
49
        n1.dx += xDist * factor
 
50
        n1.dy += yDist * factor
 
51
        n2.dx -= xDist * factor
 
52
        n2.dy -= yDist * factor
 
53
 
 
54
# Gravity repulsion function.  For some reason, gravity was included
 
55
# within the linRepulsion function in the original gephi java code,
 
56
# which doesn't make any sense (considering a. gravity is unrelated to
 
57
# nodes repelling each other, and b. gravity is actually an
 
58
# attraction).
 
59
def linGravity(n, g, coefficient=0):
 
60
    xDist = n.x
 
61
    yDist = n.y
 
62
    distance = sqrt(xDist*xDist + yDist*yDist)
 
63
 
 
64
    if distance > 0:
 
65
        factor = coefficient * n.mass * g / distance
 
66
        n.dx -= xDist * factor
 
67
        n.dy -= yDist * factor
 
68
 
 
69
# Strong gravity force function.  `n` should be a node, and `g`
 
70
# should be a constant by which to apply the force.
 
71
def strongGravity(n, g, coefficient=0):
 
72
    xDist = n.x
 
73
    yDist = n.y
 
74
 
 
75
    if xDist != 0 and yDist != 0:
 
76
        factor = coefficient * n.mass * g
 
77
        n.dx -= xDist * factor
 
78
        n.dy -= yDist * factor
 
79
 
 
80
# Attraction function.  `n1` and `n2` should be nodes.  This will
 
81
# adjust the dx and dy values of `n1` (and optionally `n2`).  It does
 
82
# not return anything.
 
83
def linAttraction(n1, n2, e, coefficient=0):
 
84
    xDist = n1.x - n2.x
 
85
    yDist = n1.y - n2.y
 
86
    factor = -coefficient * e
 
87
    n1.dx += xDist * factor
 
88
    n1.dy += yDist * factor
 
89
    n2.dx -= xDist * factor
 
90
    n2.dy -= yDist * factor
 
91
 
 
92
def apply_repulsion(nodes, coefficient):
 
93
    for i in range(0, len(nodes)):
 
94
        for j in range(0, i):
 
95
            linRepulsion(nodes[i], nodes[j], coefficient)
 
96
 
 
97
def apply_gravity(nodes, gravity, scalingRatio):
 
98
    for i in range(0, len(nodes)):
 
99
        linGravity(nodes[i], gravity/scalingRatio, scalingRatio)
 
100
 
 
101
def apply_attraction(nodes, edges, coefficient, edgeWeightInfluence):
 
102
    # Optimization, since usually edgeWeightInfluence is 0 or 1, and pow is slow
 
103
    if edgeWeightInfluence == 0:
 
104
        for edge in edges:
 
105
            linAttraction(nodes[edge.node1], nodes[edge.node2], 1, coefficient)
 
106
    elif edgeWeightInfluence == 1:
 
107
        for edge in edges:
 
108
            linAttraction(nodes[edge.node1], nodes[edge.node2], 1, edge.weight)
 
109
    else:
 
110
        for edge in edges:
 
111
            linAttraction(nodes[edge.node1], nodes[edge.node2], pow(edge.weight, edgeWeightInfluence), coefficient)
 
112
 
 
113
 
 
114
try:
 
115
    import cython
 
116
    if cython.compiled:
 
117
        print("Using compiled version of cython")
 
118
    else:
 
119
        print("Warning: uncompiled.  Compile with cython for extra speed")
 
120
except:
 
121
    print("No cython detected")