Adapt to handle scipy.optimize.leastsq for scipy >= 0.8.0.
[hooke.git] / hooke / util / fit.py
index 0a97135aeba2a31cc3b76965ddbeca5bf563daf8..6267341849f450e757ba1d64193898e3b5c4903c 100644 (file)
@@ -1,4 +1,4 @@
-# Copyright (C) 2010 W. Trevor King <wking@drexel.edu>
+# Copyright (C) 2010-2011 W. Trevor King <wking@drexel.edu>
 #
 # This file is part of Hooke.
 #
 #
 # This file is part of Hooke.
 #
 """
 
 from numpy import arange, ndarray
 """
 
 from numpy import arange, ndarray
+from scipy import __version__ as _scipy_version
 from scipy.optimize import leastsq
 from scipy.optimize import leastsq
-import scipy.optimize
+
+_strings = _scipy_version.split('.')
+# Don't convert third string to an integer in case of (for example) '0.7.2rc3'
+_SCIPY_VERSION = (int(_strings[0]), int(_strings[1]), _strings[2])
+del _strings
 
 
 class PoorFit (ValueError):
 
 
 class PoorFit (ValueError):
@@ -128,6 +133,25 @@ class ModelFitter (object):
     7.000
     >>> print '%.3f' % offset
     -32.890
     7.000
     >>> print '%.3f' % offset
     -32.890
+
+    Test single-parameter models:
+
+    >>> class SingleParameterModel (LinearModel):
+    ...     '''Simple linear model.
+    ...     '''
+    ...     def model(self, params):
+    ...         return super(SingleParameterModel, self).model([params[0], 0.])
+    ...     def guess_initial_params(self, outqueue=None):
+    ...         return super(SingleParameterModel, self
+    ...             ).guess_initial_params(outqueue)[:1]
+    ...     def guess_scale(self, params, outqueue=None):
+    ...         return super(SingleParameterModel, self
+    ...             ).guess_scale([params[0], 0.], outqueue)[:1]
+    >>> data = 20*numpy.sin(arange(1000)) + 7.*arange(1000)
+    >>> m = SingleParameterModel(data)
+    >>> slope, = m.fit(outqueue=outqueue)
+    >>> print '%.3f' % slope
+    7.000
     """
     def __init__(self, *args, **kwargs):
         self.set_data(*args, **kwargs)
     """
     def __init__(self, *args, **kwargs):
         self.set_data(*args, **kwargs)
@@ -233,13 +257,13 @@ class ModelFitter (object):
         params,cov,info,mesg,ier = leastsq(
             func=self.residual, x0=active_params, full_output=True,
             diag=scale, **kwargs)
         params,cov,info,mesg,ier = leastsq(
             func=self.residual, x0=active_params, full_output=True,
             diag=scale, **kwargs)
+        if len(initial_params) == 1 and _SCIPY_VERSION < (0, 8, '0'):
+            # params is a float for scipy < 0.8.0.  Convert to list.
+            params = [params]
         if self._rescale == True:
             active_params = params
         if self._rescale == True:
             active_params = params
-            if len(initial_params) == 1:  # params is a float
-                params = params * self._param_scale_factors[0]
-            else:
-                params = [p*s for p,s in zip(params,
-                                             self._param_scale_factors)]
+            params = [p*s for p,s in zip(params,
+                                         self._param_scale_factors)]
         else:
             active_params = params
         if outqueue != None:
         else:
             active_params = params
         if outqueue != None: