Commit before breaking everything
[matches/honours.git] / research / transmission_spectroscopy / 24-1-12 / simulation.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

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