research/transmission_spectroscopy/simulator/old.py

   1 import sys\r
   2 import os\r
   3 import random\r
   4 import pygame\r
   5 \r
   6 scale = 20 #Scale in pixels per angstrom\r
   7 \r
   8 mass = {"proton": 10**4, "electron":1} #Masses normalised to electron mass\r
   9 charge = {"proton":+1, "electron":-1} #Charges normalised to fundamental charge\r
  10 k = 10**-3 #Coulomb constant\r
  11 A = 10**-3 #Repulsion constant\r
  12 m = 6 #Inverse power for the repulsion law (Fr = A/r^m)\r
  13 \r
  14 class particle():\r
  15    def __init__(self, mass, charge, r):\r
  16       """Create a particle"""\r
  17       self.mass = mass\r
  18       self.charge = charge\r
  19       self.r = r\r
  20       self.v = [0,0]\r
  21       self.f = [0,0]\r
  22 \r
  23    def step(self):\r
  24       """Do a single step"""\r
  25       #Update velocity\r
  26       if self.mass != 0:\r
  27          for i in range(0,2):\r
  28             self.v[i] += self.f[i] / self.mass\r
  29             \r
  30       #Update position \r
  31       for i in range(0,2):\r
  32          self.r[i] += self.v[i]\r
  33 \r
  34       #Reflect off walls\r
  35       if self.r[0] <= 0 or self.r[0] >= 640:\r
  36          self.v[0] = -self.v[0]\r
  37       if self.r[1] <= 0 or self.r[1] >= 480:\r
  38          self.v[1] = -self.v[1]\r
  39          \r
  40       \r
  41 \r
  42       \r
  43       \r
  44          \r
  45       #sys.stdout.write("Position of " + str(self) + " is " + str(self.r) + "\n")\r
  46                      \r
  47 \r
  48    def calcForce(a,b):\r
  49       r = [0,0]\r
  50       for i in range(0,2):\r
  51          r[i] = b.r[i] - a.r[i]\r
  52       if r[0] == 0 and r[1] == 0:\r
  53          return [0,0]\r
  54       #if r[0] > 10**-10 or r[1] > 10**-10:\r
  55        #  return [0,0]\r
  56       f = [0,0]\r
  57       #Coulomb force\r
  58 \r
  59       distance = (r[0]**2 + r[1]**2)**(1/2)\r
  60       #sys.stderr.write(str(r[i] * (k * a.charge * b.charge)) + " ")\r
  61       #sys.stderr.write("Distance is " + str(distance) + "\n")\r
  62       for i in range(0,2):\r
  63          #Print out calculation\r
  64          #sys.stderr.write("f["+str(i)+"] -= " + str(r[i]) + " * (" + str(k) + " * " + str(a.charge) + " * " + str(b.charge) + " / (" + str(r[0]) + "^2 + " + str(r[1]) + "^2)^(3/2)\n")\r
  65          f[i] -= r[i] * (k * a.charge * b.charge / (distance**3))\r
  66       #Repulsive force\r
  67       #if distance > 10:\r
  68       for i in range(0,2):\r
  69          f[i] -= r[i] * (A / ((distance)**(m+1)))\r
  70       #else:\r
  71        #  for i in range(0,2):\r
  72         #    f[] -= r[i] * (A / ((distance-10)**(m+1)))\r
  73       return [f[0], f[1]]\r
  74 \r
  75    def draw(self, window):\r
  76       \r
  77       if self.charge == 1:\r
  78          colour = pygame.Color(255,0,0,1)\r
  79          size = 4\r
  80       elif self.charge == -1:\r
  81          colour = pygame.Color(0,0,255,1)\r
  82          size = 4\r
  83       elif self.charge == 0:\r
  84          colour = pygame.Color(255,255,255,1)\r
  85          size = 4\r
  86       else:\r
  87          colour = pygame.Color(0,255,0,1)\r
  88          size = 4\r
  89 \r
  90       pygame.draw.circle(window,colour,[int(self.r[0] * scale), int(self.r[1] * scale)],size,0)\r
  91       #pygame.draw.line(window,colour,[int((self.r[0] - self.v[0]) / scale), int((self.r[1] - self.v[1]) / scale)], [int(self.r[0] / scale), int(self.r[1] / scale)])\r
  92       \r
  93       \r
  94 \r
  95    \r
  96 class universe():\r
  97    def __init__(self):\r
  98       self.particles = []\r
  99       for i in range(0,0):\r
 100          self.particles.append(particle(mass["proton"],charge["proton"], universe.getRandomPosition()))\r
 101          self.particles.append(particle(mass["proton"] + 2*mass["electron"],charge["electron"],universe.getRandomPosition()))\r
 102 \r
 103    def step(self):\r
 104       for a in self.particles:\r
 105          a.f = [0,0]\r
 106          for b in self.particles:\r
 107             if b != a:\r
 108                fAB = particle.calcForce(a,b)\r
 109                for i in range(0,2):\r
 110                   a.f[i] += fAB[i]\r
 111 \r
 112          a.step()\r
 113 \r
 114    def getRandomPosition():\r
 115       return [float(random.randint(0,640) / scale), float(random.randint(0,480) / scale)]\r
 116 \r
 117    def draw(self, window):\r
 118       for a in self.particles:\r
 119          a.draw(window)\r
 120 \r
 121 if __name__ == "__main__":\r
 122    #create the screen\r
 123    window = pygame.display.set_mode((640, 480))\r
 124    theUniverse = universe()\r
 125 \r
 126    run = True\r
 127    while run:\r
 128       #Clear screen\r
 129       \r
 130      \r
 131       #Step\r
 132       #sys.stdout.write("Stepping...\n")\r
 133       theUniverse.step()\r
 134       #Draw\r
 135       window.fill(pygame.Color(0,0,0,0))\r
 136       theUniverse.draw(window)\r
 137       #Update draw events\r
 138       pygame.display.flip()\r
 139       #Wait a bit\r
 140       #pygame.time.wait(1)\r
 141       #Input handling\r
 142       for event in pygame.event.get(): \r
 143          if event.type == pygame.QUIT: \r
 144              pygame.quit()\r
 145              run = False\r
 146          elif event.type == pygame.MOUSEBUTTONDOWN:\r
 147              mousestate = pygame.mouse.get_pressed()\r
 148              mousepos = pygame.mouse.get_pos()\r
 149              if mousestate[0] and mousestate[2] == False:\r
 150                 theUniverse.particles.append(particle(mass["proton"],charge["proton"], [mousepos[0]/scale,mousepos[1]/scale]))\r
 151              elif mousestate[2] and mousestate[0] == False:\r
 152                 theUniverse.particles.append(particle(mass["electron"],charge["electron"], [mousepos[0]/scale,mousepos[1]/scale]))\r
 153              break\r
 154          else: \r
 155              #print(event)\r
 156             pass\r
 157             \r