path: root/examples/lock_eit.py

import logging
# this should be done before justpy is called or log formatter does not work
logging.basicConfig(format='%(asctime)s %(levelname)8s %(name)s: %(message)s', datefmt='%m/%d/%Y %H:%M:%S')

import justpy as jp
import pyvisa
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

import qolab.tsdb as tsdb
import asyncio
import time

import qolab.gui.web as gui
from qolab.hardware import BasicInstrument
from qolab.hardware.scope import SDS1104X
from qolab.hardware.rf_generator import AgilentE8257D
from qolab.hardware.lockin import SRS_SR865A
from qolab.feedback import PID
from qolab.data import TraceSetSameX, TraceXY, Trace

l = logging.getLogger('qolab.gui.web')
l.setLevel(logging.INFO)
l = logging.getLogger('qolab.tsdb')
l.setLevel(logging.INFO)
logger = logging.getLogger('Magnetometer')
logger.setLevel(logging.INFO)

def getConfig(apparatus):
    config = apparatus.config.copy()
    ai = apparatus.instruments
    for n, i in ai.items():
        config[n]=i.getConfig()
    return config

def initLog(extra_tags={}):
    errorTrace = Trace("error")
    errorTrace.config['unit']='V'
    timeTrace = Trace("time")
    timeTrace.config['unit']='S'
    timeTrace.config['type'] = 'timestamp'
    errorLog = TraceXY("error")
    errorLog.x = timeTrace
    errorLog.y = errorTrace

    freqTrace = Trace("frequency")
    freqTrace.config['unit']='Hz'
    freqLog = TraceXY("frequency")
    freqLog.x = timeTrace
    freqLog.y = freqTrace

    feedbackTrace = Trace("feedback")
    feedbackTrace.config['unit']='Hz'
    feedbackLog = TraceXY("feedback")
    feedbackLog.x = timeTrace
    feedbackLog.y = feedbackTrace

    eitTrace = Trace("eit")
    eitTrace.config['unit']='V'
    eitLog = TraceXY("eit")
    eitLog.x = timeTrace
    eitLog.y = eitTrace

    rfPoutTrace = Trace("rfPout")
    rfPoutTrace.config['unit']='dBm'
    rfPoutLog = TraceXY("rfPout")
    rfPoutLog.x = timeTrace
    rfPoutLog.y = rfPoutTrace


    log = TraceSetSameX("timelog")
    log.addTrace(errorLog)
    log.addTrace(freqLog)
    log.addTrace(feedbackLog)
    log.addTrace(eitLog)
    log.addTrace(rfPoutLog)
    log.config['tags']['apparatus']=getConfig(apparatus)
    log.config['tags'].update(extra_tags)
    return log

async def feedbackLoop(apparatus, nsteps):
    # while True:
    for i in range(0,nsteps):
        adjustRFandLog(apparatus)
        await asyncio.sleep(0.1)
    apparatus.runStatus = False

def adjustRFandLog(apparatus):
    timenow = time.time()

    ai = apparatus.instruments
    pid = ai['pid']
    scope_fdbk = ai['scope.feedback']
    log   = apparatus.gui_log.traces

    err, eit = apparatus.getErrAndEit()
    fdbck = pid.feedback(err)

    freq0 = apparatus.getRFreq()
    apparatus.setFeefback(fdbck)

    rfPout = apparatus.getRFAmplitude()

    apparatus.getBCurrent(); # this automatically logs value to TSDB

    log.addPointToTrace(timenow)
    log.addPointToTrace(err, "error")
    log.addPointToTrace(freq0, "frequency")
    log.addPointToTrace(fdbck, "feedback")
    log.addPointToTrace(rfPout, "rfPout")
    log.addPointToTrace(eit, "eit")

async def initial_lock_to_eit(apparatus, extra_tags={}):
    rfPout = apparatus.getRFAmplitude()
    apparatus.state = f'Initial lock RF power {rfPout} dBm'
    logger.info(apparatus.state)
    update_webpage(apparatus=apparatus)

    ai = apparatus.instruments
    apparatus.gui_log.setTraces( initLog(extra_tags=extra_tags) )
    ai['pid'].reset()
    ai['pid'].setEnable(True)
    apparatus.runStatus = True
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=50)
            )

async def calibratingLockin(apparatus, extra_tags={}):
    ai = apparatus.instruments
    rfPout = apparatus.getRFAmplitude()
    apparatus.state = f'Calibrating lockin response at RF power {rfPout} dBm'
    logger.info(apparatus.state)
    update_webpage(apparatus=apparatus)
    ai['pid'].setEnable(False)
    ai['pid'].reset()
    apparatus.gui_log.setTraces( initLog(extra_tags=extra_tags) )
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=20)
            )

    fr0 = ai['rfgen'].getFreqFixed()
    df = 10
    ai['pid'].setEnable(False)
    ai['pid'].reset()
    apparatus.setRFreq(fr0+df)
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=20)
            )
    apparatus.setRFreq(fr0)

    log   = apparatus.gui_log.traces
    trE=log.getTrace('error')

    e1 = np.mean(trE.y.values[0:20])
    e2 = np.mean(trE.y.values[20:])
    dE = e2-e1
    slope = dE/df
    apparatus.setLockinSlope(slope)

    logger.info(f'dE = {dE}')
    logger.info(f'lockin error signal slope = {slope} V/Hz')
    ai['error_signal_response_to_eit_detuning'].conversion_factor = slope

    return float(slope)

async def responseToChangeOfBfieldControlVoltage(apparatus, extra_tags={}):
    ai = apparatus.instruments
    rfPout = apparatus.getRFAmplitude()
    apparatus.state = f'Calibrating lockin response to change of Bfield control voltage, {rfPout} dBm'
    logger.info(apparatus.state)
    update_webpage(apparatus=apparatus)
    ai['pid'].setEnable(True)
    ai['pid'].reset()
    apparatus.gui_log.setTraces( initLog(extra_tags=extra_tags) )
    dV = 0.01
    ai['lockin'].AuxOut1=dV

    # initial lock
    ai['pid'].setEnable(True)
    ai['pid'].reset()
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=20)
            )

    # this is for the money lock
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=30)
            )

    fr0 = ai['rfgen'].getFreqFixed()

    # this is for the money lock
    ai['lockin'].AuxOut1=0
    ai['pid'].setEnable(True)
    ai['pid'].reset()
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=30)
            )
    fr1 = ai['rfgen'].getFreqFixed()

    log   = apparatus.gui_log.traces
    trE=log.getTrace('error')

    df = fr0 - fr1
    slope = df/dV

    logger.info(f'df = {df}')
    logger.info(f'response to magnetic control voltage = {slope} Hz/V')
    ai['B_control_voltage_to_eit_detuning'].conversion_factor = slope

    # relock to default state
    ai['lockin'].AuxOut1=0
    ai['pid'].setEnable(True)
    ai['pid'].reset()
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=20)
            )

    return float(slope)

async def record_magnetometer_noise(apparatus):
    ai = apparatus.instruments
    rfPout = apparatus.getRFAmplitude()
    ai['pid'].setEnable(False)

    apparatus.state = f'Record magnetometer noise,  RF power {rfPout} dBm'
    logger.info(apparatus.state)

    scope = ai['scope.monitor']
    old_config = scope.getConfig()
    tperdiv=1
    scope.setTimePerDiv(tperdiv)
    scope.setRoll(True)

    await asyncio.sleep(tperdiv*14 + 1)
    errorTr=scope.getTrace(1, maxRequiredPoints=10000)
    errorTr.config['label'] = 'error vs time'
    errorTr.y.config['label'] = 'error'
    log = errorTr
    log.config['tags']['apparatus']=getConfig(apparatus)

    apparatus.gui_log.setTraces( log )
    update_webpage(apparatus=apparatus)

    logger.info('Restoring scope settings')
    scope.setConfig(old_config)

async def longTermLock(apparatus, extra_tags={}):
    ai = apparatus.instruments
    rfPout = apparatus.getRFAmplitude()
    apparatus.state = f'Long term lock RF power {rfPout} dBm'
    logger.info(apparatus.state)

    ai = apparatus.instruments
    rfPout = apparatus.getRFAmplitude()
    apparatus.runStatus = True
    ai['pid'].setEnable(True)

    apparatus.gui_log.setTraces( initLog(extra_tags=extra_tags) )
    update_webpage(apparatus=apparatus)

    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=10000000)
            )
    return res

async def sweepRFPower(apparatus, extra_tags={}, rfPowerList=[]):
    apparatus.gui_log.setTraces( initLog(extra_tags=extra_tags) )
    ai = apparatus.instruments
    # for p in np.linspace(-10,10, 110):
    ai['pid'].setEnable(True)
    for p in rfPowerList:
        ai['rfgen'].setRFAmplitude(p)
        rfPout = ai['rfgen'].getRFAmplitude()
        apparatus.state = f'lock with RF power {rfPout} dBm'
        logger.info(apparatus.state)
        update_webpage(apparatus=apparatus)
        ai['pid'].reset()
        apparatus.runStatus = True
        res = await asyncio.gather(
                feedbackLoop(apparatus, nsteps=50)
                )
    return res


freqZero = 6_834_686_400
freqDeltaMp2 = 6_835_385_570
dfB = freqDeltaMp2 - freqZero
freqDeltaMm2 = freqZero - dfB;

async def main():
    app_nickname=apparatus.config['DeviceNickname']

    task_wp_update_loop = asyncio.create_task(update_webpage_loop(update_interval=1))

    apparatus.gui_log = gui.QOLTimeLog(a=wp)
    apparatus.gui_log.save_controls.getNextDataFile = apparatus.getNextDataFile

    logger.info("Accessing hardware")
    rm = pyvisa.ResourceManager()
    instr_scope=rm.open_resource('TCPIP::192.168.0.61::INSTR')
    scope_fdbk = SDS1104X(instr_scope, device_nickname='.'.join([app_nickname, 'scope.feedback']), tsdb_ingester=tsdb_ingester)
    instr_scope=rm.open_resource('TCPIP::192.168.0.62::INSTR')
    scope_mon = SDS1104X(instr_scope, device_nickname='.'.join([app_nickname, 'scope.monitor']), tsdb_ingester=tsdb_ingester)
    instr_rfgen=rm.open_resource('TCPIP::192.168.0.114::INSTR')
    rfgen = AgilentE8257D(instr_rfgen, device_nickname='.'.join([app_nickname, 'rfgen']), tsdb_ingester=tsdb_ingester)
    instr_lockin=rm.open_resource('TCPIP::192.168.0.51::INSTR')
    lockin = SRS_SR865A(instr_lockin, device_nickname='.'.join([app_nickname, 'lockin']), tsdb_ingester=tsdb_ingester)
    # print('------ Header start -------------')
    # print(str.join('\n', scope.getHeader()))
    # print(str.join('\n', rfgen.getHeader()))
    # print('------ Header ends  -------------')
    # ch1  = scope.getTrace(1)
    # traces = scope.getAllTraces()
    # pid = PID(100,400,0, sign=-1); # good for dm=-2 resonance
    pid = PID(50,200,0, sign=-1, device_nickname='.'.join([app_nickname, 'pid']), tsdb_ingester=tsdb_ingester)

    apparatus.instruments={}
    ai = apparatus.instruments
    ai['rfgen'] = rfgen
    ai['lockin']   = lockin
    ai['pid']   = pid
    ai['scope.feedback'] = scope_fdbk
    scope_fdbk.setRoll(False)
    scope_fdbk.setRun(True)
    scope_fdbk.setTimePerDiv(0.0005)
    scope_fdbk.setChanVoltsPerDiv(1, 0.02)
    scope_fdbk.setChanVoltsPerDiv(3, 0.02)
    ai['scope.monitor'] = scope_mon
    scope_mon.setRoll(False)
    scope_mon.setRun(True)
    scope_mon.setTimePerDiv(0.0005)
    scope_mon.setChanVoltsPerDiv(2, 0.002)
    scope_mon.setChanVoltageOffset(2, -0.0726)
    apparatus.state = None
    apparatus.runStatus = False

    error_signal_response_to_eit_detuning = BasicInstrument(device_nickname='.'.join([app_nickname, 'error_signal_response_to_eit_detuning']), tsdb_ingester=tsdb_ingester)
    error_signal_response_to_eit_detuning.deviceProperties = {'conversion_factor', 'unit'}
    error_signal_response_to_eit_detuning.conversion_factor = None
    error_signal_response_to_eit_detuning.unit = 'V_per_Hz'
    ai['error_signal_response_to_eit_detuning']=error_signal_response_to_eit_detuning

    B_control_voltage_to_eit_detuning = BasicInstrument(device_nickname='.'.join([app_nickname, 'B_control_voltage_to_eit_detuning']), tsdb_ingester=tsdb_ingester)
    B_control_voltage_to_eit_detuning.deviceProperties = {'conversion_factor', 'unit'}
    B_control_voltage_to_eit_detuning.conversion_factor = None
    B_control_voltage_to_eit_detuning.unit = 'Hz_per_V'
    ai['B_control_voltage_to_eit_detuning'] = B_control_voltage_to_eit_detuning


    SweepSpan = 10000
    ai['rfgen'].stopFrequencySweep()
    # ai['rfgen'].setFreqFixed(freqDeltaMm2)
    # ai['rfgen'].setFreqFixed(freqZero)
    apparatus.setRFreq(freqDeltaMp2)
    # apparatus.setRFreq(freqZero)

    rfPstart = -10 
    rfPstop = 10
    rfPowerList = np.linspace(rfPstart, rfPstop, 11)
    rfPower0 = rfPstart
    rfPower0 = 3.4 
    apparatus.setRFAmplitude(rfPower0)

    d=getConfig(apparatus)
    instruments_config=gui.QOLDictionary(a=wp, name='Instruments configs', container=d)
    extra_tags={}

    await initial_lock_to_eit(apparatus, extra_tags=extra_tags)

    lockin_slope= await calibratingLockin(apparatus, extra_tags=extra_tags)
    instruments_config.container = getConfig(apparatus)
    instruments_config.display_container_dictionary()

    B_response= await responseToChangeOfBfieldControlVoltage(apparatus, extra_tags=extra_tags)
    instruments_config.container = getConfig(apparatus)
    instruments_config.display_container_dictionary()

    # await sweepRFPower(apparatus, extra_tags=extra_tags, rfPowerList=rfPowerList)

    # await record_magnetometer_noise(apparatus)


    await longTermLock(apparatus, extra_tags=extra_tags)

    apparatus.state = 'Done working with hardware'
    logger.info(apparatus.state)
    update_webpage(apparatus=apparatus)

    apparatus.gui_log.stop_tasks()
    task_wp_update_loop.cancel()
    logger.info("exiting main loop")
    update_webpage(apparatus=apparatus)
    return(apparatus)

def update_webpage(byWhom=None, apparatus=None):
    timestr = time.strftime("%a, %d %b %Y, %H:%M:%S", time.localtime())
    clock_upd.text = f'Last update at {timestr}' 
    try:
        status_line.text = apparatus.state
    except:
        pass
    jp.run_task(wp.update())

async def update_webpage_loop(update_interval=1):
    while True:
        update_webpage()
        await asyncio.sleep(update_interval)

async def getPage():
    return wp

async def  jp_startup():
    jp.run_task(main())


class Apparatus(BasicInstrument):
    def __init__(self, *args, **kwds):
        super().__init__(*args, **kwds)
        self.setLockinSlope(0)

    @BasicInstrument.tsdb_append
    def getBCurrent(self):
        scope = self.instruments['scope.monitor']
        ch=scope.getTrace(2, decimate=False)
        bCurrent = np.mean(ch.y.values); # mV -> mA
        return bCurrent

    def getErrAndEit(self):
        scope = self.instruments['scope.feedback']
        # we need to freeze scope since this two values are taken on the same scope
        trg_mode = scope.getTriggerMode()
        scope.setTriggerMode('STOP')
        err = self.getErr()
        eit = self.getEIT()
        scope.setTriggerMode(trg_mode)
        return(err,eit)



    @BasicInstrument.tsdb_append
    def getErr(self):
        scope = self.instruments['scope.feedback']
        ch1=scope.getTrace(1, decimate=False)
        err = np.mean(ch1.y.values)
        return err

    @BasicInstrument.tsdb_append
    def getEIT(self):
        scope = self.instruments['scope.feedback']
        ch3=scope.getTrace(3, decimate=False)
        eit = np.mean(ch3.y.values)
        return eit

    @BasicInstrument.tsdb_append
    def setFeefback(self, val):
        freq0 = self.getRFreq()
        self.setRFreq(freq0+val)

    @BasicInstrument.tsdb_append
    def getRFreq(self):
        rfgen = self.instruments['rfgen']
        return rfgen.getFreqFixed()

    @BasicInstrument.tsdb_append
    def setRFreq(self,val):
        rfgen = self.instruments['rfgen']
        rfgen.setFreqFixed(val)

    @BasicInstrument.tsdb_append
    def getRFAmplitude(self):
        rfgen = self.instruments['rfgen']
        return rfgen.getRFAmplitude()

    @BasicInstrument.tsdb_append
    def setRFAmplitude(self, val):
        rfgen = self.instruments['rfgen']
        return rfgen.setRFAmplitude(val)

    @BasicInstrument.tsdb_append
    def setLockinSlope(self, val):
        self.lockin_slope=val


if __name__ == '__main__': 
    logger.info('========== Start up ===========')
    tsdb_client = tsdb.Client('influx', 'http://lumus.physics.wm.edu:8428', database='qolab')
    tsdb_ingester = tsdb.Ingester(tsdb_client, batch=11, measurement_prefix='VAMPIRE.HighPower')

    wp = jp.WebPage(delete_flag=False)
    apparatus = Apparatus(tsdb_ingester=tsdb_ingester, device_nickname='magnetometer',)
    apparatus.config['Device type'] = 'QOL VAMPIRE HighPower magnetometer'
    apparatus.config['Device model'] = 'v0.1'
    apparatus.config['FnamePrefix'] = 'magnetometer'
    apparatus.config['SavePath'] = '/mnt/qol_grp_data/data.VAMPIRE.HighPower'
    # apparatus.config['SavePath'] = './data'

    d=jp.Div(text='Magnetometer log', a=wp, classes='text-white bg-blue-500 text-center text-xl')

    div_status = jp.Div(classes='text-xl flex m0 p-1 space-x-4 bg-gray-300 font-mono', a=wp)
    clock_upd = jp.Div(text='Clock Loading...', classes='text-xl bg-gray-400', a=div_status)
    status_line = jp.Div(text='Status Loading...',classes='text-xl', a=div_status)

    # mpl.use("TkAgg")
    # apparatus = asyncio.run(main())
    jp.justpy(getPage, startup=jp_startup)