pyerrors.fits

  1import gc
  2from collections.abc import Sequence
  3import warnings
  4import numpy as np
  5import autograd.numpy as anp
  6import scipy.optimize
  7import scipy.stats
  8import matplotlib.pyplot as plt
  9from matplotlib import gridspec
 10from scipy.odr import ODR, Model, RealData
 11import iminuit
 12from autograd import jacobian as auto_jacobian
 13from autograd import hessian as auto_hessian
 14from autograd import elementwise_grad as egrad
 15from numdifftools import Jacobian as num_jacobian
 16from numdifftools import Hessian as num_hessian
 17from .obs import Obs, derived_observable, covariance, cov_Obs, invert_corr_cov_cholesky
 18
 19
 20class Fit_result(Sequence):
 21    """Represents fit results.
 22
 23    Attributes
 24    ----------
 25    fit_parameters : list
 26        results for the individual fit parameters,
 27        also accessible via indices.
 28    chisquare_by_dof : float
 29        reduced chisquare.
 30    p_value : float
 31        p-value of the fit
 32    t2_p_value : float
 33        Hotelling t-squared p-value for correlated fits.
 34    """
 35
 36    def __init__(self):
 37        self.fit_parameters = None
 38
 39    def __getitem__(self, idx):
 40        return self.fit_parameters[idx]
 41
 42    def __len__(self):
 43        return len(self.fit_parameters)
 44
 45    def gamma_method(self, **kwargs):
 46        """Apply the gamma method to all fit parameters"""
 47        [o.gamma_method(**kwargs) for o in self.fit_parameters]
 48
 49    gm = gamma_method
 50
 51    def __str__(self):
 52        my_str = 'Goodness of fit:\n'
 53        if hasattr(self, 'chisquare_by_dof'):
 54            my_str += '\u03C7\u00b2/d.o.f. = ' + f'{self.chisquare_by_dof:2.6f}' + '\n'
 55        elif hasattr(self, 'residual_variance'):
 56            my_str += 'residual variance = ' + f'{self.residual_variance:2.6f}' + '\n'
 57        if hasattr(self, 'chisquare_by_expected_chisquare'):
 58            my_str += '\u03C7\u00b2/\u03C7\u00b2exp  = ' + f'{self.chisquare_by_expected_chisquare:2.6f}' + '\n'
 59        if hasattr(self, 'p_value'):
 60            my_str += 'p-value   = ' + f'{self.p_value:2.4f}' + '\n'
 61        if hasattr(self, 't2_p_value'):
 62            my_str += 't\u00B2p-value = ' + f'{self.t2_p_value:2.4f}' + '\n'
 63        my_str += 'Fit parameters:\n'
 64        for i_par, par in enumerate(self.fit_parameters):
 65            my_str += str(i_par) + '\t' + ' ' * int(par >= 0) + str(par).rjust(int(par < 0.0)) + '\n'
 66        return my_str
 67
 68    def __repr__(self):
 69        m = max(map(len, list(self.__dict__.keys()))) + 1
 70        return '\n'.join([key.rjust(m) + ': ' + repr(value) for key, value in sorted(self.__dict__.items())])
 71
 72
 73def least_squares(x, y, func, priors=None, silent=False, **kwargs):
 74    r'''Performs a non-linear fit to y = func(x).
 75        ```
 76
 77    Parameters
 78    ----------
 79    For an uncombined fit:
 80
 81    x : list
 82        list of floats.
 83    y : list
 84        list of Obs.
 85    func : object
 86        fit function, has to be of the form
 87
 88        ```python
 89        import autograd.numpy as anp
 90
 91        def func(a, x):
 92            return a[0] + a[1] * x + a[2] * anp.sinh(x)
 93        ```
 94
 95        For multiple x values func can be of the form
 96
 97        ```python
 98        def func(a, x):
 99            (x1, x2) = x
100            return a[0] * x1 ** 2 + a[1] * x2
101        ```
102        It is important that all numpy functions refer to autograd.numpy, otherwise the differentiation
103        will not work.
104
105    OR For a combined fit:
106
107    x : dict
108        dict of lists.
109    y : dict
110        dict of lists of Obs.
111    funcs : dict
112        dict of objects
113        fit functions have to be of the form (here a[0] is the common fit parameter)
114        ```python
115        import autograd.numpy as anp
116        funcs = {"a": func_a,
117                "b": func_b}
118
119        def func_a(a, x):
120            return a[1] * anp.exp(-a[0] * x)
121
122        def func_b(a, x):
123            return a[2] * anp.exp(-a[0] * x)
124
125        It is important that all numpy functions refer to autograd.numpy, otherwise the differentiation
126        will not work.
127
128    priors : dict or list, optional
129        priors can either be a dictionary with integer keys and the corresponding priors as values or
130        a list with an entry for every parameter in the fit. The entries can either be
131        Obs (e.g. results from a previous fit) or strings containing a value and an error formatted like
132        0.548(23), 500(40) or 0.5(0.4)
133    silent : bool, optional
134        If True all output to the console is omitted (default False).
135    initial_guess : list
136        can provide an initial guess for the input parameters. Relevant for
137        non-linear fits with many parameters. In case of correlated fits the guess is used to perform
138        an uncorrelated fit which then serves as guess for the correlated fit.
139    method : str, optional
140        can be used to choose an alternative method for the minimization of chisquare.
141        The possible methods are the ones which can be used for scipy.optimize.minimize and
142        migrad of iminuit. If no method is specified, Levenberg–Marquardt is used.
143        Reliable alternatives are migrad, Powell and Nelder-Mead.
144    tol: float, optional
145        can be used (only for combined fits and methods other than Levenberg–Marquardt) to set the tolerance for convergence
146        to a different value to either speed up convergence at the cost of a larger error on the fitted parameters (and possibly
147        invalid estimates for parameter uncertainties) or smaller values to get more accurate parameter values
148        The stopping criterion depends on the method, e.g. migrad: edm_max = 0.002 * tol * errordef (EDM criterion: edm < edm_max)
149    correlated_fit : bool
150        If True, use the full inverse covariance matrix in the definition of the chisquare cost function.
151        For details about how the covariance matrix is estimated see `pyerrors.obs.covariance`.
152        In practice the correlation matrix is Cholesky decomposed and inverted (instead of the covariance matrix).
153        This procedure should be numerically more stable as the correlation matrix is typically better conditioned (Jacobi preconditioning).
154    inv_chol_cov_matrix [array,list], optional
155        array: shape = (number of y values) X (number of y values)
156        list:   for an uncombined fit: [""]
157                for a combined fit: list of keys belonging to the corr_matrix saved in the array, must be the same as the keys of the y dict in alphabetical order
158        If correlated_fit=True is set as well, can provide an inverse covariance matrix (y errors, dy_f included!) of your own choosing for a correlated fit.
159        The matrix must be a lower triangular matrix constructed from a Cholesky decomposition: The function invert_corr_cov_cholesky(corr, inverrdiag) can be
160        used to construct it from a correlation matrix (corr) and the errors dy_f of the data points (inverrdiag = np.diag(1 / np.asarray(dy_f))). For the correct
161        ordering the correlation matrix (corr) can be sorted via the function sort_corr(corr, kl, yd) where kl is the list of keys and yd the y dict.
162    expected_chisquare : bool
163        If True estimates the expected chisquare which is
164        corrected by effects caused by correlated input data (default False).
165    resplot : bool
166        If True, a plot which displays fit, data and residuals is generated (default False).
167    qqplot : bool
168        If True, a quantile-quantile plot of the fit result is generated (default False).
169    num_grad : bool
170        Use numerical differentation instead of automatic differentiation to perform the error propagation (default False).
171    n_parms : int, optional
172        Number of fit parameters. Overrides automatic detection of parameter count.
173        Useful when autodetection fails. Must match the length of initial_guess or priors (if provided).
174
175    Returns
176    -------
177    output : Fit_result
178        Parameters and information on the fitted result.
179    Examples
180    ------
181    >>> # Example of a correlated (correlated_fit = True, inv_chol_cov_matrix handed over) combined fit, based on a randomly generated data set
182    >>> import numpy as np
183    >>> from scipy.stats import norm
184    >>> from scipy.linalg import cholesky
185    >>> import pyerrors as pe
186    >>> # generating the random data set
187    >>> num_samples = 400
188    >>> N = 3
189    >>> x = np.arange(N)
190    >>> x1 = norm.rvs(size=(N, num_samples)) # generate random numbers
191    >>> x2 = norm.rvs(size=(N, num_samples)) # generate random numbers
192    >>> r = r1 = r2 = np.zeros((N, N))
193    >>> y = {}
194    >>> for i in range(N):
195    >>>    for j in range(N):
196    >>>        r[i, j] = np.exp(-0.8 * np.fabs(i - j)) # element in correlation matrix
197    >>> errl = np.sqrt([3.4, 2.5, 3.6]) # set y errors
198    >>> for i in range(N):
199    >>>    for j in range(N):
200    >>>        r[i, j] *= errl[i] * errl[j] # element in covariance matrix
201    >>> c = cholesky(r, lower=True)
202    >>> y = {'a': np.dot(c, x1), 'b': np.dot(c, x2)} # generate y data with the covariance matrix defined
203    >>> # random data set has been generated, now the dictionaries and the inverse covariance matrix to be handed over are built
204    >>> x_dict = {}
205    >>> y_dict = {}
206    >>> chol_inv_dict = {}
207    >>> data = []
208    >>> for key in y.keys():
209    >>>    x_dict[key] = x
210    >>>    for i in range(N):
211    >>>        data.append(pe.Obs([[i + 1 + o for o in y[key][i]]], ['ens'])) # generate y Obs from the y data
212    >>>    [o.gamma_method() for o in data]
213    >>>    corr = pe.covariance(data, correlation=True)
214    >>>    inverrdiag = np.diag(1 / np.asarray([o.dvalue for o in data]))
215    >>>    chol_inv = pe.obs.invert_corr_cov_cholesky(corr, inverrdiag) # gives form of the inverse covariance matrix needed for the combined correlated fit below
216    >>> y_dict = {'a': data[:3], 'b': data[3:]}
217    >>> # common fit parameter p[0] in combined fit
218    >>> def fit1(p, x):
219    >>>    return p[0] + p[1] * x
220    >>> def fit2(p, x):
221    >>>    return p[0] + p[2] * x
222    >>> fitf_dict = {'a': fit1, 'b':fit2}
223    >>> fitp_inv_cov_combined_fit = pe.least_squares(x_dict,y_dict, fitf_dict, correlated_fit = True, inv_chol_cov_matrix = [chol_inv,['a','b']])
224    Fit with 3 parameters
225    Method: Levenberg-Marquardt
226    `ftol` termination condition is satisfied.
227    chisquare/d.o.f.: 0.5388013574561786 # random
228    fit parameters [1.11897846 0.96361162 0.92325319] # random
229
230    '''
231    output = Fit_result()
232
233    if (isinstance(x, dict) and isinstance(y, dict) and isinstance(func, dict)):
234        xd = {key: anp.asarray(x[key]) for key in x}
235        yd = y
236        funcd = func
237        output.fit_function = func
238    elif (isinstance(x, dict) or isinstance(y, dict) or isinstance(func, dict)):
239        raise TypeError("All arguments have to be dictionaries in order to perform a combined fit.")
240    else:
241        x = np.asarray(x)
242        xd = {"": x}
243        yd = {"": y}
244        funcd = {"": func}
245        output.fit_function = func
246
247    if kwargs.get('num_grad') is True:
248        jacobian = num_jacobian
249        hessian = num_hessian
250    else:
251        jacobian = auto_jacobian
252        hessian = auto_hessian
253
254    key_ls = sorted(list(xd.keys()))
255
256    if sorted(list(yd.keys())) != key_ls:
257        raise ValueError('x and y dictionaries do not contain the same keys.')
258
259    if sorted(list(funcd.keys())) != key_ls:
260        raise ValueError('x and func dictionaries do not contain the same keys.')
261
262    x_all = np.concatenate([np.array(xd[key]).transpose() for key in key_ls]).transpose()
263    y_all = np.concatenate([np.array(yd[key]) for key in key_ls])
264
265    y_f = [o.value for o in y_all]
266    dy_f = [o.dvalue for o in y_all]
267
268    if len(x_all.shape) > 2:
269        raise ValueError("Unknown format for x values")
270
271    if np.any(np.asarray(dy_f) <= 0.0):
272        raise Exception("No y errors available, run the gamma method first.")
273
274    # number of fit parameters
275    if 'n_parms' in kwargs:
276        n_parms = kwargs.get('n_parms')
277        if not isinstance(n_parms, int):
278            raise TypeError(
279                f"'n_parms' must be an integer, got {n_parms!r} "
280                f"of type {type(n_parms).__name__}."
281            )
282        if n_parms <= 0:
283            raise ValueError(
284                f"'n_parms' must be a positive integer, got {n_parms}."
285            )
286    else:
287        n_parms_ls = []
288        for key in key_ls:
289            if not callable(funcd[key]):
290                raise TypeError('func (key=' + key + ') is not a function.')
291            if np.asarray(xd[key]).shape[-1] != len(yd[key]):
292                raise ValueError('x and y input (key=' + key + ') do not have the same length')
293            for n_loc in range(100):
294                try:
295                    funcd[key](np.arange(n_loc), x_all.T[0])
296                except TypeError:
297                    continue
298                except IndexError:
299                    continue
300                else:
301                    break
302            else:
303                raise RuntimeError("Fit function (key=" + key + ") is not valid.")
304            n_parms_ls.append(n_loc)
305
306        n_parms = max(n_parms_ls)
307
308    if len(key_ls) > 1:
309        for key in key_ls:
310            if np.asarray(yd[key]).shape != funcd[key](np.arange(n_parms), xd[key]).shape:
311                raise ValueError(f"Fit function {key} returns the wrong shape ({funcd[key](np.arange(n_parms), xd[key]).shape} instead of {np.asarray(yd[key]).shape})\nIf the fit function is just a constant you could try adding x*0 to get the correct shape.")
312
313    if not silent:
314        print('Fit with', n_parms, 'parameter' + 's' * (n_parms > 1))
315
316    if priors is not None:
317        if isinstance(priors, (list, np.ndarray)):
318            if n_parms != len(priors):
319                raise ValueError("'priors' does not have the correct length.")
320
321            loc_priors = []
322            for i_n, i_prior in enumerate(priors):
323                loc_priors.append(_construct_prior_obs(i_prior, i_n))
324
325            prior_mask = np.arange(len(priors))
326            output.priors = loc_priors
327
328        elif isinstance(priors, dict):
329            loc_priors = []
330            prior_mask = []
331            output.priors = {}
332            for pos, prior in priors.items():
333                if isinstance(pos, int):
334                    prior_mask.append(pos)
335                else:
336                    raise TypeError("Prior position needs to be an integer.")
337                loc_priors.append(_construct_prior_obs(prior, pos))
338
339                output.priors[pos] = loc_priors[-1]
340            if max(prior_mask) >= n_parms:
341                raise ValueError("Prior position out of range.")
342        else:
343            raise TypeError("Unkown type for `priors`.")
344
345        p_f = [o.value for o in loc_priors]
346        dp_f = [o.dvalue for o in loc_priors]
347        if np.any(np.asarray(dp_f) <= 0.0):
348            raise Exception("No prior errors available, run the gamma method first.")
349    else:
350        p_f = dp_f = np.array([])
351        prior_mask = []
352        loc_priors = []
353
354    if 'initial_guess' in kwargs:
355        x0 = kwargs.get('initial_guess')
356        if len(x0) != n_parms:
357            raise ValueError('Initial guess does not have the correct length: %d vs. %d' % (len(x0), n_parms))
358    else:
359        x0 = [0.1] * n_parms
360
361    if priors is None:
362        def general_chisqfunc_uncorr(p, ivars, pr):
363            model = anp.concatenate([anp.array(funcd[key](p, xd[key])).reshape(-1) for key in key_ls])
364            return (ivars - model) / dy_f
365    else:
366        def general_chisqfunc_uncorr(p, ivars, pr):
367            model = anp.concatenate([anp.array(funcd[key](p, xd[key])).reshape(-1) for key in key_ls])
368            return anp.concatenate(((ivars - model) / dy_f, (p[prior_mask] - pr) / dp_f))
369
370    def chisqfunc_uncorr(p):
371        return anp.sum(general_chisqfunc_uncorr(p, y_f, p_f) ** 2)
372
373    if kwargs.get('correlated_fit') is True:
374        if 'inv_chol_cov_matrix' in kwargs:
375            chol_inv = kwargs.get('inv_chol_cov_matrix')
376            if (chol_inv[0].shape[0] != len(dy_f)):
377                raise TypeError('The number of columns of the inverse covariance matrix handed over needs to be equal to the number of y errors.')
378            if (chol_inv[0].shape[0] != chol_inv[0].shape[1]):
379                raise TypeError('The inverse covariance matrix handed over needs to have the same number of rows as columns.')
380            if (chol_inv[1] != key_ls):
381                raise ValueError('The keys of inverse covariance matrix are not the same or do not appear in the same order as the x and y values.')
382            chol_inv = chol_inv[0]
383            if np.any(np.diag(chol_inv) <= 0) or (not np.all(chol_inv == np.tril(chol_inv))):
384                raise ValueError('The inverse covariance matrix inv_chol_cov_matrix[0] has to be a lower triangular matrix constructed from a Cholesky decomposition.')
385        else:
386            corr = covariance(y_all, correlation=True, **kwargs)
387            inverrdiag = np.diag(1 / np.asarray(dy_f))
388            chol_inv = invert_corr_cov_cholesky(corr, inverrdiag)
389
390        def general_chisqfunc(p, ivars, pr):
391            model = anp.concatenate([anp.array(funcd[key](p, xd[key])).reshape(-1) for key in key_ls])
392            return anp.concatenate((anp.dot(chol_inv, (ivars - model)), (p[prior_mask] - pr) / dp_f))
393
394        def chisqfunc(p):
395            return anp.sum(general_chisqfunc(p, y_f, p_f) ** 2)
396    else:
397        general_chisqfunc = general_chisqfunc_uncorr
398        chisqfunc = chisqfunc_uncorr
399
400    output.method = kwargs.get('method', 'Levenberg-Marquardt')
401    if not silent:
402        print('Method:', output.method)
403
404    if output.method != 'Levenberg-Marquardt':
405        if output.method == 'migrad':
406            tolerance = 1e-4  # default value of 1e-1 set by iminuit can be problematic
407            if 'tol' in kwargs:
408                tolerance = kwargs.get('tol')
409            fit_result = iminuit.minimize(chisqfunc_uncorr, x0, tol=tolerance)  # Stopping criterion 0.002 * tol * errordef
410            if kwargs.get('correlated_fit') is True:
411                fit_result = iminuit.minimize(chisqfunc, fit_result.x, tol=tolerance)
412            output.iterations = fit_result.nfev
413        else:
414            tolerance = 1e-12
415            if 'tol' in kwargs:
416                tolerance = kwargs.get('tol')
417            fit_result = scipy.optimize.minimize(chisqfunc_uncorr, x0, method=kwargs.get('method'), tol=tolerance)
418            if kwargs.get('correlated_fit') is True:
419                fit_result = scipy.optimize.minimize(chisqfunc, fit_result.x, method=kwargs.get('method'), tol=tolerance)
420            output.iterations = fit_result.nit
421
422        chisquare = fit_result.fun
423
424    else:
425        if 'tol' in kwargs:
426            print('tol cannot be set for Levenberg-Marquardt')
427
428        def chisqfunc_residuals_uncorr(p):
429            return general_chisqfunc_uncorr(p, y_f, p_f)
430
431        fit_result = scipy.optimize.least_squares(chisqfunc_residuals_uncorr, x0, method='lm', ftol=1e-15, gtol=1e-15, xtol=1e-15)
432        if kwargs.get('correlated_fit') is True:
433            def chisqfunc_residuals(p):
434                return general_chisqfunc(p, y_f, p_f)
435
436            fit_result = scipy.optimize.least_squares(chisqfunc_residuals, fit_result.x, method='lm', ftol=1e-15, gtol=1e-15, xtol=1e-15)
437
438        chisquare = np.sum(fit_result.fun ** 2)
439        assert np.isclose(chisquare, chisqfunc(fit_result.x), atol=1e-14)
440
441        output.iterations = fit_result.nfev
442
443    if not fit_result.success:
444        raise Exception('The minimization procedure did not converge.')
445
446    output.chisquare = chisquare
447    output.dof = y_all.shape[-1] - n_parms + len(loc_priors)
448    output.p_value = 1 - scipy.stats.chi2.cdf(output.chisquare, output.dof)
449    if output.dof > 0:
450        output.chisquare_by_dof = output.chisquare / output.dof
451    else:
452        output.chisquare_by_dof = float('nan')
453
454    output.message = fit_result.message
455    if not silent:
456        print(fit_result.message)
457        print('chisquare/d.o.f.:', output.chisquare_by_dof)
458        print('fit parameters', fit_result.x)
459
460    def prepare_hat_matrix():
461        hat_vector = []
462        for key in key_ls:
463            if (len(xd[key]) != 0):
464                hat_vector.append(jacobian(funcd[key])(fit_result.x, xd[key]))
465        hat_vector = [item for sublist in hat_vector for item in sublist]
466        return hat_vector
467
468    if kwargs.get('expected_chisquare') is True:
469        if kwargs.get('correlated_fit') is not True:
470            W = np.diag(1 / np.asarray(dy_f))
471            cov = covariance(y_all)
472            hat_vector = prepare_hat_matrix()
473            A = W @ hat_vector
474            P_phi = A @ np.linalg.pinv(A.T @ A) @ A.T
475            expected_chisquare = np.trace((np.identity(y_all.shape[-1]) - P_phi) @ W @ cov @ W)
476            output.chisquare_by_expected_chisquare = output.chisquare / expected_chisquare
477            if not silent:
478                print('chisquare/expected_chisquare:', output.chisquare_by_expected_chisquare)
479
480    fitp = fit_result.x
481
482    try:
483        hess = hessian(chisqfunc)(fitp)
484    except TypeError:
485        raise Exception("It is required to use autograd.numpy instead of numpy within fit functions, see the documentation for details.") from None
486
487    len_y = len(y_f)
488
489    def chisqfunc_compact(d):
490        return anp.sum(general_chisqfunc(d[:n_parms], d[n_parms: n_parms + len_y], d[n_parms + len_y:]) ** 2)
491
492    jac_jac_y = hessian(chisqfunc_compact)(np.concatenate((fitp, y_f, p_f)))
493
494    # Compute hess^{-1} @ jac_jac_y[:n_parms + m, n_parms + m:] using LAPACK dgesv
495    try:
496        deriv_y = -scipy.linalg.solve(hess, jac_jac_y[:n_parms, n_parms:])
497    except np.linalg.LinAlgError:
498        raise Exception("Cannot invert hessian matrix.")
499
500    result = []
501    for i in range(n_parms):
502        result.append(derived_observable(lambda x_all, **kwargs: (x_all[0] + np.finfo(np.float64).eps) / (y_all[0].value + np.finfo(np.float64).eps) * fitp[i], list(y_all) + loc_priors, man_grad=list(deriv_y[i])))
503
504    output.fit_parameters = result
505
506    # Hotelling t-squared p-value for correlated fits.
507    if kwargs.get('correlated_fit') is True:
508        n_cov = np.min(np.vectorize(lambda x_all: x_all.N)(y_all))
509        output.t2_p_value = 1 - scipy.stats.f.cdf((n_cov - output.dof) / (output.dof * (n_cov - 1)) * output.chisquare,
510                                                  output.dof, n_cov - output.dof)
511
512    if kwargs.get('resplot') is True:
513        for key in key_ls:
514            residual_plot(xd[key], yd[key], funcd[key], result, title=key)
515
516    if kwargs.get('qqplot') is True:
517        for key in key_ls:
518            qqplot(xd[key], yd[key], funcd[key], result, title=key)
519
520    return output
521
522
523def total_least_squares(x, y, func, silent=False, **kwargs):
524    r'''Performs a non-linear fit to y = func(x) and returns a list of Obs corresponding to the fit parameters.
525
526    Parameters
527    ----------
528    x : list
529        list of Obs, or a tuple of lists of Obs
530    y : list
531        list of Obs. The dvalues of the Obs are used as x- and yerror for the fit.
532    func : object
533        func has to be of the form
534
535        ```python
536        import autograd.numpy as anp
537
538        def func(a, x):
539            return a[0] + a[1] * x + a[2] * anp.sinh(x)
540        ```
541
542        For multiple x values func can be of the form
543
544        ```python
545        def func(a, x):
546            (x1, x2) = x
547            return a[0] * x1 ** 2 + a[1] * x2
548        ```
549
550        It is important that all numpy functions refer to autograd.numpy, otherwise the differentiation
551        will not work.
552    silent : bool, optional
553        If True all output to the console is omitted (default False).
554    initial_guess : list
555        can provide an initial guess for the input parameters. Relevant for non-linear
556        fits with many parameters.
557    expected_chisquare : bool
558        If True prints the expected chisquare which is
559        corrected by effects caused by correlated input data.
560        This can take a while as the full correlation matrix
561        has to be calculated (default False).
562    num_grad : bool
563        Use numerical differentiation instead of automatic differentiation to perform the error propagation (default False).
564    n_parms : int, optional
565        Number of fit parameters. Overrides automatic detection of parameter count.
566        Useful when autodetection fails. Must match the length of initial_guess (if provided).
567
568    Notes
569    -----
570    Based on the orthogonal distance regression module of scipy.
571
572    Returns
573    -------
574    output : Fit_result
575        Parameters and information on the fitted result.
576    '''
577
578    output = Fit_result()
579
580    output.fit_function = func
581
582    x = np.array(x)
583
584    x_shape = x.shape
585
586    if kwargs.get('num_grad') is True:
587        jacobian = num_jacobian
588        hessian = num_hessian
589    else:
590        jacobian = auto_jacobian
591        hessian = auto_hessian
592
593    if not callable(func):
594        raise TypeError('func has to be a function.')
595
596    if 'n_parms' in kwargs:
597        n_parms = kwargs.get('n_parms')
598        if not isinstance(n_parms, int):
599            raise TypeError(
600                f"'n_parms' must be an integer, got {n_parms!r} "
601                f"of type {type(n_parms).__name__}."
602            )
603        if n_parms <= 0:
604            raise ValueError(
605                f"'n_parms' must be a positive integer, got {n_parms}."
606            )
607    else:
608        for i in range(100):
609            try:
610                func(np.arange(i), x.T[0])
611            except TypeError:
612                continue
613            except IndexError:
614                continue
615            else:
616                break
617        else:
618            raise RuntimeError("Fit function is not valid.")
619
620        n_parms = i
621
622    if not silent:
623        print('Fit with', n_parms, 'parameter' + 's' * (n_parms > 1))
624
625    x_f = np.vectorize(lambda o: o.value)(x)
626    dx_f = np.vectorize(lambda o: o.dvalue)(x)
627    y_f = np.array([o.value for o in y])
628    dy_f = np.array([o.dvalue for o in y])
629
630    if np.any(np.asarray(dx_f) <= 0.0):
631        raise Exception('No x errors available, run the gamma method first.')
632
633    if np.any(np.asarray(dy_f) <= 0.0):
634        raise Exception('No y errors available, run the gamma method first.')
635
636    if 'initial_guess' in kwargs:
637        x0 = kwargs.get('initial_guess')
638        if len(x0) != n_parms:
639            raise Exception('Initial guess does not have the correct length: %d vs. %d' % (len(x0), n_parms))
640    else:
641        x0 = [1] * n_parms
642
643    data = RealData(x_f, y_f, sx=dx_f, sy=dy_f)
644    model = Model(func)
645    odr = ODR(data, model, x0, partol=np.finfo(np.float64).eps)
646    odr.set_job(fit_type=0, deriv=1)
647    out = odr.run()
648
649    output.residual_variance = out.res_var
650
651    output.method = 'ODR'
652
653    output.message = out.stopreason
654
655    output.xplus = out.xplus
656
657    if not silent:
658        print('Method: ODR')
659        print(*out.stopreason)
660        print('Residual variance:', output.residual_variance)
661
662    if out.info > 3:
663        raise Exception('The minimization procedure did not converge.')
664
665    m = x_f.size
666
667    def odr_chisquare(p):
668        model = func(p[:n_parms], p[n_parms:].reshape(x_shape))
669        chisq = anp.sum(((y_f - model) / dy_f) ** 2) + anp.sum(((x_f - p[n_parms:].reshape(x_shape)) / dx_f) ** 2)
670        return chisq
671
672    if kwargs.get('expected_chisquare') is True:
673        W = np.diag(1 / np.asarray(np.concatenate((dy_f.ravel(), dx_f.ravel()))))
674
675        if kwargs.get('covariance') is not None:
676            cov = kwargs.get('covariance')
677        else:
678            cov = covariance(np.concatenate((y, x.ravel())))
679
680        number_of_x_parameters = int(m / x_f.shape[-1])
681
682        old_jac = jacobian(func)(out.beta, out.xplus)
683        fused_row1 = np.concatenate((old_jac, np.concatenate((number_of_x_parameters * [np.zeros(old_jac.shape)]), axis=0)))
684        fused_row2 = np.concatenate((jacobian(lambda x, y: func(y, x))(out.xplus, out.beta).reshape(x_f.shape[-1], x_f.shape[-1] * number_of_x_parameters), np.identity(number_of_x_parameters * old_jac.shape[0])))
685        new_jac = np.concatenate((fused_row1, fused_row2), axis=1)
686
687        A = W @ new_jac
688        P_phi = A @ np.linalg.pinv(A.T @ A) @ A.T
689        expected_chisquare = np.trace((np.identity(P_phi.shape[0]) - P_phi) @ W @ cov @ W)
690        if expected_chisquare <= 0.0:
691            warnings.warn("Negative expected_chisquare.", RuntimeWarning)
692            expected_chisquare = np.abs(expected_chisquare)
693        output.chisquare_by_expected_chisquare = odr_chisquare(np.concatenate((out.beta, out.xplus.ravel()))) / expected_chisquare
694        if not silent:
695            print('chisquare/expected_chisquare:',
696                  output.chisquare_by_expected_chisquare)
697
698    fitp = out.beta
699    try:
700        hess = hessian(odr_chisquare)(np.concatenate((fitp, out.xplus.ravel())))
701    except TypeError:
702        raise Exception("It is required to use autograd.numpy instead of numpy within fit functions, see the documentation for details.") from None
703
704    def odr_chisquare_compact_x(d):
705        model = func(d[:n_parms], d[n_parms:n_parms + m].reshape(x_shape))
706        chisq = anp.sum(((y_f - model) / dy_f) ** 2) + anp.sum(((d[n_parms + m:].reshape(x_shape) - d[n_parms:n_parms + m].reshape(x_shape)) / dx_f) ** 2)
707        return chisq
708
709    jac_jac_x = hessian(odr_chisquare_compact_x)(np.concatenate((fitp, out.xplus.ravel(), x_f.ravel())))
710
711    # Compute hess^{-1} @ jac_jac_x[:n_parms + m, n_parms + m:] using LAPACK dgesv
712    try:
713        deriv_x = -scipy.linalg.solve(hess, jac_jac_x[:n_parms + m, n_parms + m:])
714    except np.linalg.LinAlgError:
715        raise Exception("Cannot invert hessian matrix.")
716
717    def odr_chisquare_compact_y(d):
718        model = func(d[:n_parms], d[n_parms:n_parms + m].reshape(x_shape))
719        chisq = anp.sum(((d[n_parms + m:] - model) / dy_f) ** 2) + anp.sum(((x_f - d[n_parms:n_parms + m].reshape(x_shape)) / dx_f) ** 2)
720        return chisq
721
722    jac_jac_y = hessian(odr_chisquare_compact_y)(np.concatenate((fitp, out.xplus.ravel(), y_f)))
723
724    # Compute hess^{-1} @ jac_jac_y[:n_parms + m, n_parms + m:] using LAPACK dgesv
725    try:
726        deriv_y = -scipy.linalg.solve(hess, jac_jac_y[:n_parms + m, n_parms + m:])
727    except np.linalg.LinAlgError:
728        raise Exception("Cannot invert hessian matrix.")
729
730    result = []
731    for i in range(n_parms):
732        result.append(derived_observable(lambda my_var, **kwargs: (my_var[0] + np.finfo(np.float64).eps) / (x.ravel()[0].value + np.finfo(np.float64).eps) * out.beta[i], list(x.ravel()) + list(y), man_grad=list(deriv_x[i]) + list(deriv_y[i])))
733
734    output.fit_parameters = result
735
736    output.odr_chisquare = odr_chisquare(np.concatenate((out.beta, out.xplus.ravel())))
737    output.dof = x.shape[-1] - n_parms
738    output.p_value = 1 - scipy.stats.chi2.cdf(output.odr_chisquare, output.dof)
739
740    return output
741
742
743def fit_lin(x, y, **kwargs):
744    """Performs a linear fit to y = n + m * x and returns two Obs n, m.
745
746    Parameters
747    ----------
748    x : list
749        Can either be a list of floats in which case no xerror is assumed, or
750        a list of Obs, where the dvalues of the Obs are used as xerror for the fit.
751    y : list
752        List of Obs, the dvalues of the Obs are used as yerror for the fit.
753
754    Returns
755    -------
756    fit_parameters : list[Obs]
757        LIist of fitted observables.
758    """
759
760    def f(a, x):
761        y = a[0] + a[1] * x
762        return y
763
764    if all(isinstance(n, Obs) for n in x):
765        out = total_least_squares(x, y, f, **kwargs)
766        return out.fit_parameters
767    elif all(isinstance(n, float) or isinstance(n, int) for n in x) or isinstance(x, np.ndarray):
768        out = least_squares(x, y, f, **kwargs)
769        return out.fit_parameters
770    else:
771        raise TypeError('Unsupported types for x')
772
773
774def qqplot(x, o_y, func, p, title=""):
775    """Generates a quantile-quantile plot of the fit result which can be used to
776       check if the residuals of the fit are gaussian distributed.
777
778    Returns
779    -------
780    None
781    """
782
783    residuals = []
784    for i_x, i_y in zip(x, o_y):
785        residuals.append((i_y - func(p, i_x)) / i_y.dvalue)
786    residuals = sorted(residuals)
787    my_y = [o.value for o in residuals]
788    probplot = scipy.stats.probplot(my_y)
789    my_x = probplot[0][0]
790    plt.figure(figsize=(8, 8 / 1.618))
791    plt.errorbar(my_x, my_y, fmt='o')
792    fit_start = my_x[0]
793    fit_stop = my_x[-1]
794    samples = np.arange(fit_start, fit_stop, 0.01)
795    plt.plot(samples, samples, 'k--', zorder=11, label='Standard normal distribution')
796    plt.plot(samples, probplot[1][0] * samples + probplot[1][1], zorder=10, label='Least squares fit, r=' + str(np.around(probplot[1][2], 3)), marker='', ls='-')
797
798    plt.xlabel('Theoretical quantiles')
799    plt.ylabel('Ordered Values')
800    plt.legend(title=title)
801    plt.draw()
802
803
804def residual_plot(x, y, func, fit_res, title=""):
805    """Generates a plot which compares the fit to the data and displays the corresponding residuals
806
807    For uncorrelated data the residuals are expected to be distributed ~N(0,1).
808
809    Returns
810    -------
811    None
812    """
813    sorted_x = sorted(x)
814    xstart = sorted_x[0] - 0.5 * (sorted_x[1] - sorted_x[0])
815    xstop = sorted_x[-1] + 0.5 * (sorted_x[-1] - sorted_x[-2])
816    x_samples = np.arange(xstart, xstop + 0.01, 0.01)
817
818    plt.figure(figsize=(8, 8 / 1.618))
819    gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1], wspace=0.0, hspace=0.0)
820    ax0 = plt.subplot(gs[0])
821    ax0.errorbar(x, [o.value for o in y], yerr=[o.dvalue for o in y], ls='none', fmt='o', capsize=3, markersize=5, label='Data')
822    ax0.plot(x_samples, func([o.value for o in fit_res], x_samples), label='Fit', zorder=10, ls='-', ms=0)
823    ax0.set_xticklabels([])
824    ax0.set_xlim([xstart, xstop])
825    ax0.set_xticklabels([])
826    ax0.legend(title=title)
827
828    residuals = (np.asarray([o.value for o in y]) - func([o.value for o in fit_res], np.asarray(x))) / np.asarray([o.dvalue for o in y])
829    ax1 = plt.subplot(gs[1])
830    ax1.plot(x, residuals, 'ko', ls='none', markersize=5)
831    ax1.tick_params(direction='out')
832    ax1.tick_params(axis="x", bottom=True, top=True, labelbottom=True)
833    ax1.axhline(y=0.0, ls='--', color='k', marker=" ")
834    ax1.fill_between(x_samples, -1.0, 1.0, alpha=0.1, facecolor='k')
835    ax1.set_xlim([xstart, xstop])
836    ax1.set_ylabel('Residuals')
837    plt.subplots_adjust(wspace=None, hspace=None)
838    plt.draw()
839
840
841def error_band(x, func, beta):
842    """Calculate the error band for an array of sample values x, for given fit function func with optimized parameters beta.
843
844    Returns
845    -------
846    err : np.array(Obs)
847        Error band for an array of sample values x
848    """
849    cov = covariance(beta)
850    if np.any(np.abs(cov - cov.T) > 1000 * np.finfo(np.float64).eps):
851        warnings.warn("Covariance matrix is not symmetric within floating point precision", RuntimeWarning)
852
853    deriv = []
854    for i, item in enumerate(x):
855        deriv.append(np.array(egrad(func)([o.value for o in beta], item)))
856
857    err = []
858    for i, item in enumerate(x):
859        err.append(np.sqrt(deriv[i] @ cov @ deriv[i]))
860    err = np.array(err)
861
862    return err
863
864
865def ks_test(objects=None):
866    """Performs a Kolmogorov–Smirnov test for the p-values of all fit object.
867
868    Parameters
869    ----------
870    objects : list
871        List of fit results to include in the analysis (optional).
872
873    Returns
874    -------
875    None
876    """
877
878    if objects is None:
879        obs_list = []
880        for obj in gc.get_objects():
881            if isinstance(obj, Fit_result):
882                obs_list.append(obj)
883    else:
884        obs_list = objects
885
886    p_values = [o.p_value for o in obs_list]
887
888    bins = len(p_values)
889    x = np.arange(0, 1.001, 0.001)
890    plt.plot(x, x, 'k', zorder=1)
891    plt.xlim(0, 1)
892    plt.ylim(0, 1)
893    plt.xlabel('p-value')
894    plt.ylabel('Cumulative probability')
895    plt.title(str(bins) + ' p-values')
896
897    n = np.arange(1, bins + 1) / np.float64(bins)
898    Xs = np.sort(p_values)
899    plt.step(Xs, n)
900    diffs = n - Xs
901    loc_max_diff = np.argmax(np.abs(diffs))
902    loc = Xs[loc_max_diff]
903    plt.annotate('', xy=(loc, loc), xytext=(loc, loc + diffs[loc_max_diff]), arrowprops=dict(arrowstyle='<->', shrinkA=0, shrinkB=0))
904    plt.draw()
905
906    print(scipy.stats.kstest(p_values, 'uniform'))
907
908
909def _extract_val_and_dval(string):
910    split_string = string.split('(')
911    if '.' in split_string[0] and '.' not in split_string[1][:-1]:
912        factor = 10 ** -len(split_string[0].partition('.')[2])
913    else:
914        factor = 1
915    return float(split_string[0]), float(split_string[1][:-1]) * factor
916
917
918def _construct_prior_obs(i_prior, i_n):
919    if isinstance(i_prior, Obs):
920        return i_prior
921    elif isinstance(i_prior, str):
922        loc_val, loc_dval = _extract_val_and_dval(i_prior)
923        return cov_Obs(loc_val, loc_dval ** 2, '#prior' + str(i_n) + f"_{np.random.randint(2147483647):010d}")
924    else:
925        raise TypeError("Prior entries need to be 'Obs' or 'str'.")