<commit_before>from astroML.time_series import multiterm_periodogram
from scipy.interpolate import UnivariateSpline as interpolate
import numpy as np        

from .base import PeriodicModeler

class LombScargle(PeriodicModeler):
    """Class implementing Lomb-Scargle fitting"""
    def __init__(self, frequencies=None, min_period=None, max_period=None, 
            num_periods=None):
        if frequencies is not None:
            self.frequencies = frequencies
            self.periods = 2*np.pi/self.frequencies
        else:
            self.periods = np.logspace(np.log10(min_period),
                    np.log_10(max_period), num_periods)
            self.frequencies = 2*np.pi/self.periods
        self.power = None
        self.power_fn = None
    
    def fit(self, t, y, dy=None, filts=None):
        """Provide data for the fit"""
        self.powers = multiterm_periodogram(t, y, dy, self.frequencies,
                                            nterms=1)
        self.power_fn = interpolate(self.periods, self.powers, s=0, k=0)

    def score(self, period):
        """Compute the score for a period or array of periods"""
        if self.power_fn is None:
            raise RuntimeError("Need to fit data first.")
        else:
            return self.power_fn(period)
<commit_msg>Added period_search override to lomb_scargle<commit_after>from astroML.time_series import multiterm_periodogram
from scipy.interpolate import UnivariateSpline as interpolate
import numpy as np

from .base import PeriodicModeler

class LombScargle(PeriodicModeler):
    """Class implementing Lomb-Scargle fitting"""
    def __init__(self, frequencies=None, min_period=None, max_period=None, 
            num_periods=None):
        if frequencies is not None:
            self.frequencies = frequencies
            self.periods = 2*np.pi/self.frequencies
        else:
            self.periods = np.logspace(np.log10(min_period),
                    np.log_10(max_period), num_periods)
            self.frequencies = 2*np.pi/self.periods
        self.power = None
        self.power_fn = None
    
    def fit(self, t, y, dy=None, filts=None):
        """Provide data for the fit"""
        self.powers = multiterm_periodogram(t, y, dy, self.frequencies,
                                            nterms=1)
        self.power_fn = interpolate(self.periods, self.powers, s=0, k=0)

    def score(self, period):
        """Compute the score for a period or array of periods"""
        if self.power_fn is None:
            raise RuntimeError("Need to fit data first.")
        else:
            return self.power_fn(period)
    
    def period_search(self, pmin=None, pmax=None, resolution=None, **kwargs):
        """Find the best period for the model"""
        if pmin is None: 
            pmin = min(self.periods)
        if pmax is None:
            pmax = max(self.periods)
        if resolution is None:
            resolution = len(self.periods))
        return super(LombScargle).period_search(pmin, pmax, 
                resolution=resolution, **kwargs)
{'fraction_non_alphanumeric': 0.09111691704768125, 'fraction_numerical': 0.005878510777269758, 'mean_word_length': 2.293548387096774, '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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6583775', 'n_tokens_mistral': 931, 'n_tokens_neox': 894, 'n_words': 236}