seizure_modeling.py 文件源码

def get_model_class(method):
    """
    Returns the class associated with a method string.
    :param method: A string describing the method to use.
    :return: A class corresponding to the method.
    """
    if method == 'logistic':
        return sklearn.linear_model.LogisticRegression
    elif method == 'svm':
        return sklearn.svm.SVC
    elif method == 'mirowski-svm':
        return sklearn.svm.SVC
    elif method == 'sgd':
        return sklearn.linear_model.SGDClassifier
    elif method == 'random-forest':
        return sklearn.ensemble.RandomForestClassifier
    elif method == 'nearest-centroid':
        return sklearn.neighbors.NearestCentroid
    elif method == 'knn':
        return sklearn.neighbors.KNeighborsClassifier
    elif method == 'bagging':
        return sklearn.ensemble.BaggingClassifier
    else:
        raise NotImplementedError("Method {} is not supported".format(method))