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