pyqtgraph/examples/verlet_chain/relax.py

71 lines
2.1 KiB
Python
Raw Normal View History

import ctypes
import os
so = os.path.join(os.path.dirname(__file__), 'maths.so')
try:
lib = ctypes.CDLL(so)
COMPILED = True
except OSError:
COMPILED = False
if COMPILED:
lib.relax.argtypes = [
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_int,
ctypes.c_int,
]
def relax(pos, links, mrel1, mrel2, lengths, push, pull, iters):
nlinks = links.shape[0]
lib.relax(pos.ctypes, links.ctypes, mrel1.ctypes, mrel2.ctypes, lengths.ctypes, push.ctypes, pull.ctypes, nlinks, iters)
else:
def relax(pos, links, mrel1, mrel2, lengths, push, pull, iters):
lengths2 = lengths**2
for i in range(iters):
#p1 = links[:, 0]
#p2 = links[:, 1]
#x1 = pos[p1]
#x2 = pos[p2]
#dx = x2 - x1
#dist = (dx**2).sum(axis=1)**0.5
#mask = (npush & (dist < lengths)) | (npull & (dist > lengths))
##dist[mask] = lengths[mask]
#change = (lengths-dist) / dist
#change[mask] = 0
#dx *= change[:, np.newaxis]
#print dx
##pos[p1] -= mrel2 * dx
##pos[p2] += mrel1 * dx
#for j in range(links.shape[0]):
#pos[links[j,0]] -= mrel2[j] * dx[j]
#pos[links[j,1]] += mrel1[j] * dx[j]
for l in range(links.shape[0]):
p1, p2 = links[l];
x1 = pos[p1]
x2 = pos[p2]
dx = x2 - x1
dist2 = (dx**2).sum()
if (push[l] and dist2 < lengths2[l]) or (pull[l] and dist2 > lengths2[l]):
dist = dist2 ** 0.5
change = (lengths[l]-dist) / dist
dx *= change
pos[p1] -= mrel2[l] * dx
pos[p2] += mrel1[l] * dx