Chain fountain simulation

chain.py

# Chain simulation
import sys
import time
import math
import numpy as np
import pygame
from pygame.locals import *

def LinkForce(pos, l1, l2):
	
	totalForce = np.zeros((3,))
	dpos = pos[l1, :] - pos[l2, :]
	dposLen = np.pow(np.sum(np.pow(dpos, 2.0)), 0.5)
	dposNorm = dpos.copy()
	if dposLen > 1e-9:		
		dposNorm /= dposLen

	linkLen = 10.0
	dlen = dposLen - linkLen
	linkForce = dlen * 0.2

	return -dposNorm * linkForce


def LinkFriction(pos, vel, l1, l2):
	dpos = pos[l1, :] - pos[l2, :]
	dposLen = np.pow(np.sum(np.pow(dpos, 2.0)), 0.5)
	dposNorm = dpos.copy()
	if dposLen > 1e-9:		
		dposNorm /= dposLen

	dvel = vel[l1, :] - vel[l2, :]
	d = dvel.dot(-dposNorm)
	dv = d * -0.05
	
	vel[l1, :] += dv 
	vel[l2, :] -= dv # Transfer momentum


def LinkFriction2(pos, vel, isActive, i):

	if i-1 >= 0 and i+1 < pos.shape[0] and isActive[i] and isActive[i-1] and isActive[i+1]:
		#LinkFriction(pos, vel, i+1, i-1)
		LinkFriction(pos, vel, i, i-1)
		LinkFriction(pos, vel, i, i+1)
	else:
		if i-1 >= 0 and isActive[i] and isActive[i-1]:
			LinkFriction(pos, vel, i, i-1)
		if i+1 < pos.shape[0] and isActive[i] and isActive[i+1]:
			LinkFriction(pos, vel, i, i+1)


if __name__=="__main__":

	pygame.init()

	screen = pygame.display.set_mode((1024,1024))

	color1 = pygame.Color(255, 0, 0)
	color2 = pygame.Color(0, 255, 0)
	color3 = pygame.Color(0, 0, 255)
	colors = [color1, color2, color3] 

	numNodes = 150
	pos = np.zeros((numNodes, 2))
	vel = np.zeros((numNodes, 2))
	prevPos = pos.copy()
	notInTray = np.zeros((numNodes,), dtype=np.uint8)
	isActive = np.zeros((numNodes,), dtype=np.uint8)
	mass = 1.0
	earthMass = 1000.0
	earthVel = 0.0

	isActive[0] = 1

	for i in range(1, 10):
		pos[i, 0] = 200
		pos[i, 1] = 500 + i * 10.0

	cursor = 200.0
	cursorDirection = 1
	for i in range(10, pos.shape[0]):
		cursor += 10.0 * cursorDirection
		if cursor >= 220:
			cursorDirection = -1
		if cursor <= 180:
			cursorDirection = 1
		
		pos[i, 0] = cursor
		pos[i, 1] = 600
		isActive[i] = 1

	for i in range(1, 10):
		notInTray[i] = 1
		isActive[i] = 1

	frameNum = 0
	colourOffset = 0

	while True:

		t = time.time()

		screen.fill("black")

		if frameNum < 30:
			pos[0, 0] = 200
			pos[0, 1] = 500 - 10.0 * frameNum
			# gets to (200, 200) on frame 30
		elif frameNum < 60:
			pos[0, 0] = 300 - 100.0 * math.cos((frameNum-30) * math.pi * 1.0 / 30.0)
			pos[0, 1] = 200 - 100.0 * math.sin((frameNum-30) * math.pi * 1.0 / 30.0)
			#  gets to (400, 200) on frame 60
		elif frameNum < 200:
			pos[0, 0] = 400
			pos[0, 1] = 200 + 10.0 * (frameNum - 60)
		else:
			notInTray[0] = 1

		vel[0, :] = pos[0, :] - prevPos[0, :]
		forceOnEarth = 0.0

		for i in range(pos.shape[0]):

			if i == 0 and frameNum < 200:
				continue

			totalForce = np.zeros((2,))
			
			# Link forces
			if i-1 >= 0 and isActive[i] and isActive[i-1]:
				df = LinkForce(pos, i, i-1)
				totalForce += df
			if i+1 < pos.shape[0] and isActive[i] and isActive[i+1]:
				df = LinkForce(pos, i, i+1)
				totalForce += df

			# Fiction
			#totalForce += -0.001 * vel[i, :]

			# Gravity
			totalForce[1] += mass * 0.05
			forceOnEarth -= mass * 0.05

			accel = totalForce / mass

			vel[i, :] += accel

			#if i == 0:
			#	print (totalForce)
			
		earthAccel = forceOnEarth / earthMass
		earthVel += earthAccel

		# Smooth velocity along chain
		for j in range(1):
		
			for i in range(pos.shape[0]):
				LinkFriction2(pos, vel, isActive, i)

			for i in range(pos.shape[0]-1, -1, -1):
				LinkFriction2(pos, vel, isActive, i)


		prevPos[:, :] = pos[:, :]

		pos += vel

		#print (vel[0, :])

		# Check momentum
		mom = np.zeros((2,))
		for i in range(pos.shape[0]):
			mom += vel[i, :] * mass * isActive[i]
		mom[1] += earthVel * earthMass
		#if frameNum > 200:
		#	print (frameNum, mom)

		# Check for when a node is picked up from tray
		for i in range(1, pos.shape[0]):

			if not isActive[i]: continue
			
			if not notInTray[i]:
				if pos[i, 1] > 600.0:
					pos[i, 1] = 600.0
					vel[i, 1] = 0.0
				
				if pos[i, 1] < 590.0:
					notInTray[i] = 1

		if frameNum > 200:
			dropIndex = None
			for i in range(pos.shape[0]):
				if pos[i, 1] < 5000:
					dropIndex = i
					break

			if dropIndex:
				pos = pos[dropIndex:, :]
				vel = vel[dropIndex:, :]
				prevPos = prevPos[dropIndex:, :]
				notInTray = notInTray[dropIndex:]
				isActive = isActive[dropIndex:]

				colourOffset += dropIndex
				colourOffset = colourOffset % len(colors)

			countInTray = 0
			for i in range(pos.shape[0]):
				countInTray += 1 - notInTray[i]

			if countInTray < 20:
				numToAdd = 10

				pos = np.vstack((pos, np.zeros((numToAdd, 2))))
				vel = np.vstack((vel, np.zeros((numToAdd, 2))))
				prevPos = np.vstack((prevPos, np.zeros((numToAdd, 2))))
				notInTray = np.hstack((notInTray, np.zeros((numToAdd,), dtype=np.uint8)))
				isActive = np.hstack((isActive, np.ones((numToAdd,), dtype=np.uint8)))

				for i in range(pos.shape[0]-numToAdd, pos.shape[0]):
					cursor += 10.0 * cursorDirection
					if cursor >= 220:
						cursorDirection = -1
					if cursor <= 180:
						cursorDirection = 1

					pos[i, 0] = cursor
					pos[i, 1] = 600
					isActive[i] = 1

		scale = 2.0

		for i in range(1, pos.shape[0]):
			pygame.draw.line(screen, colors[0], pos[i-1, :] / scale, pos[i, :] / scale, 2)

		for i in range(pos.shape[0]):
			pygame.draw.circle(screen, colors[(i+colourOffset) % 3], pos[i, :] / scale, 4.0)

		for event in pygame.event.get():
			if event.type == QUIT:
				pygame.quit()
				sys.exit()

		frameNum += 1

		pygame.display.update()
		time.sleep(0.02)