Aluode/PerceptionLabPortable
0
1import collections2import os3import sys4from time import perf_counter5 6import numpy as np7 8import pyqtgraph as pg9from pyqtgraph import configfile10from pyqtgraph.parametertree import Parameter, ParameterTree11from pyqtgraph.parametertree import types as pTypes12from pyqtgraph.Qt import QtCore, QtGui, QtWidgets13 14 15class RelativityGUI(QtWidgets.QWidget):16 def __init__(self):17 QtWidgets.QWidget.__init__(self)18 19 self.animations = []20 self.animTimer = QtCore.QTimer()21 self.animTimer.timeout.connect(self.stepAnimation)22 self.animTime = 023 self.animDt = .01624 self.lastAnimTime = 025 26 self.setupGUI()27 28 self.objectGroup = ObjectGroupParam()29 30 self.params = Parameter.create(name='params', type='group', children=[31 dict(name='Load Preset..', type='list', limits=[]),32 #dict(name='Unit System', type='list', limits=['', 'MKS']),33 dict(name='Duration', type='float', value=10.0, step=0.1, limits=[0.1, None]),34 dict(name='Reference Frame', type='list', limits=[]),35 dict(name='Animate', type='bool', value=True),36 dict(name='Animation Speed', type='float', value=1.0, dec=True, step=0.1, limits=[0.0001, None]),37 dict(name='Recalculate Worldlines', type='action'),38 dict(name='Save', type='action'),39 dict(name='Load', type='action'),40 self.objectGroup,41 ])42 self.tree.setParameters(self.params, showTop=False)43 self.params.param('Recalculate Worldlines').sigActivated.connect(self.recalculate)44 self.params.param('Save').sigActivated.connect(self.save)45 self.params.param('Load').sigActivated.connect(self.load)46 self.params.param('Load Preset..').sigValueChanged.connect(self.loadPreset)47 self.params.sigTreeStateChanged.connect(self.treeChanged)48 49 ## read list of preset configs50 presetDir = os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), 'presets')51 if os.path.exists(presetDir):52 presets = [os.path.splitext(p)[0] for p in os.listdir(presetDir)]53 self.params.param('Load Preset..').setLimits(['']+presets)54 55 56 57 58 def setupGUI(self):59 self.layout = QtWidgets.QVBoxLayout()60 self.layout.setContentsMargins(0,0,0,0)61 self.setLayout(self.layout)62 self.splitter = QtWidgets.QSplitter()63 self.splitter.setOrientation(QtCore.Qt.Orientation.Horizontal)64 self.layout.addWidget(self.splitter)65 66 self.tree = ParameterTree(showHeader=False)67 self.splitter.addWidget(self.tree)68 69 self.splitter2 = QtWidgets.QSplitter()70 self.splitter2.setOrientation(QtCore.Qt.Orientation.Vertical)71 self.splitter.addWidget(self.splitter2)72 73 self.worldlinePlots = pg.GraphicsLayoutWidget()74 self.splitter2.addWidget(self.worldlinePlots)75 76 self.animationPlots = pg.GraphicsLayoutWidget()77 self.splitter2.addWidget(self.animationPlots)78 79 self.splitter2.setSizes([int(self.height()*0.8), int(self.height()*0.2)])80 81 self.inertWorldlinePlot = self.worldlinePlots.addPlot()82 self.refWorldlinePlot = self.worldlinePlots.addPlot()83 84 self.inertAnimationPlot = self.animationPlots.addPlot()85 self.inertAnimationPlot.setAspectLocked(1)86 self.refAnimationPlot = self.animationPlots.addPlot()87 self.refAnimationPlot.setAspectLocked(1)88 89 self.inertAnimationPlot.setXLink(self.inertWorldlinePlot)90 self.refAnimationPlot.setXLink(self.refWorldlinePlot)91 92 def recalculate(self):93 ## build 2 sets of clocks94 clocks1 = collections.OrderedDict()95 clocks2 = collections.OrderedDict()96 for cl in self.params.param('Objects'):97 clocks1.update(cl.buildClocks())98 clocks2.update(cl.buildClocks())99 100 ## Inertial simulation101 dt = self.animDt * self.params['Animation Speed']102 sim1 = Simulation(clocks1, ref=None, duration=self.params['Duration'], dt=dt)103 sim1.run()104 sim1.plot(self.inertWorldlinePlot)105 self.inertWorldlinePlot.autoRange(padding=0.1)106 107 ## reference simulation108 ref = self.params['Reference Frame']109 dur = clocks1[ref].refData['pt'][-1] ## decide how long to run the reference simulation110 sim2 = Simulation(clocks2, ref=clocks2[ref], duration=dur, dt=dt)111 sim2.run()112 sim2.plot(self.refWorldlinePlot)113 self.refWorldlinePlot.autoRange(padding=0.1)114 115 116 ## create animations117 self.refAnimationPlot.clear()118 self.inertAnimationPlot.clear()119 self.animTime = 0120 121 self.animations = [Animation(sim1), Animation(sim2)]122 self.inertAnimationPlot.addItem(self.animations[0])123 self.refAnimationPlot.addItem(self.animations[1])124 125 ## create lines representing all that is visible to a particular reference126 #self.inertSpaceline = Spaceline(sim1, ref)127 #self.refSpaceline = Spaceline(sim2)128 self.inertWorldlinePlot.addItem(self.animations[0].items[ref].spaceline())129 self.refWorldlinePlot.addItem(self.animations[1].items[ref].spaceline())130 131 132 133 134 def setAnimation(self, a):135 if a:136 self.lastAnimTime = perf_counter()137 self.animTimer.start(int(self.animDt*1000))138 else:139 self.animTimer.stop()140 141 def stepAnimation(self):142 now = perf_counter()143 dt = (now-self.lastAnimTime) * self.params['Animation Speed']144 self.lastAnimTime = now145 self.animTime += dt146 if self.animTime > self.params['Duration']:147 self.animTime = 0148 for a in self.animations:149 a.restart()150 151 for a in self.animations:152 a.stepTo(self.animTime)153 154 155 def treeChanged(self, *args):156 clocks = []157 for c in self.params.param('Objects'):158 clocks.extend(c.clockNames())159 #for param, change, data in args[1]:160 #if change == 'childAdded':161 self.params.param('Reference Frame').setLimits(clocks)162 self.setAnimation(self.params['Animate'])163 164 def save(self):165 filename, _ = QtWidgets.QFileDialog.getSaveFileName(self, "Save State..", "untitled.cfg", "Config Files (*.cfg)")166 if not filename:167 return168 state = self.params.saveState()169 configfile.writeConfigFile(state, filename)170 171 def load(self):172 filename, _ = QtWidgets.QFileDialog.getOpenFileName(self, "Save State..", "", "Config Files (*.cfg)")173 if not filename:174 return175 state = configfile.readConfigFile(filename)176 self.loadState(state)177 178 def loadPreset(self, param, preset):179 if preset == '':180 return181 path = os.path.abspath(os.path.dirname(__file__))182 fn = os.path.join(path, 'presets', preset+".cfg")183 state = configfile.readConfigFile(fn)184 self.loadState(state)185 186 def loadState(self, state):187 if 'Load Preset..' in state['children']:188 del state['children']['Load Preset..']['limits']189 del state['children']['Load Preset..']['value']190 self.params.param('Objects').clearChildren()191 self.params.restoreState(state, removeChildren=False)192 self.recalculate()193 194 195class ObjectGroupParam(pTypes.GroupParameter):196 def __init__(self):197 pTypes.GroupParameter.__init__(self, name="Objects", addText="Add New..", addList=['Clock', 'Grid'])198 199 def addNew(self, typ):200 if typ == 'Clock':201 self.addChild(ClockParam())202 elif typ == 'Grid':203 self.addChild(GridParam())204 205class ClockParam(pTypes.GroupParameter):206 def __init__(self, **kwds):207 defs = dict(name="Clock", autoIncrementName=True, renamable=True, removable=True, children=[208 dict(name='Initial Position', type='float', value=0.0, step=0.1),209 #dict(name='V0', type='float', value=0.0, step=0.1),210 AccelerationGroup(),211 212 dict(name='Rest Mass', type='float', value=1.0, step=0.1, limits=[1e-9, None]),213 dict(name='Color', type='color', value=(100,100,150)),214 dict(name='Size', type='float', value=0.5),215 dict(name='Vertical Position', type='float', value=0.0, step=0.1),216 ])217 #defs.update(kwds)218 pTypes.GroupParameter.__init__(self, **defs)219 self.restoreState(kwds, removeChildren=False)220 221 def buildClocks(self):222 x0 = self['Initial Position']223 y0 = self['Vertical Position']224 color = self['Color']225 m = self['Rest Mass']226 size = self['Size']227 prog = self.param('Acceleration').generate()228 c = Clock(x0=x0, m0=m, y0=y0, color=color, prog=prog, size=size)229 return {self.name(): c}230 231 def clockNames(self):232 return [self.name()]233 234pTypes.registerParameterType('Clock', ClockParam)235 236class GridParam(pTypes.GroupParameter):237 def __init__(self, **kwds):238 defs = dict(name="Grid", autoIncrementName=True, renamable=True, removable=True, children=[239 dict(name='Number of Clocks', type='int', value=5, limits=[1, None]),240 dict(name='Spacing', type='float', value=1.0, step=0.1),241 ClockParam(name='ClockTemplate'),242 ])243 #defs.update(kwds)244 pTypes.GroupParameter.__init__(self, **defs)245 self.restoreState(kwds, removeChildren=False)246 247 def buildClocks(self):248 clocks = {}249 template = self.param('ClockTemplate')250 spacing = self['Spacing']251 for i in range(self['Number of Clocks']):252 c = list(template.buildClocks().values())[0]253 c.x0 += i * spacing254 clocks[self.name() + '%02d' % i] = c255 return clocks256 257 def clockNames(self):258 return [self.name() + '%02d' % i for i in range(self['Number of Clocks'])]259 260pTypes.registerParameterType('Grid', GridParam)261 262class AccelerationGroup(pTypes.GroupParameter):263 def __init__(self, **kwds):264 defs = dict(name="Acceleration", addText="Add Command..")265 pTypes.GroupParameter.__init__(self, **defs)266 self.restoreState(kwds, removeChildren=False)267 268 def addNew(self):269 nextTime = 0.0270 if self.hasChildren():271 nextTime = self.children()[-1]['Proper Time'] + 1272 self.addChild(Parameter.create(name='Command', autoIncrementName=True, type=None, renamable=True, removable=True, children=[273 dict(name='Proper Time', type='float', value=nextTime),274 dict(name='Acceleration', type='float', value=0.0, step=0.1),275 ]))276 277 def generate(self):278 prog = []279 for cmd in self:280 prog.append((cmd['Proper Time'], cmd['Acceleration']))281 return prog 282 283pTypes.registerParameterType('AccelerationGroup', AccelerationGroup)284 285 286class Clock(object):287 nClocks = 0288 289 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):290 Clock.nClocks += 1291 self.pen = pg.mkPen(color)292 self.brush = pg.mkBrush(color)293 self.y0 = y0294 self.x0 = x0295 self.v0 = v0296 self.m0 = m0297 self.t0 = t0298 self.prog = prog299 self.size = size300 301 def init(self, nPts):302 ## Keep records of object from inertial frame as well as reference frame303 self.inertData = np.empty(nPts, dtype=[('x', float), ('t', float), ('v', float), ('pt', float), ('m', float), ('f', float)])304 self.refData = np.empty(nPts, dtype=[('x', float), ('t', float), ('v', float), ('pt', float), ('m', float), ('f', float)])305 306 ## Inertial frame variables307 self.x = self.x0308 self.v = self.v0309 self.m = self.m0310 self.t = 0.0 ## reference clock always starts at 0311 self.pt = self.t0 ## proper time starts at t0312 313 ## reference frame variables314 self.refx = None315 self.refv = None316 self.refm = None317 self.reft = None318 319 self.recordFrame(0)320 321 def recordFrame(self, i):322 f = self.force()323 self.inertData[i] = (self.x, self.t, self.v, self.pt, self.m, f)324 self.refData[i] = (self.refx, self.reft, self.refv, self.pt, self.refm, f)325 326 def force(self, t=None):327 if len(self.prog) == 0:328 return 0.0329 if t is None:330 t = self.pt331 332 ret = 0.0333 for t1,f in self.prog:334 if t >= t1:335 ret = f336 return ret337 338 def acceleration(self, t=None):339 return self.force(t) / self.m0340 341 def accelLimits(self):342 ## return the proper time values which bound the current acceleration command343 if len(self.prog) == 0:344 return -np.inf, np.inf345 t = self.pt346 ind = -1347 for i, v in enumerate(self.prog):348 t1,f = v349 if t >= t1:350 ind = i351 352 if ind == -1:353 return -np.inf, self.prog[0][0]354 elif ind == len(self.prog)-1:355 return self.prog[-1][0], np.inf356 else:357 return self.prog[ind][0], self.prog[ind+1][0]358 359 360 def getCurve(self, ref=True):361 362 if ref is False:363 data = self.inertData364 else:365 data = self.refData[1:]366 367 x = data['x']368 y = data['t']369 370 curve = pg.PlotCurveItem(x=x, y=y, pen=self.pen)371 #x = self.data['x'] - ref.data['x']372 #y = self.data['t']373 374 step = 1.0375 #mod = self.data['pt'] % step376 #inds = np.argwhere(abs(mod[1:] - mod[:-1]) > step*0.9)377 inds = [0]378 pt = data['pt']379 for i in range(1,len(pt)):380 diff = pt[i] - pt[inds[-1]]381 if abs(diff) >= step:382 inds.append(i)383 inds = np.array(inds)384 385 #t = self.data['t'][inds]386 #x = self.data['x'][inds] 387 pts = []388 for i in inds:389 x = data['x'][i]390 y = data['t'][i]391 if i+1 < len(data):392 dpt = data['pt'][i+1]-data['pt'][i]393 dt = data['t'][i+1]-data['t'][i]394 else:395 dpt = 1396 397 if dpt > 0:398 c = pg.mkBrush((0,0,0))399 else:400 c = pg.mkBrush((200,200,200))401 pts.append({'pos': (x, y), 'brush': c})402 403 points = pg.ScatterPlotItem(pts, pen=self.pen, size=7)404 405 return curve, points406 407 408class Simulation:409 def __init__(self, clocks, ref, duration, dt):410 self.clocks = clocks411 self.ref = ref412 self.duration = duration413 self.dt = dt414 415 @staticmethod416 def hypTStep(dt, v0, x0, tau0, g):417 ## Hyperbolic step. 418 ## If an object has proper acceleration g and starts at position x0 with speed v0 and proper time tau0419 ## as seen from an inertial frame, then return the new v, x, tau after time dt has elapsed.420 if g == 0:421 return v0, x0 + v0*dt, tau0 + dt * (1. - v0**2)**0.5422 v02 = v0**2423 g2 = g**2424 425 tinit = v0 / (g * (1 - v02)**0.5)426 427 B = (1 + (g2 * (dt+tinit)**2))**0.5428 429 v1 = g * (dt+tinit) / B430 431 dtau = (np.arcsinh(g * (dt+tinit)) - np.arcsinh(g * tinit)) / g432 433 tau1 = tau0 + dtau434 435 x1 = x0 + (1.0 / g) * ( B - 1. / (1.-v02)**0.5 )436 437 return v1, x1, tau1438 439 440 @staticmethod441 def tStep(dt, v0, x0, tau0, g):442 ## Linear step.443 ## Probably not as accurate as hyperbolic step, but certainly much faster.444 gamma = (1. - v0**2)**-0.5445 dtau = dt / gamma446 return v0 + dtau * g, x0 + v0*dt, tau0 + dtau447 448 @staticmethod449 def tauStep(dtau, v0, x0, t0, g):450 ## linear step in proper time of clock.451 ## If an object has proper acceleration g and starts at position x0 with speed v0 at time t0452 ## as seen from an inertial frame, then return the new v, x, t after proper time dtau has elapsed.453 454 455 ## Compute how much t will change given a proper-time step of dtau456 gamma = (1. - v0**2)**-0.5457 if g == 0:458 dt = dtau * gamma459 else:460 v0g = v0 * gamma461 dt = (np.sinh(dtau * g + np.arcsinh(v0g)) - v0g) / g462 463 #return v0 + dtau * g, x0 + v0*dt, t0 + dt464 v1, x1, t1 = Simulation.hypTStep(dt, v0, x0, t0, g)465 return v1, x1, t0+dt466 467 @staticmethod468 def hypIntersect(x0r, t0r, vr, x0, t0, v0, g):469 ## given a reference clock (seen from inertial frame) has rx, rt, and rv,470 ## and another clock starts at x0, t0, and v0, with acceleration g,471 ## compute the intersection time of the object clock's hyperbolic path with 472 ## the reference plane.473 474 ## I'm sure we can simplify this...475 476 if g == 0: ## no acceleration, path is linear (and hyperbola is undefined)477 #(-t0r + t0 v0 vr - vr x0 + vr x0r)/(-1 + v0 vr)478 479 t = (-t0r + t0 *v0 *vr - vr *x0 + vr *x0r)/(-1 + v0 *vr)480 return t481 482 gamma = (1.0-v0**2)**-0.5483 sel = (1 if g>0 else 0) + (1 if vr<0 else 0)484 sel = sel%2485 if sel == 0:486 #(1/(g^2 (-1 + vr^2)))(-g^2 t0r + g gamma vr + g^2 t0 vr^2 - 487 #g gamma v0 vr^2 - g^2 vr x0 + 488 #g^2 vr x0r + \[Sqrt](g^2 vr^2 (1 + gamma^2 (v0 - vr)^2 - vr^2 + 489 #2 g gamma (v0 - vr) (-t0 + t0r + vr (x0 - x0r)) + 490 #g^2 (t0 - t0r + vr (-x0 + x0r))^2)))491 492 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)))493 494 else:495 496 #-(1/(g^2 (-1 + vr^2)))(g^2 t0r - g gamma vr - g^2 t0 vr^2 + 497 #g gamma v0 vr^2 + g^2 vr x0 - 498 #g^2 vr x0r + \[Sqrt](g^2 vr^2 (1 + gamma^2 (v0 - vr)^2 - vr^2 + 499 #2 g gamma (v0 - vr) (-t0 + t0r + vr (x0 - x0r)) + 500 #g^2 (t0 - t0r + vr (-x0 + x0r))^2)))501 502 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)))503 return t504 505 def run(self):506 nPts = int(self.duration/self.dt)+1507 for cl in self.clocks.values():508 cl.init(nPts)509 510 if self.ref is None:511 self.runInertial(nPts)512 else:513 self.runReference(nPts)514 515 def runInertial(self, nPts):516 clocks = self.clocks517 dt = self.dt518 tVals = np.linspace(0, dt*(nPts-1), nPts)519 for cl in self.clocks.values():520 for i in range(1,nPts):521 nextT = tVals[i]522 while True:523 tau1, tau2 = cl.accelLimits()524 x = cl.x525 v = cl.v526 tau = cl.pt527 g = cl.acceleration()528 529 v1, x1, tau1 = self.hypTStep(dt, v, x, tau, g)530 if tau1 > tau2:531 dtau = tau2-tau532 cl.v, cl.x, cl.t = self.tauStep(dtau, v, x, cl.t, g)533 cl.pt = tau2534 else:535 cl.v, cl.x, cl.pt = v1, x1, tau1536 cl.t += dt537 538 if cl.t >= nextT:539 cl.refx = cl.x540 cl.refv = cl.v541 cl.reft = cl.t542 cl.recordFrame(i)543 break544 545 546 def runReference(self, nPts):547 clocks = self.clocks548 ref = self.ref549 dt = self.dt550 dur = self.duration551 552 ## make sure reference clock is not present in the list of clocks--this will be handled separately.553 clocks = clocks.copy()554 for k,v in clocks.items():555 if v is ref:556 del clocks[k]557 break558 559 ref.refx = 0560 ref.refv = 0561 ref.refm = ref.m0562 563 ## These are the set of proper times (in the reference frame) that will be simulated564 ptVals = np.linspace(ref.pt, ref.pt + dt*(nPts-1), nPts)565 566 for i in range(1,nPts):567 568 ## step reference clock ahead one time step in its proper time569 nextPt = ptVals[i] ## this is where (when) we want to end up570 while True:571 tau1, tau2 = ref.accelLimits()572 dtau = min(nextPt-ref.pt, tau2-ref.pt) ## do not step past the next command boundary573 g = ref.acceleration()574 v, x, t = Simulation.tauStep(dtau, ref.v, ref.x, ref.t, g)575 ref.pt += dtau576 ref.v = v577 ref.x = x578 ref.t = t579 ref.reft = ref.pt580 if ref.pt >= nextPt:581 break582 #else:583 #print "Stepped to", tau2, "instead of", nextPt584 ref.recordFrame(i)585 586 ## determine plane visible to reference clock587 ## this plane goes through the point ref.x, ref.t and has slope = ref.v588 589 590 ## update all other clocks591 for cl in clocks.values():592 while True:593 g = cl.acceleration()594 tau1, tau2 = cl.accelLimits()595 ##Given current position / speed of clock, determine where it will intersect reference plane596 #t1 = (ref.v * (cl.x - cl.v * cl.t) + (ref.t - ref.v * ref.x)) / (1. - cl.v)597 t1 = Simulation.hypIntersect(ref.x, ref.t, ref.v, cl.x, cl.t, cl.v, g)598 dt1 = t1 - cl.t599 600 ## advance clock by correct time step601 v, x, tau = Simulation.hypTStep(dt1, cl.v, cl.x, cl.pt, g)602 603 ## check to see whether we have gone past an acceleration command boundary.604 ## if so, we must instead advance the clock to the boundary and start again605 if tau < tau1:606 dtau = tau1 - cl.pt607 cl.v, cl.x, cl.t = Simulation.tauStep(dtau, cl.v, cl.x, cl.t, g)608 cl.pt = tau1-0.000001 609 continue610 if tau > tau2:611 dtau = tau2 - cl.pt612 cl.v, cl.x, cl.t = Simulation.tauStep(dtau, cl.v, cl.x, cl.t, g)613 cl.pt = tau2614 continue615 616 ## Otherwise, record the new values and exit the loop617 cl.v = v618 cl.x = x619 cl.pt = tau620 cl.t = t1621 cl.m = None622 break623 624 ## transform position into reference frame625 x = cl.x - ref.x626 t = cl.t - ref.t627 gamma = (1.0 - ref.v**2) ** -0.5628 vg = -ref.v * gamma629 630 cl.refx = gamma * (x - ref.v * t)631 cl.reft = ref.pt # + gamma * (t - ref.v * x) # this term belongs here, but it should always be equal to 0.632 cl.refv = (cl.v - ref.v) / (1.0 - cl.v * ref.v)633 cl.refm = None634 cl.recordFrame(i)635 636 t += dt637 638 def plot(self, plot):639 plot.clear()640 for cl in self.clocks.values():641 c, p = cl.getCurve()642 plot.addItem(c)643 plot.addItem(p)644 645class Animation(pg.ItemGroup):646 def __init__(self, sim):647 pg.ItemGroup.__init__(self)648 self.sim = sim649 self.clocks = sim.clocks650 651 self.items = {}652 for name, cl in self.clocks.items():653 item = ClockItem(cl)654 self.addItem(item)655 self.items[name] = item656 657 def restart(self):658 for cl in self.items.values():659 cl.reset()660 661 def stepTo(self, t):662 for i in self.items.values():663 i.stepTo(t)664 665 666class ClockItem(pg.ItemGroup):667 def __init__(self, clock):668 pg.ItemGroup.__init__(self)669 self.size = clock.size670 self.item = QtWidgets.QGraphicsEllipseItem(QtCore.QRectF(0, 0, self.size, self.size))671 tr = QtGui.QTransform.fromTranslate(-self.size*0.5, -self.size*0.5)672 self.item.setTransform(tr)673 self.item.setPen(pg.mkPen(100,100,100))674 self.item.setBrush(clock.brush)675 self.hand = QtWidgets.QGraphicsLineItem(0, 0, 0, self.size*0.5)676 self.hand.setPen(pg.mkPen('w'))677 self.hand.setZValue(10)678 self.flare = QtWidgets.QGraphicsPolygonItem(QtGui.QPolygonF([679 QtCore.QPointF(0, -self.size*0.25),680 QtCore.QPointF(0, self.size*0.25),681 QtCore.QPointF(self.size*1.5, 0),682 QtCore.QPointF(0, -self.size*0.25),683 ]))684 self.flare.setPen(pg.mkPen('y'))685 self.flare.setBrush(pg.mkBrush(255,150,0))686 self.flare.setZValue(-10)687 self.addItem(self.hand)688 self.addItem(self.item)689 self.addItem(self.flare)690 691 self.clock = clock692 self.i = 1693 694 self._spaceline = None695 696 697 def spaceline(self):698 if self._spaceline is None:699 self._spaceline = pg.InfiniteLine()700 self._spaceline.setPen(self.clock.pen)701 return self._spaceline702 703 def stepTo(self, t):704 data = self.clock.refData705 706 while self.i < len(data)-1 and data['t'][self.i] < t:707 self.i += 1708 while self.i > 1 and data['t'][self.i-1] >= t:709 self.i -= 1710 711 self.setPos(data['x'][self.i], self.clock.y0)712 713 t = data['pt'][self.i]714 self.hand.setRotation(-0.25 * t * 360.)715 716 v = data['v'][self.i]717 gam = (1.0 - v**2)**0.5718 self.setTransform(QtGui.QTransform.fromScale(gam, 1.0))719 720 f = data['f'][self.i]721 tr = QtGui.QTransform()722 if f < 0:723 tr.translate(self.size*0.4, 0)724 else:725 tr.translate(-self.size*0.4, 0)726 727 tr.scale(-f * (0.5+np.random.random()*0.1), 1.0)728 self.flare.setTransform(tr)729 730 if self._spaceline is not None:731 self._spaceline.setPos(pg.Point(data['x'][self.i], data['t'][self.i]))732 self._spaceline.setAngle(data['v'][self.i] * 45.)733 734 735 def reset(self):736 self.i = 1737 738 739#class Spaceline(pg.InfiniteLine):740 #def __init__(self, sim, frame):741 #self.sim = sim742 #self.frame = frame743 #pg.InfiniteLine.__init__(self)744 #self.setPen(sim.clocks[frame].pen)745 746 #def stepTo(self, t):747 #self.setAngle(0)748 749 #pass750 751if __name__ == '__main__':752 app = pg.mkQApp()753 #import pyqtgraph.console754 #cw = pyqtgraph.console.ConsoleWidget()755 #cw.show()756 #cw.catchNextException()757 win = RelativityGUI()758 win.setWindowTitle("Relativity!")759 win.show()760 win.resize(1100,700)761 762 pg.exec()763 