Positronium! ...ish

#!/usr/bin/python

#
# Honors project for PHYS 4700 - Electricity & Magnetism (milestone 13)
# 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 print_function, division
from visual import *
import random
import time
import gc
import inspect
import pprint

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.05  # simulation time per step in seconds
frame_time = 1/4  # minimum real time per step in seconds when continuously stepping
#drag_time = 0.1   # sleep time for drag position updates


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 = 8.9875517873681764 * 10**9 # N*m*m/C/C

    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 ):
        angle = math.pi/6
        start = rotate( (length/2,0,0), math.pi/2-angle, (0,0,1) )
        side = (count-1)/2
        next = -start/side
        position = start + pos # backward for type safety
        gap = rotate( (width/2,0,0), math.pi*2-angle, (0,0,1) )
        while 0 < count:
            self.electron( position-gap )
            self.positron( position+gap )
            position = position + next
            count -= 1

    def circle( self, count, radius, pos, angle ):
        inc = math.pi/count
        while 0 < count:
            offset = rotate( (radius,0,0), angle, (0,0,1) )
            self.positron( offset+pos )
            angle += inc
            offset = rotate( (radius,0,0), angle, (0,0,1) )
            self.electron( offset+pos )
            angle += inc
            count -= 1

    def positronium( self, radius, angle, pos=(0,0,0), clockwise = false ):
        offset = rotate( (radius,0,0), angle, (0,0,1) )
        p = self.positron( offset+pos )
        n = self.electron( -offset+pos )
        f = Particle.force_electric( p, n )
        d = 1
        if clockwise:
            d = -d
        p.acceleration = f/p.mass
        a = p.acceleration * sqrt( p.acceleration.mag * radius )
        p.velocity = rotate( a, d*math.pi/2, (0,0,1) )
        n.acceleration = -f/n.mass
        a = n.acceleration * sqrt( n.acceleration.mag * radius )
        n.velocity = rotate( a, d*math.pi/2, (0,0,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 = 10 ** 10
    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 + electric.mag * GridPoint.arrow_scale
        if electric.mag > self.arena.grid_max:
            electric.mag = 0
        self.electric = electric
        self.in_sync = false

    def sync( self ):
        if not self.in_sync:
            self.avatar.axis = self.electric
            self.in_sync = true

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

    @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 )

    @classmethod
    def force_electric( self, p1, p2 ):
        range = p1.position - p2.position
        return Arena.coulomb_constant * range * p1.charge * p2.charge / (range.mag**3)

    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 __str__( self ):
        return "{ Particle "+str(id(self))+": "+str(self.mass)+"kg "+str(self.charge)+"C @"\
               +str(self.position)+" & "+str(id(self.avatar))+"=>"+str(id(self.avatar.position))+" }"

    def sync( self ): # synchronize avatar
        if not self.in_sync:
            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:
                force = Particle.force_electric( self, particle )
                electric += force
        self.acceleration = electric / self.mass
        if self.acceleration.mag > self.arena.accel_limit:
            self.acceleration.mag = self.arena.accel_limit
        self.in_sync = false #avatar is outdated

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

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

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

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 = 1 # proportional length of velocity & acceleration arrows
    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.update()

    def __str__( self ):
        return "{ Avatar "+str(id(self))+": "+str(id(self.position))+" "+str(id(self.velocity))+" "\
               +str(id(self.acceleration))+" }"

    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

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



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) ) )
#arena.circle( 1, arena_size/20, (0,0,0), 2*math.pi*random.random() )
#arena.circle( ceil(arena_grid/3), arena_size/6, (0,0,0), 2*math.pi*random.random() )
arena.positronium( arena_size/7.85, 2*math.pi*random.random() )

negative = true
target = None
event = None
step = false
changed = true
drag = false # dragging a particle
run = false # stepping continuously (~dragging motion)

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

while true:
    if changed:
        tickstart = time.time()
        if step:
            #print( "calculating step" )
            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 * 1.35 )
        #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
        #print( "refresh complete" )
    #else:
        #print( "refresh skipped" )
    if 0 == scene.mouse.events and (drag or run):
        #print( "dragging..." )
        #time.sleep( drag_time )
        newpos = vector( scene.mouse.pos[0], scene.mouse.pos[1], 0 )
    else:
        #print( "dequeing 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:  "+str(id(event.pick))+" selects "+str(target) )
        elif "motion" == role:
            #print( "press:  "+str(id(event.pick))+" steps once" )
            step = true
            changed = true
            run = true
        elif None != target:
            #print( "press:  "+str(id(event.pick))+" replaces "+str(target)+" at "+srt(newpos) )
            target.move( newpos )
            arena.add( target )
            negative = target.charge > 0
            changed = true
            drag = true
        elif negative:
            target = arena.electron( newpos )
            #print( "press:  "+str(id(event.pick))+" creates "+str(target) )
            negative = false
            changed = true
            drag = true
        else:
            target = arena.positron( newpos )
            #print( "press:  "+str(id(event.pick))+" creates "+str(target) )
            negative = true
            changed = true
            drag = true
    elif event.drag:
        if None != target:
            #print( "drag:  "+str(id(event.pick))+" moves "+str(target)+" to "+str(newpos) )
            target.move( newpos )
            changed = true
            drag = true
        elif run:
            #print( "drag:  "+str(id(event.pick))+" steps again" )
            step = true
            changed = true
        else: # how would this happen?
            #print( "drag:  "+str(id(event.pick))+" remains idle\t\t(I sense a disturbance...)" )
            changed = false
    elif event.release:
        if None != target:
            if false == drag:
                #print( "release:  "+str(id(event.pick))+" removes "+str(target) )
                arena.remove( target )
                target = None
                changed = true
            else:
                #print( "release:  "+str(id(event.pick))+" drops "+str(target)+" at "+str(newpos) )
                target.move( newpos )
                negative = target.charge > 0
                target = None
                drag = false
                changed = true
        elif run:
            #print( "release:  "+str(id(event.pick))+" stops stepping" )
            step = false
            run = false
        else:
            #print( "release:  "+str(id(event.pick))+" remains idle" )
            changed = false
    else:
        print( "Unknown Event!", pprint.pprint( inspect.getmembers( event ) ) )
        changed = false