Repository: scikit-learn/scikit-learn
File: sklearn/semi_supervised/tests/test_label_propagation.py (Python)

Code before:
)

SPARSE_TYPES = ("sparse_csr", "sparse_csc", "sparse_csr_array", "sparse_csc_array")
CONSTRUCTOR_TYPES = ("array",) + SPARSE_TYPES

ESTIMATORS = [
    (label_propagation.LabelPropagation, {"kernel": "rbf"}),
    (label_propagation.LabelPropagation, {"kernel": "knn", "n_neighbors": 2}),
    (
        label_propagation.LabelPropagation,
        {"kernel": lambda x, y: rbf_kernel(x, y, gamma=20)},
    ),
    (label_propagation.LabelSpreading, {"kernel": "rbf"}),
    (label_propagation.LabelSpreading, {"kernel": "knn", "n_neighbors": 2}),
    (
        label_propagation.LabelSpreading,
        {"kernel": lambda x, y: rbf_kernel(x, y, gamma=20)},
    ),
]

LP_ESTIMATORS = [
    est for est in ESTIMATORS if isinstance(est, label_propagation.LabelPropagation)
]


@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_fit_transduction(global_dtype, Estimator, parameters):
    samples = np.asarray([[1.0, 0.0], [0.0, 2.0], [1.0, 3.0]], dtype=global_dtype)
    labels = [0, 1, -1]
    clf = Estimator(**parameters).fit(samples, labels)
    assert clf.transduction_[2] == 1


@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_distribution(global_dtype, Estimator, parameters):
    if parameters["kernel"] == "knn":
        pytest.skip(
            "Unstable test for this configuration: changes in k-NN ordering break it."
        )
    samples = np.asarray([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=global_dtype)
    labels = [0, 1, -1]
    clf = Estimator(**parameters).fit(samples, labels)
    assert_allclose(clf.label_distributions_[2], [0.5, 0.5], atol=1e-2)


@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_predict(global_dtype, Estimator, parameters):
    samples = np.asarray([[1.0, 0.0], [0.0, 2.0], [1.0, 3.0]], dtype=global_dtype)
    labels = [0, 1, -1]
    clf = Estimator(**parameters).fit(samples, labels)
    assert_array_equal(clf.predict([[0.5, 2.5]]), np.array([1]))

Review comment (suggestion):
```suggestion
LP_ESTIMATORS = [
    (klass, params)
    for (klass, params) in ESTIMATORS
    if klass == label_propagation.LabelPropagation
]
```

Code after:
)

SPARSE_TYPES = ("sparse_csr", "sparse_csc", "sparse_csr_array", "sparse_csc_array")
CONSTRUCTOR_TYPES = ("array",) + SPARSE_TYPES

ESTIMATORS = [
    (label_propagation.LabelPropagation, {"kernel": "rbf"}),
    (label_propagation.LabelPropagation, {"kernel": "knn", "n_neighbors": 2}),
    (
        label_propagation.LabelPropagation,
        {"kernel": lambda x, y: rbf_kernel(x, y, gamma=20)},
    ),
    (label_propagation.LabelSpreading, {"kernel": "rbf"}),
    (label_propagation.LabelSpreading, {"kernel": "knn", "n_neighbors": 2}),
    (
        label_propagation.LabelSpreading,
        {"kernel": lambda x, y: rbf_kernel(x, y, gamma=20)},
    ),
]

LP_ESTIMATORS = [
    (klass, params)
    for (klass, params) in ESTIMATORS
    if klass == label_propagation.LabelPropagation
]


@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_fit_transduction(global_dtype, Estimator, parameters):
    samples = np.asarray([[1.0, 0.0], [0.0, 2.0], [1.0, 3.0]], dtype=global_dtype)
    labels = [0, 1, -1]
    clf = Estimator(**parameters).fit(samples, labels)
    assert clf.transduction_[2] == 1


@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_distribution(global_dtype, Estimator, parameters):
    if parameters["kernel"] == "knn":
        pytest.skip(
            "Unstable test for this configuration: changes in k-NN ordering break it."
        )
    samples = np.asarray([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=global_dtype)
    labels = [0, 1, -1]
    clf = Estimator(**parameters).fit(samples, labels)
    assert_allclose(clf.label_distributions_[2], [0.5, 0.5], atol=1e-2)


@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_predict(global_dtype, Estimator, parameters):
    samples = np.asarray([[1.0, 0.0], [0.0, 2.0], [1.0, 3.0]], dtype=global_dtype)
    labels = [0, 1, -1]