Adjust fit-parameter handling for the polymer_fit plugin.
[hooke.git] / hooke / plugin / polymer_fit.py
index 2ea503c58af398687463711674397e23b4008537..a4b65074bcd3e5ef57edb9e8b0e12941ab794a35 100644 (file)
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 #
-# Copyright (C) 2008-2010 Alberto Gomez-Casado
+# Copyright (C) 2008-2011 Alberto Gomez-Casado
 #                         Fabrizio Benedetti
 #                         Massimo Sandal <devicerandom@gmail.com>
 #                         W. Trevor King <wking@drexel.edu>
@@ -37,13 +37,12 @@ from scipy.optimize import newton
 from ..command import Command, Argument, Success, Failure
 from ..config import Setting
 from ..curve import Data
-from ..plugin import Plugin, argument_to_setting
 from ..util.callback import is_iterable
 from ..util.fit import PoorFit, ModelFitter
 from ..util.peak import Peak
 from ..util.si import join_data_label, split_data_label
-from .curve import CurveArgument
-from .vclamp import scale
+from . import Plugin, argument_to_setting
+from .curve import ColumnAccessCommand, ColumnAddingCommand
 
 
 kB = 1.3806503e-23  # Joules/Kelvin
@@ -245,7 +244,7 @@ class FJC (ModelFitter):
 
     >>> info['Kuhn length (m)'] = 2*a
     >>> model = FJC(d_data, info=info, rescale=True)
-    >>> L = model.fit(outqueue=outqueue)
+    >>> L, = model.fit(outqueue=outqueue)
     >>> fit_info = outqueue.get(block=False)
     >>> print L  # doctest: +ELLIPSIS
     3.199...e-08
@@ -341,11 +340,9 @@ class FJC (ModelFitter):
 
     def fit(self, *args, **kwargs):
         params = super(FJC, self).fit(*args, **kwargs)
-        if is_iterable(params):
-            params[0] = self.L(params[0])  # convert Lp -> L
+        params[0] = self.L(params[0])  # convert Lp -> L
+        if len(params) > 1:
             params[1] = abs(params[1])  # take the absolute value of `a`
-        else:  # params is a float
-            params = self.L(params)  # convert Lp -> L
         return params
 
     def guess_initial_params(self, outqueue=None):
@@ -508,7 +505,7 @@ class FJC_PEG (ModelFitter):
 
     >>> info['Kuhn length (m)'] = 2*kwargs['a']
     >>> model = FJC_PEG(d_data, info=info, rescale=True)
-    >>> N = model.fit(outqueue=outqueue)
+    >>> N, = model.fit(outqueue=outqueue)
     >>> fit_info = outqueue.get(block=False)
     >>> print N  # doctest: +ELLIPSIS
     96.931...
@@ -614,11 +611,9 @@ class FJC_PEG (ModelFitter):
 
     def fit(self, *args, **kwargs):
         params = super(FJC_PEG, self).fit(*args, **kwargs)
-        if is_iterable(params):
-            params[0] = self.L(params[0])  # convert Nr -> N
+        params[0] = self.L(params[0])  # convert Nr -> N
+        if len(params) > 1:
             params[1] = abs(params[1])  # take the absolute value of `a`
-        else:  # params is a float
-            params = self.L(params)  # convert Nr -> N
         return params
 
     def guess_initial_params(self, outqueue=None):
@@ -727,7 +722,7 @@ class WLC (ModelFitter):
 
     >>> info['persistence length (m)'] = 2*p
     >>> model = WLC(d_data, info=info, rescale=True)
-    >>> L = model.fit(outqueue=outqueue)
+    >>> L, = model.fit(outqueue=outqueue)
     >>> fit_info = outqueue.get(block=False)
     >>> print L  # doctest: +ELLIPSIS
     3.318...e-08
@@ -823,11 +818,9 @@ class WLC (ModelFitter):
 
     def fit(self, *args, **kwargs):
         params = super(WLC, self).fit(*args, **kwargs)
-        if is_iterable(params):
-            params[0] = self.L(params[0])  # convert Lp -> L
+        params[0] = self.L(params[0])  # convert Lp -> L
+        if len(params) > 1:
             params[1] = abs(params[1])  # take the absolute value of `p`
-        else:  # params is a float
-            params = self.L(params)  # convert Lp -> L
         return params
 
     def guess_initial_params(self, outqueue=None):
@@ -888,18 +881,16 @@ trans-trans-gauche (ttg) PEG states in units of kBT.
             self._settings.append(argument_to_setting(
                     self.setting_section, argument))
             argument.default = None # if argument isn't given, use the config.
-        self._arguments.extend([  # Non-configurable arguments
-                Argument(name='input distance column', type='string',
-                         default='cantilever adjusted extension (m)',
-                         help="""
+        self._input_columns = [
+            ('distance column', 'cantilever adjusted extension (m)',
+             """
 Name of the column to use as the surface position input.
 """.strip()),
-                Argument(name='input deflection column', type='string',
-                         default='deflection (N)',
-                         help="""
+            ('deflection column', 'deflection (N)',
+             """
 Name of the column to use as the deflection input.
 """.strip()),
-                ])
+            ]
         self._commands = [
             PolymerFitCommand(self), PolymerFitPeaksCommand(self),
             TranslateFlatPeaksCommand(self),
@@ -912,7 +903,7 @@ Name of the column to use as the deflection input.
         return self._settings
 
 
-class PolymerFitCommand (Command):
+class PolymerFitCommand (ColumnAddingCommand):
     """Polymer model (WLC, FJC, etc.) fitting.
 
     Fits an entropic elasticity function to a given chunk of the
@@ -926,62 +917,53 @@ class PolymerFitCommand (Command):
     def __init__(self, plugin):
         super(PolymerFitCommand, self).__init__(
             name='polymer fit',
+            columns=plugin._input_columns,
+            new_columns=[
+                ('output tension column', 'polymer tension',
+                 """
+Name of the column (without units) to use as the polymer tension output.
+""".strip()),
+                ],
             arguments=[
-                CurveArgument,
-                Argument(name='block', aliases=['set'], type='int', default=0,
+                Argument(name='fit parameters info name', type='string',
+                         default='polymer fit',
                          help="""
-Data block for which the fit should be calculated.  For an
-approach/retract force curve, `0` selects the approaching curve and
-`1` selects the retracting curve.
+Name (without units) for storing the fit parameters in the `.info` dictionary.
 """.strip()),
                 Argument(name='bounds', type='point', optional=False, count=2,
                          help="""
 Indicies of points bounding the selected data.
 """.strip()),
-                ] + plugin._arguments + [
-                Argument(name='output tension column', type='string',
-                         default='polymer tension',
-                         help="""
-Name of the column (without units) to use as the polymer tension output.
-""".strip()),
-                Argument(name='fit parameters info name', type='string',
-                         default='surface deflection offset',
-                         help="""
-Name (without units) for storing the fit parameters in the `.info` dictionary.
-""".strip()),
-                ],
+                ] + plugin._arguments,
             help=self.__doc__, plugin=plugin)
 
     def _run(self, hooke, inqueue, outqueue, params):
-        for key,value in params.items():
-            if value == None:  # Use configured default value.
-                params[key] = self.plugin.config[key]
-        scale(hooke, params['curve'], params['block'])  # TODO: is autoscaling a good idea? (explicit is better than implicit)
-        data = params['curve'].data[params['block']]
-        # HACK? rely on params['curve'] being bound to the local hooke
-        # playlist (i.e. not a copy, as you would get by passing a
-        # curve through the queue).  Ugh.  Stupid queues.  As an
-        # alternative, we could pass lookup information through the
-        # queue...
+        params = self._setup_params(hooke, params)
+        data = self._block(hooke, params)
         model = params['polymer model']
-        new = Data((data.shape[0], data.shape[1]+1), dtype=data.dtype)
-        new.info = copy.deepcopy(data.info)
-        new[:,:-1] = data
-        new.info['columns'].append(
-            join_data_label(params['output tension column'], 'N'))
-        z_data = data[:,data.info['columns'].index(
-                params['input distance column'])]
-        d_data = data[:,data.info['columns'].index(
-                params['input deflection column'])]
+        dist_data = self._get_column(
+            hooke=hooke, params=params, column_name='distance column')
+        def_data = self._get_column(
+            hooke=hooke, params=params, column_name='deflection column')
         start,stop = params['bounds']
         tension_data,ps = self.fit_polymer_model(
-            params, z_data, d_data, start, stop, outqueue)
-        new.info[params['fit parameters info name']] = ps
-        new.info[params['fit parameters info name']]['model'] = model
-        new[:,-1] = tension_data
-        params['curve'].data[params['block']] = new
+            params, dist_data, def_data, start, stop, outqueue)
+        data.info[params['fit parameters info name']] = ps
+        data.info[params['fit parameters info name']]['model'] = model
+        self._set_column(hooke=hooke, params=params,
+                         column_name='output tension column',
+                         values=tension_data)
+
+    def _setup_params(self, hooke, params):
+        for key,value in params.items():
+            if value == None:  # Use configured default value.
+                params[key] = self.plugin.config[key]
+        name,def_unit = split_data_label(params['deflection column'])
+        params['output tension column'] = join_data_label(
+            params['output tension column'], def_unit)
+        return params
 
-    def fit_polymer_model(self, params, z_data, d_data, start, stop,
+    def fit_polymer_model(self, params, dist_data, def_data, start, stop,
                           outqueue=None):
         """Railyard for the `fit_*_model` family.
 
@@ -990,89 +972,90 @@ Name (without units) for storing the fit parameters in the `.info` dictionary.
         """
         fn = getattr(self, 'fit_%s_model'
                      % params['polymer model'].replace('-','_'))
-        return fn(params, z_data, d_data, start, stop, outqueue)
+        return fn(params, dist_data, def_data, start, stop, outqueue)
 
-    def fit_FJC_model(self, params, z_data, d_data, start, stop,
+    def fit_FJC_model(self, params, dist_data, def_data, start, stop,
                       outqueue=None):
         """Fit the data with :class:`FJC`.
         """
         info = {
             'temperature (K)': self.plugin.config['temperature'],
-            'x data (m)': z_data[start:stop],
+            'x data (m)': dist_data[start:stop],
             }
         if True:  # TODO: optionally free persistence length
             info['Kuhn length (m)'] = (
                 params['FJC Kuhn length'])
-        model = FJC(d_data[start:stop], info=info, rescale=True)
+        model = FJC(def_data[start:stop], info=info, rescale=True)
         queue = Queue()
         params = model.fit(outqueue=queue)
         if True:  # TODO: if Kuhn length fixed
-            params = [params]
             a = info['Kuhn length (m)']
         else:
             a = params[1]
         L = params[0]
         T = info['temperature (K)']
         fit_info = queue.get(block=False)
-        f_data = numpy.ones(z_data.shape, dtype=z_data.dtype) * numpy.nan
-        f_data[start:stop] = FJC_fn(z_data[start:stop], T=T, L=L, a=a)
+        f_data = numpy.ones(dist_data.shape, dtype=dist_data.dtype) * numpy.nan
+        f_data[start:stop] = FJC_fn(dist_data[start:stop], T=T, L=L, a=a)
         return [f_data, fit_info]
 
-    def fit_FJC_PEG_model(self, params, z_data, d_data, start, stop,
+    def fit_FJC_PEG_model(self, params, dist_data, def_data, start, stop,
                           outqueue=None):
         """Fit the data with :class:`FJC_PEG`.
         """
         info = {
             'temperature (K)': self.plugin.config['temperature'],
-            'x data (m)': z_data[start:stop],
+            'x data (m)': dist_data[start:stop],
             # TODO: more info
             }
         if True:  # TODO: optionally free persistence length
             info['Kuhn length (m)'] = (
                 params['FJC Kuhn length'])
-        model = FJC_PEG(d_data[start:stop], info=info, rescale=True)
+        model = FJC_PEG(def_data[start:stop], info=info, rescale=True)
         queue = Queue()
         params = model.fit(outqueue=queue)
         if True:  # TODO: if Kuhn length fixed
-            params = [params]
             a = info['Kuhn length (m)']
         else:
             a = params[1]
         N = params[0]
         T = info['temperature (K)']
         fit_info = queue.get(block=False)
-        f_data = numpy.ones(z_data.shape, dtype=z_data.dtype) * numpy.nan
-        f_data[start:stop] = FJC_PEG_fn(z_data[start:stop], **kwargs)
+        f_data = numpy.ones(dist_data.shape, dtype=dist_data.dtype) * numpy.nan
+        f_data[start:stop] = FJC_PEG_fn(dist_data[start:stop], **kwargs)
         return [f_data, fit_info]
 
-    def fit_WLC_model(self, params, z_data, d_data, start, stop,
+    def fit_WLC_model(self, params, dist_data, def_data, start, stop,
                       outqueue=None):
         """Fit the data with :class:`WLC`.
         """
         info = {
             'temperature (K)': self.plugin.config['temperature'],
-            'x data (m)': z_data[start:stop],
+            'x data (m)': dist_data[start:stop],
             }
         if True:  # TODO: optionally free persistence length
             info['persistence length (m)'] = (
                 params['WLC persistence length'])
-        model = WLC(d_data[start:stop], info=info, rescale=True)
+        model = WLC(def_data[start:stop], info=info, rescale=True)
         queue = Queue()
         params = model.fit(outqueue=queue)
         if True:  # TODO: if persistence length fixed
-            params = [params]
             p = info['persistence length (m)']
         else:
             p = params[1]
+            info['persistence length (m)'] = p
         L = params[0]
+        info['contour length (m)'] = L
         T = info['temperature (K)']
         fit_info = queue.get(block=False)
-        f_data = numpy.ones(z_data.shape, dtype=z_data.dtype) * numpy.nan
-        f_data[start:stop] = WLC_fn(z_data[start:stop], T=T, L=L, p=p)
-        return [f_data, fit_info]
+        info['fit'] = fit_info
+        info.pop('x data (m)')
+        f_data = numpy.ones(dist_data.shape, dtype=dist_data.dtype) * numpy.nan
+        f_data[start:stop] = WLC_fn(dist_data[start:stop], T=T, L=L, p=p)
+        return [f_data, info]
 
 
-class PolymerFitPeaksCommand (Command):
+class PolymerFitPeaksCommand (ColumnAccessCommand):
     """Polymer model (WLC, FJC, etc.) fitting.
 
     Use :class:`PolymerFitCommand` to fit the each peak in a list of
@@ -1081,14 +1064,8 @@ class PolymerFitPeaksCommand (Command):
     def __init__(self, plugin):
         super(PolymerFitPeaksCommand, self).__init__(
             name='polymer fit peaks',
+            columns=plugin._input_columns,
             arguments=[
-                CurveArgument,
-                Argument(name='block', aliases=['set'], type='int', default=0,
-                         help="""
-Data block for which the fit should be calculated.  For an
-approach/retract force curve, `0` selects the approaching curve and
-`1` selects the retracting curve.
-""".strip()),
                 Argument(name='peak info name', type='string',
                          default='polymer peaks',
                          help="""
@@ -1102,11 +1079,12 @@ Index of the selected peak in the list of peaks.  Use `None` to fit all peaks.
             help=self.__doc__, plugin=plugin)
 
     def _run(self, hooke, inqueue, outqueue, params):
-        data = params['curve'].data[params['block']]
-        fit_command = [c for c in hooke.commands
-                       if c.name=='polymer fit'][0]
+        data = self._block(hooke, params)
+        fit_command = hooke.command_by_name['polymer fit']
         inq = Queue()
         outq = Queue()
+        curve = params['curve']
+        params['curve'] = None
         p = copy.deepcopy(params)
         p['curve'] = params['curve']
         del(p['peak info name'])
@@ -1122,7 +1100,7 @@ Index of the selected peak in the list of peaks.  Use `None` to fit all peaks.
                 raise ret
 
 
-class TranslateFlatPeaksCommand (Command):
+class TranslateFlatPeaksCommand (ColumnAccessCommand):
     """Translate flat filter peaks into polymer peaks for fitting.
 
     Use :class:`~hooke.plugin.flatfilt.FlatPeaksCommand` creates a
@@ -1134,62 +1112,59 @@ class TranslateFlatPeaksCommand (Command):
     polymer loading regions.
     """
     def __init__(self, plugin):
-        plugin_arguments = [a for a in plugin._arguments
-                            if a.name in ['input distance column',
-                                          'input deflection column']]
-        arguments = [
-            CurveArgument,
-            Argument(name='block', aliases=['set'], type='int', default=0,
-                     help="""
-Data block for which the fit should be calculated.  For an
-approach/retract force curve, `0` selects the approaching curve and
-`1` selects the retracting curve.
-""".strip()),
-            ] + plugin_arguments + [
-            Argument(name='input peak info name', type='string',
-                     default='flat filter peaks',
-                     help="""
+        super(TranslateFlatPeaksCommand, self).__init__(
+            name='flat peaks to polymer peaks',
+            columns=plugin._input_columns,
+            arguments=[
+                Argument(name='input peak info name', type='string',
+                         default='flat filter peaks',
+                         help="""
 Name for the input peaks in the `.info` dictionary.
 """.strip()),
-            Argument(name='output peak info name', type='string',
-                     default='polymer peaks',
-                     help="""
+                Argument(name='output peak info name', type='string',
+                         default='polymer peaks',
+                         help="""
 Name for the ouptput peaks in the `.info` dictionary.
 """.strip()),
-            Argument(name='end offset', type='int', default=-1,
-                     help="""
+                Argument(name='end offset', type='int', default=-1,
+                         help="""
 Number of points between the end of a new peak and the start of the old.
 """.strip()),
-            Argument(name='start fraction', type='float', default=0.2,
-                     help="""
+                Argument(name='start fraction', type='float', default=0.2,
+                         help="""
 Place the start of the new peak at `start_fraction` from the end of
 the previous old peak to the end of the new peak.  Because the first
 new peak will have no previous old peak, it uses a distance of zero
 instead.
 """.strip()),
-            ]
-        super(TranslateFlatPeaksCommand, self).__init__(
-            name='flat peaks to polymer peaks',
-            arguments=arguments,
+            ] + plugin._arguments,
             help=self.__doc__, plugin=plugin)
 
     def _run(self, hooke, inqueue, outqueue, params):
-        data = params['curve'].data[params['block']]
-        z_data = data[:,data.info['columns'].index(
-                params['input distance column'])]
-        d_data = data[:,data.info['columns'].index(
-                params['input deflection column'])]
-        previous_old_stop = numpy.absolute(z_data).argmin()
+        data = self._block(hooke, params)
+        dist_data = self._get_column(
+            hooke=hooke, params=params, column_name='distance column')
+        def_data = self._get_column(
+            hooke=hooke, params=params, column_name='deflection column')
+        previous_old_stop = numpy.absolute(dist_data).argmin()
         new = []
-        for i,peak in enumerate(data.info[params['input peak info name']]):
+        try:
+            peaks = data.info[params['input peak info name']]
+        except KeyError, e:
+            raise Failure('No value for %s in %s.info (%s): %s'
+                          % (params['input peak info name'], data.info['name'],
+                             sorted(data.info.keys()), e))
+        for i,peak in enumerate(peaks):
             next_old_start = peak.index
             stop = next_old_start + params['end offset'] 
-            z_start = z_data[previous_old_stop] + params['start fraction']*(
-                z_data[stop] - z_data[previous_old_stop])
-            start = numpy.absolute(z_data - z_start).argmin()
+            dist_start = (
+                dist_data[previous_old_stop]
+                + params['start fraction']*(
+                    dist_data[stop] - dist_data[previous_old_stop]))
+            start = numpy.absolute(dist_data - dist_start).argmin()
             p = Peak('polymer peak %d' % i,
                      index=start,
-                     values=d_data[start:stop])
+                     values=def_data[start:stop])
             new.append(p)
             previous_old_stop = peak.post_index()
         data.info[params['output peak info name']] = new