Commit before breaking everything
[matches/honours.git] / 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

UCC git Repository :: git.ucc.asn.au