see http://bit.ly/SZfe3U

#!/usr/bin/python

#
# Honors project for PHYS 4700 - Electricity & Magnetism (milestone 9)
# Shad Sterling <[email protected]>
#
# Green particles have negative charge, Red particles have positive charge
# Blue arrows represent velocity, Yellow arrows represent acceleration
# Grey grid arrows are scaled electric field vectors
#
# Add particles by clicking on empty space (new particles will alternate charge)
# Move particles by drag & drop (next new particle will have opposite charge as moved particle)
# Remove particles by clicking without dragging (next new particle will have same charge)
# Advance motion by clicking on any motion arrow (velocity or acceleration)
# Continuously advance motion by dragging on a motion arrow and holding button down
#


from __future__ import division
from visual import *
import random
import os
import inspect
import pprint
import gc
import time

## - use for real units
#arena_size = 65   # arena size in meters (significantly smaller sizes are problematic)
#arena_grid = 24   # number of gridpoints across each dimension of the arena
#step_time = 0.1   # simulation time per step in seconds
#frame_time = 0.5  # minimum real time per step in seconds when continuously stepping

## - use for simplified units
arena_size = 26
arena_grid = 24
step_time = 0.5
frame_time = 0.5



class Arena:
    grid_base = 0.25     # minimum length for grid vectors as proportion of grid spacing
    grid_max = 1/4       # maximum length for grid vectors as proportion of arena dimension
    rebound_ratio = 0.1  # proportion of velocity reflected from boundary
    coulomb_constant = 1 # 8.9875517873681764 * 10**9 # N*m*m/C/C # simplified: 1

    def __init__( self, size, grid, scene ):
        self.size = size # size of arena in each dimension in meters
        self.boundary = size/2
        self.particles = []
        self.grid = []
        scene.autoscale = false
        scene.autocenter = false
        scene.range = self.boundary
        self.regrid( grid ) # setup indicator grid
        self.accel_limit = self.size

    def regrid( self, grid ):
        for point in self.grid:
            point.hide()
        self.grid_size = ceil(abs(grid)) # gridpoint count in each dimension
        self.grid_spacing = self.size/(self.grid_size+1) # distance in meters between gridpoints (no points on edges)
        self.grid_base = self.grid_spacing * Arena.grid_base
        self.grid_max = self.size * Arena.grid_max
        start = -self.boundary + self.grid_spacing
        stop = self.boundary - self.grid_spacing*1.5
        x = start
        y = self.boundary - self.grid_spacing
        i = self.grid_size ** 2 #number of gridpoints
        while i > 0:
            self.grid.append( GridPoint( self, vector( x, y, 0 ) ) )
            if x < stop:
                x += self.grid_spacing
            else:
                x = start
                y -= self.grid_spacing
            i = i - 1
        for particle in self.particles:
            show = particle.Avatar.visible
            particle.Avatar.hide()
            particle.avatar = Avatar( particle )
            if show:
                particle.Avatar.show()

    def capacitor( self, count, length, width, pos, angle ):
        c = cos( angle )
        s = sin( angle )
        count = count - 1
        next = vector( length*s/count, length*c/count, 0 )
        gap = vector( width*c/2, -width*s/2, 0 )
        pos = vector(pos) - count*next/2
        while( 0 <= count ):
            self.electron( pos-gap )
            self.positron( pos+gap )
            pos = pos + next
            count = count - 1

    def electron( self, position, velocity=(0,0,0) ):
        new = Particle.electron( self, position, velocity )
        self.add( new )
        return new

    def positron( self, position, velocity=(0,0,0) ):
        new = Particle.positron( self, position, velocity )
        self.add( new )
        return new

    def add( self, particle ):
        self.particles.append( particle )
        particle.reveal()

    def remove( self, particle ):
        particle.hide()
        self.particles.remove( particle )

    def bound( self, particle ):
        i = 0
        while i < 2:
            if particle.position[i] > self.boundary:
                particle.position[i] = self.boundary
                particle.velocity[i] = -abs( particle.velocity[i] * self.rebound_ratio )
            elif particle.position[i] < -self.boundary:
                particle.position[i] = -self.boundary
                particle.velocity[i] = abs( particle.velocity[i] * self.rebound_ratio )
            i = i + 1

    def update( self ):
        for point in self.grid:
            point.update()
        for particle in self.particles:
            particle.update() # update acceleration

    def step( self, step_time ):
        for particle in self.particles:
            particle.step( step_time ) # update position & velocity
            self.bound( particle )     # constrain to arena bounds
        self.update()

    def sync( self ):
        for particle in self.particles:
            particle.sync()
        for point in self.grid:
            point.sync()

class GridPoint:
    arrow_radius = 0.1
    arrow_scale = 1 # 10 ** 10 # simplified: 1
    color_electric = color.gray(0.7)

    def __init__( self, arena, position ):
        self.arena = arena
        self.position = position
        self.avatar = arrow( pos = position, fixedwidth = true, color = GridPoint.color_electric,
                             shaftwidth = GridPoint.arrow_radius * self.arena.grid_spacing,
                             role = "grid" )
        self.electric = vector( 0,0,0 )

    def update( self ):
        electric = vector( 0,0,0 )
        for particle in self.arena.particles:
            range = self.position - particle.position
            if 0 != particle.charge:
                electric += Arena.coulomb_constant * range * particle.charge / (range.mag**3)
        electric.mag = self.arena.grid_base + sqrt(electric.mag)*GridPoint.arrow_scale # simplified
        electric.mag = self.arena.grid_base + electric.mag * GridPoint.arrow_scale # real
        if electric.mag > self.arena.grid_max:
            electric.mag = 0
        self.electric = electric
        self.in_sync = false

    def sync( self ):
        self.avatar.axis = self.electric

class Particle:
    electron_mass = 1 # kilograms; real: 9.10938291 * 10**(-31) # kilograms # simplified: 1
    electron_charge = -1 # -1.602176565 * 10**(-19) # Coulombs # simplified: -1

    @classmethod
    def electron( self, arena, position, velocity=(0,0,0) ):
        return self( arena, Particle.electron_mass, Particle.electron_charge, position, velocity )

    @classmethod
    def positron( self, arena, position, velocity=(0,0,0) ):
        return self( arena, Particle.electron_mass, -Particle.electron_charge, position, velocity )

    def __init__( self, arena, mass, charge, position, velocity, acceleration=(0,0,0) ):
        self.arena = arena
        self.mass = mass
        self.charge = charge
        self.position = vector( position )
        self.velocity = vector( velocity )
        self.acceleration = vector( acceleration )
        self.avatar = Avatar( self )
        self.in_sync = true #true when avatar matches

    def sync( self ): # synchronize avatar
        self.avatar.update()
        self.in_sync = true

    def step( self, step_time ): # step position & velocity
        self.position += self.velocity * step_time
        self.velocity += self.acceleration * step_time
        self.in_sync = false #avatar is outdated

    def update( self ): # recalculate acceleration
        electric = vector( 0, 0, 0 )
        for particle in self.arena.particles:
            if 0 != particle.charge and self is not particle:
                range = self.position - particle.position
                force = Arena.coulomb_constant * range * self.charge * particle.charge / (range.mag**3)
                electric += force
        self.acceleration = max( self.arena.accel_limit, electric / self.mass )
        self.in_sync = false #avatar is outdated

    def move( self, new_position ):
        self.position = new_position
        self.in_sync = false #avatar is outdated

    def hide( self ):
        self.avatar.hide()

    def reveal( self ):
        self.avatar.reveal()

    def dim( self ):
        self.avatar.dim()

    def bright( self ):
        self.avatar.bright()

class Avatar:
    avatar_radius = 0.5 # for avatars of point particles
    color_negative = (0,.9,0) # negative particles are (almost) green
    color_positive = (1,.1,.1) # positive particles are (almost) red
    arrow_radius = 0.35
    arrow_scale = 3 # 1 # proportional length of velocity & acceleration arrows # simplified: 3
    color_velocity = (0,.2,1) #velocity arrows are (almost) blue
    color_acceleration = color.yellow #acceleration arrows are yellow

    @classmethod
    def color_of( self, charge ):
        if charge < 0:
            return Avatar.color_negative
        elif charge > 0:
            return Avatar.color_positive

    @classmethod
    def arrow_of( self, base, motion ): # velocity arrow
        m = vector( motion )
        m.mag = m.mag * Avatar.arrow_scale + base
        return m

    def __init__( self, particle ):
        self.radius = Avatar.avatar_radius * particle.arena.grid_spacing
        self.arrow_width = Avatar.arrow_radius * particle.arena.grid_spacing
        self.position = visual.sphere( color = Avatar.color_of( particle.charge ),
                                       radius = self.radius, role = "particle", real = particle )
        self.velocity = visual.arrow( shaftwidth = self.arrow_width, fixedwidth = true,
                                      color = Avatar.color_velocity, role = "motion" )
        self.acceleration = visual.arrow( shaftwidth = self.arrow_width, fixedwidth = true,
                                          color = Avatar.color_acceleration, role = "motion" )
        self.represents = particle
        self.represents.avatar = self
        self.update()

    def update( self ):
        self.position.pos = self.represents.position
        self.velocity.pos = self.represents.position
        self.velocity.axis = Avatar.arrow_of( self.radius, self.represents.velocity )
        self.acceleration.pos = self.represents.position
        self.acceleration.axis = Avatar.arrow_of( self.radius, self.represents.acceleration )

    def hide( self ):
        self.position.visible = false
        self.velocity.visible = false
        self.acceleration.visible = false
        self.represents.visible = false

    def reveal( self ):
        self.position.visible = true
        self.velocity.visible = true
        self.acceleration.visible = true
        self.represents.visible = true

    def dim( self ):
        self.position.opacity = 0.75
        self.velocity.opacity = 0.75
        self.acceleration.opacity = 0.75

    def bright( self ):
        self.position.opacity = 1
        self.velocity.opacity = 1
        self.acceleration.opacity = 1


arena = Arena( arena_size, arena_grid, scene )
arena.capacitor( ceil(arena_grid/3), arena_size*2/3, arena_size/2, (0,0,0), 2*math.pi*random.random() )
arena.positron( (0,0,0), rotate( (arena_size/4/Avatar.arrow_scale,0,0), 2*math.pi*random.random(), (0,0,1) ) )

negative = true
target = None
event = None
step = false
changed = true
run = false
new = false

ticksum = 0
tickcount = 0
tickmisscount = 0
tickinterval = frame_time

while true:
    if changed:
        tickstart = time.time()
        if step:
            #print "stepping..."
            arena.step( step_time )
        arena.update() #recalculate forces
        gc.collect()
        changed = false
        tickend = time.time()
        ticktime = tickend - tickstart # tick time (not including waiting)
        tickcount = tickcount + 1
        ticksum = ticksum + ticktime
        tickavg = ticksum/tickcount # average seconds per tick
        #if ticktime > tickinterval:
        #   tickmisscount += 1
        #tickmissrate = tickmisscount / tickcount
        tickinterval = max( frame_time, tickavg * 2 )
        tickwait = max( 0, tickinterval - ticktime )
        #print "tick", tickcount, "took", ticktime, "avg", tickavg, "sum", ticksum, "interval", \
        #      tickinterval, "missed", tickmisscount, "rate", tickmissrate, "wait", tickwait
        time.sleep( tickwait )
        arena.sync() #update avatars
    #else:
        #print "refresh skipped"
    if 0 == scene.mouse.events and ( run or None != target ) and None != event and event.drag:
        #print "dragging"
        newpos = vector( scene.mouse.pos[0], scene.mouse.pos[1], 0 )
    else:
        #print "awaiting event..."
        event = scene.mouse.getevent()
        newpos = vector( event.pos[0], event.pos[1], 0 )
    if event.press:
        if None != event.pick:
            role = event.pick.role
        else:
            role = None
        if "particle" == role:
            target = event.pick.real
            #print "press:  targeted", target, "at", target.position
            startpos = vector( target.position )
        elif "motion" == role:
                #print "press:  step"
                step = true
                changed = true
                run = true
        elif None != target:
                #print "press:  replace at", newpos, "from", startpos
                target.move( newpos )
                arena.add( target )
                negative = target.charge > 0
                changed = true
                new = true
        elif negative:
            #print "press:  electron at", newpos
            target = arena.electron( newpos )
            startpos = vector(target.position)
            negative = false
            changed = true
            new = true
        else:
            #print "press:  positron at", newpos
            target = arena.positron( newpos )
            startpos = vector(target.position)
            negative = true
            changed = true
            new = true
    elif event.drag:
        if None != target:
            #print "drag:  dragged to", newpos, "via", target.position, "from", startpos
            target.move( newpos )
            target.dim()
            changed = true
            drag = true
        elif run:
            #print "drag:  continue stepping"
            step = true
            changed = true
        else: # how would this happen?
            #print "drag:  no object targeted"
            changed = false
    elif event.release:
        if None != target:
            if false == new and target.position == startpos:
                #print "release: removed from", startpos
                arena.remove( target )
                changed = true
            else:
                #print "release: dropped at", newpos, "from", startpos
                target.move( newpos )
                target.bright()
                negative = target.charge > 0
                target = None
                startpos = None
                changed = true
        elif run:
            #print "release:  end stepping"
            step = false
            run = false
        else:
            #print "release:  nothing in progress"
            changed = false
    else:
        print "Unknown Event!", pprint.pprint( inspect.getmembers( event ) )
        changed = false