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 modify it
   9 # under the terms of the GNU Lesser General Public License as
  10 # published by the Free Software Foundation, either version 3 of the
  11 # License, or (at your option) any later version.
  12 #
  13 # Hooke is distributed in the hope that it will be useful, but WITHOUT
  14 # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  15 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General
  16 # 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 from ..util.si import ppSI, join_data_label, split_data_label
  36
  37
  38
  39 class CurvePlugin (Builtin):
  40     def __init__(self):
  41         super(CurvePlugin, self).__init__(name='curve')
  42         self._commands = [
  43             GetCommand(self), InfoCommand(self), DeltaCommand(self),
  44             ExportCommand(self), DifferenceCommand(self),
  45             DerivativeCommand(self), PowerSpectrumCommand(self)]
  46
  47
  48 # Define common or complicated arguments
  49
  50 def current_curve_callback(hooke, command, argument, value):
  51     if value != None:
  52         return value
  53     playlist = current_playlist_callback(hooke, command, argument, value)
  54     curve = playlist.current()
  55     if curve == None:
  56         raise Failure('No curves in %s' % playlist)
  57     return curve
  58
  59 CurveArgument = Argument(
  60     name='curve', type='curve', callback=current_curve_callback,
  61     help="""
  62 :class:`hooke.curve.Curve` to act on.  Defaults to the current curve
  63 of the current playlist.
  64 """.strip())
  65
  66
  67 # Define commands
  68
  69 class GetCommand (Command):
  70     """Return a :class:`hooke.curve.Curve`.
  71     """
  72     def __init__(self, plugin):
  73         super(GetCommand, self).__init__(
  74             name='get curve', arguments=[CurveArgument],
  75             help=self.__doc__, plugin=plugin)
  76
  77     def _run(self, hooke, inqueue, outqueue, params):
  78         outqueue.put(params['curve'])
  79
  80 class InfoCommand (Command):
  81     """Get selected information about a :class:`hooke.curve.Curve`.
  82     """
  83     def __init__(self, plugin):
  84         args = [
  85             CurveArgument,
  86             Argument(name='all', type='bool', default=False, count=1,
  87                      help='Get all curve information.'),
  88             ]
  89         self.fields = ['name', 'path', 'experiment', 'driver', 'filetype', 'note',
  90                        'blocks', 'block sizes']
  91         for field in self.fields:
  92             args.append(Argument(
  93                     name=field, type='bool', default=False, count=1,
  94                     help='Get curve %s' % field))
  95         super(InfoCommand, self).__init__(
  96             name='curve info', arguments=args,
  97             help=self.__doc__, plugin=plugin)
  98
  99     def _run(self, hooke, inqueue, outqueue, params):
 100         fields = {}
 101         for key in self.fields:
 102             fields[key] = params[key]
 103         if reduce(lambda x,y: x and y, fields.values()) == False:
 104             params['all'] = True # No specific fields set, default to 'all'
 105         if params['all'] == True:
 106             for key in self.fields:
 107                 fields[key] = True
 108         lines = []
 109         for key in self.fields:
 110             if fields[key] == True:
 111                 get = getattr(self, '_get_%s' % key.replace(' ', '_'))
 112                 lines.append('%s: %s' % (key, get(params['curve'])))
 113         outqueue.put('\n'.join(lines))
 114
 115     def _get_name(self, curve):
 116         return curve.name
 117
 118     def _get_path(self, curve):
 119         return curve.path
 120
 121     def _get_experiment(self, curve):
 122         return curve.info.get('experiment', None)
 123
 124     def _get_driver(self, curve):
 125         return curve.driver
 126
 127     def _get_filetype(self, curve):
 128         return curve.info.get('filetype', None)
 129
 130     def _get_note(self, curve):
 131         return curve.info.get('note', None)
 132
 133     def _get_blocks(self, curve):
 134         return len(curve.data)
 135
 136     def _get_block_sizes(self, curve):
 137         return [block.shape for block in curve.data]
 138
 139
 140 class DeltaCommand (Command):
 141     """Get distance information between two points.
 142
 143     With two points A and B, the returned distances are A-B.
 144     """
 145     def __init__(self, plugin):
 146         super(DeltaCommand, self).__init__(
 147             name='delta',
 148             arguments=[
 149                 CurveArgument,
 150                 Argument(name='block', type='int', default=0,
 151                     help="""
 152 Data block that points are selected from.  For an approach/retract
 153 force curve, `0` selects the approaching curve and `1` selects the
 154 retracting curve.
 155 """.strip()),
 156                 Argument(name='point', type='point', optional=False, count=2,
 157                          help="""
 158 Indicies of points bounding the selected data.
 159 """.strip()),
 160                 Argument(name='SI', type='bool', default=False,
 161                          help="""
 162 Return distances in SI notation.
 163 """.strip())
 164                 ],
 165             help=self.__doc__, plugin=plugin)
 166
 167     def _run(self, hooke, inqueue, outqueue, params):
 168         data = params['curve'].data[params['block']]
 169         As = data[params['point'][0],:]
 170         Bs = data[params['point'][1],:]
 171         ds = [A-B for A,B in zip(As, Bs)]
 172         if params['SI'] == False:
 173             out = [(name, d) for name,d in zip(data.info['columns'], ds)]
 174         else:
 175             out = []
 176             for name,d in zip(data.info['columns'], ds):
 177                 n,units = split_data_label(name)
 178                 out.append(
 179                   (n, ppSI(value=d, unit=units, decimals=2)))
 180         outqueue.put(out)
 181
 182
 183 class ExportCommand (Command):
 184     """Export a :class:`hooke.curve.Curve` data block as TAB-delimeted
 185     ASCII text.
 186
 187     A "#" prefixed header will optionally appear at the beginning of
 188     the file naming the columns.
 189     """
 190     def __init__(self, plugin):
 191         super(ExportCommand, self).__init__(
 192             name='export block',
 193             arguments=[
 194                 CurveArgument,
 195                 Argument(name='block', aliases=['set'], type='int', default=0,
 196                          help="""
 197 Data block to save.  For an approach/retract force curve, `0` selects
 198 the approaching curve and `1` selects the retracting curve.
 199 """.strip()),
 200                 Argument(name='output', type='file', default='curve.dat',
 201                          help="""
 202 File name for the output data.  Defaults to 'curve.dat'
 203 """.strip()),
 204                 Argument(name='header', type='bool', default=True,
 205                          help="""
 206 True if you want the column-naming header line.
 207 """.strip()),
 208                 ],
 209             help=self.__doc__, plugin=plugin)
 210
 211     def _run(self, hooke, inqueue, outqueue, params):
 212         data = params['curve'].data[params['block']]
 213
 214         f = open(params['output'], 'w')
 215         if params['header'] == True:
 216             f.write('# %s \n' % ('\t'.join(data.info['columns'])))
 217         numpy.savetxt(f, data, delimiter='\t')
 218         f.close()
 219
 220 class DifferenceCommand (Command):
 221     """Calculate the difference between two blocks of data.
 222     """
 223     def __init__(self, plugin):
 224         super(DifferenceCommand, self).__init__(
 225             name='block difference',
 226             arguments=[
 227                 CurveArgument,
 228                 Argument(name='block one', aliases=['set one'], type='int',
 229                          default=1,
 230                          help="""
 231 Block A in A-B.  For an approach/retract force curve, `0` selects the
 232 approaching curve and `1` selects the retracting curve.
 233 """.strip()),
 234                 Argument(name='block two', aliases=['set two'], type='int',
 235                          default=0,
 236                          help='Block B in A-B.'),
 237                 Argument(name='x column', type='int', default=0,
 238                          help="""
 239 Column of data to return as x values.
 240 """.strip()),
 241                 Argument(name='y column', type='int', default=1,
 242                          help="""
 243 Column of data block to difference.
 244 """.strip()),
 245                 ],
 246             help=self.__doc__, plugin=plugin)
 247
 248     def _run(self, hooke, inqueue, outqueue, params):
 249         a = params['curve'].data[params['block one']]
 250         b = params['curve'].data[params['block two']]
 251         assert a[:,params['x column']] == b[:,params['x column']]
 252         out = Data((a.shape[0],2), dtype=a.dtype)
 253         out[:,0] = a[:,params['x column']]
 254         out[:,1] = a[:,params['y column']] - b[:,params['y column']]
 255         outqueue.put(out)
 256
 257 class DerivativeCommand (Command):
 258     """Calculate the derivative (actually, the discrete differentiation)
 259     of a curve data block.
 260
 261     See :func:`hooke.util.calculus.derivative` for implementation
 262     details.
 263     """
 264     def __init__(self, plugin):
 265         super(DerivativeCommand, self).__init__(
 266             name='block derivative',
 267             arguments=[
 268                 CurveArgument,
 269                 Argument(name='block', aliases=['set'], type='int', default=0,
 270                          help="""
 271 Data block to differentiate.  For an approach/retract force curve, `0`
 272 selects the approaching curve and `1` selects the retracting curve.
 273 """.strip()),
 274                 Argument(name='x column', type='int', default=0,
 275                          help="""
 276 Column of data block to differentiate with respect to.
 277 """.strip()),
 278                 Argument(name='f column', type='int', default=1,
 279                          help="""
 280 Column of data block to differentiate.
 281 """.strip()),
 282                 Argument(name='weights', type='dict', default={-1:-0.5, 1:0.5},
 283                          help="""
 284 Weighting scheme dictionary for finite differencing.  Defaults to
 285 central differencing.
 286 """.strip()),
 287                 ],
 288             help=self.__doc__, plugin=plugin)
 289
 290     def _run(self, hooke, inqueue, outqueue, params):
 291         data = params['curve'].data[params['block']]
 292         outqueue.put(derivative(
 293                 block, x_col=params['x column'], f_col=params['f column'],
 294                 weights=params['weights']))
 295
 296 class PowerSpectrumCommand (Command):
 297     """Calculate the power spectrum of a data block.
 298     """
 299     def __init__(self, plugin):
 300         super(PowerSpectrumCommand, self).__init__(
 301             name='block power spectrum',
 302             arguments=[
 303                 CurveArgument,
 304                 Argument(name='block', aliases=['set'], type='int', default=0,
 305                          help="""
 306 Data block to act on.  For an approach/retract force curve, `0`
 307 selects the approaching curve and `1` selects the retracting curve.
 308 """.strip()),
 309                 Argument(name='column', type='string', optional=False,
 310                          help="""
 311 Name of the data block column containing to-be-transformed data.
 312 """.strip()),
 313                 Argument(name='bounds', type='point', optional=True, count=2,
 314                          help="""
 315 Indicies of points bounding the selected data.
 316 """.strip()),
 317                 Argument(name='freq', type='float', default=1.0,
 318                          help="""
 319 Sampling frequency.
 320 """.strip()),
 321                 Argument(name='freq units', type='string', default='Hz',
 322                          help="""
 323 Units for the sampling frequency.
 324 """.strip()),
 325                 Argument(name='chunk size', type='int', default=2048,
 326                          help="""
 327 Number of samples per chunk.  Use a power of two.
 328 """.strip()),
 329                 Argument(name='overlap', type='bool', default=False,
 330                          help="""
 331 If `True`, each chunk overlaps the previous chunk by half its length.
 332 Otherwise, the chunks are end-to-end, and not overlapping.
 333 """.strip()),
 334                 Argument(name='output block name', type='string',
 335                          default='power spectrum',
 336                          help="""
 337 Name of the new data block (without `of <source block name> <source column name>`) for storing the power spectrum.
 338 """.strip()),
 339                 ],
 340             help=self.__doc__, plugin=plugin)
 341
 342     def _run(self, hooke, inqueue, outqueue, params):
 343         data = params['curve'].data[params['block']]
 344         col = data.info['columns'].index(params['column'])
 345         d = data[:,col]
 346         if bounds != None:
 347             d = d[params['bounds'][0]:params['bounds'][1]]
 348         freq_axis,power = unitary_avg_power_spectrum(
 349             d, freq=params['freq'],
 350             chunk_size=params['chunk size'],
 351             overlap=params['overlap'])
 352         name,data_units = split_data_label(params['column'])
 353         b = Data((len(freq_axis),2), dtype=data.dtype)
 354         b.info['name'] = '%s of %s %s' % (
 355             params['output block name'], data.info['name'], params['column'])
 356         b.info['columns'] = [
 357             join_data_label('frequency axis', params['freq units']),
 358             join_data_label('power density',
 359                             '%s^2/%s' % (data_units, params['freq units'])),
 360             ]
 361         b[:,0] = freq_axis
 362         b[:,1] = power
 363         params['curve'].data.append(b)
 364         outqueue.put(b)