Fix approacing -> approaching typo in cut and curve plugins
[hooke.git] / hooke / plugin / curve.py
1 # Copyright (C) 2008-2010 Alberto Gomez-Casado
2 #                         Fabrizio Benedetti
3 #                         Massimo Sandal <devicerandom@gmail.com>
4 #                         W. Trevor King <wking@drexel.edu>
5 #
6 # This file is part of Hooke.
7 #
8 # Hooke is free software: you can redistribute it and/or
9 # modify it under the terms of the GNU Lesser General Public
10 # License as published by the Free Software Foundation, either
11 # version 3 of the License, or (at your option) any later version.
12 #
13 # Hooke is distributed in the hope that it will be useful,
14 # but WITHOUT ANY WARRANTY; without even the implied warranty of
15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
16 # GNU Lesser General Public License for more details.
17 #
18 # You should have received a copy of the GNU Lesser General Public
19 # License along with Hooke.  If not, see
20 # <http://www.gnu.org/licenses/>.
21
22 """The ``curve`` module provides :class:`CurvePlugin` and several
23 associated :class:`hooke.command.Command`\s for handling
24 :mod:`hooke.curve` classes.
25 """
26
27 import numpy
28
29 from ..command import Command, Argument, Failure
30 from ..curve import Data
31 from ..plugin import Builtin
32 from ..plugin.playlist import current_playlist_callback
33 from ..util.calculus import derivative
34 from ..util.fft import unitary_avg_power_spectrum
35
36
37 class CurvePlugin (Builtin):
38     def __init__(self):
39         super(CurvePlugin, self).__init__(name='curve')
40         self._commands = [
41             InfoCommand(self), ExportCommand(self),
42             DifferenceCommand(self), DerivativeCommand(self),
43             PowerSpectrumCommand(self)]
44
45
46 # Define common or complicated arguments
47
48 def current_curve_callback(hooke, command, argument, value):
49     if value != None:
50         return value
51     playlist = current_playlist_callback(hooke, command, argument, value)
52     curve = playlist.current()
53     if curve == None:
54         raise Failure('No curves in %s' % playlist)
55     return curve
56
57 CurveArgument = Argument(
58     name='curve', type='curve', callback=current_curve_callback,
59     help="""
60 :class:`hooke.curve.Curve` to act on.  Defaults to the current curve
61 of the current playlist.
62 """.strip())
63
64
65 # Define commands
66
67 class InfoCommand (Command):
68     """Get selected information about a :class:`hooke.curve.Curve`.
69     """
70     def __init__(self, plugin):
71         args = [
72             CurveArgument,                    
73             Argument(name='all', type='bool', default=False, count=1,
74                      help='Get all curve information.'),
75             ]
76         self.fields = ['name', 'path', 'experiment', 'driver', 'filetype', 'note',
77                        'blocks', 'block sizes']
78         for field in self.fields:
79             args.append(Argument(
80                     name=field, type='bool', default=False, count=1,
81                     help='Get curve %s' % field))
82         super(InfoCommand, self).__init__(
83             name='curve info', arguments=args,
84             help=self.__doc__, plugin=plugin)
85
86     def _run(self, hooke, inqueue, outqueue, params):
87         fields = {}
88         for key in self.fields:
89             fields[key] = params[key]
90         if reduce(lambda x,y: x and y, fields.values()) == False:
91             params['all'] = True # No specific fields set, default to 'all'
92         if params['all'] == True:
93             for key in self.fields:
94                 fields[key] = True
95         lines = []
96         for key in self.fields:
97             if fields[key] == True:
98                 get = getattr(self, '_get_%s' % key.replace(' ', '_'))
99                 lines.append('%s: %s' % (key, get(params['curve'])))
100         outqueue.put('\n'.join(lines))
101
102     def _get_name(self, curve):
103         return curve.name
104
105     def _get_path(self, curve):
106         return curve.path
107
108     def _get_experiment(self, curve):
109         return curve.info.get('experiment', None)
110
111     def _get_driver(self, curve):
112         return curve.driver
113
114     def _get_filetype(self, curve):
115         return curve.info.get('filetype', None)
116
117     def _get_note(self, curve):
118         return curve.info.get('note', None)
119                               
120     def _get_blocks(self, curve):
121         return len(curve.data)
122
123     def _get_block_sizes(self, curve):
124         return [block.shape for block in curve.data]
125
126 class ExportCommand (Command):
127     """Export a :class:`hooke.curve.Curve` data block as TAB-delimeted
128     ASCII text.
129     """
130     def __init__(self, plugin):
131         super(ExportCommand, self).__init__(
132             name='export block',
133             arguments=[
134                 CurveArgument,
135                 Argument(name='block', aliases=['set'], type='int', default=0,
136                          help="""
137 Data block to save.  For an approach/retract force curve, `0` selects
138 the approaching curve and `1` selects the retracting curve.
139 """.strip()),
140                 Argument(name='output', type='file', default='curve.dat',
141                          help="""
142 File name for the output data.  Defaults to 'curve.dat'
143 """.strip()),
144                 ],
145             help=self.__doc__, plugin=plugin)
146
147     def _run(self, hooke, inqueue, outqueue, params):
148         data = params['curve'].data[int(params['block'])] # HACK, int() should be handled by ui
149         numpy.savetxt(params['output'], data, delimiter='\t')
150
151 class DifferenceCommand (Command):
152     """Calculate the derivative (actually, the discrete differentiation)
153     of a curve data block.
154
155     See :func:`hooke.util.calculus.derivative` for implementation
156     details.
157     """
158     def __init__(self, plugin):
159         super(DifferenceCommand, self).__init__(
160             name='block difference',
161             arguments=[
162                 CurveArgument,
163                 Argument(name='block one', aliases=['set one'], type='int',
164                          default=1,
165                          help="""
166 Block A in A-B.  For an approach/retract force curve, `0` selects the
167 approaching curve and `1` selects the retracting curve.
168 """.strip()),
169                 Argument(name='block two', aliases=['set two'], type='int',
170                          default=0,
171                          help='Block B in A-B.'),
172                 Argument(name='x column', type='int', default=0,
173                          help="""
174 Column of data block to differentiate with respect to.
175 """.strip()),
176                 Argument(name='f column', type='int', default=1,
177                          help="""
178 Column of data block to differentiate.
179 """.strip()),
180                 ],
181             help=self.__doc__, plugin=plugin)
182
183     def _run(self, hooke, inqueue, outqueue, params):
184         a = params['curve'].data[params['block one']]
185         b = params['curve'].data[params['block two']]
186         assert a[:,params['x column']] == b[:,params['x column']]
187         out = Data((a.shape[0],2), dtype=a.dtype)
188         out[:,0] = a[:,params['x column']]
189         out[:,1] = a[:,params['f column']] - b[:,params['f column']]
190         outqueue.put(out)
191
192 class DerivativeCommand (Command):
193     """Calculate the difference between two blocks of data.
194     """
195     def __init__(self, plugin):
196         super(DerivativeCommand, self).__init__(
197             name='block derivative',
198             arguments=[
199                 CurveArgument,
200                 Argument(name='block', aliases=['set'], type='int', default=0,
201                          help="""
202 Data block to differentiate.  For an approach/retract force curve, `0`
203 selects the approaching curve and `1` selects the retracting curve.
204 """.strip()),
205                 Argument(name='x column', type='int', default=0,
206                          help="""
207 Column of data block to differentiate with respect to.
208 """.strip()),
209                 Argument(name='f column', type='int', default=1,
210                          help="""
211 Column of data block to differentiate.
212 """.strip()),
213                 Argument(name='weights', type='dict', default={-1:-0.5, 1:0.5},
214                          help="""
215 Weighting scheme dictionary for finite differencing.  Defaults to
216 central differencing.
217 """.strip()),
218                 ],
219             help=self.__doc__, plugin=plugin)
220
221     def _run(self, hooke, inqueue, outqueue, params):
222         data = params['curve'].data[params['block']]
223         outqueue.put(derivative(
224                 block, x_col=params['x column'], f_col=params['f column'],
225                 weights=params['weights']))
226
227 class PowerSpectrumCommand (Command):
228     """Calculate the power spectrum of a data block.
229     """
230     def __init__(self, plugin):
231         super(PowerSpectrumCommand, self).__init__(
232             name='block power spectrum',
233             arguments=[
234                 CurveArgument,
235                 Argument(name='block', aliases=['set'], type='int', default=0,
236                          help="""
237 Data block to act on.  For an approach/retract force curve, `0`
238 selects the approaching curve and `1` selects the retracting curve.
239 """.strip()),
240                 Argument(name='f column', type='int', default=1,
241                          help="""
242 Column of data block to differentiate with respect to.
243 """.strip()),
244                 Argument(name='freq', type='float', default=1.0,
245                          help="""
246 Sampling frequency.
247 """.strip()),
248                 Argument(name='chunk size', type='int', default=2048,
249                          help="""
250 Number of samples per chunk.  Use a power of two.
251 """.strip()),
252                 Argument(name='overlap', type='bool', default=False,
253                          help="""
254 If `True`, each chunk overlaps the previous chunk by half its length.
255 Otherwise, the chunks are end-to-end, and not overlapping.
256 """.strip()),
257                 ],
258             help=self.__doc__, plugin=plugin)
259
260     def _run(self, hooke, inqueue, outqueue, params):
261         data = params['curve'].data[params['block']]
262         outqueue.put(unitary_avg_power_spectrum(
263                 data[:,params['f column']], freq=params['freq'],
264                 chunk_size=params['chunk size'],
265                 overlap=params['overlap']))