<commit_before>import numpy as np
from sklearn.exceptions import NotFittedError

from mushroom_rl.core import Serializable


class Ensemble(Serializable):
    """
    This class is used to create an ensemble of regressors.

    """
    def __init__(self, model, n_models, prediction='mean', **params):
        """
        Constructor.

        Args:
            approximator (class): the model class to approximate the
                Q-function.
            n_models (int): number of regressors in the ensemble;
            prediction (str, ['mean', 'sum', 'min', 'max']): the type of
                prediction to make;
            **params: parameters dictionary to create each regressor.

        """
        self._model = list()
        self._prediction = prediction

        for _ in range(n_models):
            self._model.append(model(**params))

        self._add_save_attr(
            _model=self._get_serialization_method(model)
        )

    def fit(self, *z, idx=None, **fit_params):
        """
        Fit the ``idx``-th model of the ensemble if ``idx`` is provided, every
        model otherwise.

        Args:
            *z: a list containing the inputs to use to predict with each
                regressor of the ensemble;
            idx (int, None): index of the model to fit;
            **fit_params: other params.

        """
        if idx is None:
            for i in range(len(self)):
                self[i].fit(*z, **fit_params)
        else:
            self[idx].fit(*z, **fit_params)

    def predict(self, *z, idx=None, prediction=None, compute_variance=False,
                **predict_params):
        """
        Predict.

        Args:
            *z: a list containing the inputs to use to predict with each
                regressor of the ensemble;
            idx (int, None): index of the model to use for prediction;
            prediction (str, None): the type of prediction to make. When
                provided, it overrides the ``prediction`` class attribute;
            compute_variance (bool, False): whether to compute the variance
                of the prediction or not;
            **predict_params: other parameters used by the predict method
                the regressor.

        Returns:
            The predictions of the model.

        """
        if idx is None:
            predictions = list()
            for i in range(len(self._model)):
                try:
                    predictions.append(self[i].predict(*z, **predict_params))
                except NotFittedError:
                    pass

            if len(predictions) == 0:
                raise NotFittedError

            prediction = self._prediction if prediction is None else prediction

            if prediction == 'mean':
                results = np.mean(predictions, axis=0)
            elif prediction == 'sum':
                results = np.sum(predictions, axis=0)
            elif prediction == 'min':
                results = np.min(predictions, axis=0)
            elif prediction == 'max':
                results = np.max(predictions, axis=0)
            else:
                raise ValueError
            if compute_variance:
                results = [results, np.var(predictions, ddof=1, axis=0)]
        else:
            try:
                results = self[idx].predict(*z, **predict_params)
            except NotFittedError:
                raise NotFittedError

        return results

    def reset(self):
        """
        Reset the model parameters.

        """
        try:
            for m in self.model:
                m.reset()
        except AttributeError:
            raise NotImplementedError('Attempt to reset weights of a'
                                      ' non-parametric regressor.')<commit_msg>Ensemble regressor can now handle list of indexes for prediction.<commit_after>import numpy as np
from sklearn.exceptions import NotFittedError

from mushroom_rl.core import Serializable


class Ensemble(Serializable):
    """
    This class is used to create an ensemble of regressors.

    """
    def __init__(self, model, n_models, prediction='mean', **params):
        """
        Constructor.

        Args:
            approximator (class): the model class to approximate the
                Q-function.
            n_models (int): number of regressors in the ensemble;
            prediction (str, ['mean', 'sum', 'min', 'max']): the type of
                prediction to make;
            **params: parameters dictionary to create each regressor.

        """
        self._model = list()
        self._prediction = prediction

        for _ in range(n_models):
            self._model.append(model(**params))

        self._add_save_attr(
            _model=self._get_serialization_method(model)
        )

    def fit(self, *z, idx=None, **fit_params):
        """
        Fit the ``idx``-th model of the ensemble if ``idx`` is provided, every
        model otherwise.

        Args:
            *z: a list containing the inputs to use to predict with each
                regressor of the ensemble;
            idx (int, None): index of the model to fit;
            **fit_params: other params.

        """
        if idx is None:
            for i in range(len(self)):
                self[i].fit(*z, **fit_params)
        else:
            self[idx].fit(*z, **fit_params)

    def predict(self, *z, idx=None, prediction=None, compute_variance=False,
                **predict_params):
        """
        Predict.

        Args:
            *z: a list containing the inputs to use to predict with each
                regressor of the ensemble;
            idx (int, None): index of the model to use for prediction;
            prediction (str, None): the type of prediction to make. When
                provided, it overrides the ``prediction`` class attribute;
            compute_variance (bool, False): whether to compute the variance
                of the prediction or not;
            **predict_params: other parameters used by the predict method
                the regressor.

        Returns:
            The predictions of the model.

        """
        if idx is None:
            idx = [x for x in range(len(self))]

        if isinstance(idx, int):
            try:
                results = self[idx].predict(*z, **predict_params)
            except NotFittedError:
                raise NotFittedError
        else:
            predictions = list()
            for i in idx:
                try:
                    predictions.append(self[i].predict(*z, **predict_params))
                except NotFittedError:
                    pass

            if len(predictions) == 0:
                raise NotFittedError

            prediction = self._prediction if prediction is None else prediction

            print(np.array(predictions).shape)
            if prediction == 'mean':
                results = np.mean(predictions, axis=0)
            elif prediction == 'sum':
                results = np.sum(predictions, axis=0)
            elif prediction == 'min':
                results = np.min(predictions, axis=0)
            elif prediction == 'max':
                results = np.max(predictions, axis=0)
            else:
                raise ValueError
            if compute_variance:
                results = [results, np.var(predictions, ddof=1, axis=0)]

        return results

    def reset(self):
        """
        Reset the model parameters.

        """
        try:
            for m in self.model:
                m.reset()
        except AttributeError:
            raise NotImplementedError('Attempt to reset weights of a'
                                      ' non-parametric regressor.')
{'fraction_non_alphanumeric': 0.08776011373917539, 'fraction_numerical': 0.0018094868812201111, 'mean_word_length': 1.6194989844278944, '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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1559221', 'n_tokens_mistral': 1987, 'n_tokens_neox': 1886, 'n_words': 654}