hooke/util/fit.py

   1 # Copyright (C) 2010 W. Trevor King <wking@drexel.edu>
   2 #
   3 # This file is part of Hooke.
   4 #
   5 # Hooke is free software: you can redistribute it and/or modify it
   6 # under the terms of the GNU Lesser General Public License as
   7 # published by the Free Software Foundation, either version 3 of the
   8 # License, or (at your option) any later version.
   9 #
  10 # Hooke is distributed in the hope that it will be useful, but WITHOUT
  11 # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  12 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General
  13 # Public License for more details.
  14 #
  15 # You should have received a copy of the GNU Lesser General Public
  16 # License along with Hooke.  If not, see
  17 # <http://www.gnu.org/licenses/>.
  18
  19 """Provide :class:`ModelFitter` to make arbitrary model fitting easy.
  20 """
  21
  22 from numpy import arange, ndarray
  23 from scipy.optimize import leastsq
  24
  25
  26 class PoorFit (ValueError):
  27     pass
  28
  29 class ModelFitter (object):
  30     """A convenient wrapper around :func:`scipy.optimize.leastsq`.
  31
  32     Parameters
  33     ----------
  34     d_data : array_like
  35         Deflection data to be analyzed for the contact position.
  36     info :
  37         Store any extra information useful inside your overridden
  38         methods.
  39     rescale : boolean
  40         Rescale parameters so the scale of each is 1.0.  Also rescale
  41         the data so data.std() == 1.0.
  42
  43     Examples
  44     --------
  45
  46     >>> from pprint import pprint
  47     >>> from Queue import Queue
  48     >>> import numpy
  49
  50     You'll want to subclass `ModelFitter`, overriding at least
  51     `.model` and potentially the parameter and scale guessing
  52     methods.
  53
  54     >>> class LinearModel (ModelFitter):
  55     ...     '''Simple linear model.
  56     ...
  57     ...     Levenberg-Marquardt is not how you want to solve this problem
  58     ...     for real systems, but it's a simple test case.
  59     ...     '''
  60     ...     def model(self, params):
  61     ...         '''A linear model.
  62     ...
  63     ...         Notes
  64     ...         -----
  65     ...         .. math:: y = p_0 x + p_1
  66     ...         '''
  67     ...         p = params  # convenient alias
  68     ...         self._model_data[:] = p[0]*arange(len(self._data)) + p[1]
  69     ...         return self._model_data
  70     ...     def guess_initial_params(self, outqueue=None):
  71     ...         return [float(self._data[-1] - self._data[0])/len(self._data),
  72     ...                 self._data[0]]
  73     ...     def guess_scale(self, params, outqueue=None):
  74     ...         slope_scale = params[0]/10.
  75     ...         if slope_scale == 0:  # data is expected to be flat
  76     ...             slope_scale = float(self._data.max()-self._data.min())/len(self._data)
  77     ...             if slope_scale == 0:  # data is completely flat
  78     ...                 slope_scale = 1.
  79     ...         offset_scale = self._data.std()/10.0
  80     ...         if offset_scale == 0:  # data is completely flat
  81     ...             offset_scale = 1.
  82     ...         return [slope_scale, offset_scale]
  83     >>> data = 20*numpy.sin(arange(1000)) + 7.*arange(1000) - 33.0
  84     >>> m = LinearModel(data)
  85     >>> outqueue = Queue()
  86     >>> slope,offset = m.fit(outqueue=outqueue)
  87     >>> info = outqueue.get()
  88     >>> pprint(info)  # doctest: +ELLIPSIS, +REPORT_UDIFF
  89     {'active fitted parameters': array([  6.999..., -32.889...]),
  90      'active parameters': array([  6.999..., -32.889...]),
  91      'active scale': [0.699..., 202.071...],
  92      'convergence flag': 2,
  93      'covariance matrix': array([[  1.199...e-08,  -5.993...e-06],
  94            [ -5.993...e-06,   3.994...e-03]]),
  95      'data scale factor': 1.0,
  96      'fitted parameters': array([  6.999..., -32.889...]),
  97      'info': {'fjac': array([[...]]),
  98               'fvec': array([...]),
  99               'ipvt': array([1, 2]),
 100               'nfev': 7,
 101               'qtf': array([  2.851...e-07,   1.992...e-06])},
 102      'initial parameters': [6.992..., -33.0],
 103      'message': 'The relative error between two consecutive iterates is at most 0.000...',
 104      'rescaled': False,
 105      'scale': [0.699..., 202.071...]}
 106
 107     We round the outputs to protect the doctest against differences in
 108     machine rounding during computation.  We expect the values to be close
 109     to the input settings (slope 7, offset -33).
 110
 111     >>> print '%.3f' % slope
 112     7.000
 113     >>> print '%.3f' % offset
 114     -32.890
 115
 116     The offset is a bit off because, the range is not a multiple of
 117     :math:`2\pi`.
 118
 119     We could also use rescaled parameters:
 120
 121     >>> m = LinearModel(data, rescale=True)
 122     >>> outqueue = Queue()
 123     >>> slope,offset = m.fit(outqueue=outqueue)
 124     >>> print '%.3f' % slope
 125     7.000
 126     >>> print '%.3f' % offset
 127     -32.890
 128     """
 129     def __init__(self, *args, **kwargs):
 130         self.set_data(*args, **kwargs)
 131
 132     def set_data(self, data, info=None, rescale=False):
 133         self._data = data
 134         self._model_data = ndarray(shape=data.shape, dtype=data.dtype)
 135         self.info = info
 136         self._rescale = rescale
 137         if rescale == True:
 138             self._data_scale_factor = data.std()
 139         else:
 140             self._data_scale_factor = 1.0
 141
 142     def model(self, params):
 143         p = params  # convenient alias
 144         self._model_data[:] = arange(len(self._data))
 145         raise NotImplementedError
 146
 147     def guess_initial_params(self, outqueue=None):
 148         return []
 149
 150     def guess_scale(self, params, outqueue=None):
 151         return None
 152
 153     def residual(self, params):
 154         if self._rescale == True:
 155             params = [p*s for p,s in zip(params, self._param_scale_factors)]
 156         residual = self._data - self.model(params)
 157         if self._rescale == True or False:
 158             residual /= self._data_scale_factor
 159         return residual
 160
 161     def fit(self, initial_params=None, scale=None, outqueue=None, **kwargs):
 162         """
 163         Parameters
 164         ----------
 165         initial_params : iterable or None
 166             Initial parameter values for residual minimization.  If
 167             `None`, they are estimated from the data using
 168             :meth:`guess_initial_params`.
 169         scale : iterable or None
 170             Parameter length scales for residual minimization.  If
 171             `None`, they are estimated from the data using
 172             :meth:`guess_scale`.
 173         outqueue : Queue or None
 174             If given, will be used to output the data and fitted model
 175             for user verification.
 176         kwargs :
 177             Any additional arguments are passed through to `leastsq`.
 178         """
 179         if initial_params == None:
 180             initial_params = self.guess_initial_params(outqueue)
 181         if scale == None:
 182             scale = self.guess_scale(initial_params, outqueue)
 183         assert min(scale) > 0, scale
 184         if self._rescale == True:
 185             self._param_scale_factors = scale
 186             active_scale = [1.0 for s in scale]
 187             active_params = [p/s for p,s in zip(initial_params,
 188                                                 self._param_scale_factors)]
 189         else:
 190             active_scale = scale
 191             active_params = initial_params
 192         params,cov,info,mesg,ier = leastsq(
 193             func=self.residual, x0=active_params, full_output=True,
 194             diag=active_scale, **kwargs)
 195         if self._rescale == True:
 196             active_params = params
 197             params = [p*s for p,s in zip(params, self._param_scale_factors)]
 198         else:
 199             active_params = params
 200         if outqueue != None:
 201             outqueue.put({
 202                     'rescaled': self._rescale,
 203                     'initial parameters': initial_params,
 204                     'active parameters': active_params,
 205                     'scale': scale,
 206                     'active scale': active_scale,
 207                     'data scale factor': self._data_scale_factor,
 208                     'active fitted parameters': active_params,
 209                     'fitted parameters': params,
 210                     'covariance matrix': cov,
 211                     'info': info,
 212                     'message': mesg,
 213                     'convergence flag': ier,
 214                     })
 215         return params