<commit_before>"""
TODO:
-
"""

# Third-party
import numpy as np
from numpy.polynomial.polynomial import polyval
from celerite.modeling import Model
from celerite import terms, GP

# Project
from comoving_rv.log import logger

__all__ = ['MeanModel', 'GPModel']

class MeanModel(Model):

    def __init__(self, n_bg_coef, **p0):
        self._n_bg_coef = n_bg_coef
        self.parameter_names = (["a{}".format(i) for i in range(n_bg_coef)])
        super(MeanModel, self).__init__(**p0)
<commit_msg>ad fit all lines func<commit_after># Third-party
import numpy as np
from numpy.polynomial.polynomial import polyval
from celerite.modeling import Model
from celerite import terms, GP

# Project
from comoving_rv.log import logger

__all__ = ['fit_all_lines', 'MeanModel', 'GPModel']

def fit_all_lines(pixels, flux, flux_ivar, line_waves, line_pixels, half_width=5):
    """
    Given a wavelength guess (a list of rough pixel-wavelength correspondences),
    measure more precise line centroids to then fit for a wavelength
    calibratrion function.

    Parameters
    ----------
    pixels : array_like
        Pixel array for a 1D comp. lamp spectrum.
    flux : array_like
        Flux array for a 1D comp. lamp spectrum.
    flux_ivar : array_like
        Inverse-variance array for a 1D comp. lamp spectrum.
    line_waves : array_like
        List of wavelength values for emission lines in a comp. lamp spectrum.
    line_pixels : array_like
        List of pixel centroids for emission lines in a comp. lamp spectrum.
    half_width : int
        Number of pixels on either side of line
    """

    _idx = np.argsort(line_waves)
    wvln = np.array(line_waves)[_idx]
    pixl = np.array(line_pixels)[_idx]

    fit_centroids = []

    for pix_ctr,wave in zip(pixl, wvln):

        logger.debug("Fitting line at predicted pix={:.2f}, λ={:.2f}"
                     .format(pix_ctr, wave))

        # indices for region around line
        i1 = int(np.floor(pix_ctr-half_width))
        i2 = int(np.ceil(pix_ctr+half_width)) + 1

        # recenter window
        i0 = i1 + flux[i1:i2].argmax()
        i1 = int(np.floor(i0-half_width))
        i2 = int(np.ceil(i0+half_width)) + 1

        _pixl = pixels[i1:i2]
        _flux = flux[i1:i2]
        _ivar = flux_ivar[i1:i2]

        # instead of doing anything fancy (e.g., fitting a profile), just
        # estimate the mean...
        x0 = np.sum(_pixl * _flux**2 * _ivar) / np.sum(_flux**2 * _ivar)
        fit_centroids.append(x0)

    return np.array(fit_centroids)

class MeanModel(Model):

    def __init__(self, n_bg_coef, **p0):
        self._n_bg_coef = n_bg_coef
        self.parameter_names = (["a{}".format(i) for i in range(n_bg_coef)])
        super(MeanModel, self).__init__(**p0)

{'fraction_non_alphanumeric': 0.10663764961915125, 'fraction_numerical': 0.011606819006166122, 'mean_word_length': 2.8735955056179776, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4872225', 'n_tokens_mistral': 972, 'n_tokens_neox': 940, 'n_words': 259}