Advertisement
Guest User

Untitled

a guest
Mar 18th, 2015
397
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 28.85 KB | None | 0 0
  1. # -*- coding: utf-8 -*-
  2. # Copyright (C) 2014 nineties
  3. # $Id: hchem.py 2014-08-12 10:47:50 nineties $
  4.  
  5. #= A clone of Tim Hutton's artificial chemistry simulator. =
  6.  
  7. import numpy as np
  8. import numpy.linalg as la
  9. import pygame, pickle
  10. import re, sys, os, time, math, datetime
  11. import Tkinter as tk
  12. import tkMessageBox
  13.  
  14. class HChemRule:
  15. def __init__(self, filename):
  16. self.cnt = 0
  17. self.num = None
  18. self.fill = []
  19. self.types = []
  20. self.map = {}
  21. self.color_count = 0
  22. self.colors = []
  23. self.colormap = {}
  24. self.wildcards = ['X', 'Y']
  25. self.wildstates = ['x', 'y']
  26. self.state_max = 10
  27. self.name = []
  28. self.ruleb = {} # Rules for bounded pair
  29. self.ruleu = {} # Rules for unbounded pair
  30. self.rule_texts = []
  31. self.colortable = [
  32. (0, 0, 0), (255, 0, 0), (0, 255, 0), (0, 0, 255),
  33. (255, 255, 0), (255, 0, 255), (0, 255, 255),
  34. (255, 127, 0), (255, 0, 127), (127, 255, 0), (127, 0, 255),
  35. (0, 255, 127), (0, 127, 255)
  36. ]
  37.  
  38. self.parse(filename)
  39. for t in self.types:
  40. for s in range(0, self.state_max+1):
  41. self.to_index(t, str(s))
  42.  
  43. def gen_color(self, a):
  44. self.color_count += 1
  45. return self.colortable[(self.color_count - 1) % len(self.colortable)]
  46.  
  47. def to_index(self, t, n):
  48. term = t + n
  49. if not term in self.map:
  50. self.map[term] = self.cnt
  51. self.colors.append(self.colormap[t])
  52. self.name.append(term)
  53. self.cnt += 1
  54. return self.map[term]
  55.  
  56. def is_valid_type(self, term):
  57. if term in self.map:
  58. return True
  59. else:
  60. return False
  61.  
  62. def get_index(self, term):
  63. return self.map[term]
  64.  
  65. def get_name(self, idx):
  66. return self.name[idx]
  67.  
  68. def parse_expr(self, str):
  69. if "-" in str:
  70. bnd = True
  71. M0, M1 = str.split("-")
  72. else:
  73. bnd = False
  74. M0, M1 = str.split()
  75. p = re.search(r'([a-wzA-Z]+)(\d+|[xy])', M0)
  76. q = re.search(r'([a-wzA-Z]+)(\d+|[xy])', M1)
  77. return (p.group(1), p.group(2), q.group(1), q.group(2), bnd)
  78.  
  79. def add_rule(self, L0, l0, L1, l1, lbnd, R0, r0, R1, r1, rbnd, prob):
  80. if l0 in self.wildstates and l1 in self.wildstates:
  81. if l0 == l1:
  82. for s in range(self.state_max+1):
  83. s = str(s)
  84. _l0 = s
  85. _l1 = s
  86. if r0 == l0: _r0 = s
  87. else: _r0 = r0
  88. if r1 == l0: _r1 = s
  89. else: _r1 = r1
  90. self.add_rule(L0, _l0, L1, _l1, lbnd, R0, _r0, R1, _r1, rbnd, prob)
  91. return
  92. else:
  93. for s0 in range(self.state_max+1):
  94. for s1 in range(self.state_max+1):
  95. if l0 == l1 and s0 > s1: continue
  96. s0 = str(s0)
  97. s1 = str(s1)
  98. _l0 = s0
  99. _l1 = s1
  100. if r0 == l0: _r0 = s0
  101. elif r0 == l1: _r0 = s1
  102. else: _r0 = r0
  103. if r1 == l0: _r1 = s0
  104. elif r1 == l1: _r1 = s1
  105. else: _r1 = r1
  106. self.add_rule(L0, _l0, L1, _l1, lbnd, R0, _r0, R1, _r1, rbnd, prob)
  107. return
  108. elif l0 in self.wildstates or l1 in self.wildstates:
  109. if l0 in self.wildstates:
  110. for s in range(self.state_max + 1):
  111. s = str(s)
  112. _l0 = s
  113. if r0 == l0: _r0 = s
  114. else: _r0 = r0
  115. if r1 == l0: _r1 = s
  116. else: _r1 = r1
  117. self.add_rule(L0, _l0, L1, l1, lbnd, R0, _r0, R1, _r1, rbnd, prob)
  118. return
  119. else:
  120. for s in range(self.state_max + 1):
  121. s = str(s)
  122. _l1 = s
  123. if r0 == l1: _r0 = s
  124. else: _r0 = r0
  125. if r1 == l1: _r1 = s
  126. else: _r1 = r1
  127. self.add_rule(L0, l0, L1, _l1, lbnd, R0, _r0, R1, _r1, rbnd, prob)
  128. return
  129. LL0 = self.to_index(L0, l0)
  130. LL1 = self.to_index(L1, l1)
  131. RR0 = self.to_index(R0, r0)
  132. RR1 = self.to_index(R1, r1)
  133. if lbnd:
  134. if not (LL0, LL1) in self.ruleb:
  135. self.ruleb[(LL0, LL1)] = []
  136. if LL0 != LL1:
  137. self.ruleb[(LL1, LL0)] = []
  138. self.ruleb[(LL0, LL1)].append((RR0, RR1, rbnd, prob))
  139. if LL0 != LL1:
  140. self.ruleb[(LL1, LL0)].append((RR1, RR0, rbnd, prob))
  141. else:
  142. if not (LL0, LL1) in self.ruleu:
  143. self.ruleu[(LL0, LL1)] = []
  144. if LL0 != LL1:
  145. self.ruleu[(LL1, LL0)] = []
  146. self.ruleu[(LL0, LL1)].append((RR0, RR1, rbnd, prob))
  147. if LL0 != LL1:
  148. self.ruleu[(LL1, LL0)].append((RR1, RR0, rbnd, prob))
  149.  
  150. def parse_rule(self, line):
  151. try:
  152. lhs, rhs = line.split("->")
  153. except:
  154. return
  155. self.rule_texts.append(line)
  156. prob = 1.0
  157. if ":" in rhs:
  158. rhs, p = rhs.split(":")
  159. prob = eval(p.strip())
  160. L0, l0, L1, l1, lbnd = self.parse_expr(lhs.strip())
  161. R0, r0, R1, r1, rbnd = self.parse_expr(rhs.strip())
  162. if L0 in self.wildcards and L1 in self.wildcards:
  163. if L0 == L1:
  164. for t in self.types:
  165. _L0 = t
  166. _L1 = t
  167. if R0 == L0: _R0 = t
  168. else: _R0 = R0
  169. if R1 == L0: _R1 = t
  170. else: _R1 = R1
  171. self.add_rule(_L0, l0, _L1, l1, lbnd, _R0, r0, _R1, r1, rbnd, prob)
  172. else:
  173. for t0 in self.types:
  174. for t1 in self.types:
  175. if l0 == l1 and t0 > t1: continue
  176. _L0 = t0
  177. _L1 = t1
  178. if R0 == L0: _R0 = t0
  179. elif R0 == L1: _R0 = t1
  180. else: _R0 = R0
  181. if R1 == L0: _R1 = t0
  182. elif R1 == L1: _R1 = t1
  183. else: _R1 = R1
  184. self.add_rule(_L0, l0, _L1, l1, lbnd, _R0, r0, _R1, r1, rbnd, prob)
  185. elif L0 in self.wildcards or L1 in self.wildcards:
  186. if L0 in self.wildcards:
  187. for t in self.types:
  188. _L0 = t
  189. if R0 == L0: _R0 = t
  190. else: _R0 = R0
  191. if R1 == L0: _R1 = t
  192. else: _R1 = R1
  193. self.add_rule(_L0, l0, L1, l1, lbnd, _R0, r0, _R1, r1, rbnd, prob)
  194. else:
  195. for t in self.types:
  196. _L1 = t
  197. if R0 == L1: _R0 = t
  198. else: _R0 = R0
  199. if R1 == L1: _R1 = t
  200. else: _R1 = R1
  201. self.add_rule(L0, l0, _L1, l1, lbnd, _R0, r0, _R1, r1, rbnd, prob)
  202. else:
  203. self.add_rule(L0, l0, L1, l1, lbnd, R0, r0, R1, r1, rbnd, prob)
  204.  
  205. def add_type(self, t):
  206. self.types.append(t)
  207. self.colormap[t] = self.gen_color(t)
  208.  
  209. def setup_types(self, str):
  210. lhs, rhs = str.split(":")
  211. for t in rhs.split(","):
  212. self.add_type(t.strip())
  213.  
  214. def setup_fill(self, str):
  215. lhs, rhs = str.split(":")
  216. for decl in rhs.split(","):
  217. t, p = decl.strip().split(" ")
  218. self.fill.append((t, eval(p)))
  219.  
  220. def parse(self, filename):
  221. f = open(filename, "r")
  222. while True:
  223. line = f.readline()
  224. if not line: break
  225. if line[0] == '#': continue
  226. line = line.strip()
  227. if line.find("type") == 0:
  228. self.setup_types(line)
  229. elif line.find("number of particles") == 0:
  230. self.num = int(line.split(":")[1].strip())
  231. elif line.find("state max") == 0:
  232. self.state_max = int(line.split(":")[1].strip())
  233. elif line.find("fill") == 0:
  234. self.setup_fill(line)
  235. elif "->" in line:
  236. self.parse_rule(line)
  237. f.close()
  238.  
  239. def check(self, L0, L1, bound):
  240. if bound:
  241. if not (L0, L1) in self.ruleb:
  242. return None
  243. return self.ruleb[(L0, L1)]
  244. else:
  245. if not (L0, L1) in self.ruleu:
  246. return None
  247. return self.ruleu[(L0, L1)]
  248.  
  249. class HChem:
  250. # n : number of particles
  251. # r : radious of particles
  252. # v0 : initial velocity of particles
  253. # dt : duration of one frame
  254. # k : strength of bonds
  255. # w,h : width and height of the universe
  256. # seed: random seed
  257. def __init__(self, filename, n = 1000, r = 10, v0 = None, dt = 0.1,
  258. width = 1200, height = 700, bucket_size = None, seed=None):
  259. self.rule = HChemRule(filename)
  260. if seed: np.random.seed(seed)
  261. if v0 == None: v0 = r
  262. if bucket_size == None: bucket_size = 2*r
  263.  
  264. self.n = n
  265. if self.rule.num: self.n = self.rule.num
  266. self.r = r
  267. self.dt = dt
  268. self.w = width
  269. self.h = height
  270. self.speed = 10
  271.  
  272. # Initialize positions of particles
  273. self.pos = np.zeros((n, 2))
  274. self.pos[:,0] = np.random.uniform(r, width-r, n)
  275. self.pos[:,1] = np.random.uniform(r, height-r, n)
  276.  
  277. # Initialize velocities of particles
  278. direction = np.random.uniform(0, 2*np.pi, n)
  279. self.vel = np.zeros((n, 2))
  280. self.vel[:,0] = v0*np.cos(direction)
  281. self.vel[:,1] = v0*np.sin(direction)
  282. # Initialize types
  283. self.types = np.zeros(n, dtype=int)
  284. for k in xrange(self.n):
  285. p = np.random.uniform(0, 1)
  286. q = 0
  287. for (t, r) in self.rule.fill:
  288. q += r
  289. if p < q:
  290. self.types[k] = self.rule.get_index(t)
  291. break
  292. self.stypes = np.zeros(n, dtype=object)
  293. for k in xrange(self.n):
  294. self.stypes[k] = self.rule.get_name(self.types[k])
  295.  
  296. # self.bonds[i] == list of indexes of particles which is bound to i.
  297. self.bonds = np.zeros(n, dtype=object)
  298. for i in xrange(n): self.bonds[i] = []
  299.  
  300. # Initialize buckets for compute_collision detection
  301. self.bucket_size = bucket_size
  302. self.nbx = int(math.ceil(float(width)/bucket_size))
  303. self.nby = int(math.ceil(float(height)/bucket_size))
  304. self.buckets = np.zeros((self.nbx, self.nby), dtype=object)
  305.  
  306. def bucket_index(self, x):
  307. return (min(max(int(x[0]/self.bucket_size), 0), self.nbx-1),
  308. min(max(int(x[1]/self.bucket_size), 0), self.nby-1))
  309.  
  310. def init_bucket(self):
  311. for i in xrange(self.nbx):
  312. for j in xrange(self.nby):
  313. self.buckets[i,j] = []
  314. for k in xrange(self.n):
  315. i, j = self.bucket_index(self.pos[k, :])
  316. self.buckets[i, j].append(k)
  317.  
  318. def add_impulse_from_walls(self):
  319. r = self.r
  320. for k in xrange(self.n):
  321. x = self.pos[k, 0]
  322. y = self.pos[k, 1]
  323. vx = self.vel[k, 0]
  324. vy = self.vel[k, 1]
  325. if (x < r and vx < 0) or (x > self.w-r and vx > 0):
  326. self.vel[k, 0] += -2*self.vel[k, 0]
  327. self.vel[k, 1] += 0
  328. if (y < r and vy < 0) or (y > self.h-r and vy > 0):
  329. self.vel[k, 0] += 0
  330. self.vel[k, 1] += -2*self.vel[k, 1]
  331.  
  332. def update_state_of_particle_pair(self, k, l):
  333. if l in self.bonds[k]:
  334. # bound pair
  335. rules = self.rule.check(self.types[k], self.types[l], True)
  336. if rules:
  337. for r in rules:
  338. p = r[3]
  339. if np.random.uniform(0, 1) < p:
  340. # print "apply:",
  341. # print self.rule.get_name(self.types[k]),
  342. # print "-",
  343. # print self.rule.get_name(self.types[l]),
  344. # print " -> ",
  345. # print self.rule.get_name(r[0]),
  346. # if r[2]:
  347. # print "-",
  348. # else:
  349. # print " ",
  350. # print self.rule.get_name(r[1])
  351. self.types[k] = r[0]
  352. self.types[l] = r[1]
  353. if not r[2]:
  354. self.bonds[k].remove(l)
  355. self.bonds[l].remove(k)
  356. return False
  357. return True
  358. return True
  359. else:
  360. # unbound pair
  361. rules = self.rule.check(self.types[k], self.types[l], False)
  362. if rules:
  363. for r in rules:
  364. p = r[3]
  365. if np.random.uniform(0, 1) < p:
  366. # print "apply:",
  367. # print self.rule.get_name(self.types[k]),
  368. # print " ",
  369. # print self.rule.get_name(self.types[l]),
  370. # print " -> ",
  371. # print self.rule.get_name(r[0]),
  372. # if r[2]:
  373. # print "-",
  374. # else:
  375. # print " ",
  376. # print self.rule.get_name(r[1])
  377. self.types[k] = r[0]
  378. self.types[l] = r[1]
  379. if r[2]:
  380. self.bonds[k].append(l)
  381. self.bonds[l].append(k)
  382. return True
  383. return False
  384. return False
  385.  
  386. def add_impulse_between_unbound_pair(self, k, l, rx, rv, d2):
  387. if self.update_state_of_particle_pair(k, l):
  388. return
  389. d = math.sqrt(d2)
  390. n = rx/d
  391. ldt = -n.dot(rv)
  392. self.vel[k,:] += ldt*n
  393. self.vel[l,:] -= ldt*n
  394.  
  395. def add_impulse_between_bound_pair(self, k, l, rx, rv, d2):
  396. d = math.sqrt(d2)
  397. n = rx/d
  398. c = rx.dot(rv)
  399. #ldt = -(2*c + 3*(d2-4*self.r*self.r))/(8*d2)
  400. #self.vel[k,:] += 2*ldt*rx
  401. #self.vel[l,:] -= 2*ldt*rx
  402. if (d < 2*self.r and c < 0) or (d > 2*self.r and c > 0):
  403. ldt = -n.dot(rv)
  404. self.vel[k,:] += ldt*n
  405. self.vel[l,:] -= ldt*n
  406.  
  407. def add_impulse_between_particles_sub(self, k, i, j):
  408. if i < 0 or j < 0 or i >= self.nbx or j >= self.nby: return
  409. for l in self.buckets[i, j]:
  410. if k >= l: continue
  411. rx = self.pos[k,:] - self.pos[l,:]
  412. rv = self.vel[k,:] - self.vel[l,:]
  413. if rx.dot(rv) >= 0: continue
  414. d2 = np.sum(rx*rx)
  415. if d2 > 4*self.r*self.r: continue
  416. self.add_impulse_between_unbound_pair(k, l, rx, rv, d2)
  417.  
  418. def add_impulse_between_particles(self):
  419. r = self.r
  420.  
  421. # add impulse between unbound pairs
  422. for k in xrange(self.n):
  423. i,j = self.bucket_index(self.pos[k, :])
  424. self.add_impulse_between_particles_sub(k, i-1, j)
  425. self.add_impulse_between_particles_sub(k, i-1, j-1)
  426. self.add_impulse_between_particles_sub(k, i-1 ,j+1)
  427. self.add_impulse_between_particles_sub(k, i, j-1)
  428. self.add_impulse_between_particles_sub(k, i, j)
  429. self.add_impulse_between_particles_sub(k, i, j+1)
  430. self.add_impulse_between_particles_sub(k, i+1, j-1)
  431. self.add_impulse_between_particles_sub(k, i+1, j)
  432. self.add_impulse_between_particles_sub(k, i+1, j+1)
  433.  
  434. def add_impulse_between_bound_particles(self):
  435. # add impulse between bound pairs
  436. for k in xrange(self.n):
  437. for l in self.bonds[k]:
  438. if k >= l: continue
  439. rx = self.pos[k,:] - self.pos[l,:]
  440. rv = self.vel[k,:] - self.vel[l,:]
  441. d2 = np.sum(rx*rx)
  442. self.add_impulse_between_bound_pair(k, l, rx, rv, d2)
  443.  
  444. def compute_impulse(self):
  445. self.init_bucket()
  446. self.add_impulse_from_walls()
  447. self.add_impulse_between_particles()
  448. self.add_impulse_between_bound_particles()
  449. self.pos += self.vel*self.dt
  450.  
  451. def change_speed(self, delta):
  452. self.speed += delta;
  453. if self.speed < 0:
  454. self.speed = 1
  455.  
  456. def update(self):
  457. # Update position
  458. for k in range(self.speed):
  459. self.compute_impulse()
  460.  
  461. def total_energy(self):
  462. return np.sum(self.vel*self.vel)/2
  463.  
  464. def save(self, fname, type):
  465. if type == "particles":
  466. self.save_particles(fname)
  467. else:
  468. self.save_rules(fname)
  469.  
  470. def save_particles(self, fname):
  471. try:
  472. with open(fname, "w") as f:
  473. f.write(repr(self.n)); f.write("\n")
  474. f.write(repr(self.dt)); f.write("\n")
  475. for k in xrange(self.n):
  476. f.write(self.stypes[k]); f.write(",")
  477. f.write(repr(self.pos[k,0])); f.write(",")
  478. f.write(repr(self.pos[k,1])); f.write(",")
  479. f.write(repr(self.vel[k,0])); f.write(",")
  480. f.write(repr(self.vel[k,1])); f.write("\n")
  481. for k in xrange(self.n):
  482. bonds = " ".join(map(str, self.bonds[k]))
  483. f.write(bonds); f.write("\n")
  484. except Exception(e):
  485. root = tk.Tk()
  486. root.withdraw()
  487. tkMessageBox.showerror('Error', str(e))
  488.  
  489. def load_particles(self, fname):
  490. try:
  491. with open(fname, "r") as f:
  492. n = int(f.readline())
  493. dt = float(f.readline())
  494. pos = []
  495. vel = []
  496. types = []
  497. bonds = []
  498. for k in xrange(n):
  499. t, p0, p1, v0, v1 = f.readline().strip().split(",")
  500. pos.append((float(p0), float(p1)))
  501. vel.append((float(v0), float(v1)))
  502. if not self.rule.is_valid_type(t):
  503. root = tk.Tk()
  504. root.withdraw()
  505. tkMessageBox.showerror('Error', 'Unknown type:' + t)
  506. return
  507. types.append(t)
  508. for k in xrange(n):
  509. bonds.append(map(int, f.readline().strip().split()))
  510. self.n = n
  511. self.dt = dt
  512. self.pos = np.array(pos)
  513. self.vel = np.array(vel)
  514. self.stypes = np.array(types, dtype=object)
  515. self.types = np.array(map(lambda t: self.rule.get_index(t), types), dtype=int)
  516. self.bonds = bonds
  517.  
  518. except Exception(e):
  519. root = tk.Tk()
  520. root.withdraw()
  521. tkMessageBox.showerror('Error', str(e))
  522.  
  523. def load(self, fname, type):
  524. if type == "particles":
  525. self.load_particles(fname)
  526. else:
  527. self.load_rules(fname)
  528.  
  529. def start_record(self):
  530. date = datetime.date.today()
  531. now = time.time()
  532. dirname = str(date) + str(now)
  533. os.mkdir(dirname)
  534. self.record_dir = dirname
  535.  
  536. def record(self, iteration):
  537. self.save_particles("%s/iteration-%d.dat" % (self.record_dir, iteration))
  538.  
  539. class HChemViewer:
  540. RED = (255, 0, 0)
  541. BLUE = (0, 0, 255)
  542. WHITE = (255, 255, 255)
  543. BLACK = (0, 0, 0)
  544. INFO = [
  545. "(P) play/pause, (F) stepwise, (R) record, (Q) quit, (T) show/hide particle types, (up) Speed up, (down) Speed down",
  546. "(drag) move particle, (shift + drag) bind/unbind particles, (double click) change type and state of a particle"
  547. ]
  548.  
  549. def __init__(self, sim, w = None, h = None):
  550. if w == None: w = sim.w
  551. if h == None: h = sim.h
  552.  
  553. self.sim = sim
  554. pygame.init()
  555. self.screen = pygame.display.set_mode((w, h)
  556. #, pygame.DOUBLEBUF | pygame.FULLSCREEN | pygame.HWSURFACE
  557. )
  558. pygame.display.set_caption("Tim Hutton's Artifical Chemistry")
  559. self.fontsize = 18
  560. self.font = pygame.font.SysFont(None, self.fontsize)
  561. info_texts = self.INFO + sim.rule.rule_texts
  562. self.info = map(lambda text: self.font.render(text, False, self.BLUE),
  563. info_texts)
  564.  
  565. self.speed = 10
  566.  
  567. # For events
  568. self.record = False
  569. self.shift = False
  570. self.play = False
  571. self.stepwise = False
  572. self.dragged = None
  573. self.moving = False
  574. self.binding = False
  575. self.display_types = False
  576. self.prev_lclick = time.time()
  577.  
  578. def get_clicked(self):
  579. for k in xrange(self.sim.n):
  580. d2 = np.sum((self.sim.pos[k,:] - pygame.mouse.get_pos())**2)
  581. if d2 < self.sim.r**2:
  582. return k
  583. break
  584. return None
  585.  
  586. def ask_particle(self):
  587. i = self.get_clicked()
  588. dialog = tk.Tk()
  589. dialog.title("Enter particle type and state")
  590. type = tk.StringVar(dialog)
  591. tk.Label(dialog, text = "type").pack(side=tk.LEFT)
  592. entry = tk.Entry(dialog, textvariable=type)
  593. entry.focus_force()
  594. entry.pack(side=tk.LEFT)
  595. entry.bind('<Return>', lambda evt: dialog.destroy())
  596. dialog.mainloop()
  597. return type.get()
  598.  
  599. def ask_file(self, title):
  600. print "Sono in ASK FILE"
  601. dialog = tk.Tk()
  602. print 1
  603. dialog.title(title)
  604. print 2
  605. filename = tk.StringVar(dialog)
  606. print 3
  607. savetype = tk.StringVar(dialog)
  608. print 4
  609. savetype.set("particles")
  610. print 5
  611. tk.Radiobutton(dialog, text="Particles", variable=savetype, value="particles").pack(anchor = tk.W)
  612. print 6
  613. tk.Radiobutton(dialog, text="Rules", variable=savetype, value="rules").pack(anchor = tk.W)
  614. print 7
  615. tk.Label(dialog, text = "filename").pack(side=tk.LEFT)
  616. print 8
  617. entry = tk.Entry(dialog, textvariable=filename)
  618. print 9
  619. entry.focus_force()
  620. print 10
  621. entry.pack(side=tk.LEFT)
  622. print 11
  623. entry.bind('<Return>', lambda evt: dialog.destroy())
  624. print 12
  625. dialog.mainloop()
  626. return filename.get(), savetype.get()
  627.  
  628. def check_event(self):
  629. for event in pygame.event.get():
  630. if event.type == pygame.KEYDOWN:
  631. key = pygame.key.get_pressed()
  632. if key[pygame.K_RSHIFT] or key[pygame.K_LSHIFT]:
  633. self.shift = True
  634. if key[pygame.K_r]:
  635. self.record = not self.record
  636. if self.record:
  637. self.sim.start_record()
  638. if key[pygame.K_q]:
  639. sys.exit()
  640. if key[pygame.K_p]:
  641. self.play = not self.play
  642. if key[pygame.K_UP]:
  643. self.sim.change_speed(1)
  644. if key[pygame.K_DOWN]:
  645. self.sim.change_speed(-1)
  646. if key[pygame.K_s]:
  647. fname, type = self.ask_file("Save configuration")
  648. if fname: self.sim.save(fname, type)
  649. if key[pygame.K_l]:
  650. fname, type = self.ask_file("Load configuration")
  651. if fname: self.sim.load(fname, type)
  652. if key[pygame.K_t]:
  653. self.display_types = not self.display_types
  654. if key[pygame.K_f]:
  655. self.stepwise = True
  656. self.play = True
  657. if event.type == pygame.KEYUP:
  658. key = pygame.key.get_pressed()
  659. if not key[pygame.K_RSHIFT] and not key[pygame.K_LSHIFT]:
  660. self.shift = False
  661. if not self.dragged and event.type == pygame.MOUSEBUTTONDOWN:
  662. self.play = False
  663. clicked = self.get_clicked()
  664.  
  665. # Detect double click
  666. t = time.time()
  667. double_click = False
  668. if t - self.prev_lclick < 1.0/3:
  669. double_click = True
  670. self.prev_lclick = t
  671.  
  672. if clicked and double_click:
  673. t = self.ask_particle()
  674. try:
  675. self.sim.stypes[clicked] = t
  676. self.sim.types[clicked] = self.sim.rule.get_index(t)
  677. except:
  678. pass
  679. elif clicked:
  680. self.dragged = clicked
  681. if not self.shift: self.moving = True
  682. elif self.shift: self.binding = True
  683. elif self.dragged and event.type == pygame.MOUSEMOTION:
  684. if self.moving:
  685. self.sim.pos[self.dragged,:] = pygame.mouse.get_pos()
  686. elif self.dragged and event.type == pygame.MOUSEBUTTONUP:
  687. if self.binding:
  688. clicked = self.get_clicked()
  689. if clicked and self.dragged != clicked and\
  690. not (clicked in self.sim.bonds[self.dragged]):
  691. self.sim.bonds[self.dragged].append(clicked)
  692. self.sim.bonds[clicked].append(self.dragged)
  693. elif clicked and self.dragged != clicked and\
  694. (clicked in self.sim.bonds[self.dragged]):
  695. self.sim.bonds[self.dragged].remove(clicked)
  696. self.sim.bonds[clicked].remove(self.dragged)
  697. self.moving = False
  698. self.binding = False
  699. self.dragged = None
  700. elif event.type == pygame.QUIT:
  701. sys.exit()
  702.  
  703. def loop(self):
  704. iteration = 0
  705. screen = self.screen
  706. while True:
  707. sim = self.sim
  708. n = sim.n
  709. r = sim.r
  710. if self.play:
  711. iteration += 1
  712. sim.update()
  713. if self.record:
  714. sim.record(iteration)
  715.  
  716. if self.stepwise:
  717. self.play = False
  718. self.stepwise = False
  719.  
  720. pos = sim.pos
  721.  
  722. screen.fill(self.WHITE)
  723. # Draw particles
  724. for k in xrange(n):
  725. pygame.draw.circle(screen, sim.rule.colors[sim.types[k]],
  726. (int(pos[k,0]), int(pos[k,1])), r, 1)
  727.  
  728. if self.display_types:
  729. for k in xrange(sim.n):
  730. t = sim.rule.get_name(sim.types[k])
  731. text = self.font.render(t, False, self.BLACK)
  732. rect = text.get_rect()
  733. rect.centerx = pos[k,0]
  734. rect.centery = pos[k,1]
  735. self.screen.blit(text, rect)
  736.  
  737. # Draw bonds
  738. for k in xrange(n):
  739. for l in sim.bonds[k]:
  740. if k >= l: continue
  741. pygame.draw.line(screen, self.BLACK, pos[k,:], pos[l,:])
  742.  
  743. # Other info
  744. if self.binding:
  745. pygame.draw.line(screen, self.BLACK,
  746. pos[self.dragged,:], pygame.mouse.get_pos())
  747.  
  748. y = 10
  749. for i in self.info:
  750. self.screen.blit(i, (10, y))
  751. y += self.fontsize
  752. text = self.font.render(
  753. "time = " + str(iteration*sim.dt),
  754. False, self.BLUE)
  755. self.screen.blit(text, (10, y))
  756. energy = sim.total_energy()
  757. text = self.font.render(
  758. "energy = " + str(energy),
  759. False, self.BLUE)
  760. self.screen.blit(text, (10, y + self.fontsize))
  761. self.check_event()
  762. pygame.display.update()
  763.  
  764. if __name__ == '__main__':
  765. print len(sys.argv)
  766.  
  767. if len(sys.argv) == 2:
  768. sim = HChem(sys.argv[1])
  769. HChemViewer(sim).loop()
  770. exit()
  771.  
  772. if len(sys.argv) == 3:
  773. sim = HChem(sys.argv[1])
  774. sim.load_particles(sys.argv[2])
  775. HChemViewer(sim).loop()
  776. exit()
  777.  
  778. print "Please insert a rules file."
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement