python - 改进一个简单的 spring 网络的 numpy 实现

Question

我想要一个用 numpy 编写的非常简单的 spring 系统。该系统将被定义为一个简单的 结 网络，由 links 链接。我对随着时间的推移评估系统不感兴趣，而是想从初始状态开始，更改变量(通常将 knot 移动到新位置)并解决系统直到它到达稳定 状态(最后施加的力低于给定阈值)。结没有质量，没有重力，力都来自每个链接的当前长度/初始长度。唯一的“特殊”变量是每个结都可以设置为“锚定”(不移动)。

所以我在下面编写了这个简单的求解器，并包含了一个简单的示例。跳到我的问题的最后。

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

np.set_printoptions(precision=4)
np.set_printoptions(suppress=True)
np.set_printoptions(linewidth =150)
np.set_printoptions(threshold=10)


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 lengths
        AB/=w1[:,None] # normalize link vectors
        f = (w1 - w0) # calculate force vectors
        f = f[:,None] * AB
        
        # Apply force vectors on each knot
        np.add.at(F, link0, f)
        np.subtract.at(F, link1, f)

        # 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

这里有一些测试数据来展示它是如何工作的。在这种情况下，我使用的是网格，但实际上这可以是任何类型的网络，比如有很多结的绳子，或者一堆乱七八糟的多边形……:

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

这是网格的样子:

如果我们使用这些数据运行我的代码，则不会发生任何事情，因为链接没有推送或拉伸(stretch):

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

现在让我们稍微移动中间的锚定结并求解系统:

# 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

下面是网格在被单个锚定结拉动后的样子:

问题:

我的实际用例比这个例子稍微复杂一点，而且根据我的口味，解决速度有点太慢:(100-200 节的网络在 200-300 个链接之间，在几秒钟内解决)。

如何让我的求解器函数运行得更快？我会考虑使用 Cython，但我对 C 的经验为零。非常感谢任何帮助。

Answer 1

粗略地看一下，您的方法似乎是一种明确的欠松弛类型的方法。计算每个结处的残余力，应用该力的一个因子作为位移，重复直到收敛。需要时间的是重复直到收敛。您拥有的点越多，每次迭代花费的时间就越长，但您还需要更多迭代，以便将网格一端的约束传播到另一端。

您是否考虑过隐式方法？写出各非约束节点处的残余力方程，拼成一个大矩阵，一步求解。现在，信息只需一步即可传播到整个问题。作为一个额外的好处，您构建的矩阵应该是稀疏的，scipy 有一个模块。

Wikipedia: explicit and implicit methods

---

编辑 隐式方法(大致)匹配您的问题的示例。这个解是线性的，所以它没有考虑计算位移对力的影响。您需要迭代(或使用非线性技术)来计算它。希望对您有所帮助。

#!/usr/bin/python3

import matplotlib.pyplot as pp
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import scipy as sp
import scipy.sparse
import scipy.sparse.linalg

#------------------------------------------------------------------------------#

# Generate a grid of knots
nX = 10
nY = 10
x = np.linspace(-0.5, 0.5, nX)
y = np.linspace(-0.5, 0.5, nY)
x, y = np.meshgrid(x, y)
knots = list(zip(x.flatten(), y.flatten()))

# Create links between the knots
links = []
# Horizontal links
for i in range(0, nY):
    for j in range(0, nX - 1):
        links.append((i*nX + j, i*nX + j + 1))
# Vertical links
for i in range(0, nY - 1):
    for j in range(0, nX):
        links.append((i*nX + j, (i + 1)*nX + j))

# Create constraints. This dict takes a knot index as a key and returns the
# fixed z-displacement associated with that knot.
constraints = {
    0          : 0.0,
    nX - 1     : 0.0,
    nX*(nY - 1): 0.0,
    nX*nY - 1  : 1.0,
    2*nX + 4   : 1.0,
    }

#------------------------------------------------------------------------------#

# Matrix i-coordinate, j-coordinate and value
Ai = []
Aj = []
Ax = []

# Right hand side array
B = np.zeros(len(knots))

# Loop over the links
for link in links:

    # Link geometry
    displacement = np.array([ knots[1][i] - knots[0][i] for i in range(2) ])
    distance = np.sqrt(displacement.dot(displacement))

    # For each node
    for i in range(2):

        # If it is not a constraint, add the force associated with the link to
        # the equation of the knot
        if link[i] not in constraints:

            Ai.append(link[i])
            Aj.append(link[i])
            Ax.append(-1/distance)

            Ai.append(link[i])
            Aj.append(link[not i])
            Ax.append(+1/distance)

        # If it is a constraint add a diagonal and a value
        else:

            Ai.append(link[i])
            Aj.append(link[i])
            Ax.append(+1.0)
            B[link[i]] += constraints[link[i]]

# Create the matrix and solve
A = sp.sparse.coo_matrix((Ax, (Ai, Aj))).tocsr()
X = sp.sparse.linalg.lsqr(A, B)[0]

#------------------------------------------------------------------------------#

# Plot the links
fg = pp.figure()
ax = fg.add_subplot(111, projection='3d')

for link in links:
    x = [ knots[i][0] for i in link ]
    y = [ knots[i][1] for i in link ]
    z = [ X[i] for i in link ]
    ax.plot(x, y, z)

pp.show()