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
import ast

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

from configparser import ConfigParser, ExtendedInterpolation
config = ConfigParser(interpolation=ExtendedInterpolation(),
        converters = {'list': lambda x: [i.strip() for i in x.strip().split('\n')]})
default_config = """
[RF]
central_frequency = 6.83468e9
frequency_span =  100e3
initial_frequency = 6.834e9
frequency_export_name = rfFreq

[DAQ]
# channels to grab and their meaning
ain0 = transmission
ain1 = lockin
ain2 = davll
ain3 = ain3_undefined
# commented out channels will not be processed or stored
# dac0 = dac0_undefined
# dac1 = dac1_undefined2

[Plot]
x_axis_data = rfFreq

[Plot_channels_visibility]
transmission = yes
lockin = yes
davll = yes

[Plot_channels_colors]
# color specification will be evaluated and should match PyQtGraph.mkColor(args)
# for example
# (R,G,B,Alpha) tuple: example_trace = (255,0,255,100)
# or color name: example_trace = 'g'
transmission = (20,20,20,100)
lockin = (85,170,255,100)
davll = (255,0,255,100)
ain3_undefined = (0,85,255,100)


[Save]
save_prefix = eit
data_dir = z:\data.VAMPIRE

"""

config.read_string(default_config)
# additional config
config.read("config.ini")

    
class Experiment:

    def __init__(self, root, config, args):
        self.root = root
        self.config = config
        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.newDataIsReady = False
        alpha=100
        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.buttons = {}
        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.data_dir="VAMPIRE.Data"
            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 = self.config['RF'].getfloat('central_frequency', fallback=6.83468e9)
        fSpan = self.config['RF'].getfloat('frequency_span', fallback=100e3)
        rf_f_init = self.config['RF'].getfloat('initial_frequency', fallback=6.834e9)
        self.hardware['rfGen'].setFreq(rf_f_init)
        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):
        v=sb.value()
        self.config['RF']['central_frequency'] = str(v)
        self.rfGenFunc.setCenter(v)
        pass

    def freqSpanValueChanged(self, sb):
        v=sb.value()
        self.config['RF']['frequency_span'] = str(v)
        self.rfGenFunc.setSpan(v)
        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))
        dS = Dock("Status", size=(1,2), autoOrientation=False)
        dS.setOrientation(o='horizontal')
        area.addDock(d1, 'top')
        area.addDock(dS, 'bottom', d1)
        area.addDock(d2, 'bottom', dS)
        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()
        self.vLineSweepPosition = vLine = pg.InfiniteLine(angle=90, movable=False, pen='r')
        self.dataPlot.addItem(vLine, ignoreBounds=True)

        # global buttons
        w1 = pg.LayoutWidget()
        bAutoZoom = QtGui.QPushButton('&AutoZoom')
        bAutoZoom.clicked.connect(self.autoZoom)
        self.buttons["AutoZoom"] = bAutoZoom
        bRestart = QtGui.QPushButton('&Restart')
        bRestart.clicked.connect(self.restart)
        self.buttons["Restart"] = bRestart
        bStartStopToggle = QtGui.QPushButton('&Start')
        bStartStopToggle.clicked.connect(self.start)
        self.buttons["StartStopToggle"] = bStartStopToggle
        bSave = QtGui.QPushButton('Sa&ve data')
        bSave.clicked.connect(self.saveCmd)
        self.buttons["Save"] = bSave
        bExit = QtGui.QPushButton('&Exit')
        bExit.clicked.connect(exit)
        self.buttons["Exit"] = bExit
        bSaveConfig = QtGui.QPushButton('Save Con&fig')
        bSaveConfig.clicked.connect(self.saveConfigCmd)
        self.buttons["SaveConfig"] = bSaveConfig
        w1.addWidget(bAutoZoom, row=0, col=0)
        w1.addWidget(bRestart, row=0, col=1)
        w1.addWidget(bStartStopToggle, row=0, col=2)
        w1.addWidget(bSave, row=0, col=3)
        w1.addWidget(bSaveConfig, row=0, col=4)
        w1.addWidget(bExit, row=0, col=5)
        d1.addWidget(w1)

        ## status line
        self.statusline = l = QtGui.QLabel("All ok")
        dS.addWidget(l, row=1, col=0)

        # 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 = {}
        # RF generator channels
        rf_freq_Name = self.config['RF'].get('frequency_export_name')
        self.data[rf_freq_Name] = []
        # DAQ channels
        for ch in self.config['DAQ']:
            ch_meaning = self.config['DAQ'][ch]
            self.data[ch_meaning] = []
        # special channels
        self.data['tic'] = []
        self.data['x'] = []


    def stop(self):
        self.sweeper.cmdStop()
        self.buttons["StartStopToggle"].setText("&Continue")
        self.buttons["StartStopToggle"].clicked.disconnect()
        self.buttons["StartStopToggle"].clicked.connect(self.start)

    def start(self):
        self.sweeper.cmdStart()
        self.buttons["StartStopToggle"].setText("&Pause")
        self.buttons["StartStopToggle"].clicked.disconnect()
        self.buttons["StartStopToggle"].clicked.connect(self.stop)

    def restart(self):
        self.clearData()
        self.sweeper.cmdRestart()
        self.buttons["StartStopToggle"].setText("&Pause")
        self.buttons["StartStopToggle"].clicked.disconnect()
        self.buttons["StartStopToggle"].clicked.connect(self.stop)

    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')
        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')
        msg = f"data saved to {csv_file}"
        print(msg)
        self.statusline.setText(msg)

        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 saveConfigCmd(self):
        with open('config.ini', 'w') as configfile:
            self.config.write(configfile)

    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)
        rf_freq_Name = self.config['RF'].get('frequency_export_name')
        self.data[rf_freq_Name].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)

        for ch in self.config['DAQ']:
            if ch[0:3] == 'ain':
                n=int(ch[3:])
                vIn = daq0.getInputCh(n)
                ch_meaning = self.config['DAQ'][ch]
                self.data[ch_meaning].append( vIn )

        # X-axis (i.e. independent variable)
        x=self.data[self.config['Plot']['x_axis_data']]
        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.newDataIsReady = 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 self.newDataIsReady:
            self.newDataIsReady = False
        else:
            return
            
        start = datetime.now()
        x = self.data['x']
        for name in self.config['Plot_channels_visibility']:
            if not self.config['Plot_channels_visibility'].getboolean(name):
                continue
            if name not in self.data:
                continue
            y = self.data[name]
            if name not in self.channelGraph:
                if name in self.config['Plot_channels_colors']:
                    color_str = self.config['Plot_channels_colors'][name]
                    color = ast.literal_eval(color_str)
                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)
        if len(x)>0:
            self.vLineSweepPosition.setValue(x[-1])
        # 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="z:\data.VAMPIRE"

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