VPython Lab9.

e09.py

from __future__ import division
from visual import *
import numpy as np

radS = 0.4
radH = 0.4
length = 100.
k = 0.5
dt = 0.1
n_balls = 52

scene = display(
    width=800,
    height=600,
    center=(length/2., 0, 0),
    background=color.black
)


def _init_vectors(n):
    n_shape = (n+2, 3)
    pos = np.zeros(n_shape)
    pos[:, 0] = np.linspace(0, length, n+2)
    vel = np.zeros(n_shape)
    f = np.zeros(n_shape)
    ax = pos[1:, :] - pos[:-1, :]
    return pos, vel, f, ax

position, velocity, force, axis = _init_vectors(n_balls)
velocity[1, 1] += 10.
velocity[-2, 1] -= 10
balls = [sphere(radius=radS, color=color.red) for r in position]
helices = [helix(radius=radH, thickness=0.1, color=color.yellow) for p, a in zip(position, axis)]

while 1:
    rate(100)

    for i, r in enumerate(position[:-1, :]):
        balls[i].pos = r
        helices[i].pos = r
    balls[-1].pos = position[-1, :]
    for i, r in enumerate(axis):
        helices[i].axis = r

    force[1:-1, :] = (-k) * (2 * position[1:-1, :] - position[2:, :] - position[:-2, :])

    velocity[1:-1, :] += force[1:-1, :] * dt
    position[1:-1, :] += velocity[1:-1, :] * dt
    axis[:, :] = position[1:, :] - position[:-1, :]