Moved testing/common/fit.py to pysawsim/fit.py and cleaned up.
[sawsim.git] / pysawsim / fit.py
1 # Copyright (C) 2008-2010  W. Trevor King <wking@drexel.edu>
2 #
3 # This program is free software: you can redistribute it and/or modify
4 # it under the terms of the GNU General Public License as published by
5 # the Free Software Foundation, either version 3 of the License, or
6 # (at your option) any later version.
7 #
8 # This program is distributed in the hope that it will be useful,
9 # but WITHOUT ANY WARRANTY; without even the implied warranty of
10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 # GNU General Public License for more details.
12 #
13 # You should have received a copy of the GNU General Public License
14 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
15 #
16 # The author may be contacted at <wking@drexel.edu> on the Internet, or
17 # write to Trevor King, Drexel University, Physics Dept., 3141 Chestnut St.,
18 # Philadelphia PA 19104, USA.
19
20 """Provide :class:`ModelFitter` to make arbitrary model fitting easy.
21 """
22
23 from copy import deepcopy
24
25 from numpy import arange, array, ndarray, exp, pi, sqrt
26 from scipy.optimize import leastsq
27
28
29 class ModelFitter (object):
30     """A convenient wrapper around :func:`scipy.optimize.leastsq`.
31
32     For minimizing fit error: returns the normalized rms difference
33     between a dataset and a function `y(X)`.  In order to use this
34     class, you must sublcass it and overwrite `ModelFitter.model()`.
35     It is recommended that you also overwrite
36     `ModelFitter.guess_initial_params()`, and then not pass
37     initializing params to `ModelFitter.model()`.  This helps ensure
38     good automation for unsupervised testing.  Developing good
39     heuristics for `ModelFitter.guess_initial_params()` is usually the
40     hardest part writing a theoretical histogram test.
41
42     For bonus points, you can override `ModelFitter.print_model()` and
43     `ModelFitter.print_params()`, to give your users an easy to
44     understand idea of what's going on.  It's better to be verbose and
45     clear.  Remember, this is the testing code.  It needs to be solid.
46     Save your wonderful, sneaky hacks for the main code ;).
47
48     TODO: look into :mod:`scipy.odr` as an alternative fitting
49     algorithm (minimizing the sum of squares of orthogonal distances,
50     vs. minimizing y distances).
51
52     Parameters
53     ----------
54     d_data : array_like
55         Deflection data to be analyzed for the contact position.
56     info :
57         Store any extra information useful inside your overridden
58         methods.
59     rescale : boolean
60         Rescale parameters so the guess for each is 1.0.  Also rescale
61         the data so data.std() == 1.0.
62
63     Examples
64     --------
65     >>> from pprint import pprint
66     >>> import numpy
67     >>> class LinearModel (ModelFitter):
68     ...     '''Simple linear model.
69     ...
70     ...     Levenberg-Marquardt is not how you want to solve this problem
71     ...     for real systems, but it's a simple test case.
72     ...     '''
73     ...     def model(self, params):
74     ...         '''A linear model.
75     ...
76     ...         Notes
77     ...         -----
78     ...         .. math:: y = p_0 x + p_1
79     ...         '''
80     ...         p = params  # convenient alias
81     ...         self._model_data[:] = p[0]*arange(len(self._data)) + p[1]
82     ...         return self._model_data
83     ...     def guess_initial_params(self):
84     ...         return [float(self._data[-1] - self._data[0])/len(self._data),
85     ...                 self._data[0]]
86     ...     def guess_scale(self, params):
87     ...         slope_scale = 0.1
88     ...         if params[1] == 0:
89     ...             offset_scale = 1
90     ...         else:
91     ...             offset_scale = 0.1*self._data.std()/abs(params[1])
92     ...             if offset_scale == 0:  # data is completely flat
93     ...                 offset_scale = 1.
94     ...         return [slope_scale, offset_scale]
95     ...     def model_string(self):
96     ...         return 'y = A x + B'
97     ...     def param_string(self, params):
98     ...         pstrings = []
99     ...         for name,param in zip(['A', 'B'], params):
100     ...             pstrings.append('%s=%.3f' % (name, param))
101     ...         return ', '.join(pstrings)
102     >>> data = 20*numpy.sin(arange(1000)) + 7.*arange(1000) - 33.0
103     >>> m = LinearModel(data)
104     >>> slope,offset = m.fit()
105     >>> pprint(m.fit_info)
106     ... # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE, +REPORT_UDIFF
107     {'active fitted parameters': array([  6.999..., -32.889...]),
108      'active parameters': array([  6.999..., -32.889...]),
109      'convergence flag': ...,
110      'covariance matrix': array([[  1.199...e-08,  -5.993...e-06],
111            [ -5.993...e-06,   3.994...e-03]]),
112      'data scale factor': 1.0,
113      'fitted parameters': array([  6.999..., -32.889...]),
114      'info': {'fjac': array([[...]]),
115               'fvec': array([...]),
116               'ipvt': array([1, 2]),
117               'nfev': 7,
118               'qtf': array([...])},
119      'initial parameters': [6.992..., -33.0],
120      'message': '...relative error between two consecutive iterates is at most 0.000...',
121      'rescaled': False,
122      'scale': [0.100..., 6.123...]}
123
124     We round the outputs to protect the doctest against differences in
125     machine rounding during computation.  We expect the values to be close
126     to the input settings (slope 7, offset -33).
127
128     >>> print '%.3f' % slope
129     7.000
130     >>> print '%.3f' % offset
131     -32.890
132
133     >>> print m.model_string()
134     y = A x + B
135     >>> print m.param_string([slope,offset])
136     A=7.000, B=-32.890
137     >>> print ModelFitter.param_string(m, [slope,offset])  # doctest: +ELLIPSIS
138     param_0=6.9..., param_1=-32.8...
139
140     The offset is a bit off because, the range is not a multiple of
141     :math:`2\pi`.
142
143     We could also use rescaled parameters:
144
145     >>> m = LinearModel(data, rescale=True)
146     >>> slope,offset = m.fit()
147     >>> print '%.3f' % slope
148     7.000
149     >>> print '%.3f' % offset
150     -32.890
151     """
152     def __init__(self, *args, **kwargs):
153         self.set_data(*args, **kwargs)
154
155     def set_data(self, data, info=None, rescale=False):
156         self._data = data
157         self._model_data = ndarray(shape=data.shape, dtype=data.dtype)
158         self.info = info
159         self._rescale = rescale
160         if rescale == True:
161             for x in [data.std(), data.max()-data.min(), abs(data.max()), 1.0]:
162                 if x != 0:
163                     self._data_scale_factor = x
164                     break
165         else:
166             self._data_scale_factor = 1.0
167
168     def model(self, params):
169         p = params  # convenient alias
170         self._model_data[:] = arange(len(self._data))
171         raise NotImplementedError
172
173     def model_string(self):
174         """OVERRIDE: Give the user some idea of what model we're using.
175         """
176         return 'model string not defined'
177
178     def param_string(self, params):
179         """OVERRIDE: Give the user nicely named format for a parameter vector.
180         """
181         pstrings = []
182         for i,p in enumerate(params):
183             pstrings.append('param_%d=%g' % (i,p))
184         return ', '.join(pstrings)
185
186     def guess_initial_params(self):
187         """OVERRIDE: Provide some intelligent heuristic to pick a
188         starting point for minimizing your model"""
189         raise NotImplementedError
190
191     def guess_scale(self, params):
192         """Guess the problem length scale in each parameter dimension.
193
194         Notes
195         -----
196         From the :func:`scipy.optimize.leastsq` documentation, `diag`
197         (which we refer to as `scale`, sets `factor * || diag * x||`
198         as the initial step.  If `x == 0`, then `factor` is used
199         instead (from `lmdif.f`_)::
200
201                       do 70 j = 1, n
202                         wa3(j) = diag(j)*x(j)
203                70       continue
204                       xnorm = enorm(n,wa3)
205                       delta = factor*xnorm
206                       if (delta .eq. zero) delta = factor
207   
208         For most situations then, you don't need to do anything fancy.
209         The default scaling (if you don't set a scale) is::
210
211             c        on the first iteration and if mode is 1, scale according
212             c        to the norms of the columns of the initial jacobian.
213
214         (i.e. `diag(j) = acnorm(j)`, where `acnorm(j) is the norm of the `j`th column
215         of the initial Jacobian).
216
217         .. _lmdif.f: http://www.netlib.org/minpack/lmdif.f
218         """
219         return None
220
221     def residual(self, params):
222         if self._rescale == True:
223             params = [p*s for p,s in zip(params, self._param_scale_factors)]
224         residual = self._data - self.model(params)
225         if False:  # fit debugging
226             if not hasattr(self, '_i_'):
227                 self._i_ = 0
228             self._data.tofile('data.%d' % self._i_, sep='\n')
229             self._model_data.tofile('model.%d' % self._i_, sep='\n')
230             self._i_ += 1
231         if self._rescale == True:
232             residual /= self._data_scale_factor
233         return residual
234
235     def fit(self, initial_params=None, scale=None, **kwargs):
236         """
237         Parameters
238         ----------
239         initial_params : iterable or None
240             Initial parameter values for residual minimization.  If
241             `None`, they are estimated from the data using
242             :meth:`guess_initial_params`.
243         scale : iterable or None
244             Parameter length scales for residual minimization.  If
245             `None`, they are estimated from the data using
246             :meth:`guess_scale`.
247         kwargs :
248             Any additional arguments are passed through to `leastsq`.
249         """
250         if initial_params == None:
251             initial_params = self.guess_initial_params()
252         if scale == None:
253             scale = self.guess_scale(initial_params)
254         if scale != None:
255             assert min(scale) > 0, scale
256         if self._rescale == True:
257             self._param_scale_factors = initial_params
258             for i,s in enumerate(self._param_scale_factors):
259                 if s == 0:
260                     self._param_scale_factors[i] = 1.0
261             active_params = [p/s for p,s in zip(initial_params,
262                                                 self._param_scale_factors)]
263         else:
264             active_params = initial_params
265         params,cov,info,mesg,ier = leastsq(
266             func=self.residual, x0=active_params, full_output=True,
267             diag=scale, **kwargs)
268         if self._rescale == True:
269             active_params = params
270             if len(initial_params) == 1:  # params is a float
271                 params = params * self._param_scale_factors[0]
272             else:
273                 params = [p*s for p,s in zip(params,
274                                              self._param_scale_factors)]
275         else:
276             active_params = params
277         self.fit_info = {
278             'rescaled': self._rescale,
279             'initial parameters': initial_params,
280             'active parameters': active_params,
281             'scale': scale,
282             'data scale factor': self._data_scale_factor,
283             'active fitted parameters': active_params,
284             'fitted parameters': params,
285             'covariance matrix': cov,
286             'info': info,
287             'message': mesg,
288             'convergence flag': ier,
289             }
290         return params
291
292
293 class HistogramModelFitter (ModelFitter):
294     """Adapt ModelFitter for easy Histogram() comparison.
295
296     TODO: look into :func:`scipy.optimize.fmin_powell` as way to
297     minimize a single residul number.  Useful for fitting `Histograms`
298     with their assorted residal methods.
299
300     Examples
301     --------
302     >>> from random import gauss
303     >>> from numpy import array
304     >>> from .histogram import Histogram
305
306     >>> class GaussianModel (HistogramModelFitter):
307     ...     '''Simple gaussian model.
308     ...     '''
309     ...     def model(self, params):
310     ...         '''A gaussian model.
311     ...
312     ...         Notes
313     ...         -----
314     ...         .. math:: y \\propto e^{-\\frac{(x-p_0)^2}{2 p_1^2}}
315     ...         '''
316     ...         p = params  # convenient alias
317     ...         p[1] = abs(p[1])  # cannot have negative std. dev.
318     ...         self._model_data.counts = self.info['N']/(p[1]*sqrt(2*pi)) * (
319     ...             exp(-((self._model_data.bin_centers - p[0])/p[1])**2 / 2))
320     ...         return self._model_data
321     ...     def guess_initial_params(self):
322     ...         return [self._data.mean, self._data.std_dev]
323     ...     def guess_scale(self, params):
324     ...         return [self._data.std_dev, self._data.std_dev]
325     ...     def model_string(self):
326     ...         return 'p(x) = A exp((x-mu)^2 / (2 sigma^2))'
327     ...     def param_string(self, params):
328     ...         pstrings = []
329     ...         for name,param in zip(['mu', 'sigma'], params):
330     ...             pstrings.append('%s=%.3f' % (name, param))
331     ...         return ', '.join(pstrings)
332
333     >>> mu = 1.45
334     >>> sigma = 3.14
335     >>> data = array([gauss(mu,sigma) for i in range(int(1e4))])
336     >>> h = Histogram()
337     >>> h.from_data(data, h.calculate_bin_edges(data, 1))
338     >>> m = GaussianModel(h)
339     >>> params = m.fit()
340     >>> print m.model_string()
341     p(x) = A exp((x-mu)^2 / (2 sigma^2))
342     >>> print m.param_string(params)  # doctest: +ELLIPSIS
343     mu=1.4..., sigma=3.1...
344     """
345     def set_data(self, data, info=None, rescale=False):
346         self._data = data  # Histogram() instance
347         self._model_data = deepcopy(data)
348         self._model_data.bin_edges = array(self._model_data.bin_edges)
349         self._model_data.counts = array(self._model_data.counts)
350         if info == None:
351             info = {}
352         self.info = info
353         self._rescale = rescale
354         self._data_scale_factor = 1.0  # bo need to rescale normalized hists.
355         self.info['N'] = self._data.total
356         # assume constant binwidth
357         self.info['binwidth'] = (
358             self._data.bin_edges[1] - self._data.bin_edges[0])
359         self._model_data.bin_centers = (
360             self._model_data.bin_edges[:-1] + self.info['binwidth']/2)
361
362     def residual(self, params):
363         if self._rescale == True:
364             params = [p*s for p,s in zip(params, self._param_scale_factors)]
365         residual = self._data.counts - self.model(params).counts
366         if True:  # fit debugging
367             if not hasattr(self, '_i_'):
368                 self._i_ = 0
369             self._data.to_stream(open('hist-data.%d' % self._i_, 'w'))
370             self._model_data.headings = [str(x) for x in params]
371             self._model_data.to_stream(open('hist-model.%d' % self._i_, 'w'))
372             self._i_ += 1
373         if self._rescale == True:
374             residual /= self._data_scale_factor
375         return residual