<commit_before>
'''
Basic models for decomposing large-scale stacked profiles
'''

import numpy as np
from astropy.modeling import models, fitting
from scipy.interpolate import InterpolatedUnivariateSpline
from scipy.special import erf
from functools import partial


def fit_2gaussian(vels, spectrum):
    '''
    Fit a 2 Gaussian model with the means tied.

    .. todo:: Allow passing sigma for weighting in the fit.

    '''

    max_vel = vels[np.argmax(spectrum)]
    specmax = spectrum.max()

    # Estimate the inner width from the HWHM
    sigma_est = find_hwhm(vels, spectrum)[0]
    # Use this as the narrow estimate, and let the wide guess be 3x

    g_HI_init = models.Gaussian1D(amplitude=0.75 * specmax, mean=max_vel,
                                  stddev=sigma_est) +  \
        models.Gaussian1D(amplitude=0.25 * specmax, mean=max_vel,
                          stddev=3 * sigma_est)

    # Force to the same mean
    def tie_mean(mod):
        return mod.mean_0

    g_HI_init.mean_1.tied = tie_mean

    fit_g = fitting.LevMarLSQFitter()

    g_HI = fit_g(g_HI_init, vels, spectrum)

    # The covariance matrix is hidden away... tricksy
    cov = fit_g.fit_info['param_cov']
    parnames = [n for n in g_HI.param_names if n not in ['mean_1']]
    parvals = [v for (n, v) in zip(g_HI.param_names, g_HI.parameters)
               if n in parnames]

    if cov is not None:
        chan_width = np.diff(vels[:2])[0]
        parerrs = []
        for name, var in zip(parnames, np.diag(cov)):
            if "mean" in name or "stddev" in name:
                # Add the finite channel width in quadrature
                stderr = np.sqrt(var + (0.5 * chan_width)**2)
            else:
                stderr = np.sqrt(var)

            parerrs.append(stderr)
    else:
        parerrs = [np.NaN] * len(parnames)

    return parvals, parerrs, cov, parnames, g_HI


def fit_gaussian(vels, spectrum, p0=None, sigma=None):
    '''
    Fit a Gaussian model.

    Parameters
    ----------
    sigma : `~astropy.units.Quantity`, optional
        Pass a single value to be used for all points or an array of values
        equal to the size of vels and spectrum. These are used to define
        weights in the fit, and the weights as defined as 1 / sigma, as is
        used in `~scipy.optimize.curve_fit`.

    '''

    if p0 is None:
        max_vel = vels[np.argmax(spectrum)]
        specmax = spectrum.max()

        # Estimate the inner width from the HWHM
        sigma_est = find_hwhm(vels, spectrum)[0]
    else:
        specmax, max_vel, sigma_est = p0

    if sigma is not None:
        if hasattr(sigma, 'size'):
            if sigma.size > 1:
                if sigma.size != vels.size:
                    raise ValueError("sigma must match the data shape when "
                                     "multiple values are given.")
                weights = 1 / np.abs(sigma)
            else:
                weights = np.array([np.abs(sigma)] * len(vels))
        else:
            weights = np.array([np.abs(sigma)] * len(vels))
    else:
        weights = None

    # Use this as the narrow estimate, and let the wide guess be 3x

    g_init = models.Gaussian1D(amplitude=specmax, mean=max_vel,
                               stddev=sigma_est)

    fit_g = fitting.LevMarLSQFitter()

    g_fit = fit_g(g_init, vels, spectrum, weights=weights)

    # The covariance matrix is hidden away... tricksy
    cov = fit_g.fit_info['param_cov']
    if cov is None:
        cov = np.zeros((3, 3)) * np.NaN
    parnames = g_fit.param_names
    parvals = g_fit.parameters

    # Sometimes the width is negative
    parvals[-1] = np.abs(parvals[-1])

    if cov is not None:
        chan_width = np.diff(vels[:2])[0]
        parerrs = []
        for name, var in zip(parnames, np.diag(cov)):
            if "mean" in name or "stddev" in name:
                # Add the finite channel width in quadrature
                stderr = np.sqrt(var + (0.5 * chan_width)**2)
            else:<commit_msg>1 / sigma<commit_after>
'''
Basic models for decomposing large-scale stacked profiles
'''

import numpy as np
from astropy.modeling import models, fitting
from scipy.interpolate import InterpolatedUnivariateSpline
from scipy.special import erf
from functools import partial


def fit_2gaussian(vels, spectrum):
    '''
    Fit a 2 Gaussian model with the means tied.

    .. todo:: Allow passing sigma for weighting in the fit.

    '''

    max_vel = vels[np.argmax(spectrum)]
    specmax = spectrum.max()

    # Estimate the inner width from the HWHM
    sigma_est = find_hwhm(vels, spectrum)[0]
    # Use this as the narrow estimate, and let the wide guess be 3x

    g_HI_init = models.Gaussian1D(amplitude=0.75 * specmax, mean=max_vel,
                                  stddev=sigma_est) +  \
        models.Gaussian1D(amplitude=0.25 * specmax, mean=max_vel,
                          stddev=3 * sigma_est)

    # Force to the same mean
    def tie_mean(mod):
        return mod.mean_0

    g_HI_init.mean_1.tied = tie_mean

    fit_g = fitting.LevMarLSQFitter()

    g_HI = fit_g(g_HI_init, vels, spectrum)

    # The covariance matrix is hidden away... tricksy
    cov = fit_g.fit_info['param_cov']
    parnames = [n for n in g_HI.param_names if n not in ['mean_1']]
    parvals = [v for (n, v) in zip(g_HI.param_names, g_HI.parameters)
               if n in parnames]

    if cov is not None:
        chan_width = np.diff(vels[:2])[0]
        parerrs = []
        for name, var in zip(parnames, np.diag(cov)):
            if "mean" in name or "stddev" in name:
                # Add the finite channel width in quadrature
                stderr = np.sqrt(var + (0.5 * chan_width)**2)
            else:
                stderr = np.sqrt(var)

            parerrs.append(stderr)
    else:
        parerrs = [np.NaN] * len(parnames)

    return parvals, parerrs, cov, parnames, g_HI


def fit_gaussian(vels, spectrum, p0=None, sigma=None):
    '''
    Fit a Gaussian model.

    Parameters
    ----------
    sigma : `~astropy.units.Quantity`, optional
        Pass a single value to be used for all points or an array of values
        equal to the size of vels and spectrum. These are used to define
        weights in the fit, and the weights as defined as 1 / sigma, as is
        used in `~scipy.optimize.curve_fit`.

    '''

    if p0 is None:
        max_vel = vels[np.argmax(spectrum)]
        specmax = spectrum.max()

        # Estimate the inner width from the HWHM
        sigma_est = find_hwhm(vels, spectrum)[0]
    else:
        specmax, max_vel, sigma_est = p0

    if sigma is not None:
        if hasattr(sigma, 'size'):
            if sigma.size > 1:
                if sigma.size != vels.size:
                    raise ValueError("sigma must match the data shape when "
                                     "multiple values are given.")
                weights = 1 / np.abs(sigma)
            else:
                weights = 1 / np.array([np.abs(sigma)] * len(vels))
        else:
            weights = 1 / np.array([np.abs(sigma)] * len(vels))
    else:
        weights = None

    # Use this as the narrow estimate, and let the wide guess be 3x

    g_init = models.Gaussian1D(amplitude=specmax, mean=max_vel,
                               stddev=sigma_est)

    fit_g = fitting.LevMarLSQFitter()

    g_fit = fit_g(g_init, vels, spectrum, weights=weights)

    # The covariance matrix is hidden away... tricksy
    cov = fit_g.fit_info['param_cov']
    if cov is None:
        cov = np.zeros((3, 3)) * np.NaN
    parnames = g_fit.param_names
    parvals = g_fit.parameters

    # Sometimes the width is negative
    parvals[-1] = np.abs(parvals[-1])

    if cov is not None:
        chan_width = np.diff(vels[:2])[0]
        parerrs = []
        for name, var in zip(parnames, np.diag(cov)):
            if "mean" in name or "stddev" in name:
                # Add the finite channel width in quadrature
                stderr = np.sqrt(var + (0.5 * chan_width)**2)
            else:
{'fraction_non_alphanumeric': 0.09808075772681954, 'fraction_numerical': 0.010094715852442673, 'mean_word_length': 2.0466970387243735, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1356196', 'n_tokens_mistral': 2672, 'n_tokens_neox': 2492, 'n_words': 782}