77906fc7a2
Add pure-python integrator to verlet chain example
71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
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
|