3D Solution for Connected Springs

springs.py

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

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

def mag(x, axis=None):
    return np.sqrt(np.sum(np.square(x), axis))

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

# Generate a grid of knots
nX = 5
nY = 5
nZ = 5
x = np.linspace(0, 1, nX)
y = np.linspace(0, 1, nY)
z = np.linspace(0, 1, nZ)
x, y, z = np.meshgrid(x, y, z, indexing='ij')
knots = list(np.array(knot) for knot in zip(x.flatten(), y.flatten(), z.flatten()))

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

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

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

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

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

# Loop over the links
for link in links:

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

        # for each dimension
        for j in range(3):

            ai = 3*link[i] + j
            aj = 3*link[not i] + j
        
            # 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(ai)
                Aj.append(ai)
                Ax.append(-1)
                Ai.append(ai)
                Aj.append(aj)
                Ax.append(+1)
                B[ai] += knots[link[not i]][j] - knots[link[i]][j]

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

                Ai.append(ai)
                Aj.append(ai)
                Ax.append(+1)
                B[ai] += constraints[link[i]][j]

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

# Put the knots at the new positions
knots = np.reshape(X, (len(knots), 3))

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

# 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 = [ knots[i][2] for i in link ]
    ax.plot(x, y, z)
for key, value in constraints.items():
    ax.plot([value[0]], [value[1]], [value[2]], 'o')

pp.show()