utils.py 文件源码

def retrieve_features(best_estimator):
    """Retrieve selected features from any estimator.

    In case it has the 'get_support' method, use it.
    Else, if it has a 'coef_' attribute, assume it's a linear model and the
    features correspond to the indices of the coefficients != 0
    """
    if hasattr(best_estimator, 'get_support'):
        return np.nonzero(best_estimator.get_support())[0]
    elif hasattr(best_estimator, 'coef_'):
        # print best_estimator.coef_
        if best_estimator.coef_.ndim > 1 and 1 not in best_estimator.coef_.shape:
            sel_feats = []
            for dim in range(best_estimator.coef_.ndim):
                sel_feats += np.nonzero(
                    best_estimator.coef_[dim])[0].ravel().tolist()
            return np.unique(sel_feats)
        return np.nonzero(best_estimator.coef_.flatten())[0]
    else:
        # Raise an error
        raise AttributeError('The best_estimator object does not have '
                             'neither the `coef_` attribute nor the '
                             '`get_support` method')