Fast Python spring network solver

Question

I wanted a very simple spring system written in Python. The system would be defined as a simple network of knots, linked by links using the following rules:

• A knot is a massless connection between links. Each knot is only affected by the push/pull forces it receives from the links it is connected to (no gravity, viscosity etc.)... Its only attribute is to be anchored or not, where an anchored knot affects the system by its movement. An unanchored knot can also affect the system if being moved, but it will be pulled back by the resulting push/pull forces.

• A link is a connection between 2 knots. It has no mass, and it applies force on the knots connected to each end derived from the difference between its current length and its initial length.

• The system takes for input the initial position of each knot, each knot's anchored state, a list of links (presented as arrays of knot indices), and each link's initial lengths. The system then begins iterating over the network by adding up all the forces affecting each knot, adjusting the knots to a their new position (dampened for stability), and keep iterating until an iteration count limit is reached, or the highest force applied at any given iteration is below a given threshold. I don't care to solve over time, I don't need velocity, all I want is the final "relaxed" position of each knot.

Leveraging NumPy's vectorization, I came up with this code:

import numpy as np
from numpy.core.umath_tests import inner1d

def solver(kPos, kAnchor, link0, link1, w0, cycles=1000, precision=0.001, dampening=0.1, debug=False):
    """
    kPos       : vector array - knot position
    kAnchor    : float array  - knot's anchor state, 0 = moves freely, 1 = anchored (not moving)
    link0      : int array    - array of links connecting each knot. each index corresponds to a knot
    link1      : int array    - array of links connecting each knot. each index corresponds to a knot
    w0         : float array  - initial link length
    cycles     : int          - eval stops when n cycles reached
    precision  : float        - eval stops when highest applied force is below this value
    dampening  : float        - keeps system stable during each iteration
    """

    kPos        = np.asarray(kPos)
    pos         = np.array(kPos) # copy of kPos
    kAnchor     = 1-np.clip(np.asarray(kAnchor).astype(float),0,1)[:,None]
    link0       = np.asarray(link0).astype(int)
    link1       = np.asarray(link1).astype(int)
    w0          = np.asarray(w0).astype(float)

    F = np.zeros(pos.shape)
    i = 0

    for i in xrange(cycles):

        # Init force applied per knot
        F = np.zeros(pos.shape)

        # Calculate forces
        AB = pos[link1] - pos[link0] # get link vectors between knots
        w1 = np.sqrt(inner1d(AB,AB)) # get link current lengths
        AB/=w1[:,None] # normalize link vectors
        f = (w1 - w0) # calculate force
        f = f[:,None] * AB # calculate force vector

        # Apply force vectors on each knot
        np.add.at(F, link0, f)      # F[link0] += f*AB
        np.subtract.at(F, link1, f) # F[link1] -= f*AB

        # Update point positions       
        pos += F * dampening * kAnchor

        # If the maximum force applied is below our precision criteria, exit
        if np.amax(F) < precision:
            break

    # Debug info
    if debug:
        print 'Iterations: %s'%i
        print 'Max Force:  %s'%np.amax(F)

    return pos

Here's some test data to show how it works. In this case I'm using a grid, but in reality the network can be of any shape or form:

import cProfile

# Create a 5x5 3D knot grid
z = np.linspace(-0.5, 0.5, 5)
x = np.linspace(-0.5, 0.5, 5)[::-1]
x,z = np.meshgrid(x,z)
kPos = np.array([np.array(thing) for thing in zip(x.flatten(), z.flatten())])
kPos = np.insert(kPos, 1, 0, axis=1)
'''
array([[-0.5 ,  0.  ,  0.5 ],
       [-0.25,  0.  ,  0.5 ],
       [ 0.  ,  0.  ,  0.5 ],
       ..., 
       [ 0.  ,  0.  , -0.5 ],
       [ 0.25,  0.  , -0.5 ],
       [ 0.5 ,  0.  , -0.5 ]])
'''


# Define the links connecting each knots
link0 = [0,1,2,3,5,6,7,8,10,11,12,13,15,16,17,18,20,21,22,23,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
link1 = [1,2,3,4,6,7,8,9,11,12,13,14,16,17,18,19,21,22,23,24,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24]
AB    = kPos[link0]-kPos[link1]
w0    = np.sqrt(inner1d(AB,AB)) # this is a square grid, each link's initial length will be 0.25

# Set the anchor states
kAnchor = np.zeros(len(kPos)) # All knots will be free floating
kAnchor[12] = 1 # Middle knot will be anchored

This is what the grid looks like:

If we run my code using this data, nothing will happen since the links aren't pushing or pulling:

print np.allclose(kPos,solver(kPos, kAnchor, link0, link1, w0, debug=True))
# Returns True
# Iterations: 0
# Max Force:  0.0

Now let's move that middle anchored knot up a bit and solve the system:

# Move the center knot up a little
kPos[12] = np.array([0,0.3,0])

# eval the system
new = solver(kPos, kAnchor, link0, link1, w0, debug=True) # positions will have moved
#Iterations: 102
#Max Force:  0.000976603249133

# Rerun with cProfile to see how fast it runs
cProfile.run('solver(kPos, kAnchor, link0, link1, w0)')
# 520 function calls in 0.008 seconds

And here's what the grid looks like after being pulled by that single anchored knot:

Question

This grid example solves plenty fast, but my actual use cases are a little more complex than this example (~100-200 knots, ~300-500 links) and solve too slow. I'm looking for ways to make this faster. I did try to get rid of the square root calls at every iteration, which didn't make much of a difference.

I did try to cythonize my program (see below) in hopes to improve performance, but with my very limited C knowledge the final performance was far worse.

If anyone can show me ways to speed up its current pure Python implementation, or could show me ways to properly cythonize my function, I would be very grateful.

Here's my very naive Cython version:

cdef extern from "math.h":
    double sqrt(double m)


import numpy as np
cimport numpy as np
cimport cython

ctypedef np.float64_t DTYPE_f
ctypedef np.int64_t DTYPE_i

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)


def eval(np.ndarray[DTYPE_f, ndim=2] kPos,
                np.ndarray[DTYPE_f, ndim=1] kAnchor,
                np.ndarray[DTYPE_i, ndim=1] link0,
                np.ndarray[DTYPE_i, ndim=1] link1,
                np.ndarray[DTYPE_f, ndim=1] w0,
                int cycles=100,
                precision=0.001,
                dampening=0.1,
                debug=False):


    cdef Py_ssize_t i, j, k, nKnots, nLinks
    cdef double w1
    cdef double f
    cdef double maxF

    nKnots = len(kPos)
    nLinks = len(link0)
    cdef np.ndarray[DTYPE_f, ndim=2] pos = np.array(kPos)
    cdef np.ndarray[DTYPE_f, ndim=2] F = np.zeros((nKnots,3))
    cdef np.ndarray[DTYPE_f, ndim=1] AB

    for i in range(cycles):
        F = np.zeros((nKnots,3))


        # Calculate forces
        for j in range(nLinks):
            AB = pos[link1[j]] - pos[link0[j]]
            w1 = sqrt(AB[0]**2 + AB[1]**2 + AB[2]**2)

            if w1 > 0:
                AB/=w1
                f = w1 - w0[j]
                F[link0[j]] += f * AB
                F[link1[j]] -= f * AB

        # Update point positions
        for j in range(nKnots):
            if kAnchor[j] < 1:
                w1 = sqrt(F[j][0]**2 + F[j][1]**2 + F[j][2]**2)

                if maxF < w1:
                    maxF = w1

                pos[j] = pos[j] + F[j] * dampening * (1-kAnchor[j])


        if maxF < precision:
            break


    # Debug info
    if debug:
        print 'Iterations: %s'%i
        print 'Max Force:  %s'%maxF


    return pos

Answer 1

Starting from Jaime's solution, I noticed that np.bincount was being called in each loop iteration. This function seems to compute the mapping between links and knots, which is something that doesn't change in any loop iteration, at least as I understand your program. Thus I wondered if things would get factor by eliminating this repeated function call.

One way to do that is to represent the mapping between links and knots as a connectivity matrix. This matrix takes up more memory than the simple lists link0 and link1, especially if dense, but if that tradeoff is acceptable, it would allow just using a single np.dot() call in each loop, rather than two different np.bincount calls in each iteration. Effectively Jaime's code is a kind of sparse dot product, while this code uses the dense dot product.

A nice bonus of this approach is that the code in the loop becomes much simpler.

def faster_solver(kPos, kAnchor, link0, link1, w0, cycles=1000, precision=0.001, dampening=0.1, debug=False):
    ### <CODE AS BEFORE> ###
    num_links = len(link0)

    # form links data into a connectivity matrix
    cnxn_mat = np.zeros((knots, num_links), dtype = 'int')
    for ix in range(num_links):
        cnxn_mat[link0[ix], ix] = 1
        cnxn_mat[link1[ix], ix] = -1
    ### <...> ###

    for i in xrange(cycles):

        # Init force applied per knot
        F = np.zeros(pos.shape)

        # Calculate forces
        AB = pos[link1] - pos[link0] # get link vectors between knots
        w1 = inner1d(AB, AB) ** 0.5 # get link current lengths
        AB /= w1[:, None] # normalize link vectors
        f = (w1 - w0) # calculate force
        f = f[:, None] * AB # calculate force vector

        F = np.dot(cnxn_mat, f)

        # Update point positions       
        ### <CONTINUE W/ PREVIOUS CODE>

I did actually time this version vs. the old version using the small test case in the OP, and there was an improvement, but it wasn't very dramatic. My untested guess is that my approach will be much more efficient as the problem size increases, if the density of links between knots is large. The caveat is that much larger amounts of memory will likely be required. (If instead there are a tremendous number of knots but relatively few links, Jaime's sparse approach might well be faster.)

# Create a 5x5 3D knot grid
z = np.linspace(-0.5, 0.5, 5)
x = np.linspace(-0.5, 0.5, 5)[::-1]
x,z = np.meshgrid(x,z)
kPos = np.array([np.array(thing) for thing in zip(x.flatten(), z.flatten())])
kPos = np.insert(kPos, 1, 0, axis=1)

# Define the links connecting each knots
link0 = [0,1,2,3,5,6,7,8,10,11,12,13,15,16,17,18,20,21,22,23,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
link1 = [1,2,3,4,6,7,8,9,11,12,13,14,16,17,18,19,21,22,23,24,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24]

AB    = kPos[link0]-kPos[link1]
w0    = np.sqrt(inner1d(AB,AB)) # this is a square grid, each link's initial length will be 0.25

# Set the anchor states
kAnchor = np.zeros(len(kPos)) # All knots will be free floating
kAnchor[12] = 1 # Middle knot will be anchored

# Move the center knot up a little
kPos[12] = np.array([0,0.3,0])

# Compare both methods
fnew, new = (
            faster_solver(kPos, kAnchor, link0, link1, w0),
            solver(kPos, kAnchor, link0, link1, w0)
             )
assert np.all(np.isclose(fnew, new))

# Time both methods
%timeit fnew = faster_solver(kPos, kAnchor, link0, link1, w0) # positions will have moved
%timeit new = solver(kPos, kAnchor, link0, link1, w0) # positions will have moved

100 loops, best of 3: 2.86 ms per loop
100 loops, best of 3: 3.25 ms per loop

UPDATE:

After a nice comment from the OP with new data, it seems that at larger systems my proposed improvement didn't help at all. With my method, the larger system would require creating a dense 5000-by-2000 matrix. However if I understand the problem then this matrix with 10,000,000 elements will have only 4000 nonzero values. Sparsification would seem to be in order. The question is which of the various sparse formats supported by SciPy would be best. I did some empirical testing and found that for matrices in which every column has two entries at random rows, the "CSR" and "COO" formats seemed to be best.

# construct a random sparse matrix in dense format that has entries in every column, but very few in rows
np.random.seed(1)
n_rows, n_cols = (2000, 5000)

col_idx = np.concatenate((xrange(n_cols), xrange(n_cols)))
row_idx = np.random.randint(n_rows, size = 2*n_cols)

data = np.concatenate(([1]*n_cols, [-1]*n_cols))

random_mat = sparse.coo_matrix((data, (row_idx, col_idx)), dtype = 'float')

mat_coo = random_mat
mat_csr = random_mat.tocsr()
mat_full = random_mat.todense()
arr_full = random_mat.toarray()

# a random dense vector
v = np.random.rand(n_cols)

# timings
print "FULL MATRIX:"
%timeit _ = np.dot(mat_full, v)
print "\n"

print "FULL ARRAY:"
%timeit _ = np.dot(arr_full, v)
print "\n"

print "COO:"
%timeit _ = mat_coo.dot(v)
print "\n"

print "CSR:"
%timeit _ = mat_csr.dot(v)
print "\n"

print "bincount:"
link0, link1 = row_idx[0:n_cols], row_idx[n_cols:]

ndim = 3
out = np.zeros(shape = (n_rows, ndim))
def my_loop():
    ndim = 3
    for dim in xrange(ndim):
        out[:, dim] = np.bincount(link0, weights = v, minlength = n_rows) - np.bincount(link1, weights = v, minlength = n_rows)
    return out
%timeit my_loop()




FULL MATRIX:
100 loops, best of 3: 4.67 ms per loop


FULL ARRAY:
100 loops, best of 3: 4.66 ms per loop


COO:
The slowest run took 5.91 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 19.5 µs per loop


CSR:
10000 loops, best of 3: 23.3 µs per loop


bincount:
10000 loops, best of 3: 77.4 µs per loop

Minor other note: using dtype = 'int in my original faster_solver code was a mistake. Although it results in a better (lower memory, more precise) representation of the connectivity matrix, eventually this matrix will be multiplied by a float vector, f, which requires that numpy convert all the elements to float anyway. It's (slightly) more efficient to instantiate the matrix with a dtype = 'float' for this reason.

Answer 2

You should profile your code, to figure out what exactly is it that is slowing your code down. It's hard to tell without some actual measurements, but my bet is on your calls to np.add.at and np.minus.at, as the .at() method is notoriously (very) slow. For some operations there is really no alternative, but for addition/subtraction you can use np.bincount. The transformed code would look something like:

knots, ndim = kPos.shape
...
for dim in range(ndim):
    # Apply force vectors on each dimension of each knot
    F[:, dim] = (np.bincount(link0, weights=f[:, dim], minlength=knots) -
                 np.bincount(link1, weights=f[:, dim], minlength=knots))

As I said, you'll need to test it, but I wouldn't be surprised if this little change makes your code 5-10x faster.

Answer 3

To compute the link current lengths, maybe numpy.linalg.norm can be useful.

w1 = np.sqrt(inner1d(AB,AB)) # get link current lengths

Becomes (no sure about the axis value)

w1 = np.linalg.norm(AB, axis=0)