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
   6 # modify it under the terms of the GNU Lesser General Public
   7 # License as published by the Free Software Foundation, either
   8 # version 3 of the License, or (at your option) any later version.
   9 #
  10 # Hooke is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU Lesser General 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
  40     Examples
  41     --------
  42
  43     >>> from pprint import pprint
  44     >>> from Queue import Queue
  45     >>> import numpy
  46
  47     You'll want to subclass `ModelFitter`, overriding at least
  48     `.model` and potentially the parameter and scale guessing
  49     methods.
  50
  51     >>> class LinearModel (ModelFitter):
  52     ...     '''Simple linear model.
  53     ...
  54     ...     Levenberg-Marquardt is not how you want to solve this problem
  55     ...     for real systems, but it's a simple test case.
  56     ...     '''
  57     ...     def model(self, params):
  58     ...         '''A linear model.
  59     ...
  60     ...         Notes
  61     ...         -----
  62     ...         .. math:: y = p_0 x + p_1
  63     ...         '''
  64     ...         p = params  # convenient alias
  65     ...         self._model_data[:] = p[0]*arange(len(self._data)) + p[1]
  66     ...         return self._model_data
  67     ...     def guess_initial_params(self, outqueue=None):
  68     ...         return [float(self._data[-1] - self._data[0])/len(self._data),
  69     ...                 self._data[0]]
  70     ...     def guess_scale(self, params, outqueue=None):
  71     ...         slope_scale = params[0]/10.
  72     ...         if slope_scale == 0:  # data is expected to be flat
  73     ...             slope_scale = float(self._data.max()-self._data.min())/len(self._data)
  74     ...             if slope_scale == 0:  # data is completely flat
  75     ...                 slope_scale = 1.
  76     ...         offset_scale = self._data.std()/10.0
  77     ...         if offset_scale == 0:  # data is completely flat
  78     ...             offset_scale = 1.
  79     ...         return [slope_scale, offset_scale]
  80     >>> data = 20*numpy.sin(arange(1000)) + 7.*arange(1000) - 33.0
  81     >>> m = LinearModel(data)
  82     >>> outqueue = Queue()
  83     >>> slope,offset = m.fit(outqueue=outqueue)
  84     >>> info = outqueue.get()
  85     >>> pprint(info)  # doctest: +ELLIPSIS, +REPORT_UDIFF
  86     {'convergence flag': 2,
  87      'covariance matrix': array([[  1.199...e-08,  -5.993...e-06],
  88            [ -5.993...e-06,   3.994...e-03]]),
  89      'fitted parameters': array([  6.999..., -32.889...]),
  90      'info': {'fjac': array([[...]]),
  91               'fvec': array([...]),
  92               'ipvt': array([1, 2]),
  93               'nfev': 7,
  94               'qtf': array([  2.851...e-07,   1.992...e-06])},
  95      'initial parameters': [6.992..., -33.0],
  96      'message': 'The relative error between two consecutive iterates is at most 0.000...',
  97      'scale': [0.699..., 202.071...]}
  98
  99     We round the outputs to protect the doctest against differences in
 100     machine rounding during computation.  We expect the values to be close
 101     to the input settings (slope 7, offset -33).
 102
 103     >>> print '%.3f' % slope
 104     7.000
 105     >>> print '%.3f' % offset
 106     -32.890
 107
 108     The offset is a bit off because, the range is not a multiple of
 109     :math:`2\pi`.
 110     """
 111     def __init__(self, data, info=None):
 112         self.set_data(data, info)
 113
 114     def set_data(self, data, info=None):
 115         self._data = data
 116         self._model_data = ndarray(shape=data.shape, dtype=data.dtype)
 117         self.info = info
 118
 119     def model(self, params):
 120         p = params  # convenient alias
 121         self._model_data[:] = arange(len(self._data))
 122         raise NotImplementedError
 123
 124     def guess_initial_params(self, outqueue=None):
 125         return []
 126
 127     def guess_scale(self, params, outqueue=None):
 128         return None
 129
 130     def residual(self, params):
 131         return self._data - self.model(params)
 132
 133     def fit(self, initial_params=None, scale=None, outqueue=None, **kwargs):
 134         """
 135         Parameters
 136         ----------
 137         initial_params : iterable or None
 138             Initial parameter values for residual minimization.  If
 139             `None`, they are estimated from the data using
 140             :meth:`guess_initial_params`.
 141         scale : iterable or None
 142             Parameter length scales for residual minimization.  If
 143             `None`, they are estimated from the data using
 144             :meth:`guess_scale`.
 145         outqueue : Queue or None
 146             If given, will be used to output the data and fitted model
 147             for user verification.
 148         kwargs :
 149             Any additional arguments are passed through to `leastsq`.
 150         """
 151         if initial_params == None:
 152             initial_params = self.guess_initial_params(outqueue)
 153         if scale == None:
 154             scale = self.guess_scale(initial_params, outqueue)
 155         assert min(scale) > 0, scale
 156         params,cov,info,mesg,ier = leastsq(
 157             func=self.residual, x0=initial_params, full_output=True,
 158             diag=scale, **kwargs)
 159         if outqueue != None:
 160             outqueue.put({
 161                     'initial parameters': initial_params,
 162                     'scale': scale,
 163                     'fitted parameters': params,
 164                     'covariance matrix': cov,
 165                     'info': info,
 166                     'message': mesg,
 167                     'convergence flag': ier,
 168                     })
 169         return params