改进简单弹簧网络的 numpy 实现

Question

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

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

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

这是网格的样子：

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

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

---

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

#!/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()