Standardize matplotlib rendering on figure.canvas.draw() and figure.show().
[pypiezo.git] / pypiezo / wiggle.py
1 # Copyright (C) 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 for wiggling a the piezo near the surface.
18
19 This helps you detect interference between the laser bouncing off the
20 tip and the surface.  The wiggling continues until you stop it, which
21 gives you feedback while you reposition the laser to minimize the
22 interference.  One day we'll all have fancy AFMs with software control
23 over the laser alignment, but until then, we've got this module
24 helping you with your thumbscrews.
25 """
26
27 import time as _time
28
29 import numpy as _numpy
30
31 try:
32     import h5py as _h5py
33     import h5config as _h5config
34     from h5config.storage.hdf5 import HDF5_Storage as _HDF5_Storage
35     from h5config.storage.hdf5 import h5_create_group as _h5_create_group
36 except ImportError, e:
37     _h5py = None
38     _h5py_import_error = e
39
40 try:
41     import matplotlib as _matplotlib
42     import matplotlib.pyplot as _matplotlib_pyplot
43 except (ImportError, RuntimeError), e:
44     _matplotlib = None
45     _matplotlib_import_error = e
46
47 from curses_check_for_keypress import CheckForKeypress as _CheckForKeypress
48
49 from . import LOG as _LOG
50 from . import base as _base
51 from . import package_config as _package_config
52
53
54 def _setup_config(piezo, config):
55     if config['wavelength'] and config['amplitude']:
56         log_string = \
57             'use either laser_wavelength or amplitude, but not both'
58         _LOG.warn(log_string)
59
60     if None in (config['amplitude'], config['offset']):
61         output_axis = piezo.axis_by_name(config['axis'])
62         maxdata = output_axis.axis_channel.get_maxdata()
63         midpoint = int(maxdata/2)
64         if config['offset'] is None:
65             offset = midpoint
66             _LOG.debug(('generated offset for interference wiggle: {}'
67                         ).format(config['offset']))
68         if config['amplitude'] is None:
69             if config['offset'] <= midpoint:
70                 max_amplitude = int(config['offset'])
71             else:
72                 max_amplitude = int(maxdata-config['offset'])
73             offset_meters = _base.convert_bits_to_meters(
74                 output_axis.config, config['offset'])
75             if config['wavelength'] is None:
76                 config['amplitude'] = 0.5*max_amplitude
77             else:
78                 bit_wavelength = _base.convert_meters_to_bits(
79                     output_axis.config,
80                     offset_meters + config['wavelength']
81                     ) - config['offset']
82                 config['amplitude'] = 2*bit_wavelength
83             _LOG.debug(('generated amplitude for interference wiggle: {}'
84                         ).format(config['amplitude']))
85             if config['amplitude'] > max_amplitude:
86                 raise ValueError(
87                     'no room for a two wavelength wiggle ({} > {})'.format(
88                         config['amplitude'], max_amplitude))
89
90 def _construct_output(piezo, config):
91     n = config['samples']  # samples in a single oscillation
92     out = (config['amplitude']
93            * _numpy.sin(_numpy.arange(2*n)*2*_numpy.pi/n)
94            + config['offset'])
95     # 2*n for 2 periods, so you can judge precision
96     out = out.reshape((len(out), 1)).astype(
97         piezo.channel_dtype(config['axis'], direction='output'))
98     return out
99
100 def _setup_datafile(filename, group, piezo, config, output):
101     if not _h5py:
102         raise _h5py_import_error
103     output_axis = afm.piezo.axis_by_name(config['axis'])
104     input_channel = afm.piezo.input_channel_by_name(config['input'])
105     with _h5py.File(filename, 'w') as f:
106         cwg = _h5_create_group(f, group)
107         storage = _HDF5_Storage()
108         for config_,key in [
109             (config, 'config/wiggle'),
110             (output.axis.config,
111              'config/{}/axis'.format(config['axis'])),
112             (input_channel.config,
113              'config/{}/channel'.format(config['input']))]:
114             if config_ is None:
115                 continue
116             config_cwg = _h5_create_group(cwg, key)
117             storage.save(config=config, group=config_cwg)
118         cwg['wiggle/raw/{}'.format(config['axis'])] = output
119
120 def _update_datafile(filename, group, config, cycle, data):
121     timestamp = ('{0}-{1:02d}-{2:02d}T{3:02d}-{4:02d}-{5:02d}'
122                  ).format(*_time.localtime())
123     with _h5py.File(filename, 'a') as f:
124         wiggle_group = _h5_create_group(f, group)
125         cwg = _h5_create_group(
126             wiggle_group, 'wiggle/{}'.format(cycle))
127         cwg['time'] = timestamp
128         cwg['raw/{}'.format(config['input'])] = data
129
130 def _setup_plot(piezo, config, output):
131     if not _matplotlib:
132         raise _matplotlib_import_error
133     _matplotlib.interactive(True)
134     figure = _matplotlib_pyplot.figure()
135     axes = figure.add_subplot(1, 1, 1)
136     axes.hold(False)
137     timestamp = _time.strftime('%H%M%S')
138     axes.set_title('wiggle for interference %s' % timestamp)
139     plot = axes.plot(output, output, 'b.-')
140     figure.canvas.draw()
141     figure.show()
142     if not _matplotlib.is_interactive():
143         _matplotlib_pyplot.show()
144     return (figure, axes, plot)
145
146 def _update_plot(figure, axes, plot, cycle, data):
147     plot[0].set_ydata(data[:,0])
148     axes.set_ylim([data.min(), data.max()])
149     _matplotlib_pyplot.draw()
150
151 def _run_wiggles(piezo, config, figure, axes, plot, output, filename=None,
152                  group='/', keypress_test_mode=False):
153     scan_frequency = config['frequency'] * config['samples']
154     cycle = 0
155     c = _CheckForKeypress(test_mode=keypress_test_mode)
156     while c.input() == None:
157         # input will need processing for multi-segment AFMs...
158         data = piezo.ramp(
159             output, scan_frequency, output_names=[config['axis']],
160             input_names=[config['input']])
161         _LOG.debug('completed a wiggle round')
162         if filename:
163             _update_datafile(
164                 filename=filename, group=group, config=config,
165                 cycle=cycle, data=data)
166         if plot:
167             _update_plot(
168                 figure=figure, axes=axes, plot=plot, cycle=cycle, data=data)
169         cycle += 1
170
171 def wiggle_for_interference(
172     piezo, config, plot=True, filename=None, group='/',
173     keypress_test_mode=False):
174     """Output a sine wave and measure interference.
175
176     With a poorly focused or aligned laser, leaked laser light
177     reflecting off the surface may interfere with the light
178     reflected off the cantilever, causing distance-dependent
179     interference with a period roughly half the laser's
180     wavelength.  This method wiggles the cantilever near the
181     surface and monitors the magnitude of deflection oscillation,
182     allowing the operator to adjust the laser alignment in order
183     to minimize the interference.
184
185     Modern commercial AFMs with computer-aligned lasers must do
186     something like this automatically.
187     """
188     if _package_config['matplotlib']:
189         plot = True
190     _setup_config(piezo=piezo, config=config)
191     _LOG.debug('oscillate for interference wiggle ({})'.format(config))
192     output = _construct_output(piezo=piezo, config=config)
193
194     if filename:
195         _setup_datafile(
196             filename=filename, group=group, piezo=piezo, config=config,
197             output=output)
198     if plot:
199         interactive = _matplotlib.is_interactive()
200         figure,axes,plot_ = _setup_plot(
201             piezo=piezo, config=config, output=output)
202     else:
203         figure = axes = plot_ = None
204     _run_wiggles(
205         piezo=piezo, config=config, figure=figure, axes=axes, plot=plot_,
206         output=output, filename=filename, group=group,
207         keypress_test_mode=keypress_test_mode)
208     if plot:
209         _matplotlib.interactive(interactive)
210     piezo.last_output[config['axis']] = output[-1,0]
211     _LOG.debug('interference wiggle complete')