Q3D-Calibration/qdx/fit.py

import numpy as np
from astropy.modeling import models, fitting

from .utils import get_hist


def fit_line(model, x, y):
    """
    Line fitting kx+b

    Parameters
    ----------
    model : astropy.modeling.Model
        fit model
    x : array
        x data, can have two columns representing two independent variables
    y : array
        y data

    Returns
    -------
    fitted_model : astropy fitting model
        fitting gaussian model
    """
    fitter = fitting.LevMarLSQFitter()
    if np.ndim(x) == 1:
        fitted_model = fitter(model, x, y)
    else:
        fitted_model = fitter(model, x[:, 0], x[:, 1], y)

    return fitted_model


def fit_hist_gaussian(x, step=1):
    """
    Gaussian fitting is performed on the histogram

    Parameters
    ----------
    x : array
        data point
    step : int, optional
        Minimum bin width. The bin width is an integer multiple of step.

    Returns
    -------
    fitted_model : astropy fitting model
        fitting gaussian model
    """
    fitter = fitting.LMLSQFitter()

    count, center = get_hist(x, step=step)
    model = models.Gaussian1D(amplitude=count.max(), mean=x.mean(), stddev=x.std())
    fitted_model = fitter(model, center, count)

    return fitted_model