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 from ..command import Command, Argument, Failure
  28 from ..curve import Data
  29 from ..plugin import Builtin
  30 from ..plugin.playlist import current_playlist_callback
  31 from ..util.calculus import derivative
  32 from ..util.fft import unitary_avg_power_spectrum
  33
  34
  35 class CurvePlugin (Builtin):
  36     def __init__(self):
  37         super(CurvePlugin, self).__init__(name='curve')
  38         self._commands = [
  39             InfoCommand(), ExportCommand(),
  40             DifferenceCommand(), DerivativeCommand(),
  41             PowerSpectrumCommand()]
  42         self._setup_commands()
  43
  44
  45 # Define common or complicated arguments
  46
  47 def current_curve_callback(hooke, command, argument, value):
  48     if value != None:
  49         return value
  50     playlist = current_playlist_callback(hooke, command, argument, value)
  51     curve = playlist.current()
  52     if curve == None:
  53         raise Failure('No curves in %s' % playlist)
  54     return curve
  55
  56 CurveArgument = Argument(
  57     name='curve', type='curve', callback=current_curve_callback,
  58     help="""
  59 :class:`hooke.curve.Curve` to act on.  Defaults to the current curve
  60 of the current playlist.
  61 """.strip())
  62
  63
  64 # Define commands
  65
  66 class InfoCommand (Command):
  67     """Get selected information about a :class:`hooke.curve.Curve`.
  68     """
  69     def __init__(self):
  70         args = [
  71             CurveArgument,
  72             Argument(name='all', type='bool', default=False, count=1,
  73                      help='Get all curve information.'),
  74             ]
  75         self.fields = ['name', 'path', 'experiment', 'driver', 'filetype', 'note',
  76                        'blocks', 'block sizes']
  77         for field in self.fields:
  78             args.append(Argument(
  79                     name=field, type='bool', default=False, count=1,
  80                     help='Get curve %s' % field))
  81         super(InfoCommand, self).__init__(
  82             name='curve info', arguments=args, help=self.__doc__)
  83
  84     def _run(self, hooke, inqueue, outqueue, params):
  85         fields = {}
  86         for key in self.fields:
  87             fields[key] = params[key]
  88         if reduce(lambda x,y: x and y, fields.values()) == False:
  89             params['all'] = True # No specific fields set, default to 'all'
  90         if params['all'] == True:
  91             for key in self.fields:
  92                 fields[key] = True
  93         lines = []
  94         for key in self.fields:
  95             if fields[key] == True:
  96                 get = getattr(self, '_get_%s' % key.replace(' ', '_'))
  97                 lines.append('%s: %s' % (key, get(params['curve'])))
  98         outqueue.put('\n'.join(lines))
  99
 100     def _get_name(self, curve):
 101         return curve.name
 102
 103     def _get_path(self, curve):
 104         return curve.path
 105
 106     def _get_experiment(self, curve):
 107         return curve.info.get('experiment', None)
 108
 109     def _get_driver(self, curve):
 110         return curve.driver
 111
 112     def _get_filetype(self, curve):
 113         return curve.info.get('filetype', None)
 114
 115     def _get_note(self, curve):
 116         return curve.info.get('note', None)
 117
 118     def _get_blocks(self, curve):
 119         return len(curve.data)
 120
 121     def _get_block_sizes(self, curve):
 122         return [block.shape for block in curve.data]
 123
 124 class ExportCommand (Command):
 125     """Export a :class:`hooke.curve.Curve` data block as TAB-delimeted
 126     ASCII text.
 127     """
 128     def __init__(self):
 129         super(ExportCommand, self).__init__(
 130             name='export block',
 131             arguments=[
 132                 CurveArgument,
 133                 Argument(name='block', aliases=['set'], type='int', default=0,
 134                          help="""
 135 Data block to save.  For an approach/retract force curve, `0` selects
 136 the approacing curve and `1` selects the retracting curve.
 137 """.strip()),
 138                 Argument(name='output', type='file', default='curve.dat',
 139                          help="""
 140 File name for the output data.  Defaults to 'curve.dat'
 141 """.strip()),
 142                 ],
 143             help=self.__doc__)
 144
 145     def _run(self, hooke, inqueue, outqueue, params):
 146         data = params['curve'].data[params['block']]
 147         f = open(params['output'], 'w')
 148         data.tofile(f, sep='\t')
 149         f.close()
 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):
 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 approacing 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__)
 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):
 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 approacing 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__)
 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):
 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 approacing 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__)
 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']))