Only call figure.show if it exits.
[pypiezo.git] / pypiezo / surface.py
1 # Copyright (C) 2011-2012 W. Trevor King <wking@tremily.us>
2 #
3 # This file is part of pypiezo.
4 #
5 # pypiezo is free software: you can redistribute it and/or modify it under the
6 # terms of the GNU General Public License as published by the Free Software
7 # Foundation, either version 3 of the License, or (at your option) any later
8 # version.
9 #
10 # pypiezo is distributed in the hope that it will be useful, but WITHOUT ANY
11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
12 # A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
13 #
14 # You should have received a copy of the GNU General Public License along with
15 # pypiezo.  If not, see <http://www.gnu.org/licenses/>.
16
17 """Utilities detecting the position of the sample surface.
18
19 TODO: doctests
20 """
21
22 SLEEP_DURING_SURF_POS_AQUISITION = False # doesn't help
23
24 #from time import sleep as _sleep
25 import numpy as _numpy
26 from scipy.optimize import leastsq as _leastsq
27
28 try:
29     import matplotlib as _matplotlib
30     import matplotlib.pyplot as _matplotlib_pyplot
31     import time as _time  # for timestamping lines on plots
32 except (ImportError, RuntimeError), e:
33     _matplotlib = None
34     _matplotlib_import_error = e
35
36 from . import LOG as _LOG
37 from . import package_config as _package_config
38 from . import base as _base
39
40
41 class SurfaceError (Exception):
42     pass
43
44
45 class PoorFit (SurfaceError):
46     pass
47
48
49 class PoorGuess (PoorFit):
50     pass
51
52
53 class FlatFit (PoorFit):
54     "Raised for slope in the non-contact region, or no slope in contact region"
55     def __init__(self, left_slope, right_slope):
56         self.left_slope = left_slope
57         self.right_slope = right_slope
58         msg = 'slopes not sufficiently different: %g and %g' % (
59             left_slope, right_slope)
60         super(FlatFit, self).__init__(msg)
61
62
63 class EdgeKink (PoorFit):
64     def __init__(self, kink, edge, window):
65         self.kink = kink
66         self.edge = edge
67         self.window = window
68         msg = 'no kink (kink %d not within %d of %d)' % (kink, window, edge)
69         super(EdgeKink, self).__init__(self, msg)
70
71
72 def _linspace(*args, **kwargs):
73     dtype = kwargs.pop('dtype')
74     out = _numpy.linspace(*args, **kwargs)
75     return out.reshape((len(out), 1)).astype(dtype)
76
77 def _get_min_max_positions(piezo, axis_name, min_position=None,
78                            max_position=None):
79     output_axis = piezo.axis_by_name(axis_name)
80     if min_position is None:
81         min_position = _base.convert_volts_to_bits(
82             output_axis.config['channel'],
83             output_axis.config['minimum'])
84     if max_position is None:
85         max_position = _base.convert_volts_to_bits(
86             output_axis.config['channel'],
87             output_axis.config['maximum'])
88     return (min_position, max_position)
89
90 def get_surface_position_data(piezo, axis_name, max_deflection, steps=2000,
91                               frequency=10e3, min_position=None,
92                               max_position=None):
93     "Measure the distance to the surface"
94     _LOG.info('get surface position')
95     orig_position = piezo.last_output[axis_name]
96     # fully retract the piezo
97     min_position,max_position = _get_min_max_positions(
98         piezo, axis_name, min_position=min_position, max_position=max_position)
99     _LOG.info('retract the piezo to %d' % min_position)
100     dtype = piezo.channel_dtype(axis_name, direction='output')
101     out = _linspace(orig_position, min_position, steps, dtype=dtype)
102     out = out.reshape((len(out), 1)).astype(
103         piezo.channel_dtype(axis_name, direction='output'))
104     ramp_kwargs = {
105         'frequency': frequency,
106         'output_names': [axis_name],
107         'input_names': ['deflection'],
108         }
109     ret1 = piezo.named_ramp(data=out, **ramp_kwargs)
110     # locate high deflection position
111     _LOG.info('approach until there is dangerous deflection (> %d)'
112                % max_deflection)
113     if SLEEP_DURING_SURF_POS_AQUISITION == True:
114         _sleep(.2) # sleeping briefly seems to reduce bounce?
115     mtpod = piezo.move_to_pos_or_def(
116         axis_name=axis_name, position=max_position, deflection=max_deflection,
117         step=(max_position-min_position)/steps, return_data=True)
118     high_contact_pos = piezo.last_output[axis_name]
119     # fully retract the piezo again
120     _LOG.info('retract the piezo to %d again' % min_position)
121     if SLEEP_DURING_SURF_POS_AQUISITION == True:
122         _sleep(.2)
123     out = _linspace(high_contact_pos, min_position, steps, dtype=dtype)
124     ret2 = piezo.named_ramp(data=out, **ramp_kwargs)
125     # scan to the high contact position
126     _LOG.info('ramp in to the deflected position %d' % high_contact_pos)
127     if SLEEP_DURING_SURF_POS_AQUISITION == True:
128         _sleep(.2)
129     out = _linspace(min_position, high_contact_pos, steps, dtype=dtype)
130     data = piezo.named_ramp(data=out, **ramp_kwargs)
131     if SLEEP_DURING_SURF_POS_AQUISITION == True:
132         _sleep(.2)
133     # return to the original position
134     _LOG.info('return to the original position %d' % orig_position)
135     out = _linspace(high_contact_pos, orig_position, steps, dtype=dtype)
136     ret3 = piezo.named_ramp(data=out, **ramp_kwargs)
137     return {'ret1':ret1, 'mtpod':mtpod, 'ret2':ret2,
138             'approach':data, 'ret3':ret3}
139
140 def bilinear(x, params):
141     """bilinear fit for surface bumps.  Model has two linear regimes
142     which meet at x=kink_position and have independent slopes.
143
144     `x` should be a `numpy.ndarray`.
145     """
146     left_offset,left_slope,kink_position,right_slope = params
147     left_mask = x < kink_position
148     right_mask = x >= kink_position # = not left_mask
149     left_y = left_offset + left_slope*x
150     right_y = (left_offset + left_slope*kink_position
151                + right_slope*(x-kink_position))
152     return left_mask * left_y + right_mask * right_y
153
154 def analyze_surface_position_data(
155     ddict, min_slope_ratio=10.0, kink_window=None,
156     return_all_parameters=False):
157     """
158
159     min_slope_ratio : float
160         Minimum ratio between the non-contact "left" slope and the
161         contact "right" slope.
162     kink_window : int (in bits) or None
163         Raise `EdgeKink` if the kink is within `kink_window` of the
164         minimum or maximum `z` position during the approach.  If
165         `None`, a default value of 2% of the approach range is used.
166     """
167     # uses ddict["approach"] for analysis
168     # the others are just along to be plotted
169     _LOG.info('analyze surface position data')
170
171     data = ddict['approach']
172     # analyze data, using bilinear model
173     # y = p0 + p1 x                for x <= p2
174     #   = p0 + p1 p2 + p3 (x-p2)   for x >= p2
175     dump_before_index = 0 # 25 # HACK!!
176     # Generate a reasonable guess...
177     start_pos = data['z'].min()
178     final_pos = data['z'].max()
179     if final_pos == start_pos:
180         raise SurfaceError('cannot fit a single approach step (too close?)')
181     start_def = data['deflection'].min()
182     final_def = data['deflection'].max()
183     # start_def and start_pos are probably for 2 different points
184     _LOG.info('min deflection {}, max deflection {}'.format(
185             start_def, final_def))
186     _LOG.info('min position {}, max position {}'.format(start_pos, final_pos))
187
188     left_offset   = start_def
189     left_slope    = 0
190     kink_position = (final_pos+start_pos)/2.0
191     right_slope   = 2.0*(final_def-start_def)/(final_pos-start_pos)
192     pstart = [left_offset, left_slope, kink_position, right_slope]
193     _LOG.info('guessed params: %s' % pstart)
194
195     offset_scale = (final_pos - start_pos)/100
196     left_slope_scale = right_slope/10
197     kink_scale = (final_pos-start_pos)/100
198     right_slope_scale = right_slope
199     scale = [offset_scale, left_slope_scale, kink_scale, right_slope_scale]
200     _LOG.info('guessed scale: %s' % scale)
201
202     def residual(p, y, x):
203         Y = bilinear(x, p)
204         return Y-y
205     params,cov,info,mesg,ier = _leastsq(
206         residual, pstart,
207         args=(data["deflection"][dump_before_index:],
208               data["z"][dump_before_index:]),
209         full_output=True, maxfev=10000)
210     _LOG.info('best fit parameters: %s' % (params,))
211
212     if _package_config['matplotlib']:
213         if not _matplotlib:
214             raise _matplotlib_import_error
215         figure = _matplotlib_pyplot.figure()
216         axes = figure.add_subplot(1, 1, 1)
217         axes.hold(True)
218         timestamp = _time.strftime('%H%M%S')
219         axes.set_title('surf_pos %5g %5g %5g %5g' % tuple(params))
220         def plot_dict(d, label):
221             _pylab.plot(d["z"], d["deflection"], '.',label=label)
222         for n,name in [('ret1', 'first retract'),
223                        ('mtpod', 'single step in'),
224                        ('ret2', 'second retract'),
225                        ('approach', 'main approach'),
226                        ('ret3', 'return to start')]:
227             axes.plot(ddict[n]['z'], ddict[n]['deflection'], label=name)
228         def fit_fn(x, params):
229             if x <= params[2]:
230                 return params[0] + params[1]*x
231             else:
232                 return (params[0] + params[1]*params[2]
233                         + params[3]*(x-params[2]))
234         axes.plot([start_pos, params[2], final_pos],
235                   [fit_fn(start_pos, params), fit_fn(params[2], params),
236                    fit_fn(final_pos, params)], '-',label='fit')
237         #_pylab.legend(loc='best')
238         figure.canvas.draw()
239         if hasattr(figure, 'show'):
240             figure.show()
241         if not _matplotlib.is_interactive():
242             _matplotlib_pyplot.show()
243
244     # check that the fit is reasonable
245     # params[1] is slope in non-contact region
246     # params[2] is kink position
247     # params[3] is slope in contact region
248     if kink_window is None:
249         kink_window = int(0.02*(final_pos-start_pos))
250
251     if abs(params[1]*min_slope_ratio) > abs(params[3]):
252         raise FlatFit(left_slope=params[1], right_slope=params[3])
253     if params[2] < start_pos+kink_window:
254         raise EdgeKink(kink=params[2], edge=start_pos, window=kink_window)
255     if params[2] > final_pos-kink_window:
256         raise EdgeKink(kink=params[2], edge=final_pos, window=kink_window)
257     _LOG.info('surface position %s' % params[2])
258     if return_all_parameters:
259         return params
260     return params[2]
261
262 def get_surface_position(piezo, axis_name, max_deflection, **kwargs):
263     ddict = get_surface_position_data(piezo, axis_name, max_deflection)
264     return analyze_surface_position_data(ddict, **kwargs)