<commit_before>""" Make a single model class to rule them all """
from .dataset import Dataset
from .train_model import train_model as _train_model
from .infer_labels import infer_labels
from .spectral_model import diagnostics as _diagnostics
from .helpers.triangle import corner
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from copy import deepcopy

plt.rc('text', usetex=True)
plt.rc('font', family='serif')

LARGE = 200.
SMALL = 1. / LARGE

class CannonModel(object):
    def __init__(self, dataset, order):
        if not isinstance(dataset, Dataset):
            txt = 'Expecting a Dataset instance, got {0}'
            raise TypeError(txt.format(type(dataset)))
        self.dataset = dataset
        self._model = None
        self.order = order


    @property
    def model(self):
        """ return the model definition or raise an error if not trained """
        if self._model is None:
            raise RuntimeError('Model not trained')
        else:
            return self._model


    def train(self, *args, **kwargs):
        """ Train the model """
        self._model = _train_model(self.dataset)


    def diagnostics(self):
        """Run a set of diagnostics on the model.

        Plot the 0th order coefficients as the baseline spectrum.
        Overplot the continuum pixels.

        Plot each label's leading coefficient as a function of wavelength.
        Color-code by label.

        Histogram of the chi squareds of the fits.
        Dotted line corresponding to DOF = npixels - nlabels
        """
        _model_diagnostics(self.dataset, self.model)


    def infer_labels(self, dataset):
        """
        Uses the model to solve for labels of the test set.

        Parameters
        ----------
        dataset: Dataset
            dataset that needs label inference

        Returns
        -------
        output of infer_labels
        """
        return infer_labels(self.model, dataset)


    def draw_spectra(self, dataset):
        """
        Parameters
        ----------
        dataset: Dataset
            dataset that needs label inference

        Returns
        -------
        cannon_set: Dataset
            same dataset as the input value with updated fluxes and variances
        """
        coeffs_all, covs, scatters, red_chisqs, pivots, label_vector = self.model
        nstars = len(dataset.test_SNR)
        cannon_fluxes = np.zeros(dataset.test_flux.shape)
        cannon_ivars = np.zeros(dataset.test_ivar.shape)
        for i in range(nstars):
            x = label_vector[:,i,:]
            spec_fit = np.einsum('ij, ij->i', x, coeffs_all)
            cannon_fluxes[i,:] = spec_fit
            bad = dataset.test_ivar[i,:] == SMALL
            cannon_ivars[i,:][~bad] = 1. / scatters[~bad] ** 2
        cannon_set = deepcopy(dataset)
        cannon_set.test_flux = cannon_fluxes
        cannon_set.test_ivar = cannon_ivars
        return cannon_set


    def split_array(self, array, num):
        """ split an array into a certain number of segments

        Parameters
        ----------
        array: numpy ndarray
            array to be split

        num: int
            number of elements to split array into

        Returns
        ------
        out: list 
            split array
        """
        avg = len(array) / float(num)
        out = []
        last = 0.0
        while last < len(array):
            out.append(array[int(last):int(last+avg)])
            last += avg
        return out


    def diagnostics(
            self, baseline_spec_plot_name = "baseline_spec_with_cont_pix",
            leading_coeffs_plot_name = "leading_coeffs.png",
            chisq_dist_plot_name = "modelfit_chisqs.png"):
        """Run a set of diagnostics on the model.

        Plot the 0th order coefficients as the baseline spectrum.
        Overplot the continuum pixels.

        Plot each label's leading coefficient as a function of wavelength.
        Color-code by label.

        Histogram of the chi squareds of the fits.
        Dotted line corresponding to DOF = npixels - nlabels

        Parameters
        ----------
        (optional) baseline_spec_plot_name: str
            plot name
        (optional) leading_coeffs_plot_name: str
            plot name
        (optional) chisq_dist_plot_name: str
            plot name
        """
        dataset = self.dataset
        model = self.model
        contmask = dataset.contmask
        lams = dataset.wl
        label_names = dataset.get_plotting_labels()
        coeffs_all, covs, scatters, chisqs, pivots, label_vector = model
        npixels = len(lams)
        nlabels = len(pivots)

        if contmask != None:
            # Baseline spectrum with continuum<commit_msg>warning about comparison to None being elementwise in the future<commit_after>""" Make a single model class to rule them all """
from .dataset import Dataset
from .train_model import train_model as _train_model
from .infer_labels import infer_labels
from .spectral_model import diagnostics as _diagnostics
from .helpers.triangle import corner
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from copy import deepcopy

plt.rc('text', usetex=True)
plt.rc('font', family='serif')

LARGE = 200.
SMALL = 1. / LARGE

class CannonModel(object):
    def __init__(self, dataset, order):
        if not isinstance(dataset, Dataset):
            txt = 'Expecting a Dataset instance, got {0}'
            raise TypeError(txt.format(type(dataset)))
        self.dataset = dataset
        self._model = None
        self.order = order


    @property
    def model(self):
        """ return the model definition or raise an error if not trained """
        if self._model is None:
            raise RuntimeError('Model not trained')
        else:
            return self._model


    def train(self, *args, **kwargs):
        """ Train the model """
        self._model = _train_model(self.dataset)


    def diagnostics(self):
        """Run a set of diagnostics on the model.

        Plot the 0th order coefficients as the baseline spectrum.
        Overplot the continuum pixels.

        Plot each label's leading coefficient as a function of wavelength.
        Color-code by label.

        Histogram of the chi squareds of the fits.
        Dotted line corresponding to DOF = npixels - nlabels
        """
        _model_diagnostics(self.dataset, self.model)


    def infer_labels(self, dataset):
        """
        Uses the model to solve for labels of the test set.

        Parameters
        ----------
        dataset: Dataset
            dataset that needs label inference

        Returns
        -------
        output of infer_labels
        """
        return infer_labels(self.model, dataset)


    def draw_spectra(self, dataset):
        """
        Parameters
        ----------
        dataset: Dataset
            dataset that needs label inference

        Returns
        -------
        cannon_set: Dataset
            same dataset as the input value with updated fluxes and variances
        """
        coeffs_all, covs, scatters, red_chisqs, pivots, label_vector = self.model
        nstars = len(dataset.test_SNR)
        cannon_fluxes = np.zeros(dataset.test_flux.shape)
        cannon_ivars = np.zeros(dataset.test_ivar.shape)
        for i in range(nstars):
            x = label_vector[:,i,:]
            spec_fit = np.einsum('ij, ij->i', x, coeffs_all)
            cannon_fluxes[i,:] = spec_fit
            bad = dataset.test_ivar[i,:] == SMALL
            cannon_ivars[i,:][~bad] = 1. / scatters[~bad] ** 2
        cannon_set = deepcopy(dataset)
        cannon_set.test_flux = cannon_fluxes
        cannon_set.test_ivar = cannon_ivars
        return cannon_set


    def split_array(self, array, num):
        """ split an array into a certain number of segments

        Parameters
        ----------
        array: numpy ndarray
            array to be split

        num: int
            number of elements to split array into

        Returns
        ------
        out: list 
            split array
        """
        avg = len(array) / float(num)
        out = []
        last = 0.0
        while last < len(array):
            out.append(array[int(last):int(last+avg)])
            last += avg
        return out


    def diagnostics(
            self, baseline_spec_plot_name = "baseline_spec_with_cont_pix",
            leading_coeffs_plot_name = "leading_coeffs.png",
            chisq_dist_plot_name = "modelfit_chisqs.png"):
        """Run a set of diagnostics on the model.

        Plot the 0th order coefficients as the baseline spectrum.
        Overplot the continuum pixels.

        Plot each label's leading coefficient as a function of wavelength.
        Color-code by label.

        Histogram of the chi squareds of the fits.
        Dotted line corresponding to DOF = npixels - nlabels

        Parameters
        ----------
        (optional) baseline_spec_plot_name: str
            plot name
        (optional) leading_coeffs_plot_name: str
            plot name
        (optional) chisq_dist_plot_name: str
            plot name
        """
        dataset = self.dataset
        model = self.model
        contmask = dataset.contmask
        lams = dataset.wl
        label_names = dataset.get_plotting_labels()
        coeffs_all, covs, scatters, chisqs, pivots, label_vector = model
        npixels = len(lams)
        nlabels = len(pivots)

        if contmask is not None:
            # Baseline spectrum with continuum
{'fraction_non_alphanumeric': 0.08502657080337606, 'fraction_numerical': 0.002292383036365531, 'mean_word_length': 1.9650911337658326, 'pattern_counts': {'":': 0, '<': 5, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6298816', 'n_tokens_mistral': 2740, 'n_tokens_neox': 2597, 'n_words': 844}