#### Damped 2d wave simulation, with mouse "kicks"

from visual import *

#### parameters: size of array of boxes, strength of
#### interactions, time step, and damping
N = 16
k = 2.
dt = 0.05
damp = 0.1

##### function to take a number and turn into a color
##### zero gives dark magenta; large positive is red, large
##### negative is blue
def httocolor(z):
    u = (tanh(z)+1.)/2.
    return (u,0,(1-u))

##### set up scene
scene.autoscale = 0
scene.center = (N/2,N/2,0)
scene.range = N*.75
scene.forward = (-1,-.5,-.7)
scene.up = (0,0,1)


#### set up a list of lists of boxes, with initial velocity (zero)
bl = []
for i in range(N):
    bl.append([])
    for j in range(N):
        bl[i].append(box(pos=(i,j,0)))
        bl[i][j].velz = sin((i+j)*0.22)*0


#### Go on simulation
while 1:
    #rate(200)
    ### update positions and velocities (only in z-direction)
    for i in range(1,N-1):
        for j in range(1,N-1):
            bl[i][j].z += dt * bl[i][j].velz
            force = k*((bl[i+1][j].z + bl[i-1][j].z
                        +bl[i][j+1].z + bl[i][j-1].z
                        ) - 4 * bl[i][j].z)
            bl[i][j].velz += force * dt
            bl[i][j].velz *= (1. - damp * dt)
            bl[i][j].color = httocolor(bl[i][j].z)
    ### process mouse clicks
    if scene.mouse.clicked:
        m = scene.mouse.getevent()
        if (m.press == 'left'):
            o = m.pickpos
            if o and o.x > 1 and o.x < N-1 and o.y > 1 and o.y < N-1:
                ox = int(o.x)
                oy = int(o.y)
                ### make a "broad" footprint on the boxes
                bl[ox][oy].z = 4.0
                bl[ox+1][oy].z = 2.0
                bl[ox][oy+1].z = 2.0
                bl[ox-1][oy].z = 2.0
                bl[ox][oy-1].z = 2.0
                bl[ox+1][oy+1].z=1.
                bl[ox-1][oy+1].z=1.
                bl[ox-1][oy-1].z=1.
                bl[ox+1][oy-1].z=1.