Packaged with setuputils.
[pypiezo.git] / piezo / z_piezo_utils.py
similarity index 75%
rename from z_piezo_utils.py
rename to piezo/z_piezo_utils.py
index f80017fcf4c6e194c1b8ad5db9497266b9876ccf..7d4880280b168ef9709d14c4ca7595278872d325 100644 (file)
@@ -1,5 +1,6 @@
 TEXT_VERBOSE = False
-PYLAB_INTERACTIVE_VERBOSE = False
+PYLAB_INTERACTIVE = True
+PYLAB_VERBOSE = True
 BASE_FIG_NUM = 60
 LOG_DATA = True
 LOG_DIR = '$DEFAULT$/z_piezo_utils'
@@ -10,13 +11,38 @@ SLEEP_DURING_SURF_POS_AQUISITION = False # doesn't help
 from numpy import array, arange, ones, zeros, linspace, uint16, float, sin, cos, pi
 from scipy.stats import linregress
 import curses_check_for_keypress
-import fit
+import scipy.optimize
 import data_logger
 import time
 
-if PYLAB_INTERACTIVE_VERBOSE :
-    from pylab import figure, plot, title, legend, hold, subplot
-        
+_flush_plot = None
+_final_flush_plot = None
+_pylab = None
+def _dummy_fn(): pass
+
+def _import_pylab():
+    """Import pylab plotting functions for when we need to plot.  This
+    function can be called multiple times, and ensures that the pylab
+    setup is only imported once.  It defines the functions
+      _flush_plot() to be called after incremental changes,
+      _final_flush_plot(), to be called at the end of any graphical processing, and
+      _pylab, a reference to the pylab module."""
+    global _pylab
+    global _flush_plot
+    global _final_flush_plot
+    if _pylab == None :
+        import pylab as _pylab
+    if _flush_plot == None :
+        if PYLAB_INTERACTIVE :
+            _flush_plot = _pylab.draw
+        else :
+            _flush_plot = _pylab.show
+    if _final_flush_plot == None :
+        if PYLAB_INTERACTIVE :
+            _final_flush_plot = _pylab.draw
+        else :
+            _final_flush_plot = _dummy_fn
+
 class error (Exception) :
     pass
 class poorFit (error) :
@@ -25,7 +51,7 @@ class tooClose (error) :
     pass
 
 def moveToPosOrDef(zpiezo, pos, defl, step, return_data=False) :
-    if return_data or LOG_DATA or PYLAB_INTERACTIVE_VERBOSE :
+    if return_data or LOG_DATA or PYLAB_VERBOSE :
         aquire_data = True
     else :
         aquire_data = False
@@ -68,12 +94,14 @@ def moveToPosOrDef(zpiezo, pos, defl, step, return_data=False) :
     
     if zpiezo.setExternalZPiezo != None :
         zpiezo.setExternalZPiezo(zpiezo.pos_out2nm(zpiezo.curPos()))
-    if PYLAB_INTERACTIVE_VERBOSE :
-        figure(BASE_FIG_NUM)
+    if PYLAB_VERBOSE :
+        _import_pylab()
+        _pylab.figure(BASE_FIG_NUM)
         timestamp = time.strftime('%H%M%S')
-        plot(pos_array, def_array, '.', label=timestamp)
-        title('step approach')
-        legend(loc='best')
+        _pylab.plot(pos_array, def_array, '.', label=timestamp)
+        _pylab.title('step approach')
+        _pylab.legend(loc='best')
+        _flush_plot()
     if LOG_DATA :
         log = data_logger.data_log(LOG_DIR, noclobber_logsubdir=False,
                                    log_name="moveToPosOrDef")
@@ -91,20 +119,20 @@ def moveToPosOrDef(zpiezo, pos, defl, step, return_data=False) :
 def wiggleForInterferenceMin(zpiezo, wig_amp=20000, wig_freq=100,
                              plotVerbose=True) :
     "Output sin to measure interference"
-    if PYLAB_INTERACTIVE_VERBOSE or plotVerbose:
-        from pylab import figure, plot, title, hold, ion, ioff, draw, axes
+    if PYLAB_VERBOSE or plotVerbose:
+        _import_pylab()
     npoints = 1000
     scanfreq = wig_freq*npoints
     out = zeros((npoints,), dtype=uint16)
     for i in range(npoints) :
         out[i] = int(sin(4*pi*i/npoints)*wig_amp+32000)
         # 4 for 2 periods
-    if PYLAB_INTERACTIVE_VERBOSE or plotVerbose:
-        figure(BASE_FIG_NUM+1)
-        a = axes()
-        title('wiggle for interference')
-        hold(False)
-        p = plot(out, out, 'b.-')
+    if PYLAB_VERBOSE or plotVerbose:
+        _pylab.figure(BASE_FIG_NUM+1)
+        a = _pylab.axes()
+        _pylab.title('wiggle for interference')
+        _pylab.hold(False)
+        p = _pylab.plot(out, out, 'b.-')
     if LOG_DATA :
         log = data_logger.data_log(LOG_DIR, noclobber_logsubdir=False,
                                    log_name="wiggle")
@@ -113,19 +141,19 @@ def wiggleForInterferenceMin(zpiezo, wig_amp=20000, wig_freq=100,
         data = zpiezo.ramp(out, scanfreq)
         dmin = data["Deflection input"].min()
         dmax = data["Deflection input"].max()
-        if PYLAB_INTERACTIVE_VERBOSE or plotVerbose:
+        if PYLAB_VERBOSE or plotVerbose:
             p[0].set_ydata(data["Deflection input"])
             a.set_ylim([dmin,dmax])
-            draw()
+            _flush_plot()
         if LOG_DATA :
             log.write_dict_of_arrays(data)
-    if PYLAB_INTERACTIVE_VERBOSE or plotVerbose:
-        hold(True)
+    if PYLAB_VERBOSE or plotVerbose:
+        _pylab.hold(True)
 
 def getSurfPosData(zpiezo, maxDefl, textVerbose=False, plotVerbose=False) :
     "Measure the distance to the surface"
-    if plotVerbose and not PYLAB_INTERACTIVE_VERBOSE :
-        from pylab import figure, plot, title, legend, hold, subplot
+    if plotVerbose and not PYLAB_VERBOSE :
+        _import_pylab()
     origPos = zpiezo.curPos()
     # fully retract the piezo
     if TEXT_VERBOSE or textVerbose :
@@ -150,16 +178,16 @@ def getSurfPosData(zpiezo, maxDefl, textVerbose=False, plotVerbose=False) :
         ret3=zpiezo.ramp(out, 10e3)
         del(out)
         zpiezo.updateInputs()
-        if PYLAB_INTERACTIVE_VERBOSE or plotVerbose :
-            figure(BASE_FIG_NUM+2)
+        if PYLAB_VERBOSE or plotVerbose :
+            _pylab.figure(BASE_FIG_NUM+2)
             def plot_dict(d, label) :
                 print d.keys()
-                plot(d["Z piezo output"], d["Deflection input"], '.',label=label)
+                _pylab.plot(d["Z piezo output"], d["Deflection input"], '.',label=label)
             plot_dict(ret1, 'First retract')
             plot_dict(mtpod, 'Single step in')
             plot_dict(ret3, 'Return to start')
-            legend(loc='best')
-            title('Too close')
+            _pylab.legend(loc='best')
+            _pylab.title('Too close')
         raise tooClose, "Way too close!\n(highContactPos %d <= origPos %d)" % (highContactPos, origPos) 
     zpiezo.updateInputs()
     # fully retract the piezo again
@@ -230,17 +258,30 @@ def analyzeSurfPosData(ddict, textVerbose=False, plotVerbose=False, retAllParams
     if TEXT_VERBOSE or textVerbose :
         print "start Def %d, final def %d, start pos %d, final pos %d" % (startDef, finalDef, startPos, finalPos)
         print "Guessed params ", pstart
-    params = fit.fit(data=[data["Z piezo output"][dump_before_index:], data["Deflection input"][dump_before_index:]],
-                     model='bilinear', start=pstart, scale=scale)
+    def bilinear(x, params):
+        left_offset,left_slope,kink_position,right_slope = params
+        if x < kink_position:
+            y = left_offset + left_slope*x
+        else:
+            y = left_offset + left_slope*kink_position \
+                + right_slope*(x-kink_position)
+    def residual(p, y, x):
+        Y = bilinear(x, p)
+        return Y-y
+    leastsq = scipy.optimize.leastsq
+    params,cov,info,mesg,ier = leastsq(
+        residual, pstart,
+        args=(data["Deflection input"][dump_before_index:],
+              data["Z piezo output"][dump_before_index:]),
+        full_output=True, maxfev=10000)
     if TEXT_VERBOSE or textVerbose :
         print "Best fit parameters: ", params
 
-    if PYLAB_INTERACTIVE_VERBOSE or plotVerbose :
-        if plotVerbose and not PYLAB_INTERACTIVE_VERBOSE :
-            from pylab import figure, plot, title, legend, hold, subplot
-        figure(BASE_FIG_NUM+3)
+    if PYLAB_VERBOSE or plotVerbose :
+        _import_pylab()
+        _pylab.figure(BASE_FIG_NUM+3)
         def plot_dict(d, label) :
-            plot(d["Z piezo output"], d["Deflection input"], '.',label=label)
+            _pylab.plot(d["Z piezo output"], d["Deflection input"], '.',label=label)
         try :
             plot_dict(ddict['ret1'], 'First retract')
             plot_dict(ddict['mtpod'], 'Single step in')
@@ -257,12 +298,13 @@ def analyzeSurfPosData(ddict, textVerbose=False, plotVerbose=False, retAllParams
                 return params[0] + params[1]*x
             else :
                 return params[0] + params[1]*params[2] + params[3]*(x-params[2])
-        plot([startPos, params[2], finalPos],
-             [fit_fn(startPos, params), fit_fn(params[2], params), fit_fn(finalPos, params)],
-             '-',label='fit')
-        title('surf_pos %5g %5g %5g %5g' % (params[0], params[1], params[2], params[3]))
-        legend(loc='best')
-
+        _pylab.plot([startPos, params[2], finalPos],
+                    [fit_fn(startPos, params), fit_fn(params[2], params), fit_fn(finalPos, params)],
+                    '-',label='fit')
+        _pylab.title('surf_pos %5g %5g %5g %5g' % (params[0], params[1], params[2], params[3]))
+        _pylab.legend(loc='best')
+        _flush_plot()
+    
     # check that the fit is reasonable
     # params[1] is slope in non-contact region
     # params[3] is slope in contact region
@@ -307,15 +349,16 @@ def test_smoothness(zp, plotVerbose=True) :
         outdata.append(out)
     if plotVerbose :
         from pylab import figure, plot, title, legend, hold, subplot        
-    if PYLAB_INTERACTIVE_VERBOSE or plotVerbose :
-        figure(BASE_FIG_NUM+4)
+    if PYLAB_VERBOSE or plotVerbose :
+        _import_pylab()
+        _pylab.figure(BASE_FIG_NUM+4)
         for i in range(10) :
-            plot(indata[i]['Z piezo output'],
-                 indata[i]['Deflection input'], '+--', label='in')
-            plot(outdata[i]['Z piezo output'],
-                 outdata[i]['Deflection input'], '.-', label='out')
-        title('test smoothness (step in, ramp out)')
-        #legend(loc='best')
+            _pylab.plot(indata[i]['Z piezo output'],
+                        indata[i]['Deflection input'], '+--', label='in')
+            _pylab.plot(outdata[i]['Z piezo output'],
+                        outdata[i]['Deflection input'], '.-', label='out')
+        _pylab.title('test smoothness (step in, ramp out)')
+        #_pylab.legend(loc='best')
     
 def test() :
     import z_piezo
@@ -327,3 +370,5 @@ def test() :
     if TEXT_VERBOSE :
         print "Surface at %g nm", pos
     print "success"
+    if PYLAB_VERBOSE and _final_flush_plot != None:
+        _final_flush_plot()