Source code for schrodinger.application.matsci.mlearn.base

"""
Classes and functions to deal with ML features.

Copyright Schrodinger, LLC. All rights reserved."""

from sklearn.base import BaseEstimator
from sklearn.base import TransformerMixin


class BaseFeaturizer(BaseEstimator, TransformerMixin):
    """
    Class that MUST be inherited to create sklearn Model.
    """

    # See http://scikit-learn.org/stable/modules/generated/sklearn.base.TransformerMixin.html
    # for more documentation

    def fit(self, data, data_y=None):
        """
        Fit and return self. Anything that evaluates properties related to the
        passed data should go here. For example, compute physical properties
        of a stucture and save them as class property, to be used in the
        transform method.

        :type data: numpy array of shape [n_samples, n_features]
        :param data: Training set

        :type data_y: numpy array of shape [n_samples]
        :param data_y: Target values

        :rtype: BaseFeaturizer
        :return: self object with fitted data
        """

        return self

    def transform(self, data):
        """
        Get numerical features. Must be implemented by a child class.

        :type data: numpy array of shape [n_samples, n_features]
        :param data: Training set

        :rtype: numpy array of shape [n_samples, n_features_new]
        :return: Transformed array
        """

        raise NotImplementedError