path: root/funcGenerator.py

import time  # For sleep, clock, time and perf_counter
from datetime import datetime, timedelta
import numpy as np

try:
    from pyqtgraph.Qt import QtCore
except:
    pass

class SinGen:
    def __init__(self, ampl=1, offset=0, phase=0, sweeper=None):
        self.ampl = ampl
        self.phase = phase
        self.sweeper = sweeper
        self.offset = offset

    def getValue(self, sweeper=None):
        if sweeper is None and self.sweeper is None:
            print("Error: generator needs sweeper")
            return 0
        if sweeper is None:
            sweeper = self.sweeper
        return self.ampl * np.sin( 2*np.pi*sweeper.getRelPos()) + self.offset

class RampGen:
    # monotonically goes from start to stop, at final point move to start
    def __init__(self, start=0, stop=0, sweeper=None):
        self.start = start
        self.stop = stop
        self.sweeper = sweeper

    def setCenterAndSpan(self, center, span):
        self.start = center-span/2
        self.stop  = center+span/2

    def setCenter(self, center):
        span = self.stop - self.start;
        self.setCenterAndSpan(center, span)

    def getCenter(self):
        return (self.stop + self.start)/2

    def setSpan(self, span):
        center = (self.stop + self.start)/2
        self.setCenterAndSpan(center, span)

    def getSpan(self):
        return self.stop - self.start;

    def getValue(self, sweeper=None):
        if sweeper is None and self.sweeper is None:
            print("Error: generator needs sweeper")
            return 0
        if sweeper is None:
            sweeper = self.sweeper
        return self.start +  sweeper.getRelPos()*(self.stop - self.start)

class SawGen(RampGen):
    # monotonically goes from start to stop, once reaches stop goes back to start
    # does it with duty_cycle,
    # i.e duty_cycle*period we go up, the other time we go down
    def __init__(self, duty_cycle=0.5, **kwds):
        self.duty_cycle = duty_cycle;
        super().__init__(**kwds)

    def getValue(self, sweeper=None):
        if sweeper is None and self.sweeper is None:
            print("Error: generator needs sweeper")
            return 0
        if sweeper is None:
            sweeper = self.sweeper
        if sweeper.getRelPos() < self.duty_cycle:
            return self.start + sweeper.getRelPos()/self.duty_cycle*(self.stop - self.start)
        return self.start + (1-sweeper.getRelPos())/(1-self.duty_cycle)*(self.stop - self.start)

class TriangleGen(SawGen):
    # monotonically goes from start to stop, once reaches stop goes back to start
    # does it with 50% duty cycle, i.e half of period we go up, the other half down
    def __init__(self, **kwds):
        super().__init__(duty_cycle=0.5, **kwds)

class PulseGen:
    # produce ampl for the first half a period and 0 for the other half
    def __init__(self, ampl=1, sweeper=None):
        self.ampl = ampl
        self.sweeper = sweeper

    def getValue(self, sweeper=None):
        if sweeper is None and self.sweeper is None:
            print("Error: generator needs sweeper")
            return 0
        if sweeper is None:
            sweeper = self.sweeper
        if sweeper.getRelPos() < 0.5:
            return self.ampl
        return 0

class Sweeper:
    # main clock sweeper for any function generator
    def __init__(self, widget, Npoints, SweepTime, onTicCallbacks=[]):
        # walk from start to stop with Npoints
        # cnt = 1 corresponds to start
        # cnt = Npoints corresponds to stop
        # variables like relVar are relative to the start of the period
        self.cnt = 0 # onTic will increase it right away
        self.widget = widget
        self.Npoints = Npoints
        self.start = 1
        self.stop = self.Npoints
        self.SweepTime = SweepTime
        self.onTicCallbacks = onTicCallbacks
        self.isOn = False
        self.isRestart = True
        self.isTicRunning = False

        self.span = self.stop - self.start
        self.center = (self.stop + self.start)/2
        self.dPos = self.span/(self.Npoints-1)
        self.dT = self.SweepTime/(self.Npoints-1)
        self.dTmS = round(self.dT*1000) # dT in milliseconds

    def reset(self):
        self.cnt = 0 # onTic will increase it right away
        self.isRestart = False
        self.startTime = datetime.now()

    def onTic(self):
        start = datetime.now()
        deadline = start +  timedelta(milliseconds=self.dTmS)

        if not self.isOn:
            self.isTicRunning = False
            return
        self.isTicRunning = True
        if self.isRestart:
            self.reset()
            self.isRestart = False
        self.incr()
        for cb in self.onTicCallbacks:
            cb(self)
        stop = datetime.now()
        self.isTicRunning = False
        if stop > deadline:
            runTime = (stop-start).seconds + float((stop-start).microseconds)/1000000
            print("Overrun: Callbacks took %s seconds instead of %s" % (runTime, self.dTmS/1000) )
            self.after(0, self.onTic)
        
        idleTime_mS = round((deadline-stop).seconds * 1000 + (deadline-stop).microseconds/1000)
        # print("Will idle for %s" % (idleTime_mS) )
        self.after(idleTime_mS, self.onTic)

    def after(self, idleTime_mS, cmnd):
        # this function modelled after Tk.after
        # it is execute 'cmnd' after idle time given in mS
        # idleTime_mS = 0 means execute immediately
        try:
            # Tk style
            self.widget.after(idleTime_mS, cmnd)
            return
        except AttributeError:
            pass

        try:
            # PyQt style
            t = QtCore.QTimer()
            t.singleShot( idleTime_mS, cmnd )
            return
        except NameError:
            pass

        raise Exception("Something went horribly wrong: we have no mechanism for Sweeper.after()")

            
    def cmdRestart(self):
        self.cnt = 0
        if self.isOn:
            self.isRestart = True
            return
        self.reset()
        self.isOn = True
        self.onTic()

    def cmdStart(self):
        self.isOn = True
        self.onTic()

    def cmdStop(self):
        self.isOn = False

    def incr(self):
        self.cnt += 1
        self.updPos()

    def updPos(self):
        self.relCnt = 1 + ((self.cnt-1) % self.Npoints)
        self.pos = self.start + self.dPos * (self.relCnt - 1)
        self.relPos = (self.pos-self.start)/self.span

    def getCnt(self):
        return self.cnt

    def getRelCnt(self):
        return self.relCnt

    def getPos(self):
        return self.pos

    def getRelPos(self):
        return self.relPos

def testOnTicTk(sweeper):
    print( sweeper.getPos() )
    if sweeper.getPos() == sweeper.Npoints:
        print("Tk sweeper is done")
        sweeper.cmdStop()
        sweeper.widget.destroy()
        sweeper.widget.quit()

def testSweeperTk():
    print("test Sweeper with Tk")
    root=Tk()
    root.geometry("800x600")
    root.withdraw(); # do not show window

    Np = 10
    SweepTime = Np
    print(f'Test sweeper: you should see a sequence of {Np} numbers updating about every {SweepTime/Np} seconds')
    sweeper = Sweeper(root, Npoints=Np, SweepTime=SweepTime, onTicCallbacks=[testOnTicTk])
    sweeper.cmdStart()

    root.mainloop()

def testOnTicPyQtGraph(sweeper):
    print( sweeper.getPos() )
    if sweeper.getPos() == sweeper.Npoints:
        print("PyQtGraph sweeper is done")
        sweeper.cmdStop()
        sweeper.widget.exit()

def testSweeperPyQtGraph():
    print("test Sweeper with PyQtGraph")
    app = QtGui.QApplication([])

    Np = 10
    SweepTime = Np
    print(f'Test sweeper: you should see a sequence of {Np} numbers updating about every {SweepTime/Np} seconds')
    sweeper = Sweeper(app, Npoints=Np, SweepTime=SweepTime, onTicCallbacks=[testOnTicPyQtGraph])
    sweeper.cmdStart()

    app.exec()

if __name__ == '__main__': 
    try:
        from tkinter import Tk
        testSweeperTk()
    except ModuleNotFoundError:
        print("Tk module not found, skipping testSweeperTk")

    try:
        from pyqtgraph.Qt import QtGui, QtCore
        testSweeperPyQtGraph()
    except ModuleNotFoundError:
        print("pyqtgraph.Qt module not found, skipping testSweeperPyQtGraph")
    print("Done with tests")