research/transmission_spectroscopy/simulator/simulation.py

   1 """\r
   2    Drude model electron gas simulation\r
   3    @author Sam Moore\r
   4    @email [email protected]\r
   5 """\r
   6 \r
   7 import sys\r
   8 import os\r
   9 import random\r
  10 import pygame #(Slow, but fast to code)\r
  11 \r
  12 \r
  13 scale = 20 #Scale in pixels (per angstrom)\r
  14 \r
  15 \r
  16 \r
  17 mass = {"proton": 10**4, "neutron":10**4, "electron":1} #Masses normalised to electron mass\r
  18 charge = {"proton":+1, "electron":-1, "neutron":0} #Charges normalised to fundamental charge\r
  19 particleType = {"proton":"hadron", "neutron":"hadron", "electron":"lepton"}\r
  20 particleSize = hadronSize = 10**(-5) #Size of a hadron (in angstroms)\r
  21 k = 10**-2 #Coulomb constant (Supposed to be 10**-8 but I increased to speed things up)\r
  22 A = 10**-2 #Repulsion constant (Seems to work best when A = k. Completely arbitrarily chosen)\r
  23 m = 4 #Inverse power for the repulsion law (Fr = A/r^m). Completely arbitrarily chosen.\r
  24 \r
  25 \r
  26 \r
  27 class particle():\r
  28    """Class to represent a single classical particle"""\r
  29    def __init__(self, mass, charge, r):\r
  30       """Create a particle"""\r
  31       self.mass = mass\r
  32       self.charge = charge\r
  33       self.r = r #Particle position (in angstrom)\r
  34       self.v = [0,0] #Particle velocity\r
  35       self.f = [0,0] #Net force acting on particle (note: recalculated every step)\r
  36 \r
  37    def step(self):\r
  38       """Do a single step"""\r
  39       #Update velocity\r
  40       if self.mass != 0:\r
  41          for i in range(0,2):\r
  42             self.v[i] += self.f[i] / self.mass\r
  43             \r
  44       #Update position \r
  45       for i in range(0,2):\r
  46          self.r[i] += self.v[i]\r
  47 \r
  48       #Reflect off walls\r
  49       #Note: Doesn't work if particle velocity is very big\r
  50       if self.r[0] <= 0 or self.r[0] >= 640:\r
  51          self.v[0] = -self.v[0]\r
  52       if self.r[1] <= 0 or self.r[1] >= 480:\r
  53          self.v[1] = -self.v[1]\r
  54   \r
  55    def calcForce(a,b):\r
  56       """Calculate the force between particles a and b (acting on a)"""\r
  57       r = [0,0] #Relative displacement of b to a\r
  58       for i in range(0,2):\r
  59          r[i] = b.r[i] - a.r[i]\r
  60       \r
  61       if r[0] == 0 and r[1] == 0: #To prevent divide by zero... \r
  62          return [0,0]\r
  63       \r
  64       distance = (r[0]**2 + r[1]**2)**(1/2) \r
  65       \r
  66       f = [0,0] #Calculated net force\r
  67 \r
  68       #Apply Coulomb force\r
  69       for i in range(0,2):\r
  70          f[i] -= r[i] * (k * a.charge * b.charge / (distance**3))\r
  71          \r
  72       #Apply generic Repulsive force\r
  73       for i in range(0,2):\r
  74          f[i] -= r[i] * (A / ((distance)**(m+1)))\r
  75       return [f[0], f[1]] #Return the resultant force\r
  76 \r
  77    def draw(self, window):\r
  78       """Draw the particle"""   \r
  79       if self.charge == 1:\r
  80          colour = pygame.Color(255,0,0,1) #Protons are red\r
  81          size = 4\r
  82       elif self.charge == -1:\r
  83          colour = pygame.Color(0,0,255,1) #Electrons are blue\r
  84          size = 4\r
  85       elif self.charge == 0:\r
  86          colour = pygame.Color(255,255,255,1) #Neutrons are white\r
  87          size = 4\r
  88       else:\r
  89          colour = pygame.Color(0,255,0,1) #And I love you\r
  90          size = 4\r
  91       #(I don't really love you. Whoever you are).\r
  92 \r
  93       #Draw the most amazingly realistic depiction of a particle\r
  94       #The result of many years of research.\r
  95       pygame.draw.circle(window,colour,[int(self.r[0] * scale), int(self.r[1] * scale)],size,0)\r
  96       \r
  97 \r
  98 \r
  99    \r
 100 class universe():\r
 101    """Class to simulate THE UNIVERSE"""\r
 102    def __init__(self):\r
 103       self.particles = [] #List of particles\r
 104       \r
 105    def step(self):\r
 106       """Do a single step"""\r
 107 \r
 108       #Calculate force pairings\r
 109       for a in self.particles:\r
 110          a.f = [0,0]\r
 111          for b in self.particles:\r
 112             if b != a:\r
 113                fAB = particle.calcForce(a,b)\r
 114                for i in range(0,2):\r
 115                   a.f[i] += fAB[i]\r
 116       #Perform steps\r
 117       #for a in self.particles:\r
 118          a.step()\r
 119 \r
 120    def draw(self, window):\r
 121       for a in self.particles:\r
 122          a.draw(window)\r
 123 \r
 124 def runSimulation():\r
 125    #create the screen\r
 126    window = pygame.display.set_mode((640, 480))\r
 127    theUniverse = universe()\r
 128 \r
 129    run = True\r
 130    while run:\r
 131       #Step\r
 132       theUniverse.step()\r
 133       #Draw\r
 134       window.fill(pygame.Color(0,0,0,0))\r
 135       theUniverse.draw(window)\r
 136       #Update draw events\r
 137       pygame.display.flip()\r
 138       #Wait a bit\r
 139       #pygame.time.wait(1)\r
 140       #Input handling\r
 141       \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] == True:\r
 150                 theUniverse.particles.append(particle(mass["proton"],charge["proton"], [mousepos[0]/scale,mousepos[1]/scale]))\r
 151              elif mousestate[1] == True:\r
 152                 theUniverse.particles.append(particle(mass["neutron"], charge["neutron"], [mousepos[0]/scale,mousepos[1]/scale]))\r
 153              elif mousestate[2] == True:\r
 154                 theUniverse.particles.append(particle(mass["electron"],charge["electron"], [mousepos[0]/scale,mousepos[1]/scale]))\r
 155              break\r
 156          else: \r
 157             print(event)\r
 158             pass\r
 159 \r
 160 \r
 161 \r
 162 if __name__ == "__main__":\r
 163    runSimulation()\r
 164 \r
 165 \r
 166 \r
 167             \r