1# -*- coding: utf-8 -*- 2import numpy as np 3import collections 4import sys, os 5import pyqtgraph as pg 6from pyqtgraph.Qt import QtGui, QtCore 7from pyqtgraph.parametertree import Parameter, ParameterTree 8from pyqtgraph.parametertree import types as pTypes 9import pyqtgraph.configfile 10 11from time import perf_counter 12 13 14class RelativityGUI(QtGui.QWidget): 15 def __init__(self): 16 QtGui.QWidget.__init__(self) 17 18 self.animations = [] 19 self.animTimer = QtCore.QTimer() 20 self.animTimer.timeout.connect(self.stepAnimation) 21 self.animTime = 0 22 self.animDt = .016 23 self.lastAnimTime = 0 24 25 self.setupGUI() 26 27 self.objectGroup = ObjectGroupParam() 28 29 self.params = Parameter.create(name='params', type='group', children=[ 30 dict(name='Load Preset..', type='list', limits=[]), 31 #dict(name='Unit System', type='list', limits=['', 'MKS']), 32 dict(name='Duration', type='float', value=10.0, step=0.1, limits=[0.1, None]), 33 dict(name='Reference Frame', type='list', limits=[]), 34 dict(name='Animate', type='bool', value=True), 35 dict(name='Animation Speed', type='float', value=1.0, dec=True, step=0.1, limits=[0.0001, None]), 36 dict(name='Recalculate Worldlines', type='action'), 37 dict(name='Save', type='action'), 38 dict(name='Load', type='action'), 39 self.objectGroup, 40 ]) 41 self.tree.setParameters(self.params, showTop=False) 42 self.params.param('Recalculate Worldlines').sigActivated.connect(self.recalculate) 43 self.params.param('Save').sigActivated.connect(self.save) 44 self.params.param('Load').sigActivated.connect(self.load) 45 self.params.param('Load Preset..').sigValueChanged.connect(self.loadPreset) 46 self.params.sigTreeStateChanged.connect(self.treeChanged) 47 48 ## read list of preset configs 49 presetDir = os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), 'presets') 50 if os.path.exists(presetDir): 51 presets = [os.path.splitext(p)[0] for p in os.listdir(presetDir)] 52 self.params.param('Load Preset..').setLimits(['']+presets) 53 54 55 56 57 def setupGUI(self): 58 self.layout = QtGui.QVBoxLayout() 59 self.layout.setContentsMargins(0,0,0,0) 60 self.setLayout(self.layout) 61 self.splitter = QtGui.QSplitter() 62 self.splitter.setOrientation(QtCore.Qt.Orientation.Horizontal) 63 self.layout.addWidget(self.splitter) 64 65 self.tree = ParameterTree(showHeader=False) 66 self.splitter.addWidget(self.tree) 67 68 self.splitter2 = QtGui.QSplitter() 69 self.splitter2.setOrientation(QtCore.Qt.Orientation.Vertical) 70 self.splitter.addWidget(self.splitter2) 71 72 self.worldlinePlots = pg.GraphicsLayoutWidget() 73 self.splitter2.addWidget(self.worldlinePlots) 74 75 self.animationPlots = pg.GraphicsLayoutWidget() 76 self.splitter2.addWidget(self.animationPlots) 77 78 self.splitter2.setSizes([int(self.height()*0.8), int(self.height()*0.2)]) 79 80 self.inertWorldlinePlot = self.worldlinePlots.addPlot() 81 self.refWorldlinePlot = self.worldlinePlots.addPlot() 82 83 self.inertAnimationPlot = self.animationPlots.addPlot() 84 self.inertAnimationPlot.setAspectLocked(1) 85 self.refAnimationPlot = self.animationPlots.addPlot() 86 self.refAnimationPlot.setAspectLocked(1) 87 88 self.inertAnimationPlot.setXLink(self.inertWorldlinePlot) 89 self.refAnimationPlot.setXLink(self.refWorldlinePlot) 90 91 def recalculate(self): 92 ## build 2 sets of clocks 93 clocks1 = collections.OrderedDict() 94 clocks2 = collections.OrderedDict() 95 for cl in self.params.param('Objects'): 96 clocks1.update(cl.buildClocks()) 97 clocks2.update(cl.buildClocks()) 98 99 ## Inertial simulation 100 dt = self.animDt * self.params['Animation Speed'] 101 sim1 = Simulation(clocks1, ref=None, duration=self.params['Duration'], dt=dt) 102 sim1.run() 103 sim1.plot(self.inertWorldlinePlot) 104 self.inertWorldlinePlot.autoRange(padding=0.1) 105 106 ## reference simulation 107 ref = self.params['Reference Frame'] 108 dur = clocks1[ref].refData['pt'][-1] ## decide how long to run the reference simulation 109 sim2 = Simulation(clocks2, ref=clocks2[ref], duration=dur, dt=dt) 110 sim2.run() 111 sim2.plot(self.refWorldlinePlot) 112 self.refWorldlinePlot.autoRange(padding=0.1) 113 114 115 ## create animations 116 self.refAnimationPlot.clear() 117 self.inertAnimationPlot.clear() 118 self.animTime = 0 119 120 self.animations = [Animation(sim1), Animation(sim2)] 121 self.inertAnimationPlot.addItem(self.animations[0]) 122 self.refAnimationPlot.addItem(self.animations[1]) 123 124 ## create lines representing all that is visible to a particular reference 125 #self.inertSpaceline = Spaceline(sim1, ref) 126 #self.refSpaceline = Spaceline(sim2) 127 self.inertWorldlinePlot.addItem(self.animations[0].items[ref].spaceline()) 128 self.refWorldlinePlot.addItem(self.animations[1].items[ref].spaceline()) 129 130 131 132 133 def setAnimation(self, a): 134 if a: 135 self.lastAnimTime = perf_counter() 136 self.animTimer.start(int(self.animDt*1000)) 137 else: 138 self.animTimer.stop() 139 140 def stepAnimation(self): 141 now = perf_counter() 142 dt = (now-self.lastAnimTime) * self.params['Animation Speed'] 143 self.lastAnimTime = now 144 self.animTime += dt 145 if self.animTime > self.params['Duration']: 146 self.animTime = 0 147 for a in self.animations: 148 a.restart() 149 150 for a in self.animations: 151 a.stepTo(self.animTime) 152 153 154 def treeChanged(self, *args): 155 clocks = [] 156 for c in self.params.param('Objects'): 157 clocks.extend(c.clockNames()) 158 #for param, change, data in args[1]: 159 #if change == 'childAdded': 160 self.params.param('Reference Frame').setLimits(clocks) 161 self.setAnimation(self.params['Animate']) 162 163 def save(self): 164 filename = pg.QtGui.QFileDialog.getSaveFileName(self, "Save State..", "untitled.cfg", "Config Files (*.cfg)") 165 if isinstance(filename, tuple): 166 filename = filename[0] # Qt4/5 API difference 167 if filename == '': 168 return 169 state = self.params.saveState() 170 pg.configfile.writeConfigFile(state, str(filename)) 171 172 def load(self): 173 filename = pg.QtGui.QFileDialog.getOpenFileName(self, "Save State..", "", "Config Files (*.cfg)") 174 if isinstance(filename, tuple): 175 filename = filename[0] # Qt4/5 API difference 176 if filename == '': 177 return 178 state = pg.configfile.readConfigFile(str(filename)) 179 self.loadState(state) 180 181 def loadPreset(self, param, preset): 182 if preset == '': 183 return 184 path = os.path.abspath(os.path.dirname(__file__)) 185 fn = os.path.join(path, 'presets', preset+".cfg") 186 state = pg.configfile.readConfigFile(fn) 187 self.loadState(state) 188 189 def loadState(self, state): 190 if 'Load Preset..' in state['children']: 191 del state['children']['Load Preset..']['limits'] 192 del state['children']['Load Preset..']['value'] 193 self.params.param('Objects').clearChildren() 194 self.params.restoreState(state, removeChildren=False) 195 self.recalculate() 196 197 198class ObjectGroupParam(pTypes.GroupParameter): 199 def __init__(self): 200 pTypes.GroupParameter.__init__(self, name="Objects", addText="Add New..", addList=['Clock', 'Grid']) 201 202 def addNew(self, typ): 203 if typ == 'Clock': 204 self.addChild(ClockParam()) 205 elif typ == 'Grid': 206 self.addChild(GridParam()) 207 208class ClockParam(pTypes.GroupParameter): 209 def __init__(self, **kwds): 210 defs = dict(name="Clock", autoIncrementName=True, renamable=True, removable=True, children=[ 211 dict(name='Initial Position', type='float', value=0.0, step=0.1), 212 #dict(name='V0', type='float', value=0.0, step=0.1), 213 AccelerationGroup(), 214 215 dict(name='Rest Mass', type='float', value=1.0, step=0.1, limits=[1e-9, None]), 216 dict(name='Color', type='color', value=(100,100,150)), 217 dict(name='Size', type='float', value=0.5), 218 dict(name='Vertical Position', type='float', value=0.0, step=0.1), 219 ]) 220 #defs.update(kwds) 221 pTypes.GroupParameter.__init__(self, **defs) 222 self.restoreState(kwds, removeChildren=False) 223 224 def buildClocks(self): 225 x0 = self['Initial Position'] 226 y0 = self['Vertical Position'] 227 color = self['Color'] 228 m = self['Rest Mass'] 229 size = self['Size'] 230 prog = self.param('Acceleration').generate() 231 c = Clock(x0=x0, m0=m, y0=y0, color=color, prog=prog, size=size) 232 return {self.name(): c} 233 234 def clockNames(self): 235 return [self.name()] 236 237pTypes.registerParameterType('Clock', ClockParam) 238 239class GridParam(pTypes.GroupParameter): 240 def __init__(self, **kwds): 241 defs = dict(name="Grid", autoIncrementName=True, renamable=True, removable=True, children=[ 242 dict(name='Number of Clocks', type='int', value=5, limits=[1, None]), 243 dict(name='Spacing', type='float', value=1.0, step=0.1), 244 ClockParam(name='ClockTemplate'), 245 ]) 246 #defs.update(kwds) 247 pTypes.GroupParameter.__init__(self, **defs) 248 self.restoreState(kwds, removeChildren=False) 249 250 def buildClocks(self): 251 clocks = {} 252 template = self.param('ClockTemplate') 253 spacing = self['Spacing'] 254 for i in range(self['Number of Clocks']): 255 c = list(template.buildClocks().values())[0] 256 c.x0 += i * spacing 257 clocks[self.name() + '%02d' % i] = c 258 return clocks 259 260 def clockNames(self): 261 return [self.name() + '%02d' % i for i in range(self['Number of Clocks'])] 262 263pTypes.registerParameterType('Grid', GridParam) 264 265class AccelerationGroup(pTypes.GroupParameter): 266 def __init__(self, **kwds): 267 defs = dict(name="Acceleration", addText="Add Command..") 268 pTypes.GroupParameter.__init__(self, **defs) 269 self.restoreState(kwds, removeChildren=False) 270 271 def addNew(self): 272 nextTime = 0.0 273 if self.hasChildren(): 274 nextTime = self.children()[-1]['Proper Time'] + 1 275 self.addChild(Parameter.create(name='Command', autoIncrementName=True, type=None, renamable=True, removable=True, children=[ 276 dict(name='Proper Time', type='float', value=nextTime), 277 dict(name='Acceleration', type='float', value=0.0, step=0.1), 278 ])) 279 280 def generate(self): 281 prog = [] 282 for cmd in self: 283 prog.append((cmd['Proper Time'], cmd['Acceleration'])) 284 return prog 285 286pTypes.registerParameterType('AccelerationGroup', AccelerationGroup) 287 288 289class Clock(object): 290 nClocks = 0 291 292 def __init__(self, x0=0.0, y0=0.0, m0=1.0, v0=0.0, t0=0.0, color=None, prog=None, size=0.5): 293 Clock.nClocks += 1 294 self.pen = pg.mkPen(color) 295 self.brush = pg.mkBrush(color) 296 self.y0 = y0 297 self.x0 = x0 298 self.v0 = v0 299 self.m0 = m0 300 self.t0 = t0 301 self.prog = prog 302 self.size = size 303 304 def init(self, nPts): 305 ## Keep records of object from inertial frame as well as reference frame 306 self.inertData = np.empty(nPts, dtype=[('x', float), ('t', float), ('v', float), ('pt', float), ('m', float), ('f', float)]) 307 self.refData = np.empty(nPts, dtype=[('x', float), ('t', float), ('v', float), ('pt', float), ('m', float), ('f', float)]) 308 309 ## Inertial frame variables 310 self.x = self.x0 311 self.v = self.v0 312 self.m = self.m0 313 self.t = 0.0 ## reference clock always starts at 0 314 self.pt = self.t0 ## proper time starts at t0 315 316 ## reference frame variables 317 self.refx = None 318 self.refv = None 319 self.refm = None 320 self.reft = None 321 322 self.recordFrame(0) 323 324 def recordFrame(self, i): 325 f = self.force() 326 self.inertData[i] = (self.x, self.t, self.v, self.pt, self.m, f) 327 self.refData[i] = (self.refx, self.reft, self.refv, self.pt, self.refm, f) 328 329 def force(self, t=None): 330 if len(self.prog) == 0: 331 return 0.0 332 if t is None: 333 t = self.pt 334 335 ret = 0.0 336 for t1,f in self.prog: 337 if t >= t1: 338 ret = f 339 return ret 340 341 def acceleration(self, t=None): 342 return self.force(t) / self.m0 343 344 def accelLimits(self): 345 ## return the proper time values which bound the current acceleration command 346 if len(self.prog) == 0: 347 return -np.inf, np.inf 348 t = self.pt 349 ind = -1 350 for i, v in enumerate(self.prog): 351 t1,f = v 352 if t >= t1: 353 ind = i 354 355 if ind == -1: 356 return -np.inf, self.prog[0][0] 357 elif ind == len(self.prog)-1: 358 return self.prog[-1][0], np.inf 359 else: 360 return self.prog[ind][0], self.prog[ind+1][0] 361 362 363 def getCurve(self, ref=True): 364 365 if ref is False: 366 data = self.inertData 367 else: 368 data = self.refData[1:] 369 370 x = data['x'] 371 y = data['t'] 372 373 curve = pg.PlotCurveItem(x=x, y=y, pen=self.pen) 374 #x = self.data['x'] - ref.data['x'] 375 #y = self.data['t'] 376 377 step = 1.0 378 #mod = self.data['pt'] % step 379 #inds = np.argwhere(abs(mod[1:] - mod[:-1]) > step*0.9) 380 inds = [0] 381 pt = data['pt'] 382 for i in range(1,len(pt)): 383 diff = pt[i] - pt[inds[-1]] 384 if abs(diff) >= step: 385 inds.append(i) 386 inds = np.array(inds) 387 388 #t = self.data['t'][inds] 389 #x = self.data['x'][inds] 390 pts = [] 391 for i in inds: 392 x = data['x'][i] 393 y = data['t'][i] 394 if i+1 < len(data): 395 dpt = data['pt'][i+1]-data['pt'][i] 396 dt = data['t'][i+1]-data['t'][i] 397 else: 398 dpt = 1 399 400 if dpt > 0: 401 c = pg.mkBrush((0,0,0)) 402 else: 403 c = pg.mkBrush((200,200,200)) 404 pts.append({'pos': (x, y), 'brush': c}) 405 406 points = pg.ScatterPlotItem(pts, pen=self.pen, size=7) 407 408 return curve, points 409 410 411class Simulation: 412 def __init__(self, clocks, ref, duration, dt): 413 self.clocks = clocks 414 self.ref = ref 415 self.duration = duration 416 self.dt = dt 417 418 @staticmethod 419 def hypTStep(dt, v0, x0, tau0, g): 420 ## Hyperbolic step. 421 ## If an object has proper acceleration g and starts at position x0 with speed v0 and proper time tau0 422 ## as seen from an inertial frame, then return the new v, x, tau after time dt has elapsed. 423 if g == 0: 424 return v0, x0 + v0*dt, tau0 + dt * (1. - v0**2)**0.5 425 v02 = v0**2 426 g2 = g**2 427 428 tinit = v0 / (g * (1 - v02)**0.5) 429 430 B = (1 + (g2 * (dt+tinit)**2))**0.5 431 432 v1 = g * (dt+tinit) / B 433 434 dtau = (np.arcsinh(g * (dt+tinit)) - np.arcsinh(g * tinit)) / g 435 436 tau1 = tau0 + dtau 437 438 x1 = x0 + (1.0 / g) * ( B - 1. / (1.-v02)**0.5 ) 439 440 return v1, x1, tau1 441 442 443 @staticmethod 444 def tStep(dt, v0, x0, tau0, g): 445 ## Linear step. 446 ## Probably not as accurate as hyperbolic step, but certainly much faster. 447 gamma = (1. - v0**2)**-0.5 448 dtau = dt / gamma 449 return v0 + dtau * g, x0 + v0*dt, tau0 + dtau 450 451 @staticmethod 452 def tauStep(dtau, v0, x0, t0, g): 453 ## linear step in proper time of clock. 454 ## If an object has proper acceleration g and starts at position x0 with speed v0 at time t0 455 ## as seen from an inertial frame, then return the new v, x, t after proper time dtau has elapsed. 456 457 458 ## Compute how much t will change given a proper-time step of dtau 459 gamma = (1. - v0**2)**-0.5 460 if g == 0: 461 dt = dtau * gamma 462 else: 463 v0g = v0 * gamma 464 dt = (np.sinh(dtau * g + np.arcsinh(v0g)) - v0g) / g 465 466 #return v0 + dtau * g, x0 + v0*dt, t0 + dt 467 v1, x1, t1 = Simulation.hypTStep(dt, v0, x0, t0, g) 468 return v1, x1, t0+dt 469 470 @staticmethod 471 def hypIntersect(x0r, t0r, vr, x0, t0, v0, g): 472 ## given a reference clock (seen from inertial frame) has rx, rt, and rv, 473 ## and another clock starts at x0, t0, and v0, with acceleration g, 474 ## compute the intersection time of the object clock's hyperbolic path with 475 ## the reference plane. 476 477 ## I'm sure we can simplify this... 478 479 if g == 0: ## no acceleration, path is linear (and hyperbola is undefined) 480 #(-t0r + t0 v0 vr - vr x0 + vr x0r)/(-1 + v0 vr) 481 482 t = (-t0r + t0 *v0 *vr - vr *x0 + vr *x0r)/(-1 + v0 *vr) 483 return t 484 485 gamma = (1.0-v0**2)**-0.5 486 sel = (1 if g>0 else 0) + (1 if vr<0 else 0) 487 sel = sel%2 488 if sel == 0: 489 #(1/(g^2 (-1 + vr^2)))(-g^2 t0r + g gamma vr + g^2 t0 vr^2 - 490 #g gamma v0 vr^2 - g^2 vr x0 + 491 #g^2 vr x0r + \[Sqrt](g^2 vr^2 (1 + gamma^2 (v0 - vr)^2 - vr^2 + 492 #2 g gamma (v0 - vr) (-t0 + t0r + vr (x0 - x0r)) + 493 #g^2 (t0 - t0r + vr (-x0 + x0r))^2))) 494 495 t = (1./(g**2 *(-1. + vr**2)))*(-g**2 *t0r + g *gamma *vr + g**2 *t0 *vr**2 - g *gamma *v0 *vr**2 - g**2 *vr *x0 + g**2 *vr *x0r + np.sqrt(g**2 *vr**2 *(1. + gamma**2 *(v0 - vr)**2 - vr**2 + 2 *g *gamma *(v0 - vr)* (-t0 + t0r + vr *(x0 - x0r)) + g**2 *(t0 - t0r + vr* (-x0 + x0r))**2))) 496 497 else: 498 499 #-(1/(g^2 (-1 + vr^2)))(g^2 t0r - g gamma vr - g^2 t0 vr^2 + 500 #g gamma v0 vr^2 + g^2 vr x0 - 501 #g^2 vr x0r + \[Sqrt](g^2 vr^2 (1 + gamma^2 (v0 - vr)^2 - vr^2 + 502 #2 g gamma (v0 - vr) (-t0 + t0r + vr (x0 - x0r)) + 503 #g^2 (t0 - t0r + vr (-x0 + x0r))^2))) 504 505 t = -(1./(g**2 *(-1. + vr**2)))*(g**2 *t0r - g *gamma* vr - g**2 *t0 *vr**2 + g *gamma *v0 *vr**2 + g**2* vr* x0 - g**2 *vr *x0r + np.sqrt(g**2* vr**2 *(1. + gamma**2 *(v0 - vr)**2 - vr**2 + 2 *g *gamma *(v0 - vr) *(-t0 + t0r + vr *(x0 - x0r)) + g**2 *(t0 - t0r + vr *(-x0 + x0r))**2))) 506 return t 507 508 def run(self): 509 nPts = int(self.duration/self.dt)+1 510 for cl in self.clocks.values(): 511 cl.init(nPts) 512 513 if self.ref is None: 514 self.runInertial(nPts) 515 else: 516 self.runReference(nPts) 517 518 def runInertial(self, nPts): 519 clocks = self.clocks 520 dt = self.dt 521 tVals = np.linspace(0, dt*(nPts-1), nPts) 522 for cl in self.clocks.values(): 523 for i in range(1,nPts): 524 nextT = tVals[i] 525 while True: 526 tau1, tau2 = cl.accelLimits() 527 x = cl.x 528 v = cl.v 529 tau = cl.pt 530 g = cl.acceleration() 531 532 v1, x1, tau1 = self.hypTStep(dt, v, x, tau, g) 533 if tau1 > tau2: 534 dtau = tau2-tau 535 cl.v, cl.x, cl.t = self.tauStep(dtau, v, x, cl.t, g) 536 cl.pt = tau2 537 else: 538 cl.v, cl.x, cl.pt = v1, x1, tau1 539 cl.t += dt 540 541 if cl.t >= nextT: 542 cl.refx = cl.x 543 cl.refv = cl.v 544 cl.reft = cl.t 545 cl.recordFrame(i) 546 break 547 548 549 def runReference(self, nPts): 550 clocks = self.clocks 551 ref = self.ref 552 dt = self.dt 553 dur = self.duration 554 555 ## make sure reference clock is not present in the list of clocks--this will be handled separately. 556 clocks = clocks.copy() 557 for k,v in clocks.items(): 558 if v is ref: 559 del clocks[k] 560 break 561 562 ref.refx = 0 563 ref.refv = 0 564 ref.refm = ref.m0 565 566 ## These are the set of proper times (in the reference frame) that will be simulated 567 ptVals = np.linspace(ref.pt, ref.pt + dt*(nPts-1), nPts) 568 569 for i in range(1,nPts): 570 571 ## step reference clock ahead one time step in its proper time 572 nextPt = ptVals[i] ## this is where (when) we want to end up 573 while True: 574 tau1, tau2 = ref.accelLimits() 575 dtau = min(nextPt-ref.pt, tau2-ref.pt) ## do not step past the next command boundary 576 g = ref.acceleration() 577 v, x, t = Simulation.tauStep(dtau, ref.v, ref.x, ref.t, g) 578 ref.pt += dtau 579 ref.v = v 580 ref.x = x 581 ref.t = t 582 ref.reft = ref.pt 583 if ref.pt >= nextPt: 584 break 585 #else: 586 #print "Stepped to", tau2, "instead of", nextPt 587 ref.recordFrame(i) 588 589 ## determine plane visible to reference clock 590 ## this plane goes through the point ref.x, ref.t and has slope = ref.v 591 592 593 ## update all other clocks 594 for cl in clocks.values(): 595 while True: 596 g = cl.acceleration() 597 tau1, tau2 = cl.accelLimits() 598 ##Given current position / speed of clock, determine where it will intersect reference plane 599 #t1 = (ref.v * (cl.x - cl.v * cl.t) + (ref.t - ref.v * ref.x)) / (1. - cl.v) 600 t1 = Simulation.hypIntersect(ref.x, ref.t, ref.v, cl.x, cl.t, cl.v, g) 601 dt1 = t1 - cl.t 602 603 ## advance clock by correct time step 604 v, x, tau = Simulation.hypTStep(dt1, cl.v, cl.x, cl.pt, g) 605 606 ## check to see whether we have gone past an acceleration command boundary. 607 ## if so, we must instead advance the clock to the boundary and start again 608 if tau < tau1: 609 dtau = tau1 - cl.pt 610 cl.v, cl.x, cl.t = Simulation.tauStep(dtau, cl.v, cl.x, cl.t, g) 611 cl.pt = tau1-0.000001 612 continue 613 if tau > tau2: 614 dtau = tau2 - cl.pt 615 cl.v, cl.x, cl.t = Simulation.tauStep(dtau, cl.v, cl.x, cl.t, g) 616 cl.pt = tau2 617 continue 618 619 ## Otherwise, record the new values and exit the loop 620 cl.v = v 621 cl.x = x 622 cl.pt = tau 623 cl.t = t1 624 cl.m = None 625 break 626 627 ## transform position into reference frame 628 x = cl.x - ref.x 629 t = cl.t - ref.t 630 gamma = (1.0 - ref.v**2) ** -0.5 631 vg = -ref.v * gamma 632 633 cl.refx = gamma * (x - ref.v * t) 634 cl.reft = ref.pt # + gamma * (t - ref.v * x) # this term belongs here, but it should always be equal to 0. 635 cl.refv = (cl.v - ref.v) / (1.0 - cl.v * ref.v) 636 cl.refm = None 637 cl.recordFrame(i) 638 639 t += dt 640 641 def plot(self, plot): 642 plot.clear() 643 for cl in self.clocks.values(): 644 c, p = cl.getCurve() 645 plot.addItem(c) 646 plot.addItem(p) 647 648class Animation(pg.ItemGroup): 649 def __init__(self, sim): 650 pg.ItemGroup.__init__(self) 651 self.sim = sim 652 self.clocks = sim.clocks 653 654 self.items = {} 655 for name, cl in self.clocks.items(): 656 item = ClockItem(cl) 657 self.addItem(item) 658 self.items[name] = item 659 660 def restart(self): 661 for cl in self.items.values(): 662 cl.reset() 663 664 def stepTo(self, t): 665 for i in self.items.values(): 666 i.stepTo(t) 667 668 669class ClockItem(pg.ItemGroup): 670 def __init__(self, clock): 671 pg.ItemGroup.__init__(self) 672 self.size = clock.size 673 self.item = QtGui.QGraphicsEllipseItem(QtCore.QRectF(0, 0, self.size, self.size)) 674 tr = QtGui.QTransform.fromTranslate(-self.size*0.5, -self.size*0.5) 675 self.item.setTransform(tr) 676 self.item.setPen(pg.mkPen(100,100,100)) 677 self.item.setBrush(clock.brush) 678 self.hand = QtGui.QGraphicsLineItem(0, 0, 0, self.size*0.5) 679 self.hand.setPen(pg.mkPen('w')) 680 self.hand.setZValue(10) 681 self.flare = QtGui.QGraphicsPolygonItem(QtGui.QPolygonF([ 682 QtCore.QPointF(0, -self.size*0.25), 683 QtCore.QPointF(0, self.size*0.25), 684 QtCore.QPointF(self.size*1.5, 0), 685 QtCore.QPointF(0, -self.size*0.25), 686 ])) 687 self.flare.setPen(pg.mkPen('y')) 688 self.flare.setBrush(pg.mkBrush(255,150,0)) 689 self.flare.setZValue(-10) 690 self.addItem(self.hand) 691 self.addItem(self.item) 692 self.addItem(self.flare) 693 694 self.clock = clock 695 self.i = 1 696 697 self._spaceline = None 698 699 700 def spaceline(self): 701 if self._spaceline is None: 702 self._spaceline = pg.InfiniteLine() 703 self._spaceline.setPen(self.clock.pen) 704 return self._spaceline 705 706 def stepTo(self, t): 707 data = self.clock.refData 708 709 while self.i < len(data)-1 and data['t'][self.i] < t: 710 self.i += 1 711 while self.i > 1 and data['t'][self.i-1] >= t: 712 self.i -= 1 713 714 self.setPos(data['x'][self.i], self.clock.y0) 715 716 t = data['pt'][self.i] 717 self.hand.setRotation(-0.25 * t * 360.) 718 719 v = data['v'][self.i] 720 gam = (1.0 - v**2)**0.5 721 self.setTransform(QtGui.QTransform.fromScale(gam, 1.0)) 722 723 f = data['f'][self.i] 724 tr = QtGui.QTransform() 725 if f < 0: 726 tr.translate(self.size*0.4, 0) 727 else: 728 tr.translate(-self.size*0.4, 0) 729 730 tr.scale(-f * (0.5+np.random.random()*0.1), 1.0) 731 self.flare.setTransform(tr) 732 733 if self._spaceline is not None: 734 self._spaceline.setPos(pg.Point(data['x'][self.i], data['t'][self.i])) 735 self._spaceline.setAngle(data['v'][self.i] * 45.) 736 737 738 def reset(self): 739 self.i = 1 740 741 742#class Spaceline(pg.InfiniteLine): 743 #def __init__(self, sim, frame): 744 #self.sim = sim 745 #self.frame = frame 746 #pg.InfiniteLine.__init__(self) 747 #self.setPen(sim.clocks[frame].pen) 748 749 #def stepTo(self, t): 750 #self.setAngle(0) 751 752 #pass 753 754if __name__ == '__main__': 755 app = pg.mkQApp() 756 #import pyqtgraph.console 757 #cw = pyqtgraph.console.ConsoleWidget() 758 #cw.show() 759 #cw.catchNextException() 760 win = RelativityGUI() 761 win.setWindowTitle("Relativity!") 762 win.show() 763 win.resize(1100,700) 764 765 pg.exec() 766