2010.05.16 11:14:43 p.m.
Simple dynamical simulation using python Tags:

This quarter I am taking a class in numerical methods using python's numerical library, numpy. For one of our assignments, we had to write a leapfrog integrator that could handle an arbitrary number of particles connected to the origin by springs and also connected pairwise by springs. After coding up the integrator, I decided I wanted to actually see it in motion.

I stumbled across the a python library called pyglet, which wraps OpenGL in a nice python interface.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from numpy import *
import time

def step_sho(n, x, p, dt, k1=1.0, k2=2.0):
  """Timestep Simple Harmonic Oscillator using leapfrog integrator
  """
  f = (-k1 * x) - (k2 * (n*x - sum(x, 0)))
  p += f * dt
  x += p * dt

if __name__ == "__main__":
  import sys
  import pyglet

  class ho_animation:
    """A simple class to animate system of particles
    """
    def __init__(self, x, p, k1=1.0, k2=2.0):
      assert len(p) == len(x), "x and p must have same length"

      # screen size
      self.w = 640
      self.h = 480

      # spring constants
      self.k1 = k1
      self.k2 = k2

      # locations, number, dimension and momenta
      self.x = x
      self.n, self.d = x.shape
      self.p = p

      # set up timeout to perform animation
      pyglet.clock.schedule_interval(self.step, 1/60.0)

      # create a window to display animation in
      self.win = pyglet.window.Window(self.w, self.h)

      # set default transform to place origin at center of window
      pyglet.gl.glLoadIdentity()
      pyglet.gl.glTranslatef(self.w/2, self.h/2, 0)

      # add event to draw particles on screen
      @self.win.event
      def on_draw():
        self.win.clear()
        pts = self.x[:,0:2].flat #project onto xy plane and flatten
        pyglet.graphics.draw(self.n, pyglet.gl.GL_POINTS, ('v2f', pts))

    # callback for timeout
    def step(self, dt):
      step_sho(self.n, self.x, self.p, 0.01, self.k1, self.k2)

    def run(self):
      pyglet.app.run()

  if len(sys.argv) < 2 or len(sys.argv) > 4:
    print "Usage: %s <number of particles> [k1] [k2]\n" % sys.argv[0]
    exit()

  # read in options
  n = int(sys.argv[1])
  k1 = 1.0
  k2 = 1.0 / n
  if len(sys.argv) > 2:
    k1 = double(sys.argv[2])
  if len(sys.argv) > 3:
    k2 = double(sys.argv[3])

  # create initial particle locations and momenta
  # (two bunches orbiting eachother)
  x = random.randn(n,3) * 20
  x[0:n/2,0] += 100
  x[n/2+1:-1,0] -= 80
  p = random.randn(n,3) * 20
  p[0:n/2,1] += 50
  p[0:n/2,2] += 20
  p[n/2+1:-1,1] -= 50

  # start animation
  ho_animation(x,p, k1, k2).run()
0 Comments

Post a comment