1# Licensed under a 3-clause BSD style license - see LICENSE.rst
2
3"""
4This module implements classes (called Fitters) which combine optimization
5algorithms (typically from `scipy.optimize`) with statistic functions to perform
6fitting. Fitters are implemented as callable classes. In addition to the data
7to fit, the ``__call__`` method takes an instance of
8`~astropy.modeling.core.FittableModel` as input, and returns a copy of the
9model with its parameters determined by the optimizer.
10
11Optimization algorithms, called "optimizers" are implemented in
12`~astropy.modeling.optimizers` and statistic functions are in
13`~astropy.modeling.statistic`. The goal is to provide an easy to extend
14framework and allow users to easily create new fitters by combining statistics
15with optimizers.
16
17There are two exceptions to the above scheme.
18`~astropy.modeling.fitting.LinearLSQFitter` uses Numpy's `~numpy.linalg.lstsq`
19function.  `~astropy.modeling.fitting.LevMarLSQFitter` uses
20`~scipy.optimize.leastsq` which combines optimization and statistic in one
21implementation.
22"""
23# pylint: disable=invalid-name
24
25import abc
26import inspect
27import operator
28import warnings
29from importlib.metadata import entry_points
30
31from functools import reduce, wraps
32
33import numpy as np
34
35from astropy.units import Quantity
36from astropy.utils.exceptions import AstropyUserWarning
37from .utils import poly_map_domain, _combine_equivalency_dict
38from .optimizers import (SLSQP, Simplex)
39from .statistic import (leastsquare)
40from .optimizers import (DEFAULT_MAXITER, DEFAULT_EPS, DEFAULT_ACC)
41from .spline import (SplineInterpolateFitter, SplineSmoothingFitter,
42                     SplineExactKnotsFitter, SplineSplrepFitter)
43
44__all__ = ['LinearLSQFitter', 'LevMarLSQFitter', 'FittingWithOutlierRemoval',
45           'SLSQPLSQFitter', 'SimplexLSQFitter', 'JointFitter', 'Fitter',
46           "ModelLinearityError", "ModelsError"]
47
48
49# Statistic functions implemented in `astropy.modeling.statistic.py
50STATISTICS = [leastsquare]
51
52# Optimizers implemented in `astropy.modeling.optimizers.py
53OPTIMIZERS = [Simplex, SLSQP]
54
55
56class Covariance():
57    """Class for covariance matrix calculated by fitter. """
58
59    def __init__(self, cov_matrix, param_names):
60        self.cov_matrix = cov_matrix
61        self.param_names = param_names
62
63    def pprint(self, max_lines, round_val):
64        # Print and label lower triangle of covariance matrix
65        # Print rows for params up to `max_lines`, round floats to 'round_val'
66        longest_name = max([len(x) for x in self.param_names])
67        ret_str = 'parameter variances / covariances \n'
68        fstring = f'{"": <{longest_name}}| {{0}}\n'
69        for i, row in enumerate(self.cov_matrix):
70            if i <= max_lines-1:
71                param = self.param_names[i]
72                ret_str += fstring.replace(' '*len(param), param, 1).\
73                           format(repr(np.round(row[:i+1], round_val))[7:-2])
74            else:
75                ret_str += '...'
76        return(ret_str.rstrip())
77
78    def __repr__(self):
79        return(self.pprint(max_lines=10, round_val=3))
80
81    def __getitem__(self, params):
82        # index covariance matrix by parameter names or indices
83        if len(params) != 2:
84            raise ValueError('Covariance must be indexed by two values.')
85        if all(isinstance(item, str) for item in params):
86            i1, i2 = self.param_names.index(params[0]), self.param_names.index(params[1])
87        elif all(isinstance(item, int) for item in params):
88            i1, i2 = params
89        else:
90            raise TypeError('Covariance can be indexed by two parameter names or integer indices.')
91        return(self.cov_matrix[i1][i2])
92
93
94class StandardDeviations():
95    """ Class for fitting uncertainties."""
96
97    def __init__(self, cov_matrix, param_names):
98        self.param_names = param_names
99        self.stds = self._calc_stds(cov_matrix)
100
101    def _calc_stds(self, cov_matrix):
102        # sometimes scipy lstsq returns a non-sensical negative vals in the
103        # diagonals of the cov_x it computes.
104        stds = [np.sqrt(x) if x > 0 else None for x in np.diag(cov_matrix)]
105        return stds
106
107    def pprint(self, max_lines, round_val):
108        longest_name = max([len(x) for x in self.param_names])
109        ret_str = 'standard deviations\n'
110        fstring = '{0}{1}| {2}\n'
111        for i, std in enumerate(self.stds):
112            if i <= max_lines-1:
113                param = self.param_names[i]
114                ret_str += fstring.format(param,
115                                          ' ' * (longest_name - len(param)),
116                                          str(np.round(std, round_val)))
117            else:
118                ret_str += '...'
119        return(ret_str.rstrip())
120
121    def __repr__(self):
122        return(self.pprint(max_lines=10, round_val=3))
123
124    def __getitem__(self, param):
125        if isinstance(param, str):
126            i = self.param_names.index(param)
127        elif isinstance(param, int):
128            i = param
129        else:
130            raise TypeError('Standard deviation can be indexed by parameter name or integer.')
131        return(self.stds[i])
132
133
134class ModelsError(Exception):
135    """Base class for model exceptions"""
136
137
138class ModelLinearityError(ModelsError):
139    """ Raised when a non-linear model is passed to a linear fitter."""
140
141
142class UnsupportedConstraintError(ModelsError, ValueError):
143    """
144    Raised when a fitter does not support a type of constraint.
145    """
146
147
148class _FitterMeta(abc.ABCMeta):
149    """
150    Currently just provides a registry for all Fitter classes.
151    """
152
153    registry = set()
154
155    def __new__(mcls, name, bases, members):
156        cls = super().__new__(mcls, name, bases, members)
157
158        if not inspect.isabstract(cls) and not name.startswith('_'):
159            mcls.registry.add(cls)
160
161        return cls
162
163
164def fitter_unit_support(func):
165    """
166    This is a decorator that can be used to add support for dealing with
167    quantities to any __call__ method on a fitter which may not support
168    quantities itself. This is done by temporarily removing units from all
169    parameters then adding them back once the fitting has completed.
170    """
171    @wraps(func)
172    def wrapper(self, model, x, y, z=None, **kwargs):
173        equivalencies = kwargs.pop('equivalencies', None)
174
175        data_has_units = (isinstance(x, Quantity) or
176                          isinstance(y, Quantity) or
177                          isinstance(z, Quantity))
178
179        model_has_units = model._has_units
180
181        if data_has_units or model_has_units:
182
183            if model._supports_unit_fitting:
184
185                # We now combine any instance-level input equivalencies with user
186                # specified ones at call-time.
187
188                input_units_equivalencies = _combine_equivalency_dict(
189                    model.inputs, equivalencies, model.input_units_equivalencies)
190
191                # If input_units is defined, we transform the input data into those
192                # expected by the model. We hard-code the input names 'x', and 'y'
193                # here since FittableModel instances have input names ('x',) or
194                # ('x', 'y')
195
196                if model.input_units is not None:
197                    if isinstance(x, Quantity):
198                        x = x.to(model.input_units[model.inputs[0]],
199                                 equivalencies=input_units_equivalencies[model.inputs[0]])
200                    if isinstance(y, Quantity) and z is not None:
201                        y = y.to(model.input_units[model.inputs[1]],
202                                 equivalencies=input_units_equivalencies[model.inputs[1]])
203
204                # Create a dictionary mapping the real model inputs and outputs
205                # names to the data. This remapping of names must be done here, after
206                # the input data is converted to the correct units.
207                rename_data = {model.inputs[0]: x}
208                if z is not None:
209                    rename_data[model.outputs[0]] = z
210                    rename_data[model.inputs[1]] = y
211                else:
212                    rename_data[model.outputs[0]] = y
213                    rename_data['z'] = None
214
215                # We now strip away the units from the parameters, taking care to
216                # first convert any parameters to the units that correspond to the
217                # input units (to make sure that initial guesses on the parameters)
218                # are in the right unit system
219                model = model.without_units_for_data(**rename_data)
220                # We strip away the units from the input itself
221                add_back_units = False
222
223                if isinstance(x, Quantity):
224                    add_back_units = True
225                    xdata = x.value
226                else:
227                    xdata = np.asarray(x)
228
229                if isinstance(y, Quantity):
230                    add_back_units = True
231                    ydata = y.value
232                else:
233                    ydata = np.asarray(y)
234
235                if z is not None:
236                    if isinstance(z, Quantity):
237                        add_back_units = True
238                        zdata = z.value
239                    else:
240                        zdata = np.asarray(z)
241                # We run the fitting
242                if z is None:
243                    model_new = func(self, model, xdata, ydata, **kwargs)
244                else:
245                    model_new = func(self, model, xdata, ydata, zdata, **kwargs)
246
247                # And finally we add back units to the parameters
248                if add_back_units:
249                    model_new = model_new.with_units_from_data(**rename_data)
250                return model_new
251
252            else:
253
254                raise NotImplementedError("This model does not support being "
255                                          "fit to data with units.")
256
257        else:
258
259            return func(self, model, x, y, z=z, **kwargs)
260
261    return wrapper
262
263
264class Fitter(metaclass=_FitterMeta):
265    """
266    Base class for all fitters.
267
268    Parameters
269    ----------
270    optimizer : callable
271        A callable implementing an optimization algorithm
272    statistic : callable
273        Statistic function
274
275    """
276
277    supported_constraints = []
278
279    def __init__(self, optimizer, statistic):
280        if optimizer is None:
281            raise ValueError("Expected an optimizer.")
282        if statistic is None:
283            raise ValueError("Expected a statistic function.")
284        if inspect.isclass(optimizer):
285            # a callable class
286            self._opt_method = optimizer()
287        elif inspect.isfunction(optimizer):
288            self._opt_method = optimizer
289        else:
290            raise ValueError("Expected optimizer to be a callable class or a function.")
291        if inspect.isclass(statistic):
292            self._stat_method = statistic()
293        else:
294            self._stat_method = statistic
295
296    def objective_function(self, fps, *args):
297        """
298        Function to minimize.
299
300        Parameters
301        ----------
302        fps : list
303            parameters returned by the fitter
304        args : list
305            [model, [other_args], [input coordinates]]
306            other_args may include weights or any other quantities specific for
307            a statistic
308
309        Notes
310        -----
311        The list of arguments (args) is set in the `__call__` method.
312        Fitters may overwrite this method, e.g. when statistic functions
313        require other arguments.
314
315        """
316        model = args[0]
317        meas = args[-1]
318        _fitter_to_model_params(model, fps)
319        res = self._stat_method(meas, model, *args[1:-1])
320        return res
321
322    @staticmethod
323    def _add_fitting_uncertainties(*args):
324        """
325        When available, calculate and sets the parameter covariance matrix
326        (model.cov_matrix) and standard deviations (model.stds).
327        """
328        return None
329
330    @abc.abstractmethod
331    def __call__(self):
332        """
333        This method performs the actual fitting and modifies the parameter list
334        of a model.
335        Fitter subclasses should implement this method.
336        """
337
338        raise NotImplementedError("Subclasses should implement this method.")
339
340
341# TODO: I have ongoing branch elsewhere that's refactoring this module so that
342# all the fitter classes in here are Fitter subclasses.  In the meantime we
343# need to specify that _FitterMeta is its metaclass.
344class LinearLSQFitter(metaclass=_FitterMeta):
345    """
346    A class performing a linear least square fitting.
347    Uses `numpy.linalg.lstsq` to do the fitting.
348    Given a model and data, fits the model to the data and changes the
349    model's parameters. Keeps a dictionary of auxiliary fitting information.
350    Notes
351    -----
352    Note that currently LinearLSQFitter does not support compound models.
353    """
354
355    supported_constraints = ['fixed']
356    supports_masked_input = True
357
358    def __init__(self, calc_uncertainties=False):
359        self.fit_info = {'residuals': None,
360                         'rank': None,
361                         'singular_values': None,
362                         'params': None
363                         }
364        self._calc_uncertainties=calc_uncertainties
365
366    @staticmethod
367    def _is_invertible(m):
368        """Check if inverse of matrix can be obtained."""
369        if m.shape[0] != m.shape[1]:
370            return False
371        if np.linalg.matrix_rank(m) < m.shape[0]:
372            return False
373        return True
374
375    def _add_fitting_uncertainties(self, model, a, n_coeff, x, y, z=None,
376                                   resids=None):
377        """
378        Calculate and parameter covariance matrix and standard deviations
379        and set `cov_matrix` and `stds` attributes.
380        """
381        x_dot_x_prime = np.dot(a.T, a)
382        masked = False or hasattr(y, 'mask')
383
384        # check if invertible. if not, can't calc covariance.
385        if not self._is_invertible(x_dot_x_prime):
386            return(model)
387        inv_x_dot_x_prime = np.linalg.inv(x_dot_x_prime)
388
389        if z is None:  # 1D models
390            if len(model) == 1:  # single model
391                mask = None
392                if masked:
393                    mask = y.mask
394                xx = np.ma.array(x, mask=mask)
395                RSS = [(1/(xx.count()-n_coeff)) * resids]
396
397            if len(model) > 1:  # model sets
398                RSS = []   # collect sum residuals squared for each model in set
399                for j in range(len(model)):
400                    mask = None
401                    if masked:
402                        mask = y.mask[..., j].flatten()
403                    xx = np.ma.array(x, mask=mask)
404                    eval_y = model(xx, model_set_axis=False)
405                    eval_y = np.rollaxis(eval_y, model.model_set_axis)[j]
406                    RSS.append((1/(xx.count()-n_coeff)) * np.sum((y[..., j] - eval_y)**2))
407
408        else:  # 2D model
409            if len(model) == 1:
410                mask = None
411                if masked:
412                    warnings.warn('Calculation of fitting uncertainties '
413                                  'for 2D models with masked values not '
414                                  'currently supported.\n',
415                                  AstropyUserWarning)
416                    return
417                xx, yy = np.ma.array(x, mask=mask), np.ma.array(y, mask=mask)
418                # len(xx) instead of xx.count. this will break if values are masked?
419                RSS = [(1/(len(xx)-n_coeff)) * resids]
420            else:
421                RSS = []
422                for j in range(len(model)):
423                    eval_z = model(x, y, model_set_axis=False)
424                    mask = None  # need to figure out how to deal w/ masking here.
425                    if model.model_set_axis == 1:
426                        # model_set_axis passed when evaluating only refers to input shapes
427                        # so output must be reshaped for model_set_axis=1.
428                        eval_z = np.rollaxis(eval_z, 1)
429                    eval_z = eval_z[j]
430                    RSS.append([(1/(len(x)-n_coeff)) * np.sum((z[j] - eval_z)**2)])
431
432        covs = [inv_x_dot_x_prime * r for r in RSS]
433        free_param_names = [x for x in model.fixed if (model.fixed[x] is False)
434                            and (model.tied[x] is False)]
435
436        if len(covs) == 1:
437            model.cov_matrix = Covariance(covs[0], model.param_names)
438            model.stds = StandardDeviations(covs[0], free_param_names)
439        else:
440            model.cov_matrix = [Covariance(cov, model.param_names) for cov in covs]
441            model.stds = [StandardDeviations(cov, free_param_names) for cov in covs]
442
443    @staticmethod
444    def _deriv_with_constraints(model, param_indices, x=None, y=None):
445        if y is None:
446            d = np.array(model.fit_deriv(x, *model.parameters))
447        else:
448            d = np.array(model.fit_deriv(x, y, *model.parameters))
449
450        if model.col_fit_deriv:
451            return d[param_indices]
452        else:
453            return d[..., param_indices]
454
455    def _map_domain_window(self, model, x, y=None):
456        """
457        Maps domain into window for a polynomial model which has these
458        attributes.
459        """
460
461        if y is None:
462            if hasattr(model, 'domain') and model.domain is None:
463                model.domain = [x.min(), x.max()]
464            if hasattr(model, 'window') and model.window is None:
465                model.window = [-1, 1]
466            return poly_map_domain(x, model.domain, model.window)
467        else:
468            if hasattr(model, 'x_domain') and model.x_domain is None:
469                model.x_domain = [x.min(), x.max()]
470            if hasattr(model, 'y_domain') and model.y_domain is None:
471                model.y_domain = [y.min(), y.max()]
472            if hasattr(model, 'x_window') and model.x_window is None:
473                model.x_window = [-1., 1.]
474            if hasattr(model, 'y_window') and model.y_window is None:
475                model.y_window = [-1., 1.]
476
477            xnew = poly_map_domain(x, model.x_domain, model.x_window)
478            ynew = poly_map_domain(y, model.y_domain, model.y_window)
479            return xnew, ynew
480
481    @fitter_unit_support
482    def __call__(self, model, x, y, z=None, weights=None, rcond=None):
483        """
484        Fit data to this model.
485
486        Parameters
487        ----------
488        model : `~astropy.modeling.FittableModel`
489            model to fit to x, y, z
490        x : array
491            Input coordinates
492        y : array-like
493            Input coordinates
494        z : array-like, optional
495            Input coordinates.
496            If the dependent (``y`` or ``z``) coordinate values are provided
497            as a `numpy.ma.MaskedArray`, any masked points are ignored when
498            fitting. Note that model set fitting is significantly slower when
499            there are masked points (not just an empty mask), as the matrix
500            equation has to be solved for each model separately when their
501            coordinate grids differ.
502        weights : array, optional
503            Weights for fitting.
504            For data with Gaussian uncertainties, the weights should be
505            1/sigma.
506        rcond :  float, optional
507            Cut-off ratio for small singular values of ``a``.
508            Singular values are set to zero if they are smaller than ``rcond``
509            times the largest singular value of ``a``.
510        equivalencies : list or None, optional, keyword-only
511            List of *additional* equivalencies that are should be applied in
512            case x, y and/or z have units. Default is None.
513
514        Returns
515        -------
516        model_copy : `~astropy.modeling.FittableModel`
517            a copy of the input model with parameters set by the fitter
518
519        """
520
521        if not model.fittable:
522            raise ValueError("Model must be a subclass of FittableModel")
523
524        if not model.linear:
525            raise ModelLinearityError('Model is not linear in parameters, '
526                                      'linear fit methods should not be used.')
527
528        if hasattr(model, "submodel_names"):
529            raise ValueError("Model must be simple, not compound")
530
531        _validate_constraints(self.supported_constraints, model)
532
533        model_copy = model.copy()
534        model_copy.sync_constraints = False
535        _, fitparam_indices = _model_to_fit_params(model_copy)
536
537        if model_copy.n_inputs == 2 and z is None:
538            raise ValueError("Expected x, y and z for a 2 dimensional model.")
539
540        farg = _convert_input(x, y, z, n_models=len(model_copy),
541                              model_set_axis=model_copy.model_set_axis)
542
543        has_fixed = any(model_copy.fixed.values())
544
545        # This is also done by _convert_inputs, but we need it here to allow
546        # checking the array dimensionality before that gets called:
547        if weights is not None:
548            weights = np.asarray(weights, dtype=float)
549
550        if has_fixed:
551
552            # The list of fixed params is the complement of those being fitted:
553            fixparam_indices = [idx for idx in
554                                range(len(model_copy.param_names))
555                                if idx not in fitparam_indices]
556
557            # Construct matrix of user-fixed parameters that can be dotted with
558            # the corresponding fit_deriv() terms, to evaluate corrections to
559            # the dependent variable in order to fit only the remaining terms:
560            fixparams = np.asarray([getattr(model_copy,
561                                            model_copy.param_names[idx]).value
562                                    for idx in fixparam_indices])
563
564        if len(farg) == 2:
565            x, y = farg
566
567            if weights is not None:
568                # If we have separate weights for each model, apply the same
569                # conversion as for the data, otherwise check common weights
570                # as if for a single model:
571                _, weights = _convert_input(
572                    x, weights,
573                    n_models=len(model_copy) if weights.ndim == y.ndim else 1,
574                    model_set_axis=model_copy.model_set_axis
575                )
576
577            # map domain into window
578            if hasattr(model_copy, 'domain'):
579                x = self._map_domain_window(model_copy, x)
580            if has_fixed:
581                lhs = np.asarray(self._deriv_with_constraints(model_copy,
582                                                              fitparam_indices,
583                                                              x=x))
584                fixderivs = self._deriv_with_constraints(model_copy, fixparam_indices, x=x)
585            else:
586                lhs = np.asarray(model_copy.fit_deriv(x, *model_copy.parameters))
587            sum_of_implicit_terms = model_copy.sum_of_implicit_terms(x)
588            rhs = y
589        else:
590            x, y, z = farg
591
592            if weights is not None:
593                # If we have separate weights for each model, apply the same
594                # conversion as for the data, otherwise check common weights
595                # as if for a single model:
596                _, _, weights = _convert_input(
597                    x, y, weights,
598                    n_models=len(model_copy) if weights.ndim == z.ndim else 1,
599                    model_set_axis=model_copy.model_set_axis
600                )
601
602            # map domain into window
603            if hasattr(model_copy, 'x_domain'):
604                x, y = self._map_domain_window(model_copy, x, y)
605
606            if has_fixed:
607                lhs = np.asarray(self._deriv_with_constraints(model_copy,
608                                                              fitparam_indices, x=x, y=y))
609                fixderivs = self._deriv_with_constraints(model_copy,
610                                                         fixparam_indices,
611                                                         x=x, y=y)
612            else:
613                lhs = np.asanyarray(model_copy.fit_deriv(x, y, *model_copy.parameters))
614            sum_of_implicit_terms = model_copy.sum_of_implicit_terms(x, y)
615
616            if len(model_copy) > 1:
617
618                # Just to be explicit (rather than baking in False == 0):
619                model_axis = model_copy.model_set_axis or 0
620
621                if z.ndim > 2:
622                    # For higher-dimensional z, flatten all the axes except the
623                    # dimension along which models are stacked and transpose so
624                    # the model axis is *last* (I think this resolves Erik's
625                    # pending generalization from 80a6f25a):
626                    rhs = np.rollaxis(z, model_axis, z.ndim)
627                    rhs = rhs.reshape(-1, rhs.shape[-1])
628                else:
629                    # This "else" seems to handle the corner case where the
630                    # user has already flattened x/y before attempting a 2D fit
631                    # but z has a second axis for the model set. NB. This is
632                    # ~5-10x faster than using rollaxis.
633                    rhs = z.T if model_axis == 0 else z
634
635                if weights is not None:
636                    # Same for weights
637                    if weights.ndim > 2:
638                        # Separate 2D weights for each model:
639                        weights = np.rollaxis(weights, model_axis, weights.ndim)
640                        weights = weights.reshape(-1, weights.shape[-1])
641                    elif weights.ndim == z.ndim:
642                        # Separate, flattened weights for each model:
643                        weights = weights.T if model_axis == 0 else weights
644                    else:
645                        # Common weights for all the models:
646                        weights = weights.flatten()
647            else:
648                rhs = z.flatten()
649                if weights is not None:
650                    weights = weights.flatten()
651
652        # If the derivative is defined along rows (as with non-linear models)
653        if model_copy.col_fit_deriv:
654            lhs = np.asarray(lhs).T
655
656        # Some models (eg. Polynomial1D) don't flatten multi-dimensional inputs
657        # when constructing their Vandermonde matrix, which can lead to obscure
658        # failures below. Ultimately, np.linalg.lstsq can't handle >2D matrices,
659        # so just raise a slightly more informative error when this happens:
660        if np.asanyarray(lhs).ndim > 2:
661            raise ValueError('{} gives unsupported >2D derivative matrix for '
662                             'this x/y'.format(type(model_copy).__name__))
663
664        # Subtract any terms fixed by the user from (a copy of) the RHS, in
665        # order to fit the remaining terms correctly:
666        if has_fixed:
667            if model_copy.col_fit_deriv:
668                fixderivs = np.asarray(fixderivs).T  # as for lhs above
669            rhs = rhs - fixderivs.dot(fixparams)  # evaluate user-fixed terms
670
671        # Subtract any terms implicit in the model from the RHS, which, like
672        # user-fixed terms, affect the dependent variable but are not fitted:
673        if sum_of_implicit_terms is not None:
674            # If we have a model set, the extra axis must be added to
675            # sum_of_implicit_terms as its innermost dimension, to match the
676            # dimensionality of rhs after _convert_input "rolls" it as needed
677            # by np.linalg.lstsq. The vector then gets broadcast to the right
678            # number of sets (columns). This assumes all the models share the
679            # same input coordinates, as is currently the case.
680            if len(model_copy) > 1:
681                sum_of_implicit_terms = sum_of_implicit_terms[..., np.newaxis]
682            rhs = rhs - sum_of_implicit_terms
683
684        if weights is not None:
685
686            if rhs.ndim == 2:
687                if weights.shape == rhs.shape:
688                    # separate weights for multiple models case: broadcast
689                    # lhs to have more dimension (for each model)
690                    lhs = lhs[..., np.newaxis] * weights[:, np.newaxis]
691                    rhs = rhs * weights
692                else:
693                    lhs *= weights[:, np.newaxis]
694                    # Don't modify in-place in case rhs was the original
695                    # dependent variable array
696                    rhs = rhs * weights[:, np.newaxis]
697            else:
698                lhs *= weights[:, np.newaxis]
699                rhs = rhs * weights
700
701        scl = (lhs * lhs).sum(0)
702        lhs /= scl
703
704        masked = np.any(np.ma.getmask(rhs))
705        if weights is not None and not masked and np.any(np.isnan(lhs)):
706            raise ValueError('Found NaNs in the coefficient matrix, which '
707                             'should not happen and would crash the lapack '
708                             'routine. Maybe check that weights are not null.')
709
710        a = None  # need for calculating covarience
711
712        if ((masked and len(model_copy) > 1) or
713                (weights is not None and weights.ndim > 1)):
714
715            # Separate masks or weights for multiple models case: Numpy's
716            # lstsq supports multiple dimensions only for rhs, so we need to
717            # loop manually on the models. This may be fixed in the future
718            # with https://github.com/numpy/numpy/pull/15777.
719
720            # Initialize empty array of coefficients and populate it one model
721            # at a time. The shape matches the number of coefficients from the
722            # Vandermonde matrix and the number of models from the RHS:
723            lacoef = np.zeros(lhs.shape[1:2] + rhs.shape[-1:], dtype=rhs.dtype)
724
725            # Arrange the lhs as a stack of 2D matrices that we can iterate
726            # over to get the correctly-orientated lhs for each model:
727            if lhs.ndim > 2:
728                lhs_stack = np.rollaxis(lhs, -1, 0)
729            else:
730                lhs_stack = np.broadcast_to(lhs, rhs.shape[-1:] + lhs.shape)
731
732            # Loop over the models and solve for each one. By this point, the
733            # model set axis is the second of two. Transpose rather than using,
734            # say, np.moveaxis(array, -1, 0), since it's slightly faster and
735            # lstsq can't handle >2D arrays anyway. This could perhaps be
736            # optimized by collecting together models with identical masks
737            # (eg. those with no rejected points) into one operation, though it
738            # will still be relatively slow when calling lstsq repeatedly.
739            for model_lhs, model_rhs, model_lacoef in zip(lhs_stack, rhs.T, lacoef.T):
740
741                # Cull masked points on both sides of the matrix equation:
742                good = ~model_rhs.mask if masked else slice(None)
743                model_lhs = model_lhs[good]
744                model_rhs = model_rhs[good][..., np.newaxis]
745                a = model_lhs
746
747                # Solve for this model:
748                t_coef, resids, rank, sval = np.linalg.lstsq(model_lhs,
749                                                             model_rhs, rcond)
750                model_lacoef[:] = t_coef.T
751
752        else:
753
754            # If we're fitting one or more models over a common set of points,
755            # we only have to solve a single matrix equation, which is an order
756            # of magnitude faster than calling lstsq() once per model below:
757
758            good = ~rhs.mask if masked else slice(None)  # latter is a no-op
759            a = lhs[good]
760            # Solve for one or more models:
761            lacoef, resids, rank, sval = np.linalg.lstsq(lhs[good],
762                                                         rhs[good], rcond)
763
764        self.fit_info['residuals'] = resids
765        self.fit_info['rank'] = rank
766        self.fit_info['singular_values'] = sval
767
768        lacoef /= scl[:, np.newaxis] if scl.ndim < rhs.ndim else scl
769        self.fit_info['params'] = lacoef
770
771        _fitter_to_model_params(model_copy, lacoef.flatten())
772
773        # TODO: Only Polynomial models currently have an _order attribute;
774        # maybe change this to read isinstance(model, PolynomialBase)
775        if hasattr(model_copy, '_order') and len(model_copy) == 1 \
776                and not has_fixed and rank != model_copy._order:
777            warnings.warn("The fit may be poorly conditioned\n",
778                          AstropyUserWarning)
779
780        # calculate and set covariance matrix and standard devs. on model
781        if self._calc_uncertainties:
782            if len(y) > len(lacoef):
783                self._add_fitting_uncertainties(model_copy, a*scl,
784                                               len(lacoef), x, y, z, resids)
785        model_copy.sync_constraints = True
786        return model_copy
787
788
789class FittingWithOutlierRemoval:
790    """
791    This class combines an outlier removal technique with a fitting procedure.
792    Basically, given a maximum number of iterations ``niter``, outliers are
793    removed and fitting is performed for each iteration, until no new outliers
794    are found or ``niter`` is reached.
795
796    Parameters
797    ----------
798    fitter : `Fitter`
799        An instance of any Astropy fitter, i.e., LinearLSQFitter,
800        LevMarLSQFitter, SLSQPLSQFitter, SimplexLSQFitter, JointFitter. For
801        model set fitting, this must understand masked input data (as
802        indicated by the fitter class attribute ``supports_masked_input``).
803    outlier_func : callable
804        A function for outlier removal.
805        If this accepts an ``axis`` parameter like the `numpy` functions, the
806        appropriate value will be supplied automatically when fitting model
807        sets (unless overridden in ``outlier_kwargs``), to find outliers for
808        each model separately; otherwise, the same filtering must be performed
809        in a loop over models, which is almost an order of magnitude slower.
810    niter : int, optional
811        Maximum number of iterations.
812    outlier_kwargs : dict, optional
813        Keyword arguments for outlier_func.
814
815    Attributes
816    ----------
817    fit_info : dict
818        The ``fit_info`` (if any) from the last iteration of the wrapped
819        ``fitter`` during the most recent fit. An entry is also added with the
820        keyword ``niter`` that records the actual number of fitting iterations
821        performed (as opposed to the user-specified maximum).
822    """
823
824    def __init__(self, fitter, outlier_func, niter=3, **outlier_kwargs):
825        self.fitter = fitter
826        self.outlier_func = outlier_func
827        self.niter = niter
828        self.outlier_kwargs = outlier_kwargs
829        self.fit_info = {'niter': None}
830
831    def __str__(self):
832        return ("Fitter: {0}\nOutlier function: {1}\nNum. of iterations: {2}" +
833                ("\nOutlier func. args.: {3}"))\
834                .format(self.fitter.__class__.__name__,
835                        self.outlier_func.__name__, self.niter,
836                        self.outlier_kwargs)
837
838    def __repr__(self):
839        return ("{0}(fitter: {1}, outlier_func: {2}," +
840                " niter: {3}, outlier_kwargs: {4})")\
841                 .format(self.__class__.__name__,
842                         self.fitter.__class__.__name__,
843                         self.outlier_func.__name__, self.niter,
844                         self.outlier_kwargs)
845
846    def __call__(self, model, x, y, z=None, weights=None, **kwargs):
847        """
848        Parameters
849        ----------
850        model : `~astropy.modeling.FittableModel`
851            An analytic model which will be fit to the provided data.
852            This also contains the initial guess for an optimization
853            algorithm.
854        x : array-like
855            Input coordinates.
856        y : array-like
857            Data measurements (1D case) or input coordinates (2D case).
858        z : array-like, optional
859            Data measurements (2D case).
860        weights : array-like, optional
861            Weights to be passed to the fitter.
862        kwargs : dict, optional
863            Keyword arguments to be passed to the fitter.
864        Returns
865        -------
866        fitted_model : `~astropy.modeling.FittableModel`
867            Fitted model after outlier removal.
868        mask : `numpy.ndarray`
869            Boolean mask array, identifying which points were used in the final
870            fitting iteration (False) and which were found to be outliers or
871            were masked in the input (True).
872        """
873
874        # For single models, the data get filtered here at each iteration and
875        # then passed to the fitter, which is the historical behavior and
876        # works even for fitters that don't understand masked arrays. For model
877        # sets, the fitter must be able to filter masked data internally,
878        # because fitters require a single set of x/y coordinates whereas the
879        # eliminated points can vary between models. To avoid this limitation,
880        # we could fall back to looping over individual model fits, but it
881        # would likely be fiddly and involve even more overhead (and the
882        # non-linear fitters don't work with model sets anyway, as of writing).
883
884        if len(model) == 1:
885            model_set_axis = None
886        else:
887            if not hasattr(self.fitter, 'supports_masked_input') or \
888               self.fitter.supports_masked_input is not True:
889                raise ValueError("{} cannot fit model sets with masked "
890                                 "values".format(type(self.fitter).__name__))
891
892            # Fitters use their input model's model_set_axis to determine how
893            # their input data are stacked:
894            model_set_axis = model.model_set_axis
895        # Construct input coordinate tuples for fitters & models that are
896        # appropriate for the dimensionality being fitted:
897        if z is None:
898            coords = (x, )
899            data = y
900        else:
901            coords = x, y
902            data = z
903
904        # For model sets, construct a numpy-standard "axis" tuple for the
905        # outlier function, to treat each model separately (if supported):
906        if model_set_axis is not None:
907
908            if model_set_axis < 0:
909                model_set_axis += data.ndim
910
911            if 'axis' not in self.outlier_kwargs:  # allow user override
912                # This also works for False (like model instantiation):
913                self.outlier_kwargs['axis'] = tuple(
914                    n for n in range(data.ndim) if n != model_set_axis
915                )
916
917        loop = False
918
919        # Starting fit, prior to any iteration and masking:
920        fitted_model = self.fitter(model, x, y, z, weights=weights, **kwargs)
921        filtered_data = np.ma.masked_array(data)
922        if filtered_data.mask is np.ma.nomask:
923            filtered_data.mask = False
924        filtered_weights = weights
925        last_n_masked = filtered_data.mask.sum()
926        n = 0  # (allow recording no. of iterations when 0)
927
928        # Perform the iterative fitting:
929        for n in range(1, self.niter + 1):
930
931            # (Re-)evaluate the last model:
932            model_vals = fitted_model(*coords, model_set_axis=False)
933
934            # Determine the outliers:
935            if not loop:
936
937                # Pass axis parameter if outlier_func accepts it, otherwise
938                # prepare for looping over models:
939                try:
940                    filtered_data = self.outlier_func(
941                        filtered_data - model_vals, **self.outlier_kwargs
942                    )
943                # If this happens to catch an error with a parameter other
944                # than axis, the next attempt will fail accordingly:
945                except TypeError:
946                    if model_set_axis is None:
947                        raise
948                    else:
949                        self.outlier_kwargs.pop('axis', None)
950                        loop = True
951
952                        # Construct MaskedArray to hold filtered values:
953                        filtered_data = np.ma.masked_array(
954                            filtered_data,
955                            dtype=np.result_type(filtered_data, model_vals),
956                            copy=True
957                        )
958                        # Make sure the mask is an array, not just nomask:
959                        if filtered_data.mask is np.ma.nomask:
960                            filtered_data.mask = False
961
962                        # Get views transposed appropriately for iteration
963                        # over the set (handling data & mask separately due to
964                        # NumPy issue #8506):
965                        data_T = np.rollaxis(filtered_data, model_set_axis, 0)
966                        mask_T = np.rollaxis(filtered_data.mask,
967                                             model_set_axis, 0)
968
969            if loop:
970                model_vals_T = np.rollaxis(model_vals, model_set_axis, 0)
971                for row_data, row_mask, row_mod_vals in zip(data_T, mask_T,
972                                                            model_vals_T):
973                    masked_residuals = self.outlier_func(
974                        row_data - row_mod_vals, **self.outlier_kwargs
975                    )
976                    row_data.data[:] = masked_residuals.data
977                    row_mask[:] = masked_residuals.mask
978
979                # Issue speed warning after the fact, so it only shows up when
980                # the TypeError is genuinely due to the axis argument.
981                warnings.warn('outlier_func did not accept axis argument; '
982                              'reverted to slow loop over models.',
983                              AstropyUserWarning)
984
985            # Recombine newly-masked residuals with model to get masked values:
986            filtered_data += model_vals
987
988            # Re-fit the data after filtering, passing masked/unmasked values
989            # for single models / sets, respectively:
990            if model_set_axis is None:
991
992                good = ~filtered_data.mask
993
994                if weights is not None:
995                    filtered_weights = weights[good]
996
997                fitted_model = self.fitter(fitted_model,
998                                           *(c[good] for c in coords),
999                                           filtered_data.data[good],
1000                                           weights=filtered_weights, **kwargs)
1001            else:
1002                fitted_model = self.fitter(fitted_model, *coords,
1003                                           filtered_data,
1004                                           weights=filtered_weights, **kwargs)
1005
1006            # Stop iteration if the masked points are no longer changing (with
1007            # cumulative rejection we only need to compare how many there are):
1008            this_n_masked = filtered_data.mask.sum()  # (minimal overhead)
1009            if this_n_masked == last_n_masked:
1010                break
1011            last_n_masked = this_n_masked
1012
1013        self.fit_info = {'niter': n}
1014        self.fit_info.update(getattr(self.fitter, 'fit_info', {}))
1015
1016        return fitted_model, filtered_data.mask
1017
1018
1019class LevMarLSQFitter(metaclass=_FitterMeta):
1020    """
1021    Levenberg-Marquardt algorithm and least squares statistic.
1022
1023    Attributes
1024    ----------
1025    fit_info : dict
1026        The `scipy.optimize.leastsq` result for the most recent fit (see
1027        notes).
1028
1029    Notes
1030    -----
1031    The ``fit_info`` dictionary contains the values returned by
1032    `scipy.optimize.leastsq` for the most recent fit, including the values from
1033    the ``infodict`` dictionary it returns. See the `scipy.optimize.leastsq`
1034    documentation for details on the meaning of these values. Note that the
1035    ``x`` return value is *not* included (as it is instead the parameter values
1036    of the returned model).
1037    Additionally, one additional element of ``fit_info`` is computed whenever a
1038    model is fit, with the key 'param_cov'. The corresponding value is the
1039    covariance matrix of the parameters as a 2D numpy array.  The order of the
1040    matrix elements matches the order of the parameters in the fitted model
1041    (i.e., the same order as ``model.param_names``).
1042
1043    """
1044
1045    supported_constraints = ['fixed', 'tied', 'bounds']
1046    """
1047    The constraint types supported by this fitter type.
1048    """
1049
1050    def __init__(self, calc_uncertainties=False):
1051        self.fit_info = {'nfev': None,
1052                         'fvec': None,
1053                         'fjac': None,
1054                         'ipvt': None,
1055                         'qtf': None,
1056                         'message': None,
1057                         'ierr': None,
1058                         'param_jac': None,
1059                         'param_cov': None}
1060        self._calc_uncertainties=calc_uncertainties
1061        super().__init__()
1062
1063    def objective_function(self, fps, *args):
1064        """
1065        Function to minimize.
1066
1067        Parameters
1068        ----------
1069        fps : list
1070            parameters returned by the fitter
1071        args : list
1072            [model, [weights], [input coordinates]]
1073
1074        """
1075
1076        model = args[0]
1077        weights = args[1]
1078        _fitter_to_model_params(model, fps)
1079        meas = args[-1]
1080        if weights is None:
1081            return np.ravel(model(*args[2: -1]) - meas)
1082        else:
1083            return np.ravel(weights * (model(*args[2: -1]) - meas))
1084
1085    @staticmethod
1086    def _add_fitting_uncertainties(model, cov_matrix):
1087        """
1088        Set ``cov_matrix`` and ``stds`` attributes on model with parameter
1089        covariance matrix returned by ``optimize.leastsq``.
1090        """
1091
1092        free_param_names = [x for x in model.fixed if (model.fixed[x] is False)
1093                            and (model.tied[x] is False)]
1094
1095        model.cov_matrix = Covariance(cov_matrix, free_param_names)
1096        model.stds = StandardDeviations(cov_matrix, free_param_names)
1097
1098    @fitter_unit_support
1099    def __call__(self, model, x, y, z=None, weights=None,
1100                 maxiter=DEFAULT_MAXITER, acc=DEFAULT_ACC,
1101                 epsilon=DEFAULT_EPS, estimate_jacobian=False):
1102        """
1103        Fit data to this model.
1104
1105        Parameters
1106        ----------
1107        model : `~astropy.modeling.FittableModel`
1108            model to fit to x, y, z
1109        x : array
1110           input coordinates
1111        y : array
1112           input coordinates
1113        z : array, optional
1114           input coordinates
1115        weights : array, optional
1116            Weights for fitting.
1117            For data with Gaussian uncertainties, the weights should be
1118            1/sigma.
1119        maxiter : int
1120            maximum number of iterations
1121        acc : float
1122            Relative error desired in the approximate solution
1123        epsilon : float
1124            A suitable step length for the forward-difference
1125            approximation of the Jacobian (if model.fjac=None). If
1126            epsfcn is less than the machine precision, it is
1127            assumed that the relative errors in the functions are
1128            of the order of the machine precision.
1129        estimate_jacobian : bool
1130            If False (default) and if the model has a fit_deriv method,
1131            it will be used. Otherwise the Jacobian will be estimated.
1132            If True, the Jacobian will be estimated in any case.
1133        equivalencies : list or None, optional, keyword-only
1134            List of *additional* equivalencies that are should be applied in
1135            case x, y and/or z have units. Default is None.
1136
1137        Returns
1138        -------
1139        model_copy : `~astropy.modeling.FittableModel`
1140            a copy of the input model with parameters set by the fitter
1141
1142        """
1143
1144        from scipy import optimize
1145
1146        model_copy = _validate_model(model, self.supported_constraints)
1147        model_copy.sync_constraints = False
1148        farg = (model_copy, weights, ) + _convert_input(x, y, z)
1149        if model_copy.fit_deriv is None or estimate_jacobian:
1150            dfunc = None
1151        else:
1152            dfunc = self._wrap_deriv
1153        init_values, _ = _model_to_fit_params(model_copy)
1154        fitparams, cov_x, dinfo, mess, ierr = optimize.leastsq(
1155            self.objective_function, init_values, args=farg, Dfun=dfunc,
1156            col_deriv=model_copy.col_fit_deriv, maxfev=maxiter, epsfcn=epsilon,
1157            xtol=acc, full_output=True)
1158        _fitter_to_model_params(model_copy, fitparams)
1159        self.fit_info.update(dinfo)
1160        self.fit_info['cov_x'] = cov_x
1161        self.fit_info['message'] = mess
1162        self.fit_info['ierr'] = ierr
1163        if ierr not in [1, 2, 3, 4]:
1164            warnings.warn("The fit may be unsuccessful; check "
1165                          "fit_info['message'] for more information.",
1166                          AstropyUserWarning)
1167
1168        # now try to compute the true covariance matrix
1169        if (len(y) > len(init_values)) and cov_x is not None:
1170            sum_sqrs = np.sum(self.objective_function(fitparams, *farg)**2)
1171            dof = len(y) - len(init_values)
1172            self.fit_info['param_cov'] = cov_x * sum_sqrs / dof
1173        else:
1174            self.fit_info['param_cov'] = None
1175
1176        if self._calc_uncertainties is True:
1177            if self.fit_info['param_cov'] is not None:
1178                self._add_fitting_uncertainties(model_copy,
1179                                               self.fit_info['param_cov'])
1180
1181        model_copy.sync_constraints = True
1182        return model_copy
1183
1184    @staticmethod
1185    def _wrap_deriv(params, model, weights, x, y, z=None):
1186        """
1187        Wraps the method calculating the Jacobian of the function to account
1188        for model constraints.
1189        `scipy.optimize.leastsq` expects the function derivative to have the
1190        above signature (parlist, (argtuple)). In order to accommodate model
1191        constraints, instead of using p directly, we set the parameter list in
1192        this function.
1193        """
1194
1195        if weights is None:
1196            weights = 1.0
1197
1198        if any(model.fixed.values()) or any(model.tied.values()):
1199            # update the parameters with the current values from the fitter
1200            _fitter_to_model_params(model, params)
1201            if z is None:
1202                full = np.array(model.fit_deriv(x, *model.parameters))
1203                if not model.col_fit_deriv:
1204                    full_deriv = np.ravel(weights) * full.T
1205                else:
1206                    full_deriv = np.ravel(weights) * full
1207            else:
1208                full = np.array([np.ravel(_) for _ in model.fit_deriv(x, y, *model.parameters)])
1209                if not model.col_fit_deriv:
1210                    full_deriv = np.ravel(weights) * full.T
1211                else:
1212                    full_deriv = np.ravel(weights) * full
1213
1214            pars = [getattr(model, name) for name in model.param_names]
1215            fixed = [par.fixed for par in pars]
1216            tied = [par.tied for par in pars]
1217            tied = list(np.where([par.tied is not False for par in pars],
1218                                 True, tied))
1219            fix_and_tie = np.logical_or(fixed, tied)
1220            ind = np.logical_not(fix_and_tie)
1221
1222            if not model.col_fit_deriv:
1223                residues = np.asarray(full_deriv[np.nonzero(ind)]).T
1224            else:
1225                residues = full_deriv[np.nonzero(ind)]
1226
1227            return [np.ravel(_) for _ in residues]
1228        else:
1229            if z is None:
1230                try:
1231                    return np.array([np.ravel(_) for _ in np.array(weights) *
1232                                     np.array(model.fit_deriv(x, *params))])
1233                except ValueError:
1234                    return np.array([np.ravel(_) for _ in np.array(weights) *
1235                                     np.moveaxis(
1236                                         np.array(model.fit_deriv(x, *params)),
1237                                         -1, 0)]).transpose()
1238            else:
1239                if not model.col_fit_deriv:
1240                    return [np.ravel(_) for _ in
1241                            (np.ravel(weights) * np.array(model.fit_deriv(x, y, *params)).T).T]
1242                return [np.ravel(_) for _ in weights * np.array(model.fit_deriv(x, y, *params))]
1243
1244
1245class SLSQPLSQFitter(Fitter):
1246    """
1247    Sequential Least Squares Programming (SLSQP) optimization algorithm and
1248    least squares statistic.
1249
1250    Raises
1251    ------
1252    ModelLinearityError
1253        A linear model is passed to a nonlinear fitter
1254
1255    Notes
1256    ------
1257    See also the `~astropy.modeling.optimizers.SLSQP` optimizer.
1258
1259    """
1260
1261    supported_constraints = SLSQP.supported_constraints
1262
1263    def __init__(self):
1264        super().__init__(optimizer=SLSQP, statistic=leastsquare)
1265        self.fit_info = {}
1266
1267    @fitter_unit_support
1268    def __call__(self, model, x, y, z=None, weights=None, **kwargs):
1269        """
1270        Fit data to this model.
1271
1272        Parameters
1273        ----------
1274        model : `~astropy.modeling.FittableModel`
1275            model to fit to x, y, z
1276        x : array
1277            input coordinates
1278        y : array
1279            input coordinates
1280        z : array, optional
1281            input coordinates
1282        weights : array, optional
1283            Weights for fitting.
1284            For data with Gaussian uncertainties, the weights should be
1285            1/sigma.
1286        kwargs : dict
1287            optional keyword arguments to be passed to the optimizer or the statistic
1288        verblevel : int
1289            0-silent
1290            1-print summary upon completion,
1291            2-print summary after each iteration
1292        maxiter : int
1293            maximum number of iterations
1294        epsilon : float
1295            the step size for finite-difference derivative estimates
1296        acc : float
1297            Requested accuracy
1298        equivalencies : list or None, optional, keyword-only
1299            List of *additional* equivalencies that are should be applied in
1300            case x, y and/or z have units. Default is None.
1301
1302        Returns
1303        -------
1304        model_copy : `~astropy.modeling.FittableModel`
1305            a copy of the input model with parameters set by the fitter
1306
1307        """
1308
1309        model_copy = _validate_model(model, self._opt_method.supported_constraints)
1310        model_copy.sync_constraints = False
1311        farg = _convert_input(x, y, z)
1312        farg = (model_copy, weights, ) + farg
1313        init_values, _ = _model_to_fit_params(model_copy)
1314        fitparams, self.fit_info = self._opt_method(
1315            self.objective_function, init_values, farg, **kwargs)
1316        _fitter_to_model_params(model_copy, fitparams)
1317
1318        model_copy.sync_constraints = True
1319        return model_copy
1320
1321
1322class SimplexLSQFitter(Fitter):
1323    """
1324    Simplex algorithm and least squares statistic.
1325
1326    Raises
1327    ------
1328    `ModelLinearityError`
1329        A linear model is passed to a nonlinear fitter
1330
1331    """
1332
1333    supported_constraints = Simplex.supported_constraints
1334
1335    def __init__(self):
1336        super().__init__(optimizer=Simplex, statistic=leastsquare)
1337        self.fit_info = {}
1338
1339    @fitter_unit_support
1340    def __call__(self, model, x, y, z=None, weights=None, **kwargs):
1341        """
1342        Fit data to this model.
1343
1344        Parameters
1345        ----------
1346        model : `~astropy.modeling.FittableModel`
1347            model to fit to x, y, z
1348        x : array
1349            input coordinates
1350        y : array
1351            input coordinates
1352        z : array, optional
1353            input coordinates
1354        weights : array, optional
1355            Weights for fitting.
1356            For data with Gaussian uncertainties, the weights should be
1357            1/sigma.
1358        kwargs : dict
1359            optional keyword arguments to be passed to the optimizer or the statistic
1360        maxiter : int
1361            maximum number of iterations
1362        acc : float
1363            Relative error in approximate solution
1364        equivalencies : list or None, optional, keyword-only
1365            List of *additional* equivalencies that are should be applied in
1366            case x, y and/or z have units. Default is None.
1367
1368        Returns
1369        -------
1370        model_copy : `~astropy.modeling.FittableModel`
1371            a copy of the input model with parameters set by the fitter
1372
1373        """
1374
1375        model_copy = _validate_model(model,
1376                                     self._opt_method.supported_constraints)
1377        model_copy.sync_constraints = False
1378        farg = _convert_input(x, y, z)
1379        farg = (model_copy, weights, ) + farg
1380
1381        init_values, _ = _model_to_fit_params(model_copy)
1382
1383        fitparams, self.fit_info = self._opt_method(
1384            self.objective_function, init_values, farg, **kwargs)
1385        _fitter_to_model_params(model_copy, fitparams)
1386        model_copy.sync_constraints = True
1387        return model_copy
1388
1389
1390class JointFitter(metaclass=_FitterMeta):
1391    """
1392    Fit models which share a parameter.
1393    For example, fit two gaussians to two data sets but keep
1394    the FWHM the same.
1395
1396    Parameters
1397    ----------
1398    models : list
1399        a list of model instances
1400    jointparameters : list
1401        a list of joint parameters
1402    initvals : list
1403        a list of initial values
1404
1405    """
1406
1407    def __init__(self, models, jointparameters, initvals):
1408        self.models = list(models)
1409        self.initvals = list(initvals)
1410        self.jointparams = jointparameters
1411        self._verify_input()
1412        self.fitparams = self._model_to_fit_params()
1413
1414        # a list of model.n_inputs
1415        self.modeldims = [m.n_inputs for m in self.models]
1416        # sum all model dimensions
1417        self.ndim = np.sum(self.modeldims)
1418
1419    def _model_to_fit_params(self):
1420        fparams = []
1421        fparams.extend(self.initvals)
1422        for model in self.models:
1423            params = model.parameters.tolist()
1424            joint_params = self.jointparams[model]
1425            param_metrics = model._param_metrics
1426            for param_name in joint_params:
1427                slice_ = param_metrics[param_name]['slice']
1428                del params[slice_]
1429            fparams.extend(params)
1430        return fparams
1431
1432    def objective_function(self, fps, *args):
1433        """
1434        Function to minimize.
1435
1436        Parameters
1437        ----------
1438        fps : list
1439            the fitted parameters - result of an one iteration of the
1440            fitting algorithm
1441        args : dict
1442            tuple of measured and input coordinates
1443            args is always passed as a tuple from optimize.leastsq
1444
1445        """
1446
1447        lstsqargs = list(args)
1448        fitted = []
1449        fitparams = list(fps)
1450        numjp = len(self.initvals)
1451        # make a separate list of the joint fitted parameters
1452        jointfitparams = fitparams[:numjp]
1453        del fitparams[:numjp]
1454
1455        for model in self.models:
1456            joint_params = self.jointparams[model]
1457            margs = lstsqargs[:model.n_inputs + 1]
1458            del lstsqargs[:model.n_inputs + 1]
1459            # separate each model separately fitted parameters
1460            numfp = len(model._parameters) - len(joint_params)
1461            mfparams = fitparams[:numfp]
1462
1463            del fitparams[:numfp]
1464            # recreate the model parameters
1465            mparams = []
1466            param_metrics = model._param_metrics
1467            for param_name in model.param_names:
1468                if param_name in joint_params:
1469                    index = joint_params.index(param_name)
1470                    # should do this with slices in case the
1471                    # parameter is not a number
1472                    mparams.extend([jointfitparams[index]])
1473                else:
1474                    slice_ = param_metrics[param_name]['slice']
1475                    plen = slice_.stop - slice_.start
1476                    mparams.extend(mfparams[:plen])
1477                    del mfparams[:plen]
1478            modelfit = model.evaluate(margs[:-1], *mparams)
1479            fitted.extend(modelfit - margs[-1])
1480        return np.ravel(fitted)
1481
1482    def _verify_input(self):
1483        if len(self.models) <= 1:
1484            raise TypeError(f"Expected >1 models, {len(self.models)} is given")
1485        if len(self.jointparams.keys()) < 2:
1486            raise TypeError("At least two parameters are expected, "
1487                            "{} is given".format(len(self.jointparams.keys())))
1488        for j in self.jointparams.keys():
1489            if len(self.jointparams[j]) != len(self.initvals):
1490                raise TypeError("{} parameter(s) provided but {} expected".format(
1491                    len(self.jointparams[j]), len(self.initvals)))
1492
1493    def __call__(self, *args):
1494        """
1495        Fit data to these models keeping some of the parameters common to the
1496        two models.
1497        """
1498
1499        from scipy import optimize
1500
1501        if len(args) != reduce(lambda x, y: x + 1 + y + 1, self.modeldims):
1502            raise ValueError("Expected {} coordinates in args but {} provided"
1503                             .format(reduce(lambda x, y: x + 1 + y + 1,
1504                                            self.modeldims), len(args)))
1505
1506        self.fitparams[:], _ = optimize.leastsq(self.objective_function,
1507                                                self.fitparams, args=args)
1508
1509        fparams = self.fitparams[:]
1510        numjp = len(self.initvals)
1511        # make a separate list of the joint fitted parameters
1512        jointfitparams = fparams[:numjp]
1513        del fparams[:numjp]
1514
1515        for model in self.models:
1516            # extract each model's fitted parameters
1517            joint_params = self.jointparams[model]
1518            numfp = len(model._parameters) - len(joint_params)
1519            mfparams = fparams[:numfp]
1520
1521            del fparams[:numfp]
1522            # recreate the model parameters
1523            mparams = []
1524            param_metrics = model._param_metrics
1525            for param_name in model.param_names:
1526                if param_name in joint_params:
1527                    index = joint_params.index(param_name)
1528                    # should do this with slices in case the parameter
1529                    # is not a number
1530                    mparams.extend([jointfitparams[index]])
1531                else:
1532                    slice_ = param_metrics[param_name]['slice']
1533                    plen = slice_.stop - slice_.start
1534                    mparams.extend(mfparams[:plen])
1535                    del mfparams[:plen]
1536            model.parameters = np.array(mparams)
1537
1538
1539def _convert_input(x, y, z=None, n_models=1, model_set_axis=0):
1540    """Convert inputs to float arrays."""
1541
1542    x = np.asanyarray(x, dtype=float)
1543    y = np.asanyarray(y, dtype=float)
1544
1545    if z is not None:
1546        z = np.asanyarray(z, dtype=float)
1547        data_ndim, data_shape = z.ndim, z.shape
1548    else:
1549        data_ndim, data_shape = y.ndim, y.shape
1550
1551    # For compatibility with how the linear fitter code currently expects to
1552    # work, shift the dependent variable's axes to the expected locations
1553    if n_models > 1 or data_ndim > x.ndim:
1554        if (model_set_axis or 0) >= data_ndim:
1555            raise ValueError("model_set_axis out of range")
1556        if data_shape[model_set_axis] != n_models:
1557            raise ValueError(
1558                "Number of data sets (y or z array) is expected to equal "
1559                "the number of parameter sets"
1560            )
1561        if z is None:
1562            # For a 1-D model the y coordinate's model-set-axis is expected to
1563            # be last, so that its first dimension is the same length as the x
1564            # coordinates.  This is in line with the expectations of
1565            # numpy.linalg.lstsq:
1566            # https://numpy.org/doc/stable/reference/generated/numpy.linalg.lstsq.html
1567            # That is, each model should be represented by a column.  TODO:
1568            # Obviously this is a detail of np.linalg.lstsq and should be
1569            # handled specifically by any fitters that use it...
1570            y = np.rollaxis(y, model_set_axis, y.ndim)
1571            data_shape = y.shape[:-1]
1572        else:
1573            # Shape of z excluding model_set_axis
1574            data_shape = (z.shape[:model_set_axis] +
1575                          z.shape[model_set_axis + 1:])
1576
1577    if z is None:
1578        if data_shape != x.shape:
1579            raise ValueError("x and y should have the same shape")
1580        farg = (x, y)
1581    else:
1582        if not (x.shape == y.shape == data_shape):
1583            raise ValueError("x, y and z should have the same shape")
1584        farg = (x, y, z)
1585    return farg
1586
1587
1588# TODO: These utility functions are really particular to handling
1589# bounds/tied/fixed constraints for scipy.optimize optimizers that do not
1590# support them inherently; this needs to be reworked to be clear about this
1591# distinction (and the fact that these are not necessarily applicable to any
1592# arbitrary fitter--as evidenced for example by the fact that JointFitter has
1593# its own versions of these)
1594# TODO: Most of this code should be entirely rewritten; it should not be as
1595# inefficient as it is.
1596def _fitter_to_model_params(model, fps):
1597    """
1598    Constructs the full list of model parameters from the fitted and
1599    constrained parameters.
1600    """
1601
1602    _, fit_param_indices = _model_to_fit_params(model)
1603
1604    has_tied = any(model.tied.values())
1605    has_fixed = any(model.fixed.values())
1606    has_bound = any(b != (None, None) for b in model.bounds.values())
1607    parameters = model.parameters
1608
1609    if not (has_tied or has_fixed or has_bound):
1610        # We can just assign directly
1611        model.parameters = fps
1612        return
1613
1614    fit_param_indices = set(fit_param_indices)
1615    offset = 0
1616    param_metrics = model._param_metrics
1617    for idx, name in enumerate(model.param_names):
1618        if idx not in fit_param_indices:
1619            continue
1620
1621        slice_ = param_metrics[name]['slice']
1622        shape = param_metrics[name]['shape']
1623        # This is determining which range of fps (the fitted parameters) maps
1624        # to parameters of the model
1625        size = reduce(operator.mul, shape, 1)
1626
1627        values = fps[offset:offset + size]
1628
1629        # Check bounds constraints
1630        if model.bounds[name] != (None, None):
1631            _min, _max = model.bounds[name]
1632            if _min is not None:
1633                values = np.fmax(values, _min)
1634            if _max is not None:
1635                values = np.fmin(values, _max)
1636
1637        parameters[slice_] = values
1638        offset += size
1639
1640    # Update model parameters before calling ``tied`` constraints.
1641    model._array_to_parameters()
1642
1643    # This has to be done in a separate loop due to how tied parameters are
1644    # currently evaluated (the fitted parameters need to actually be *set* on
1645    # the model first, for use in evaluating the "tied" expression--it might be
1646    # better to change this at some point
1647    if has_tied:
1648        for idx, name in enumerate(model.param_names):
1649            if model.tied[name]:
1650                value = model.tied[name](model)
1651                slice_ = param_metrics[name]['slice']
1652
1653                # To handle multiple tied constraints, model parameters
1654                # need to be updated after each iteration.
1655                parameters[slice_] = value
1656                model._array_to_parameters()
1657
1658
1659def _model_to_fit_params(model):
1660    """
1661    Convert a model instance's parameter array to an array that can be used
1662    with a fitter that doesn't natively support fixed or tied parameters.
1663    In particular, it removes fixed/tied parameters from the parameter
1664    array.
1665    These may be a subset of the model parameters, if some of them are held
1666    constant or tied.
1667    """
1668
1669    fitparam_indices = list(range(len(model.param_names)))
1670    if any(model.fixed.values()) or any(model.tied.values()):
1671        params = list(model.parameters)
1672        param_metrics = model._param_metrics
1673        for idx, name in list(enumerate(model.param_names))[::-1]:
1674            if model.fixed[name] or model.tied[name]:
1675                slice_ = param_metrics[name]['slice']
1676                del params[slice_]
1677                del fitparam_indices[idx]
1678        return (np.array(params), fitparam_indices)
1679    return (model.parameters, fitparam_indices)
1680
1681
1682def _validate_constraints(supported_constraints, model):
1683    """Make sure model constraints are supported by the current fitter."""
1684
1685    message = 'Optimizer cannot handle {0} constraints.'
1686
1687    if (any(model.fixed.values()) and
1688            'fixed' not in supported_constraints):
1689        raise UnsupportedConstraintError(
1690            message.format('fixed parameter'))
1691
1692    if any(model.tied.values()) and 'tied' not in supported_constraints:
1693        raise UnsupportedConstraintError(
1694            message.format('tied parameter'))
1695
1696    if (any(tuple(b) != (None, None) for b in model.bounds.values()) and
1697            'bounds' not in supported_constraints):
1698        raise UnsupportedConstraintError(
1699            message.format('bound parameter'))
1700
1701    if model.eqcons and 'eqcons' not in supported_constraints:
1702        raise UnsupportedConstraintError(message.format('equality'))
1703
1704    if model.ineqcons and 'ineqcons' not in supported_constraints:
1705        raise UnsupportedConstraintError(message.format('inequality'))
1706
1707
1708def _validate_model(model, supported_constraints):
1709    """
1710    Check that model and fitter are compatible and return a copy of the model.
1711    """
1712
1713    if not model.fittable:
1714        raise ValueError("Model does not appear to be fittable.")
1715    if model.linear:
1716        warnings.warn('Model is linear in parameters; '
1717                      'consider using linear fitting methods.',
1718                      AstropyUserWarning)
1719    elif len(model) != 1:
1720        # for now only single data sets ca be fitted
1721        raise ValueError("Non-linear fitters can only fit "
1722                         "one data set at a time.")
1723    _validate_constraints(supported_constraints, model)
1724
1725    model_copy = model.copy()
1726    return model_copy
1727
1728
1729def populate_entry_points(entry_points):
1730    """
1731    This injects entry points into the `astropy.modeling.fitting` namespace.
1732    This provides a means of inserting a fitting routine without requirement
1733    of it being merged into astropy's core.
1734
1735    Parameters
1736    ----------
1737    entry_points : list of `~importlib.metadata.EntryPoint`
1738        entry_points are objects which encapsulate importable objects and
1739        are defined on the installation of a package.
1740
1741    Notes
1742    -----
1743    An explanation of entry points can be found `here <http://setuptools.readthedocs.io/en/latest/setuptools.html#dynamic-discovery-of-services-and-plugins>`
1744    """
1745
1746    for entry_point in entry_points:
1747        name = entry_point.name
1748        try:
1749            entry_point = entry_point.load()
1750        except Exception as e:
1751            # This stops the fitting from choking if an entry_point produces an error.
1752            warnings.warn(AstropyUserWarning(
1753                f'{type(e).__name__} error occurred in entry point {name}.'))
1754        else:
1755            if not inspect.isclass(entry_point):
1756                warnings.warn(AstropyUserWarning(
1757                    f'Modeling entry point {name} expected to be a Class.'))
1758            else:
1759                if issubclass(entry_point, Fitter):
1760                    name = entry_point.__name__
1761                    globals()[name] = entry_point
1762                    __all__.append(name)
1763                else:
1764                    warnings.warn(AstropyUserWarning(
1765                        'Modeling entry point {} expected to extend '
1766                        'astropy.modeling.Fitter' .format(name)))
1767
1768
1769def _populate_ep():
1770    # TODO: Exclusively use select when Python minversion is 3.10
1771    ep = entry_points()
1772    if hasattr(ep, 'select'):
1773        populate_entry_points(ep.select(group='astropy.modeling'))
1774    else:
1775        populate_entry_points(ep.get('astropy.modeling', []))
1776
1777
1778_populate_ep()
1779