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
    knots, ndim = kPos.shape

    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

        # 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
        #
        # Now using @Jaime's suggestion, this is a huge improvement!
        for dim in xrange(ndim):
            F[:,dim] = (np.bincount(link0, weights=f[:, dim], minlength=knots) - np.bincount(link1, weights=f[:, dim], minlength=knots))


        # 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