CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
relativity.py763 linesDownload Raw Back to relativity
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 
Aluode/PerceptionLabPortable · CoolFace