more modular lock_eit script

1 files changed, 77 insertions, 37 deletions

diff --git a/examples/lock_eit.py b/examples/lock_eit.py
index 3553baf..de87de0 100644
--- a/examples/lock_eit.py
+++ b/examples/lock_eit.py
@@ -10,9 +10,7 @@ from qolab.hardware.rf_generator import AgilentE8257D
 from qolab.feedback import PID
 from qolab.data import TraceSetSameX, TraceXY, Trace
 
-def adjustRF(pid, scope, rfgen, nsteps=1000, fname=None):
-    pid.reset()
-
+def initLog():
     errorTrace = Trace("error")
     errorTrace.config['unit']='V'
     timeTrace = Trace("time")
@@ -38,36 +36,59 @@ def adjustRF(pid, scope, rfgen, nsteps=1000, fname=None):
     log.addTrace(errorLog)
     log.addTrace(freqLog)
     log.addTrace(feedbackLog)
+    return log
 
-
+def feedbackLoop(apparatus, nsteps):
     # while True:
     for i in range(0,nsteps):
-        timenow = time.time()
-
-        ch1=scope.getTrace(1, decimate=False)
-        err = np.mean(ch1.y.values)
-        fdbck = pid.feedback(err)
-
-        freq0 = rfgen.getFreqFixed()
-        freq = freq0 + fdbck
-        rfgen.setFreqFixed(freq)
-        print(f'error = {err}, feedback = {fdbck}, request freq = {freq}')
-
-        log.addPointToTrace(timenow)
-        log.addPointToTrace(err, "error")
-        log.addPointToTrace(freq0, "frequency")
-        log.addPointToTrace(fdbck, "feedback")
+        adjustRFandLog(apparatus)
+    apparatus.runStatus = True
+
+def adjustRFandLog(apparatus):
+    timenow = time.time()
+
+    scope = apparatus.scope
+    rfgen = apparatus.rfgen
+    pid   = apparatus.pid
+    log   = apparatus.log
+
+    ch1=scope.getTrace(1, decimate=False)
+    err = np.mean(ch1.y.values)
+    fdbck = pid.feedback(err)
+
+    freq0 = rfgen.getFreqFixed()
+    freq = freq0 + fdbck
+    rfgen.setFreqFixed(freq)
+    print(f'error = {err}, feedback = {fdbck}, request freq = {freq}')
+
+    log.addPointToTrace(timenow)
+    log.addPointToTrace(err, "error")
+    log.addPointToTrace(freq0, "frequency")
+    log.addPointToTrace(fdbck, "feedback")
         
-        if fname is not None:
-            log.save(fname, item_format='.15e')
 
+def updateLogPlot(apparatus):
+    log   = apparatus.log
     plt.clf()
     log.plot()
-    return log
+    plt.pause(.0001)
 
+def updateLogPlotLoop(apparatus):
+    while True:
+        updateLogPlot(apparatus)
+        if not apparatus.runStatus:
+            break
 
+def updateLogFile(apparatus):
+    log   = apparatus.log
+    logFileName = apparatus.logFileName
+    if logFileName is not None:
+        log.save(logFileName, item_format='.15e')
 
-if __name__ == '__main__': 
+
+freqZero = 6834687190
+freqDeltaM2 = 6835387164
+def main():
     apparatus = BasicInstrument()
     apparatus.config['Device type'] = 'QOL VAMPIRE HighPower magnetometer'
     apparatus.config['Device model'] = 'v0.1'
@@ -85,30 +106,49 @@ if __name__ == '__main__':
     print(str.join('\n', scope.getHeader()))
     print(str.join('\n', rfgen.getHeader()))
     print('------ Header ends  -------------')
-    ch1  = scope.getTrace(1)
-    traces = scope.getAllTraces()
+    # ch1  = scope.getTrace(1)
+    # traces = scope.getAllTraces()
+    pid = PID(100,400,0, sign=-1)
 
-    freqZero = 6834687190
-    freqDeltaM2 = 6835387164
-    SweepSpan = 10000
+    apparatus.scope = scope
+    apparatus.rfgen = rfgen
+    apparatus.pid   = pid
+    apparatus.log   = initLog()
+    apparatus.runStatus = False
 
-    rfgen.stopFrequencySweep()
-    rfgen.setFreqFixed(freqDeltaM2+100)
 
-    pid = PID(100,400,0, sign=-1)
+
+    SweepSpan = 10000
+    apparatus.rfgen.stopFrequencySweep()
+    apparatus.rfgen.setFreqFixed(freqDeltaM2+100)
+    apparatus.logFileName = apparatus.getNextDataFile()
+    print(apparatus.logFileName)
+
 
     # initial lock
-    log=adjustRF(pid, scope, rfgen, nsteps=100)
+    apparatus.logFileName = None
+    apparatus.pid.reset()
+    apparatus.runStatus = True
+    feedbackLoop(apparatus, nsteps=10)
+    updateLogPlot(apparatus)
 
     # long term lock
-    fname = apparatus.getNextDataFile()
-    print(fname)
-    log=adjustRF(pid, scope, rfgen, nsteps=10000, fname=fname)
+    apparatus.logFileName = None
+    apparatus.pid.reset()
+    apparatus.runStatus = True
+    feedbackLoop(apparatus, nsteps=100)
+    updateLogPlot(apparatus)
+
+    # apparatus.log.save(apparatus.logFileName, item_format='.15e')
+    print(apparatus.logFileName)
+    return(apparatus)
 
-    log.save(fname, item_format='.15e')
-    print(fname)
 
 
+if __name__ == '__main__': 
+    mpl.use("TkAgg")
+    apparatus=main()
+