Repository: scikit-learn/scikit-learn
File: sklearn/ensemble/_hist_gradient_boosting/tests/test_binning.py (Python)

Code before:
    data = distinct_values[repeated_indices]
    rng.shuffle(data)
    assert_array_equal(np.unique(data), np.sort(distinct_values))

    data = data.reshape(-1, 1)

    mapper_1 = _BinMapper(n_bins=n_distinct + 1)
    binned_1 = mapper_1.fit_transform(data)
    assert_array_equal(np.unique(binned_1[:, 0]), np.arange(n_distinct))

    # Adding more bins to the mapper yields the same results (same thresholds)
    mapper_2 = _BinMapper(n_bins=min(256, n_distinct * 3) + 1)
    binned_2 = mapper_2.fit_transform(data)

    assert_allclose(mapper_1.bin_thresholds_[0], mapper_2.bin_thresholds_[0])
    assert_array_equal(binned_1, binned_2)


@pytest.mark.parametrize("n_bins", [50, None])
def test_binmapper_weighted_vs_repeated_equivalence(global_random_seed, n_bins):
    rng = np.random.RandomState(global_random_seed)

    n_samples = 200
    X = rng.randn(n_samples, 3)
    if n_bins is None:
        n_bins = np.unique(X[:, rng.randint(3)]).shape[0] + rng.randint(5)

    sw = rng.randint(0, 5, size=n_samples)
    X_repeated = np.repeat(X, sw, axis=0)

    est_weighted = _BinMapper(n_bins=n_bins).fit(X, sample_weight=sw)
    est_repeated = _BinMapper(n_bins=n_bins).fit(X_repeated, sample_weight=None)
    assert_allclose(est_weighted.bin_thresholds_, est_repeated.bin_thresholds_)

    X_trans_weighted = est_weighted.transform(X)
    X_trans_repeated = est_repeated.transform(X)
    assert_array_equal(X_trans_weighted, X_trans_repeated)


@pytest.mark.parametrize(
    "max_bins, scale, offset",
    [
        (3, 2, -1),
        (42, 1, 0),
        (255, 0.3, 42),
    ],
)
def test_bin_mapper_identity_small(max_bins, scale, offset):
    data = np.arange(max_bins).reshape(-1, 1) * scale + offset
    # max_bins is the number of bins for non-missing values
    n_bins = max_bins + 1

Review comment (suggestion):
```suggestion
        # Set n_bins to a value strictly larger than the number of unique values in at
        # least one of the data columns: in this case, the bin thresholds should
        # be set to midpoints between successive unique values.
        n_bins = np.unique(X[:, rng.randint(X.shape[1])]).shape[0] + rng.randint(5) + 1
```

Code after:
    ],
)
def test_bin_mapper_identity_repeated_values(max_bins, n_distinct, multiplier):
    data = np.array(list(range(n_distinct)) * multiplier).reshape(-1, 1)
    # max_bins is the number of bins for non-missing values
    n_bins = max_bins + 1
    binned = _BinMapper(n_bins=n_bins).fit_transform(data)
    assert_array_equal(data, binned)


@pytest.mark.parametrize("n_distinct", [2, 7, 42])
def test_bin_mapper_repeated_values_invariance(n_distinct):
    rng = np.random.RandomState(42)
    distinct_values = rng.normal(size=n_distinct)
    assert len(np.unique(distinct_values)) == n_distinct

    repeated_indices = rng.randint(low=0, high=n_distinct, size=1000)
    data = distinct_values[repeated_indices]
    rng.shuffle(data)
    assert_array_equal(np.unique(data), np.sort(distinct_values))

    data = data.reshape(-1, 1)

    mapper_1 = _BinMapper(n_bins=n_distinct + 1)
    binned_1 = mapper_1.fit_transform(data)
    assert_array_equal(np.unique(binned_1[:, 0]), np.arange(n_distinct))

    # Adding more bins to the mapper yields the same results (same thresholds)
    mapper_2 = _BinMapper(n_bins=min(256, n_distinct * 3) + 1)
    binned_2 = mapper_2.fit_transform(data)

    assert_allclose(mapper_1.bin_thresholds_[0], mapper_2.bin_thresholds_[0])
    assert_array_equal(binned_1, binned_2)


@pytest.mark.parametrize("n_bins", [50, 250])
def test_binmapper_weighted_vs_repeated_equivalence(global_random_seed, n_bins):
    rng = np.random.RandomState(global_random_seed)

    n_samples = 200
    X = rng.randn(n_samples, 3)
    sw = rng.randint(0, 5, size=n_samples)
    X_repeated = np.repeat(X, sw, axis=0)

    est_weighted = _BinMapper(n_bins=n_bins).fit(X, sample_weight=sw)
    est_repeated = _BinMapper(n_bins=n_bins).fit(X_repeated, sample_weight=None)
    assert_allclose(est_weighted.bin_thresholds_, est_repeated.bin_thresholds_)

    X_trans_weighted = est_weighted.transform(X)
    X_trans_repeated = est_repeated.transform(X)
    assert_array_equal(X_trans_weighted, X_trans_repeated)