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

Code before:
    """Check that PolynomialCountSketch results are the same for dense and sparse
    input.
    """
    ps_dense = PolynomialCountSketch(
        n_components=500, gamma=gamma, degree=degree, coef0=coef0, random_state=42
    )
    Xt_dense = ps_dense.fit_transform(X)
    Yt_dense = ps_dense.transform(Y)

    ps_sparse = PolynomialCountSketch(
        n_components=500, gamma=gamma, degree=degree, coef0=coef0, random_state=42
    )
    Xt_sparse = ps_sparse.fit_transform(csr_container(X))
    Yt_sparse = ps_sparse.transform(csr_container(Y))

    assert_allclose(Xt_dense, Xt_sparse)
    assert_allclose(Yt_dense, Yt_sparse)


def _linear_kernel(X, Y):
    if X.ndim == 1 and Y.ndim == 1:
        return X @ Y
    elif X.ndim == 2 and Y.ndim == 2:
        return X @ Y.T
    else:
        raise ValueError("Incompatible shapes for linear kernel")


@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
def test_additive_chi2_sampler(csr_container):
    # test that AdditiveChi2Sampler approximates kernel on random data

    # compute exact kernel
    # abbreviations for easier formula
    X_ = X[:, np.newaxis, :].copy()
    Y_ = Y[np.newaxis, :, :].copy()

    large_kernel = 2 * X_ * Y_ / (X_ + Y_)

    # reduce to n_samples_x x n_samples_y by summing over features
    kernel = large_kernel.sum(axis=2)

    # approximate kernel mapping
    transform = AdditiveChi2Sampler(sample_steps=3)
    X_trans = transform.fit_transform(X)
    Y_trans = transform.transform(Y)

    kernel_approx = np.dot(X_trans, Y_trans.T)

    assert_array_almost_equal(kernel, kernel_approx, 1)


Review comment (suggestion):
Wouldn't the following work in all cases?

```suggestion
    return X @ Y.T
```

Code after:
@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
@pytest.mark.parametrize("gamma", [0.1, 1.0])
@pytest.mark.parametrize("degree", [1, 2, 3])
@pytest.mark.parametrize("coef0", [0, 2.5])
def test_polynomial_count_sketch_dense_sparse(gamma, degree, coef0, csr_container):
    """Check that PolynomialCountSketch results are the same for dense and sparse
    input.
    """
    ps_dense = PolynomialCountSketch(
        n_components=500, gamma=gamma, degree=degree, coef0=coef0, random_state=42
    )
    Xt_dense = ps_dense.fit_transform(X)
    Yt_dense = ps_dense.transform(Y)

    ps_sparse = PolynomialCountSketch(
        n_components=500, gamma=gamma, degree=degree, coef0=coef0, random_state=42
    )
    Xt_sparse = ps_sparse.fit_transform(csr_container(X))
    Yt_sparse = ps_sparse.transform(csr_container(Y))

    assert_allclose(Xt_dense, Xt_sparse)
    assert_allclose(Yt_dense, Yt_sparse)


def _linear_kernel(x, y):
    return x @ y


@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
def test_additive_chi2_sampler(csr_container):
    # test that AdditiveChi2Sampler approximates kernel on random data

    # compute exact kernel
    # abbreviations for easier formula
    X_ = X[:, np.newaxis, :].copy()
    Y_ = Y[np.newaxis, :, :].copy()

    large_kernel = 2 * X_ * Y_ / (X_ + Y_)

    # reduce to n_samples_x x n_samples_y by summing over features
    kernel = large_kernel.sum(axis=2)

    # approximate kernel mapping
    transform = AdditiveChi2Sampler(sample_steps=3)
    X_trans = transform.fit_transform(X)
    Y_trans = transform.transform(Y)

    kernel_approx = np.dot(X_trans, Y_trans.T)

    assert_array_almost_equal(kernel, kernel_approx, 1)