Updated to work with newly bundled pizeo package.
[unfold-protein.git] / unfold.py
1 TEXT_VERBOSE = False
2 PYLAB_INTERACTIVE_VERBOSE = True
3 BASE_FIG_NUM = 40
4 LOG_DATA = True
5 LOG_DIR = '$DEFAULT$/unfold'
6
7 import stepper
8 import piezo.z_piezo as z_piezo
9 import piezo.z_piezo_utils as z_piezo_utils
10 import piezo.x_piezo as x_piezo
11 import temperature
12 import threading
13 import time
14 import os, os.path
15 from numpy import arange, sin, pi, isnan
16 import data_logger
17
18 if PYLAB_INTERACTIVE_VERBOSE :
19     from pylab import figure, plot, title, legend, hold, subplot, draw
20     import time # for timestamping lines on plots
21
22 EFILE = file(os.path.expanduser('~/rsrch/debug.unfold'), 'a')
23
24 class ExceptionTooClose (Exception) :
25     """
26     The piezo is too close to the surface,
27     an unfolding would extend past pzMin
28     """
29 class ExceptionTooFar (Exception) :
30     """
31     The piezo has reached pzMax without
32     reaching the desired deflection setpoint.
33     """
34
35 def plot_dict(d, label) :
36     try :
37         plot(d["Z piezo output"], d["Deflection input"], '.',label=label)
38     except KeyError, string:
39         print d.keys()
40         raise KeyError, string
41
42 class unfold_data_log (data_logger.data_log) :
43     def save(self, unfold, timestamp, CTemp, VSetpoint, nmBindPos,
44              sBindTime,
45              nmDist, nmStep, ppStep, nmPsRate, freq,
46              approach_data, unfold_data) :
47         filepath,timestamp = self.get_filename(timestamp)
48         # save the unfolding data
49         dataFile = open(filepath, "w")
50         for i in range(0, len(unfold_data["Z piezo output"])) :
51                 dataFile.write("%d\t%d\t%d\n" % (unfold_data["Z piezo output"][i],
52                                                  unfold_data["Deflection input"][i],
53                                                  unfold_data["Z piezo input"][i]))
54         dataFile.close()
55         dataFile = open(filepath+"_approach", "w")
56         for i in range(0, len(approach_data["Z piezo output"])) :
57                 dataFile.write("%d\t%d\n" % (approach_data["Z piezo output"][i],
58                                                  approach_data["Deflection input"][i]))
59         dataFile.close()
60
61         # save parameters
62         paramFile = open(filepath+"_param", "w")
63         paramFile.write("Environment\n")
64         paramFile.write("Time:\t"+timestamp+"\n")
65         if CTemp != None :
66             paramFile.write("Temperature (C):\t"+str(CTemp)+"\n")
67
68         paramFile.write("\nPiezo parameters\n")
69         paramFile.write("Z piezo sensitivity (nm/Vp):\t"+str(unfold.zp.sensitivity)+"\n")
70         paramFile.write("Z piezo gain (Vp/Vo):\t"+str(unfold.zp.gain)+"\n")
71         paramFile.write("X piezo sensitivity (nm/Vp):\t"+str(unfold.xp.xpSensitivity)+"\n")
72         paramFile.write("X piezo gain (Vp/Vo):\t"+str(unfold.xp.gain)+"\n")
73         paramFile.write("X piezo position (nm):\t"+str(unfold.xp.out2nm(unfold.xp.curPos()))+"\n")
74
75         paramFile.write("\nApproach parameters\n")
76         paramFile.write("Setpoint (V):\t"+str(VSetpoint)+"\n")
77         paramFile.write("Bind pos (nm):\t"+str(nmBindPos)+"\n")
78         
79         paramFile.write("\nBinding parameters\n")
80         paramFile.write("Bind time (s):\t"+str(sBindTime)+"\n")
81
82         paramFile.write("\nUnfolding parameters\n")
83         paramFile.write("Distance (nm):\t"+str(nmDist)+"\n")
84         paramFile.write("Step size (nm):\t"+str(nmStep)+"\n")
85         paramFile.write("Points per step:\t"+str(ppStep)+"\n")
86         paramFile.write("Unfold rate (nm/s):\t"+str(nmPsRate)+"\n")
87         paramFile.write("Sample rate (Hz):\t"+str(freq)+"\n")
88
89         paramFile.write("\nData fields:\tZ_piezo_out Deflection_in Z_piezo_in\n")
90         paramFile.close()
91
92         return timestamp
93
94 class unfold :
95     def __init__(self, controlTemp=True, dataDirectory=LOG_DIR) :
96         self.step = stepper.stepper_obj()
97         self.zp = z_piezo.z_piezo()
98         self.xp = x_piezo.x_piezo()
99         #self.zp = z_piezo.z_piezo(zp_chan = '/Dev1/ao0',
100         #                          zp_mon_chan = '/Dev1/ai1',
101         #                          def_chan = '/Dev1/ai0')
102         #self.xp = x_piezo.x_piezo(xp_chan = '/Dev1/ao1')
103         self.zp.jumpToPos(self.zp.pos_nm2out(0))
104         self.xp.jumpToPos(self.xp.nm2out(0))
105         if controlTemp == True :
106             self.T = temperature.tempController(maxCurrent=1.0)
107         else: self.T = None
108         self.log = unfold_data_log(dataDirectory, log_name="unfold")
109     def unfold(self, setpoint=None, rel_setpoint=1.0, sBindTime = 10.0,
110                nmDist=600, nmStep=0.5, ppStep=10, nmPsRate=1000,
111                dataDirectory=LOG_DIR, fileID=None) :
112         while True :
113             try :
114                 data = self.unfold_cycle(setpoint, rel_setpoint, sBindTime,
115                                          nmDist, nmStep, ppStep, nmPsRate,
116                                          dataDirectory, fileID)
117                 break
118             except ExceptionTooFar :
119                 EFILE.write('caught ExceptionTooFar\n'); EFILE.flush()
120                 try : # try for a useful surface distance...
121                     if setpoint == None : # HACK! redundant!
122                         assert rel_setpoint != None, "must have some sort of setpoint"
123                         setpoint = self.curDef() + rel_setpoint
124                         print "setpoint = %g" % setpoint
125                     EFILE.write('attempt getSurfPos\n'); EFILE.flush()
126                     surfPos = self.getSurfPos(setpoint)
127                     EFILE.write('getSurfPos succeeded\n'); EFILE.flush()
128                     print "Too far (surface at %g nm), stepping closer" % surfPos
129                 except z_piezo_utils.poorFit, string : # ... oh well, print what we know
130                     EFILE.write('getSurfPos failed\n'); EFILE.flush()
131                     print "Too far, stepping closer"
132                     print "(Fit failed with: %s)" % string
133                 EFILE.write('zero Z piezo\n'); EFILE.flush()
134                 self.zp.jumpToPos(self.zp.pos_nm2out(0))
135                 EFILE.write('step closer\n'); EFILE.flush()
136                 self.stepCloser()
137                 EFILE.write('step closer successful\n'); EFILE.flush()
138             except ExceptionTooClose :
139                 EFILE.write('caught ExceptionTooFar\n'); EFILE.flush()
140                 try : # try for a useful surface distance...
141                     if setpoint == None : # !HACK !redundant!
142                         assert rel_setpoint != None, "must have some sort of setpoint"
143                         setpoint = self.curDef() + rel_setpoint
144                         print "setpoint = %g" % setpoint
145                     EFILE.write('attempt getSurfPos\n'); EFILE.flush()
146                     surfPos = self.getSurfPos(setpoint)
147                     EFILE.write('getSurfPos succeeded\n'); EFILE.flush()
148                     print "Too close (surface at %g nm), stepping away" % surfPos
149                 except z_piezo_utils.poorFit, string : # ... oh well, print what we know
150                     EFILE.write('getSurfPos failed\n'); EFILE.flush()
151                     print "Too close, stepping away"
152                     print "(Fit failed with: %s)" % string
153                 EFILE.write('zero Z piezo\n'); EFILE.flush()
154                 self.zp.jumpToPos(self.zp.pos_nm2out(0))
155                 EFILE.write('step away\n'); EFILE.flush()
156                 self.stepAway()
157                 EFILE.write('step away successful\n'); EFILE.flush()
158             print "Too close, stepping away"
159         return data
160     def stepApproach(self, setpoint) :
161         cd = self.curDef()
162         while cd < setpoint or isnan(cd) : # sometimes during approach, we get negative nan values
163             print "deflection %g < setpoint %g.  step closer" % (cd, setpoint)
164             self.stepCloser()
165             cd = self.curDef()
166     def stepApproachRetractCurve(self, setpoint) :
167         def_array = []
168         step_array = []
169         orig_step = self.step.get_cur_pos()
170         def store_point() :
171             cd = self.curDef()
172             def_array.append(cd)
173             step_array.append(self.step.get_cur_pos())
174             return cd        
175         cd = store_point()
176         while cd < setpoint or isnan(cd) : # sometimes during approach, we get negative nan values
177             print "deflection %g < setpoint %g.  step closer" % (cd, setpoint)
178             self.step.step_rel(2)
179             cd = store_point()
180         while self.step.get_cur_pos() > orig_step :
181             self.step.step_rel(-2)
182             store_point()
183         return {"Deflection":def_array, "Stepper position":step_array}
184     def piezoApproach(self, setpoint, return_data=False) :
185         startPos = self.zp.curPos()
186         curVals,data = z_piezo_utils.moveToPosOrDef(self.zp,
187                            self.zp.zpMax,
188                            self.zp.def_V2in(setpoint),
189                            step=10,
190                            return_data=return_data)
191         if curVals["Deflection input"] < self.zp.def_V2in(setpoint) :
192             EFILE.write('Unfolding too far\n'); EFILE.flush()
193             self.zp.jumpToPos(startPos)
194             self.zp.updateInputs()
195             if PYLAB_INTERACTIVE_VERBOSE == True :
196                 figure(BASE_FIG_NUM+1)
197                 hold(False)
198                 plot_dict(data, 'Approach')
199                 hold(True)
200                 title('Unfolding too far')
201             EFILE.write('Raising ExceptionTooFar\n'); EFILE.flush()
202             raise ExceptionTooFar
203         if return_data == True :
204             return (curVals, data)
205         else :
206             return curVals
207     def bind(self, sTime=10.0) :
208         time.sleep(sTime)
209     def unfold_pull(self, nmDist=600, nmStep=0.5, ppStep=10, nmPsRate=1000) :
210         if nmStep < 0 :
211                 nmStep = -nmStep
212         numSteps = int(nmDist/nmStep+1)
213         out = [0.0]*numSteps*ppStep
214         pos = self.zp.curPos()
215         startPos = pos
216         step = int(self.zp.pos_nm2out(nmStep)-self.zp.pos_nm2out(0))
217         for i in range(0, numSteps) :
218                 for j in range(0, ppStep) :
219                         out[i*ppStep + j] = pos
220                 pos -= step
221         # points/step * step/nm * nm/s  =  points/s
222         freq = ppStep / nmStep * nmPsRate
223         print "unfolding %d points per %d steps from %d to %d at freq of %g" % (ppStep, numSteps, startPos, pos, freq) 
224         return {'freq':freq, 'data':self.zp.ramp(out, freq)}
225     def generate_refold_bind_pull_output(self, nmSurfPos=0,
226             sBindTime=2, nmUnfoldDist=160, nmRefoldDist=145, sPauseTime=2,
227             cycles=50, nmStep=0.5, ppStepUnfold=10, ppStepRefold=1,
228             nmPsRateUnfold=1000) :
229         
230         if nmStep < 0 :##
231             nmStep = -nmStep
232         numSteps = int(nmDist/nmStep+1)
233         out = [0.0]*numSteps*ppStep
234         pos = self.zp.curPos()
235         startPos = pos
236         step = int(self.zp.pos_nm2out(nmStep)-self.zp.pos_nm2out(0))
237         for i in range(0, numSteps) :
238                 for j in range(0, ppStep) :
239                         out[i*ppStep + j] = pos
240                 pos -= step
241         # points/step * step/nm * nm/s  =  points/s
242         freq = ppStep / nmStep * nmPsRate
243         print "unfolding %d points per %d steps from %d to %d at freq of %g" % (ppStep, numSteps, startPos, pos, freq) 
244         return {'freq':freq, 'data':self.zp.ramp(out, freq)}
245     def unfold_cycle(self, setpoint=None, rel_setpoint=1.0, sBindTime = 10.0,
246                      nmDist=600, nmStep=0.5, ppStep=10, nmPsRate=1000,
247                      dataDirectory=LOG_DIR, fileID=None) :
248         if setpoint == None :
249             assert rel_setpoint != None, "must have some sort of setpoint"
250             setpoint = self.curDef() + rel_setpoint
251             print "setpoint = %g" % setpoint
252         timestamp = time.strftime("%Y%m%d%H%M%S")
253         temp = None
254         if self.T != None : temp = self.T.getTemp()
255         print "approaching"
256         startPos = self.zp.curPos()
257         curVals,approach_data = self.piezoApproach(setpoint, return_data=True)
258         bindPos = curVals['Z piezo output'] # in output units
259         finalPos = bindPos - (self.zp.pos_nm2out(nmDist)-self.zp.pos_nm2out(0))
260         try :
261             self.zp._check_range(finalPos)
262         except z_piezo.outOfRange :
263             self.zp.jumpToPos(startPos)
264             self.zp.updateInputs()
265             if PYLAB_INTERACTIVE_VERBOSE == True :
266                 figure(BASE_FIG_NUM+1)
267                 hold(False)
268                 plot_dict(approach_data, 'Approach')
269                 hold(True)
270                 title('Unfolding too close')
271             raise ExceptionTooClose
272         print "binding for %.3f seconds" % sBindTime
273         self.bind(sBindTime)
274         out = self.unfold_pull(nmDist, nmStep, ppStep, nmPsRate)
275
276         if PYLAB_INTERACTIVE_VERBOSE == True :
277             figure(BASE_FIG_NUM)
278             hold(False)
279             plot_dict(approach_data, 'Approach')
280             hold(True)
281             plot_dict(out['data'], 'Unfold')
282             legend(loc='best')
283             title('Unfolding')
284             draw()
285
286         if LOG_DATA == True:
287             print "saving"
288             timestamp = self.log.save(self, timestamp, temp, setpoint,
289                                       int(self.zp.pos_out2nm(bindPos)), # don't need lots of precision...
290                                       sBindTime,
291                                       nmDist, nmStep, ppStep, nmPsRate, out['freq'],
292                                       approach_data, out['data'])
293         return out
294     def getSurfPos(self, setpoint=2, textVerbose=False, plotVerbose=False) :
295         return self.zp.pos_out2nm(z_piezo_utils.getSurfPos(self.zp, self.zp.def_V2in(setpoint), textVerbose, plotVerbose))
296     def curDef(self) :
297         self.zp.updateInputs()
298         return self.zp.def_in2V(self.zp.curDef())
299     def stepCloser(self) :
300         "Backlash-robust stepping"
301         self.step.step_rel(2) # two half-steps in full step mode
302     def stepAway(self) :
303         "Backlash-robust stepping"
304         self.step.step_rel(-120)
305         self.step.step_rel(110) # HACK, should come back 118
306     def xpWander(self) :
307         self.xp.wander()
308
309 def _test_unfold(controlTemp=True) :
310     print "Test unfold"
311     u = unfold(controlTemp=controlTemp)
312     u.unfold(setpoint=0.5, sBindTime=1.0)
313     del u
314     print "unfold successful\n"
315
316
317 class unfold_expt :
318     def __init__(self, controlTemp=True) :
319         self.u = unfold(controlTemp=controlTemp)
320         self.runUnfolds = False
321         self.bgRun = None
322         self.getExternalTempSetpoint = None
323         if self.u.T != None :
324             self.tempSetpoint = self.u.T.getTemp()
325         self.getExternalDeflectionSetpoint = None
326         self.deflectionSetpoint = 0.5
327         self.getExternalBindTime = None
328         self.bindTime = 1.0
329         self.getExternalnmDist = None
330         self.nmDist = 600.0
331         self.getExternalnmStep = None
332         self.nmStep = 0.5
333         self.getExternalppStep = None
334         self.ppStep = 10
335         self.getExternalnmPsRate = None
336         self.nmPsRate = 1000.0
337         self.getExternalDataDirectory = None
338         self.dataDirectory = LOG_DIR
339         self.plotData = None
340         self.i=0
341         self.showUnfoldingIndex = None
342     def run(self) :
343         print "running"
344         if self.bgRun != None :
345             raise Exception, "Can't run two backgrounds simultaneously"
346         self.runUnfolds = True
347         print "unfolding in the background"
348         self.bgRun = bg_run( self._run)
349     def stepApproach(self) :
350         print "approaching"
351         if self.getExternalDeflectionSetpoint != None :
352             self.deflectionSetpoint = self.getExternalDeflectionSetpoint()
353         self.u.stepApproach(self.deflectionSetpoint)
354     def _run(self) :
355         print "starting unfold loop"
356         while self.runUnfolds == True :
357             print "get unfold parameters"
358             if self.u.T != None and self.getExternalTempSetpoint != None :
359                 setpoint = self.getExternalTempSetpoint()
360                 if setpoint != self.tempSetpoint :
361                     self.u.T.setTemp(setpoint)
362                     self.tempSetpoint = setpoint
363             if self.getExternalDeflectionSetpoint != None :
364                 self.deflectionSetpoint = self.getExternalDeflectionSetpoint()
365             if self.getExternalBindTime != None :
366                 self.bindTime = self.getExternalBindTime()
367             if self.getExternalnmDist != None :
368                 self.nmDist = self.getExternalnmDist()
369             if self.getExternalnmStep != None :
370                 self.nmStep = self.getExternalnmStep()
371             if self.getExternalppStep != None :
372                 self.ppStep = self.getExternalppStep()
373             if self.getExternalnmPsRate != None :
374                 self.nmPsRate = self.getExternalnmPsRate()
375             if self.getExternalDataDirectory != None :
376                 self.dataDirectory = self.getExternalDataDirectory()
377             print "run unfold"
378             #print "setpoint  ", self.deflectionSetpoint
379             #print "bind time ", self.bindTime
380             #print "nmDist    ", self.nmDist
381             #print "nmStep    ", self.nmStep
382             #print "ppStep    ", self.ppStep
383             #print "nmPsRate  ", self.nmPsRate
384             #print "fileID    ", self.i
385             #print "dataDir   ", self.dataDirectory
386             data = self.u.unfold(setpoint=self.deflectionSetpoint,
387                                  sBindTime=self.bindTime,
388                                  nmDist=self.nmDist,
389                                  nmStep=self.nmStep,
390                                  ppStep=self.ppStep,
391                                  nmPsRate=self.nmPsRate,
392                                  fileID=str(self.i),
393                                  dataDirectory=self.dataDirectory)
394             print "plot data"
395             if self.plotData != None :
396                 self.plotData(data["data"]["Z piezo output"],
397                               data["data"]["Deflection input"])
398             if self.showUnfoldingIndex != None :
399                 self.showUnfoldingIndex(self.i)
400             self.i += 1
401         return False
402     def stop(self) :
403         if self.bgRun != None :
404             self.runUnfolds = False
405             self.bgRun.wait()
406             self.bgRun = None
407
408 def _test_unfold_expt(controlTemp=True) :
409     print "Test unfold_expt"
410     u_expt = unfold.unfold_expt(controlTemp=controlTemp)
411     u_expt.run()
412     time.sleep(100)
413     u_expt.stop()
414     print "unfold_expt working\n"
415
416 class bg_run :
417     def __init__(self, function, args=None, finishCallback=None, interruptCallback=None) :
418         print "init bg"
419         def tempFn() :
420             print "temp fn started"
421             complete = False
422             if args == None :
423                 complete = function()
424             else :
425                 complete = function(args)
426             print "temp Fn finished"
427             if finishCallback != None and complete == True :
428                 finishCallback()
429             elif interruptCallback != None and complete == False :
430                 interruptCallback()
431         print "tempFn defined"
432         self.thd = threading.Thread(target=tempFn, name="Background Fn")
433         print "starting thread"
434         self.thd.start()
435         print "thread started"
436     def wait(self) :
437         print "joining thread"
438         self.thd.join()
439         print "thread joined"
440
441 def _test_bg_run() :
442     print "Test bg_run"
443     print "test without args or callbacks"
444     def say_hello() :
445         for i in range(10) :
446             print "Hello"
447             time.sleep(0.1)
448         return True
449     bg = unfold.bg_run(say_hello)
450     time.sleep(.2)
451     print "The main thread's still going"
452     bg.wait()
453     print "The background process is done"
454     
455     print "test without args, but with callbacks"
456     def inter() :
457         print "I was interrupted"
458     def fin() :
459         print "I'm finished"
460     bg = unfold.bg_run(say_hello, finishCallback=fin, interruptCallback=inter)
461     time.sleep(.2)
462     print "The main thread's still going"
463     bg.wait()
464     print "The background process is done"
465
466     print "test with args and callbacks"
467     del say_hello
468     def say_hello(stop=[False]) :
469         for i in range(10) :
470             if stop[0] == True : return False
471             print "Hello"
472             time.sleep(0.1)
473         return True
474     stop = [False]
475     bg = unfold.bg_run(say_hello, args=stop, finishCallback=fin, interruptCallback=inter)
476     time.sleep(.2)
477     print "The main thread's still going, stop it"
478     stop[0] = True
479     bg.wait()
480     del bg
481     del stop
482     del say_hello
483     del inter
484     del fin
485     print "The background process is done"
486     
487     print "bg_run working\n"
488
489 def loop_rates(u, rates, num_loops=10, die_file=None, song=None, **kwargs):
490     """
491     loop_rates(u, rates, num_loops=10, die_file=None, song=None, **kwargs)
492     
493     Constant speed unfolding using the unfold instance u for num_loops
494     through the series of rates listed in rates.  You may set die_file
495     to a path, loop_rates() will check that location before each
496     unfolding attempt, and cleanly stop unfolding if the file exists.
497     **kwargs are passed on to u.unfold(), so a full call might look like
498     
499     loop_rates(u, rates=[20,200,2e3], num_loops=5, die_file='~/die', song='~wking/Music/system/towerclo.wav', rel_setpoint=1, nmDist=800, sBindTime=2)
500     """
501     if die_file != None:
502         die_file = os.path.expanduser(die_file)
503     if song != None:
504         song = os.path.expanduser(song)
505     for i in range(num_loops):
506         for nmPsRate in rates:
507             if die_file != None and os.path.exists(die_file):
508                 return None
509             u.unfold(nmPsRate=nmPsRate, **kwargs)
510             u.xpWander()
511     if song != None:
512         os.system("aplay '%s'" % song)
513
514 def test() :
515     _test_unfold(controlTemp=False)
516     _test_bg_run()
517     _test_unfold_expt(controlTemp=False)
518
519 if __name__ == "__main__" :
520     test()