see http://bit.ly/SZfe3U

#!/usr/bin/python

#
# Honors project for PHYS 4700 - Electricity & Magnetism
# 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


grid_span = 25    #number of field vectors across field width & height
vec_width = 0.1   #width of grid vector arrows
vec_scale = 1     #proportional length of vector arrows
vec_limit = 0.5   #don't draw vectors more than this factor of the display width
vec_halo = 0.9    #don't draw vectors within in this factor of a charge radius
move_width = 0.6  #proportion of charge radius used for velocity & acceleration arrow widths
move_scale = 3    #proportional length of velocity & acceleration arrows
move_limit = 1    #distance for calculating force must be at most this proportion of radius
rebound = 0.1     #velocity reflected from boundary
step_time = 0.5   #simulation time per frame


def electron( pos ):
    e = sphere()
    e.pos = pos
    e.radius = .5 #arbitrary?
    e.charge = -1 # 1e = -1.602176565(35) * 10**(-19) C
    e.mass = 1 # 9.10938291(40) * 10**(-31) kg
    e.color = (0,.9,0)
    e.removable = true
    e.vel = arrow( pos = e.pos, shaftwidth = e.radius*move_width, fixedwidth = true,
                   axis = vector( 0, 0, 0 ), color = (0,.2,1), stepper = true )
    e.accel = arrow( pos = e.pos, shaftwidth = e.radius*move_width, fixedwidth = true,
                     axis = vector( 0, 0, 0 ), color = color.yellow, stepper = true )
    return e

def positron( pos ):
    e = sphere()
    e.pos = pos
    e.radius = .5 #arbitrary?
    e.charge = 1 # 1e = -1.602176565(35) * 10**(-19) C
    e.mass = 1 # 9.10938291(40) * 10**(-31) kg
    e.color = (1,.1,.1)
    e.removable = true
    e.vel = arrow( pos = e.pos, shaftwidth = e.radius*move_width, fixedwidth = true,
                   axis = vector( 0, 0, 0 ), color = (.2,0,1), stepper = true )
    e.accel = arrow( pos = e.pos, shaftwidth = e.radius*move_width, fixedwidth = true,
                     axis = vector( 0, 0, 0 ), color = color.yellow, stepper = true )
    return e

def move( obj, pos ):
    obj.pos = pos
    obj.vel.pos = pos
    obj.accel.pos = pos
    return None

def hide( obj ):
    obj.visible = false
    obj.vel.visible = false
    obj.accel.visible = false

def reveal( obj ):
    obj.visible = true
    obj.vel.visible = true
    obj.accel.visible = true

def dim( obj ):
    obj.opacity = 0.75
    obj.vel.opacity = 0.75
    obj.accel.opacity = 0.75

def bright( obj ):
    obj.opacity = 1
    obj.vel.opacity = 1
    obj.accel.opacity = 1

def getvel( obj ):
    newvel = vector(obj.vel.axis)
    newvel.mag = (newvel.mag-obj.radius)/move_scale
    return newvel

def setvel( obj, newvel ):
    obj.vel.axis = vector( newvel )
    obj.vel.axis.mag = obj.vel.axis.mag*move_scale+obj.radius
    return obj.vel.axis

def getaccel( obj ):
    newaccel = vector(obj.accel.axis)
    newaccel.mag = (newaccel.mag-obj.radius)/move_scale
    return newaccel

def setaccel( obj, newaccel ):
    obj.accel.axis = vector( newaccel )
    obj.accel.axis.mag = obj.accel.axis.mag*move_scale+obj.radius
    return obj.accel.axis

def getarrow( objects, gridarrow ):
    axis = vector( 0, 0, 0 )
    for obj in objects:
        if hasattr( obj, "charge" ):
            r = gridarrow.pos - obj.pos #from arrow to obj
            if r.mag < obj.radius*vec_halo:
                axis = vector( 0, 0, 0 )
                break
            axis += r * obj.charge / (r.mag**3)
    axis.mag = 0.5 + sqrt(axis.mag)*vec_scale
    if axis.mag < gridarrow.shaftwidth or axis.mag > vec_limit*scene.range:
        axis = vector( 0, 0, 0 )
    return axis

def getforce( objects, object ):
    force = vector( 0, 0, 0 )
    for obj in objects:
        if obj is not object and hasattr( obj, "charge" ):
            r = object.pos - obj.pos
            r.mag = max( r.mag, object.radius * move_limit )
            force += r * object.charge * obj.charge / (r.mag**3)
    return force

def repos( obj, pos=None ):
    i = 0
    r = false
    if None == pos:
        p = vector( obj.pos )
    else:
        p = pos
    v = getvel( obj )
    t = vector( p )
    while i < 2:
        if abs(p[i]) > scene.range[i]:
            p[i] = sign(p[i]) * scene.range[i]
            v[i] = -rebound * v[i]
            r = true
        i = i + 1
    move( obj, p )
    setvel( obj, v )
    return r

def capacitor( 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 ):
        electron( pos-gap )
        positron( pos+gap )
        pos = pos + next
        count = count - 1
    return None

def setup( scene, size, spacing=1 ): #range is width of field arrow area
    half = ceil(abs(size/2))
    scene.autoscale = false
    scene.autocenter = false
    scene.range = half
    for obj in scene.objects:
        if hasattr( obj, "gridarrow" ):
            obj.visible = false
    for y in xrange( int(floor(-half)), int(ceil(half)), spacing ):
        for x in xrange( int(floor(-half)), int(ceil(half)), spacing ):
            a = arrow( pos = ( x, y, 0 ), shaftwidth = vec_width, fixedwidth = true,
                       color = color.gray(0.7) )
            a.gridarrow = true
            a.axis = getarrow( scene.objects, a )
    return None

def refresh( objects ):
    for obj in objects:
        if hasattr( obj, "gridarrow" ):
            obj.axis = getarrow( objects, obj )
        if hasattr( obj, "accel" ) and 0 != obj.charge:
            setaccel( obj, getforce( objects, obj ) / obj.mass )
    return None

def update( objects ):
    for obj in objects:
        if hasattr( obj, "accel" ) and 0 != obj.accel.axis.mag:
            #repos moves within bounds, handles rebound velocity; accel is updated by refresh()
            repos( obj, obj.pos + getvel(obj)*step_time )
            setvel( obj, getvel(obj)+getaccel(obj)*step_time )
    return None

setup( scene, grid_span )
capacitor( ceil(grid_span/3), grid_span*2/3, grid_span/2, (0,0,0), 2*math.pi*random.random() )
#rotate described by http://www.vpython.org/webdoc/visual/vector.html
setvel( positron( (0,0,0) ),
        rotate( (grid_span/6/move_scale,0,0), 2*math.pi*random.random(), (0,0,1) ) )
negative = true
target = None
event = None
step = false
changed = true
run = false

ticksum = 0
tickcount = 0
tickwait = 0.5

while true:
    if changed:
        #tickstart = time.time()
        if step:
            update( scene.objects )
        refresh( scene.objects )
        gc.collect()
        changed = false
        #time.sleep( tickwait )
        #tickend = time.time()
        #ticktime = tickend - tickstart #tick time, including sync wait
        #tickcount = tickcount + 1
        #ticksum = ticksum + ticktime
        #tickrate = tickcount/ticksum # ticks per second
        #tickwait = ticksum/tickcount/2 # half of the average tick time
        #print "tick", tickcount, ticktime, ticksum, 1/tickrate, tickwait, tickrate
        #visual.rate( tickrate ) #why doesn't this affect drawing?
        #print "refresh complete"
    #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:
            if hasattr( event.pick, "removable" ):
                target = event.pick
                startpos = vector(target.pos)
                changed = false
                #print "press:  targeted", startpos
                ##TODO: return to start if focus leaves while pressed
            elif hasattr( event.pick, "stepper" ):
                #print "press:  stepping"
                #update( scene.objects )
                step = true
                changed = true
                run = true
        elif None != target:
                #print "press:  replace at", newpos, "from", startpos
                move( target, newpos )
                reveal( target )
                negative = target.charge > 0
                changed = true
        else:
            if negative:
                #print "press:  electron at", newpos
                target = electron( newpos )
                startpos = vector(target.pos)
                negative = false
                changed = true
            else:
                #print "press:  positron at", newpos
                target = positron( newpos )
                startpos = vector(target.pos)
                negative = true
                changed = true
    elif event.drag:
        if None != target and target.pos != newpos:
            #print "drag:  dragged to", newpos, "via", target.pos, "from", startpos
            move( target, newpos )
            dim( target )
            changed = true
            drag = true
        elif run:
            #print "drag:  continue stepping"
            #update( scene.objects )
            step = true
            changed = true
        else:
            #print "drag:  no object targeted"
            changed = false
    elif event.release:
        if None != target:
            if target.pos == startpos:
                #print "release: removed from", startpos
                hide( target )
                changed = true
            else:
                #print "release: dropped at", newpos, "from", startpos
                move( target, newpos )
                bright( target )
                negative = target.charge > 0
                target = None
                startpos = None
                changed = true
        elif run:
            #print "release:  end stepping"
            run = false
        else:
            #print "release:  nothing in progress"
            changed = false
    else:
        print "Unknown Event!", pprint.pprint( inspect.getmembers( event ) )
        changed = false

os._exit(0)