""" Glue's fitting classes are designed to be easily subclassed for performing custom model fitting in Glue. See the guide on :ref:`writing custom fit plugins ` for help with using custom fitting utilities in Glue. """ from __future__ import absolute_import, division, print_function import numpy as np from glue.core.simpleforms import IntOption, Option __all__ = ['BaseFitter1D', 'PolynomialFitter', 'AstropyFitter1D', 'SimpleAstropyGaussianFitter', 'BasicGaussianFitter'] class BaseFitter1D(object): """ Base class for 1D fitters. This abstract class must be overwritten. """ label = "Fitter" """A short label for the fit, used by the GUI""" param_names = [] """list of parameter names that support restrictions""" def __init__(self, **params): self._constraints = {} for k, v in params.items(): if k in self.param_names: self.set_constraint(k, value=v) else: setattr(self, k, v) def plot(self, fit_result, axes, x): """ Plot the result of a fit. :param fit_result: The output from fit :param axes: The Matplotlib axes to add the fit to :param x: The values of X at which to visualize the model :returns: A list of matplotlib artists. **This is important:** plots will not be properly cleared if this isn't provided """ y = self.predict(fit_result, x) result = axes.plot(x, y, '#4daf4a', lw=3, alpha=0.8, scalex=False, scaley=False) return result def _sigma_to_weights(self, dy): if dy is not None: return 1. / np.asarray(dy) ** 2 @property def options(self): """ A dictionary of the current setting of each model hyperparameter. Hyperparameters are defined in subclasses by creating class-level :mod:`Option ` attributes. This attribute dict maps ``{hyperparameter_name: current_value}`` """ result = [] for typ in type(self).mro(): result.extend(k for k, v in typ.__dict__.items() if isinstance(v, Option)) return dict((o, getattr(self, o)) for o in result) def summarize(self, fit_result, x, y, dy=None): """ Return a textual summary of the fit. :param fit_result: The return value from :meth:`fit` :param x: The x values passed to :meth:`fit` :returns: A description of the fit result :rtype: str """ return str(fit_result) @property def constraints(self): """ A dict of the constraints on each parameter in :attr:`param_names`. Each value is itself a dict with 3 items: :key value: The default value :key fixed: True / False, indicating whether the parameter is fixed :key bounds: [min, max] or None, indicating lower/upper limits """ result = {} for p in self.param_names: result[p] = dict(value=None, fixed=False, limits=None) result[p].update(self._constraints.get(p, {})) return result def set_constraint(self, parameter_name, value=None, fixed=None, limits=None): """ Update a constraint. :param parameter_name: name of the parameter to update :type parameter_name: str :param value: Set the default value (optional) :param limits: Set the limits to[min, max] (optional) :param fixed: Set whether the parameter is fixed (optional) """ c = self._constraints.setdefault(parameter_name, {}) if value is not None: c['value'] = value if fixed is not None: c['fixed'] = fixed if limits is not None: c['limits'] = limits def build_and_fit(self, x, y, dy=None): """ Method which builds the arguments to fit, and calls that method """ x = np.asarray(x).ravel() y = np.asarray(y).ravel() if dy is not None: dy = np.asarray(dy).ravel() return self.fit(x, y, dy=dy, constraints=self.constraints, **self.options) def fit(self, x, y, dy, constraints, **options): """ Fit the model to data. *This must be overriden by a subclass.* :param x: The x values of the data :type x: :class:`numpy.ndarray` :param y: The y values of the data :type y: :class:`numpy.ndarray` :param dy: 1 sigma uncertainties on each datum (optional) :type dy: :class:`numpy.ndarray` :param constraints: The current value of :attr:`constraints` :param options: kwargs for model hyperparameters. :returns: An object representing the fit result. """ raise NotImplementedError() def predict(self, fit_result, x): """ Evaulate the model at a set of locations. **This must be overridden in a subclass.** :param fit_result: The result from the fit method :param x: Locations to evaluate model at :type x: :class:`numpy.ndarray` :returns: model(x) :rtype: :class:`numpy.ndarray` """ raise NotImplementedError() class AstropyFitter1D(BaseFitter1D): """ A base class for wrapping :mod:`astropy.modeling`. Subclasses must override :attr:`model_cls` :attr:`fitting_cls` to point to the desired Astropy :mod:`model ` and :mod:`fitter ` classes. In addition, they should override :attr:`label` with a better label, and :meth:`parameter_guesses` to generate initial guesses """ model_cls = None """class describing the model""" fitting_cls = None """class to fit the model""" label = "Base Astropy Fitter" """UI Label""" @property def param_names(self): return self.model_cls.param_names def predict(self, fit_result, x): model, _ = fit_result return model(x) def summarize(self, fit_result, x, y, dy=None): model, fitter = fit_result result = [_report_fitter(fitter), ""] pnames = list(sorted(model.param_names)) maxlen = max(map(len, pnames)) result.extend("%s = %e" % (p.ljust(maxlen), getattr(model, p).value) for p in pnames) return "\n".join(result) def fit(self, x, y, dy, constraints): m, f = self._get_model_fitter(x, y, dy, constraints) dy = self._sigma_to_weights(dy) return f(m, x, y, weights=dy), f def _get_model_fitter(self, x, y, dy, constraints): if self.model_cls is None or self.fitting_cls is None: raise NotImplementedError("Model or fitting class is unspecified.") params = dict((k, v['value']) for k, v in constraints.items()) # update unset parameters with guesses from data for k, v in self.parameter_guesses(x, y, dy).items(): if params[k] is not None or constraints[k]['fixed']: continue params[k] = v m = self.model_cls(**params) f = self.fitting_cls() for param_name, constraint in constraints.items(): param = getattr(m, param_name) if constraint['fixed']: param.fixed = True if constraint['limits']: param.min, param.max = constraint['limits'] return m, f def parameter_guesses(self, x, y, dy): """ Provide initial guesses for each model parameter. **The base implementation does nothing, and should be overridden** :param x: X - values of the data :type x: :class:`numpy.ndarray` :param y: Y - values of the data :type y: :class:`numpy.ndarray` :param dy: ncertainties on Y(assumed to be 1 sigma) :type dy: :class:`numpy.ndarray` :returns: A dict maping ``{parameter_name: value guess}`` for each parameter """ return {} def _gaussian_parameter_estimates(x, y, dy): amplitude = np.percentile(y, 95) y = np.maximum(y / y.sum(), 0) mean = (x * y).sum() stddev = np.sqrt((y * (x - mean) ** 2).sum()) return dict(mean=mean, stddev=stddev, amplitude=amplitude) class BasicGaussianFitter(BaseFitter1D): """ Fallback Gaussian fitter, for astropy < 0.3. If :mod:`astropy.modeling` is installed, this class is replaced by :class:`SimpleAstropyGaussianFitter` """ label = "Gaussian" def _errorfunc(self, params, x, y, dy): yp = self.eval(x, *params) result = (yp - y) if dy is not None: result /= dy return result @staticmethod def eval(x, amplitude, mean, stddev): return np.exp(-(x - mean) ** 2 / (2 * stddev ** 2)) * amplitude @staticmethod def fit_deriv(x, amplitude, mean, stddev): """ Gaussian1D model function derivatives. """ d_amplitude = np.exp(-0.5 / stddev ** 2 * (x - mean) ** 2) d_mean = amplitude * d_amplitude * (x - mean) / stddev ** 2 d_stddev = amplitude * d_amplitude * (x - mean) ** 2 / stddev ** 3 return [d_amplitude, d_mean, d_stddev] def fit(self, x, y, dy, constraints): from scipy import optimize init_values = _gaussian_parameter_estimates(x, y, dy) init_values = [init_values[p] for p in ['amplitude', 'mean', 'stddev']] farg = (x, y, dy) dfunc = None fitparams, status, dinfo, mess, ierr = optimize.leastsq( self._errorfunc, init_values, args=farg, Dfun=dfunc, full_output=True) return fitparams def predict(self, fit_result, x): return self.eval(x, *fit_result) def summarize(self, fit_result, x, y, dy=None): return ("amplitude = %e\n" "mean = %e\n" "stddev = %e" % tuple(fit_result)) GaussianFitter = BasicGaussianFitter try: from astropy.modeling import models, fitting class SimpleAstropyGaussianFitter(AstropyFitter1D): """ Guassian fitter using astropy.modeling. """ model_cls = models.Gaussian1D try: fitting_cls = fitting.LevMarLSQFitter except AttributeError: # astropy v0.3 fitting_cls = fitting.NonLinearLSQFitter label = "Gaussian" parameter_guesses = staticmethod(_gaussian_parameter_estimates) GaussianFitter = SimpleAstropyGaussianFitter except ImportError: pass class PolynomialFitter(BaseFitter1D): """ A polynomial model. The degree of the polynomial is specified by :attr:`degree` """ label = "Polynomial" degree = IntOption(min=0, max=5, default=3, label="Polynomial Degree") def fit(self, x, y, dy, constraints, degree=2): """ Fit a ``degree``-th order polynomial to the data. """ w = self._sigma_to_weights(dy) return np.polyfit(x, y, degree, w=w) def predict(self, fit_result, x): return np.polyval(fit_result, x) def summarize(self, fit_result, x, y, dy=None): return "Coefficients:\n" + "\n".join("%e" % coeff for coeff in fit_result.tolist()) def _report_fitter(fitter): if "nfev" in fitter.fit_info: return "Converged in %i iterations" % fitter.fit_info['nfev'] return 'Converged' __FITTERS__ = [PolynomialFitter, GaussianFitter]