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

Code before:
@pytest.mark.parametrize("accepted_sparse_type", SPARSE_TYPES)
@pytest.mark.parametrize("index_dtype", [np.int32, np.int64])
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_sparse_input_types(
    accepted_sparse_type, index_dtype, dtype, Estimator, parameters
):
    # This is non-regression test for #17085
    X = _convert_container([[1.0, 0.0], [0.0, 2.0], [1.0, 3.0]], accepted_sparse_type)
    X.data = X.data.astype(dtype, copy=False)
    X.indices = X.indices.astype(index_dtype, copy=False)
    X.indptr = X.indptr.astype(index_dtype, copy=False)
    labels = [0, 1, -1]
    clf = Estimator(**parameters).fit(X, labels)
    assert_array_equal(clf.predict([[0.5, 2.5]]), np.array([1]))


@pytest.mark.parametrize("constructor", CONSTRUCTOR_TYPES)
@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS[:2])
def test_label_propagation_build_graph_normalized(constructor, Estimator, parameters):
    # required but unused X and labels values
    X = np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 3.0]])
    labels = [0, 1, -1]

    # test normalization of an affinity_matrix
    aff_matrix = np.array([[1.0, 1.0, 0.0], [2.0, 1.0, 1.0], [0.0, 1.0, 3.0]])
    expected = np.array([[0.5, 0.5, 0.0], [0.5, 0.25, 0.25], [0.0, 0.25, 0.75]])

    def kernel_affinity_matrix(x, y=None):
        return _convert_container(aff_matrix, constructor)

    clf = Estimator(kernel=kernel_affinity_matrix).fit(X, labels)
    graph = clf._build_graph()
    assert_array_equal(getattr(graph, "toarray", lambda: graph)(), expected)


@pytest.mark.parametrize("constructor_type", CONSTRUCTOR_TYPES)
def test_convergence_speed(constructor_type):
    # This is a non-regression test for #5774
    X = _convert_container([[1.0, 0.0], [0.0, 1.0], [1.0, 2.5]], constructor_type)
    y = np.array([0, 1, -1])
    mdl = label_propagation.LabelSpreading(kernel="rbf", max_iter=5000)
    mdl.fit(X, y)

    # this should converge quickly:
    assert mdl.n_iter_ < 10
    assert_array_equal(mdl.predict(X), [0, 1, 1])


def test_convergence_warning():
    # This is a non-regression test for #5774

Review comment (suggestion):
It could be nice to add a test to check that `graph` is now properly row normalized, eg something like `assert_allclose(graph.sum(axis=1),  1)`.

Code after:
    expected[unlabelled_idx, :] = Y_u
    expected /= expected.sum(axis=1)[:, np.newaxis]

    assert_allclose(expected, clf.label_distributions_, atol=1e-4)


@pytest.mark.parametrize("accepted_sparse_type", SPARSE_TYPES)
@pytest.mark.parametrize("index_dtype", [np.int32, np.int64])
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
@pytest.mark.parametrize("Estimator, parameters", ESTIMATORS)
def test_sparse_input_types(
    accepted_sparse_type, index_dtype, dtype, Estimator, parameters
):
    # This is non-regression test for #17085
    X = _convert_container([[1.0, 0.0], [0.0, 2.0], [1.0, 3.0]], accepted_sparse_type)
    X.data = X.data.astype(dtype, copy=False)
    X.indices = X.indices.astype(index_dtype, copy=False)
    X.indptr = X.indptr.astype(index_dtype, copy=False)
    labels = [0, 1, -1]
    clf = Estimator(**parameters).fit(X, labels)
    assert_array_equal(clf.predict([[0.5, 2.5]]), np.array([1]))


@pytest.mark.parametrize("constructor", CONSTRUCTOR_TYPES)
@pytest.mark.parametrize("Estimator, parameters", LP_ESTIMATORS)
def test_label_propagation_build_graph_normalized(constructor, Estimator, parameters):
    # required but unused X and labels values
    X = np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 3.0]])
    labels = [0, 1, -1]

    # test normalization of an affinity_matrix
    aff_matrix = np.array([[1.0, 1.0, 0.0], [2.0, 1.0, 1.0], [0.0, 1.0, 3.0]])
    expected = np.array([[0.5, 0.5, 0.0], [0.5, 0.25, 0.25], [0.0, 0.25, 0.75]])

    def kernel_affinity_matrix(x, y=None):
        return _convert_container(aff_matrix, constructor)

    clf = Estimator(kernel=kernel_affinity_matrix).fit(X, labels)
    graph = clf._build_graph()
    assert_allclose(graph.sum(axis=1), 1)  # normalized rows

    if issparse(graph):
        graph = graph.toarray()
    assert_allclose(graph, expected)


@pytest.mark.parametrize("constructor_type", CONSTRUCTOR_TYPES)
def test_convergence_speed(constructor_type):
    # This is a non-regression test for #5774
    X = _convert_container([[1.0, 0.0], [0.0, 1.0], [1.0, 2.5]], constructor_type)
    y = np.array([0, 1, -1])