path: root/eitControl.py

from pyqtgraph.Qt import QtGui, QtCore
import pyqtgraph as pg
from pyqtgraph.dockarea import *
import pyqtgraph.exporters

import platform
import os
import argparse

from threading import Thread

import time  # For sleep, clock, time and perf_counter
from datetime import datetime, timedelta, date

import numpy as np
import csv

import ue9qol
from funcGenerator import Sweeper, SinGen, TriangleGen, RampGen, PulseGen
import rfGen
    
class Experiment:

    def __init__(self, root, args):
        self.root = root
        if args.save_prefix:
            self.save_prefix = args.save_prefix
        else:
            self.save_prefix = "set_the_prefix"
        if args.data_dir:
            self.data_dir = args.data_dir
        else:
            self.data_dir = "set_the_data_dir"

        self.save_cnt = 0
        self.tic = 0
        self.channelsNames2grab={'tic', 'x','rfFreq','dac0', 'dac1', 'adc0', 'adc1', 'adc2', 'adc3'}
        alpha=100
        self.channelsColor = { 'adc0': (20,20,20,alpha), 'adc1': (85,170,255,alpha), 'adc2': (255,0,255,alpha), 'adc3': (0,85,255,alpha), 'dac0': 'r', 'dac1': 'g'}
        #self.channelsNames2plot={'dac0', 'dac1', 'adc0', 'adc1', 'adc2', 'adc3'}
        if args.test:
            self.channelsNames2plot={'adc0', 'adc1', 'adc2', 'adc3'}
        else:
            self.channelsNames2plot={'adc0', 'adc1', 'adc2'}
        self.xChannelName='rfFreq' # can be also 'tic' or any of above
        self.xlabel=''
        self.channelGraph={}
        self.clearData()
        self.hardware = {}
        self.hardwareSetup(args)

        # now we ready to do gui
        self.guiSweeper = Sweeper(self.root, Npoints=10, SweepTime=1, onTicCallbacks=[self.updatePlot])
        self.guiSetup(root)
        self.guiSweeper.cmdStart()

    
    def hardwareSetup(self,args):
        if args.test:
            print("Test mode, run with fake hardware")
            self.sweeper = Sweeper(self.root, Npoints=100, SweepTime=1, onTicCallbacks=[self.onTic])
            self.hardware['LabJack'] = ue9qol.UE9qolDummy(sweeper=self.sweeper)
            self.hardware['rfGen'] = rfGen.rfGenLMX2487Dummy(port='/dev/ttyUSB0', speed=115200, timeout=1)
        else:
            self.sweeper = Sweeper(self.root, Npoints=500, SweepTime=10, onTicCallbacks=[self.onTic])
            self.hardware['LabJack'] = ue9qol.UE9qol()
            if platform.system() == 'Linux':
                rf=rfGen.rfGenLMX2487(port='/dev/ttyUSB0', speed=115200, timeout=1)
            else:
                rf=rfGen.rfGenLMX2487(port='COM5', speed=115200, timeout=1)
            self.hardware['rfGen'] = rf

        fCent = 6.83468e9
        fSpan = 100e3
        self.rfGenFunc = RampGen(start=0, stop=0, sweeper = self.sweeper)
        self.rfGenFunc.setCenter(fCent)
        self.rfGenFunc.setSpan(fSpan)

        self.funcGen = TriangleGen(0, 5, sweeper = self.sweeper)


    def centralFreqValueChanged(self, sb):
        self.rfGenFunc.setCenter(sb.value())
        pass

    def freqSpanValueChanged(self, sb):
        self.rfGenFunc.setSpan(sb.value())
        pass

    def guiSetup(self, root): 
        self.dockArea = area = DockArea()
        d1 = Dock("Global", size=(5,1))
        d2 = Dock("Data", size=(100,100))
        d3 = Dock("RF Gen", size=(1,2))
        area.addDock(d1, 'top')
        area.addDock(d2, 'bottom', d1)
        area.addDock(d3, 'bottom', d2)
        self.root.addWidget(area)

        self.dataPlot = pg.PlotWidget(name='Plot1')
        d2.addWidget(self.dataPlot)
        self.dataPlot.showGrid(x=True, y=True)
        self.dataPlot.addLegend()

        # global buttons
        w1 = pg.LayoutWidget()
        bAutoZoom = QtGui.QPushButton('AutoZoom')
        bAutoZoom.clicked.connect(self.autoZoom)
        bRestart = QtGui.QPushButton('Restart')
        bRestart.clicked.connect(self.restart)
        bStart = QtGui.QPushButton('Start')
        bStart.clicked.connect(self.start)
        bStop = QtGui.QPushButton('Stop')
        bStop.clicked.connect(self.stop)
        bSave = QtGui.QPushButton('Save data')
        bSave.clicked.connect(self.saveCmd)
        bExit = QtGui.QPushButton('Exit')
        bExit.clicked.connect(exit)
        w1.addWidget(bAutoZoom, row=0, col=0)
        w1.addWidget(bRestart, row=0, col=1)
        w1.addWidget(bStart, row=0, col=2)
        w1.addWidget(bStop, row=0, col=3)
        w1.addWidget(bSave, row=0, col=4)
        w1.addWidget(bExit, row=0, col=5)
        d1.addWidget(w1)

        # RF gen gui
        fCent=self.rfGenFunc.getCenter()
        fSpan=self.rfGenFunc.getSpan()
        spins = [
                ("Central Frequency", 
                    pg.SpinBox(value=fCent, bounds=[6.83e9, 6.84e9], suffix='Hz', siPrefix=True, step=1e3, decimals=10),
                    self.centralFreqValueChanged),
                ("Frequency Span", 
                    pg.SpinBox(value=fSpan, bounds=[1, 10e6], dec=True, step=0.5,  suffix='Hz', siPrefix=True, minStep=1),
                    self.freqSpanValueChanged)
                ]
        w3 = pg.LayoutWidget()
        d3.addWidget(w3)
        for text, spin, cb in spins:
            l=QtGui.QLabel(text)
            w3.addWidget(l)
            w3.addWidget(spin)
            spin.sigValueChanged.connect(cb)

    def clearData(self):
        self.data = {}
        for ch in self.channelsNames2grab:
            self.data[ch] = []

    def stop(self):
        self.sweeper.cmdStop()

    def start(self):
        self.sweeper.cmdStart()

    def restart(self):
        self.clearData()
        self.sweeper.cmdRestart()

    def getNewDataFileName(self, ext="csv"):
        data_dir = self.data_dir
        if not os.path.exists(data_dir):
            os.mkdir(data_dir)
        if not os.path.isdir(data_dir):
            print(f"ERROR: cannot create directory for data: {data_dir}")
            print(f"Will use current dir for storage")
            data_dir = "."

        prefix = self.save_prefix
        today = date.today()
        datestr = today.strftime("%Y%m%d")
        self.save_cnt += 1

        base_name = f"{prefix}_{datestr}_{self.save_cnt:#04}"
        file_name = f"{base_name}.{ext}"
        data_file = os.path.join(data_dir, file_name)
        if os.path.exists(data_file):
            data_file = self.getNewDataFileName(ext=ext)
        return data_file

    def saveCmd(self):
        csv_file = self.getNewDataFileName(ext='csv')
        print(f"Saving to {csv_file}")
        data = self.data
        try:
            with open(csv_file, 'w') as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow(data.keys())
                writer.writerows(zip(*data.values()))
        except IOError:
            print('I/O error')

        png_file=csv_file.replace(".csv", ".png")
        print(f"Picture saved to {png_file}")
        plt = self.dataPlot.getPlotItem()
        ex = pg.exporters.ImageExporter(plt)
        if pg.__version__ == '0.10.0':
            # Workaround for PyQtGraph version <= 0.10.0
            # see https://github.com/pyqtgraph/pyqtgraph/issues/538#issuecomment-361405356
            w= int(plt.width())
            h= int(plt.height())
            # the value in setValue need to be different from default
            # otherwise it will not be taken
            ex.parameters().param('width').setValue(w+1, blockSignal=ex.widthChanged)
            ex.parameters().param('height').setValue(h+1, blockSignal=ex.heightChanged)
            # now we set actual value
            # ex.parameters()['width'] = w
            ex.parameters().param('width').setValue(w, blockSignal=ex.widthChanged)
            ex.parameters().param('height').setValue(h, blockSignal=ex.heightChanged)
            # beware this is bad workaround!!! plot data is misplaced
        ex.export(png_file)

    def onTic(self,swp=None):
        start = datetime.now()
        if swp is None:
            swp = self.sweeper
        
        # global tic counter
        tic = self.sweeper.getCnt()
        self.data['tic'].append(tic)

        # RF generator
        rfFreq = self.rfGenFunc.getValue(swp)
        self.hardware['rfGen'].setFreq(rfFreq)
        self.data['rfFreq'].append(rfFreq)

        # DAQ
        daq0 = self.hardware['LabJack']

        # dac0
        # dac0 = self.funcGen.getValue(swp)
        dac0 = 0
        dac0 = self.funcGen.getValue(swp)
        daq0.setOutputCh(0, dac0)
        self.data['dac0'].append(dac0)

        # dac1
        # dac1 = PulseGen(ampl=5, sweeper=swp).getValue()
        dac1 = 0
        daq0.setOutputCh(1, dac1)
        self.data['dac1'].append(dac1)

        # adc0
        adc0= daq0.getInputCh(0)
        # adc0 =  SinGen(ampl=4, sweeper=swp).getValue()
        self.data['adc0'].append( adc0 )

        # adc1
        adc1= daq0.getInputCh(1)
        # adc1 =  SinGen(ampl=1, sweeper=swp).getValue()
        self.data['adc1'].append( adc1 )

        # adc2
        adc2= daq0.getInputCh(2)
        # adc2 =  SinGen(ampl=2, sweeper=swp).getValue()
        self.data['adc2'].append( adc2 )

        # adc3
        adc3= daq0.getInputCh(3)
        # adc3 =  SinGen(ampl=3, sweeper=swp).getValue()
        self.data['adc3'].append( adc3 )

        # X-axis (i.e. independent variable)
        x=self.data[self.xChannelName]
        x=np.array(x)
        fCent = self.rfGenFunc.getCenter()
        x=(x-fCent)
        self.data['x'] = x
        self.dataPlot.setLabel('bottom', 'Frequency offset', units='Hz')
        self.dataPlot.setLabel('left', 'Signal', units='V')

        self.data['newData'] = True
        stop = datetime.now()
        runTime = (stop-start).seconds + float((stop-start).microseconds)/1000000
        # print("onTic DAQ took %s seconds." % (runTime) )

    def autoZoom(self):
        self.dataPlot.autoRange()

    def updatePlot(self,swp=None):
        if 'newData' in  self.data and self.data['newData']:
            self.data['newData'] = False
        else:
            return
            
        start = datetime.now()
        x = self.data['x']
        for name in self.channelsNames2plot:
            if name not in self.data:
                continue
            y = self.data[name]
            if name not in self.channelGraph:
                if name in self.channelsColor:
                    color = self.channelsColor[name]
                else:
                    color = (255,0,0)
                self.channelGraph[name]=self.dataPlot.plot(x,y, pen=None, symbol='o', symbolPen=None, symbolBrush=color, symbolSize=5, name=name)
            else:
                self.channelGraph[name].setData(x,y)
        # centralFreqFormatted = pg.siFormat(self.fCent, suffix='Hz', precision=4)
        # showing trailing zeros is tricky
        fCent = self.rfGenFunc.getCenter()
        centralFreqFormatted = f"{fCent/1e9:.9f}"
        centralFreqFormatted = str.ljust(centralFreqFormatted, 11, '0') + " GHz"
        self.dataPlot.setTitle(f"Signals around center frequency {centralFreqFormatted}")

        stop = datetime.now()
        runTime = (stop-start).seconds + float((stop-start).microseconds)/1000000
        print("Replot took %s seconds to plot %s points per channel." % (runTime, len(self.data['x'])) )

 
if __name__ == '__main__': 
    parser = argparse.ArgumentParser(description='Perform EIT based experiment.')
    parser.add_argument('--test', '-t', action='store_true',
            help='test mode, use fake/dummy hardware')
    args = parser.parse_args()

    app = QtGui.QApplication([])
    pg.setConfigOption('background', 'w')
    pg.setConfigOption('foreground', 'k')
    mw = QtGui.QMainWindow()
    mw.setWindowTitle('pyqtgraph example: PlotWidget')
    mw.resize(800,800)
    cw = QtGui.QWidget()
    mw.setCentralWidget(cw)
    l = QtGui.QVBoxLayout()
    cw.setLayout(l)

    mw.show()

    args.save_prefix="eit"
    args.data_dir="VAMPIRE.Data"

    experiment=Experiment(l, args)
    app.exec()