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             GetCommand(self), 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 GetCommand (Command):
  68     """Return a :class:`hooke.curve.Curve`.
  69     """
  70     def __init__(self, plugin):
  71         super(GetCommand, self).__init__(
  72             name='get curve', arguments=[CurveArgument],
  73             help=self.__doc__, plugin=plugin)
  74
  75     def _run(self, hooke, inqueue, outqueue, params):
  76         outqueue.put(params['curve'])
  77
  78 class InfoCommand (Command):
  79     """Get selected information about a :class:`hooke.curve.Curve`.
  80     """
  81     def __init__(self, plugin):
  82         args = [
  83             CurveArgument,
  84             Argument(name='all', type='bool', default=False, count=1,
  85                      help='Get all curve information.'),
  86             ]
  87         self.fields = ['name', 'path', 'experiment', 'driver', 'filetype', 'note',
  88                        'blocks', 'block sizes']
  89         for field in self.fields:
  90             args.append(Argument(
  91                     name=field, type='bool', default=False, count=1,
  92                     help='Get curve %s' % field))
  93         super(InfoCommand, self).__init__(
  94             name='curve info', arguments=args,
  95             help=self.__doc__, plugin=plugin)
  96
  97     def _run(self, hooke, inqueue, outqueue, params):
  98         fields = {}
  99         for key in self.fields:
 100             fields[key] = params[key]
 101         if reduce(lambda x,y: x and y, fields.values()) == False:
 102             params['all'] = True # No specific fields set, default to 'all'
 103         if params['all'] == True:
 104             for key in self.fields:
 105                 fields[key] = True
 106         lines = []
 107         for key in self.fields:
 108             if fields[key] == True:
 109                 get = getattr(self, '_get_%s' % key.replace(' ', '_'))
 110                 lines.append('%s: %s' % (key, get(params['curve'])))
 111         outqueue.put('\n'.join(lines))
 112
 113     def _get_name(self, curve):
 114         return curve.name
 115
 116     def _get_path(self, curve):
 117         return curve.path
 118
 119     def _get_experiment(self, curve):
 120         return curve.info.get('experiment', None)
 121
 122     def _get_driver(self, curve):
 123         return curve.driver
 124
 125     def _get_filetype(self, curve):
 126         return curve.info.get('filetype', None)
 127
 128     def _get_note(self, curve):
 129         return curve.info.get('note', None)
 130
 131     def _get_blocks(self, curve):
 132         return len(curve.data)
 133
 134     def _get_block_sizes(self, curve):
 135         return [block.shape for block in curve.data]
 136
 137 class ExportCommand (Command):
 138     """Export a :class:`hooke.curve.Curve` data block as TAB-delimeted
 139     ASCII text.
 140
 141     A "#" prefixed header will optionally appear at the beginning of
 142     the file naming the columns.
 143     """
 144     def __init__(self, plugin):
 145         super(ExportCommand, self).__init__(
 146             name='export block',
 147             arguments=[
 148                 CurveArgument,
 149                 Argument(name='block', aliases=['set'], type='int', default=0,
 150                          help="""
 151 Data block to save.  For an approach/retract force curve, `0` selects
 152 the approaching curve and `1` selects the retracting curve.
 153 """.strip()),
 154                 Argument(name='output', type='file', default='curve.dat',
 155                          help="""
 156 File name for the output data.  Defaults to 'curve.dat'
 157 """.strip()),
 158                 Argument(name='header', type='bool', default=True,
 159                          help="""
 160 True if you want the column-naming header line.
 161 """.strip()),
 162                 ],
 163             help=self.__doc__, plugin=plugin)
 164
 165     def _run(self, hooke, inqueue, outqueue, params):
 166         data = params['curve'].data[int(params['block'])] # HACK, int() should be handled by ui
 167
 168         f = open(params['output'], 'w')
 169         if params['header'] == True:
 170             f.write('# %s \n' % ('\t'.join(data.info['columns'])))
 171         numpy.savetxt(f, data, delimiter='\t')
 172         f.close()
 173
 174 class DifferenceCommand (Command):
 175     """Calculate the derivative (actually, the discrete differentiation)
 176     of a curve data block.
 177
 178     See :func:`hooke.util.calculus.derivative` for implementation
 179     details.
 180     """
 181     def __init__(self, plugin):
 182         super(DifferenceCommand, self).__init__(
 183             name='block difference',
 184             arguments=[
 185                 CurveArgument,
 186                 Argument(name='block one', aliases=['set one'], type='int',
 187                          default=1,
 188                          help="""
 189 Block A in A-B.  For an approach/retract force curve, `0` selects the
 190 approaching curve and `1` selects the retracting curve.
 191 """.strip()),
 192                 Argument(name='block two', aliases=['set two'], type='int',
 193                          default=0,
 194                          help='Block B in A-B.'),
 195                 Argument(name='x column', type='int', default=0,
 196                          help="""
 197 Column of data block to differentiate with respect to.
 198 """.strip()),
 199                 Argument(name='f column', type='int', default=1,
 200                          help="""
 201 Column of data block to differentiate.
 202 """.strip()),
 203                 ],
 204             help=self.__doc__, plugin=plugin)
 205
 206     def _run(self, hooke, inqueue, outqueue, params):
 207         a = params['curve'].data[params['block one']]
 208         b = params['curve'].data[params['block two']]
 209         assert a[:,params['x column']] == b[:,params['x column']]
 210         out = Data((a.shape[0],2), dtype=a.dtype)
 211         out[:,0] = a[:,params['x column']]
 212         out[:,1] = a[:,params['f column']] - b[:,params['f column']]
 213         outqueue.put(out)
 214
 215 class DerivativeCommand (Command):
 216     """Calculate the difference between two blocks of data.
 217     """
 218     def __init__(self, plugin):
 219         super(DerivativeCommand, self).__init__(
 220             name='block derivative',
 221             arguments=[
 222                 CurveArgument,
 223                 Argument(name='block', aliases=['set'], type='int', default=0,
 224                          help="""
 225 Data block to differentiate.  For an approach/retract force curve, `0`
 226 selects the approaching curve and `1` selects the retracting curve.
 227 """.strip()),
 228                 Argument(name='x column', type='int', default=0,
 229                          help="""
 230 Column of data block to differentiate with respect to.
 231 """.strip()),
 232                 Argument(name='f column', type='int', default=1,
 233                          help="""
 234 Column of data block to differentiate.
 235 """.strip()),
 236                 Argument(name='weights', type='dict', default={-1:-0.5, 1:0.5},
 237                          help="""
 238 Weighting scheme dictionary for finite differencing.  Defaults to
 239 central differencing.
 240 """.strip()),
 241                 ],
 242             help=self.__doc__, plugin=plugin)
 243
 244     def _run(self, hooke, inqueue, outqueue, params):
 245         data = params['curve'].data[params['block']]
 246         outqueue.put(derivative(
 247                 block, x_col=params['x column'], f_col=params['f column'],
 248                 weights=params['weights']))
 249
 250 class PowerSpectrumCommand (Command):
 251     """Calculate the power spectrum of a data block.
 252     """
 253     def __init__(self, plugin):
 254         super(PowerSpectrumCommand, self).__init__(
 255             name='block power spectrum',
 256             arguments=[
 257                 CurveArgument,
 258                 Argument(name='block', aliases=['set'], type='int', default=0,
 259                          help="""
 260 Data block to act on.  For an approach/retract force curve, `0`
 261 selects the approaching curve and `1` selects the retracting curve.
 262 """.strip()),
 263                 Argument(name='f column', type='int', default=1,
 264                          help="""
 265 Column of data block to differentiate with respect to.
 266 """.strip()),
 267                 Argument(name='freq', type='float', default=1.0,
 268                          help="""
 269 Sampling frequency.
 270 """.strip()),
 271                 Argument(name='chunk size', type='int', default=2048,
 272                          help="""
 273 Number of samples per chunk.  Use a power of two.
 274 """.strip()),
 275                 Argument(name='overlap', type='bool', default=False,
 276                          help="""
 277 If `True`, each chunk overlaps the previous chunk by half its length.
 278 Otherwise, the chunks are end-to-end, and not overlapping.
 279 """.strip()),
 280                 ],
 281             help=self.__doc__, plugin=plugin)
 282
 283     def _run(self, hooke, inqueue, outqueue, params):
 284         data = params['curve'].data[params['block']]
 285         outqueue.put(unitary_avg_power_spectrum(
 286                 data[:,params['f column']], freq=params['freq'],
 287                 chunk_size=params['chunk size'],
 288                 overlap=params['overlap']))