from typing import Any, Dict, List import numpy as np import pytest from hypothesis import given, settings, strategies from scipy import sparse import xgboost as xgb from xgboost.testing import ( IteratorForTest, make_batches, make_batches_sparse, make_categorical, make_ltr, make_sparse_regression, predictor_equal, ) from xgboost.testing.data import check_inf, np_dtypes from xgboost.testing.data_iter import run_mixed_sparsity from xgboost.testing.quantile_dmatrix import ( check_categorical_strings, check_ref_quantile_cut, ) class TestQuantileDMatrix: def test_basic(self) -> None: """Checks for np array, list, tuple.""" n_samples = 234 n_features = 8 rng = np.random.default_rng() X = rng.normal(loc=0, scale=3, size=n_samples * n_features).reshape( n_samples, n_features ) y = rng.normal(0, 3, size=n_samples) Xy = xgb.QuantileDMatrix(X, y) assert Xy.num_row() == n_samples assert Xy.num_col() == n_features X = sparse.random(n_samples, n_features, density=0.1, format="csr") Xy = xgb.QuantileDMatrix(X, y) assert Xy.num_row() == n_samples assert Xy.num_col() == n_features X = sparse.random(n_samples, n_features, density=0.8, format="csr") Xy = xgb.QuantileDMatrix(X, y) assert Xy.num_row() == n_samples assert Xy.num_col() == n_features n_samples = 64 data = [] for f in range(n_samples): row = [f] * n_features data.append(row) assert np.array(data).shape == (n_samples, n_features) Xy = xgb.QuantileDMatrix(data, max_bin=256) assert Xy.num_row() == n_samples assert Xy.num_col() == n_features r = np.arange(1.0, n_samples) np.testing.assert_allclose(Xy.get_data().toarray()[1:, 0], r) def test_categorical_strings(self) -> None: check_categorical_strings("cpu") def test_error(self): from sklearn.model_selection import train_test_split rng = np.random.default_rng(1994) X, y = make_categorical( n_samples=128, n_features=2, n_categories=3, onehot=False ) reg = xgb.XGBRegressor(tree_method="hist", enable_categorical=True) w = rng.uniform(0, 1, size=y.shape[0]) X_train, X_test, y_train, y_test, w_train, w_test = train_test_split( X, y, w, random_state=1994 ) with pytest.raises(ValueError, match="sample weight"): reg.fit( X, y, sample_weight=w_train, eval_set=[(X_test, y_test)], sample_weight_eval_set=[w_test], ) with pytest.raises(ValueError, match="sample weight"): reg.fit( X_train, y_train, sample_weight=w, eval_set=[(X_test, y_test)], sample_weight_eval_set=[w_test], ) @pytest.mark.parametrize("sparsity", [0.0, 0.1, 0.8, 0.9]) def test_with_iterator(self, sparsity: float) -> None: n_samples_per_batch = 317 n_features = 8 n_batches = 7 if sparsity == 0.0: it = IteratorForTest( *make_batches(n_samples_per_batch, n_features, n_batches, False), cache=None, ) else: it = IteratorForTest( *make_batches_sparse( n_samples_per_batch, n_features, n_batches, sparsity ), cache=None, ) Xy = xgb.QuantileDMatrix(it) assert Xy.num_row() == n_samples_per_batch * n_batches assert Xy.num_col() == n_features def test_different_size(self) -> None: n_samples_per_batch = 317 n_features = 8 n_batches = 7 it = IteratorForTest( *make_batches( n_samples_per_batch, n_features, n_batches, False, vary_size=True ), cache=None, ) Xy = xgb.QuantileDMatrix(it) assert Xy.num_row() == 2429 X, y, w = it.as_arrays() Xy1 = xgb.QuantileDMatrix(X, y, weight=w) assert predictor_equal(Xy, Xy1) @pytest.mark.parametrize("sparsity", [0.0, 0.1, 0.5, 0.8, 0.9]) def test_training(self, sparsity: float) -> None: n_samples_per_batch = 317 n_features = 8 n_batches = 7 if sparsity == 0.0: it = IteratorForTest( *make_batches(n_samples_per_batch, n_features, n_batches, False), cache=None, ) else: it = IteratorForTest( *make_batches_sparse( n_samples_per_batch, n_features, n_batches, sparsity ), cache=None, ) parameters = {"tree_method": "hist", "max_bin": 256} Xy_it = xgb.QuantileDMatrix(it, max_bin=parameters["max_bin"]) from_it = xgb.train(parameters, Xy_it) X, y, w = it.as_arrays() w_it = Xy_it.get_weight() np.testing.assert_allclose(w_it, w) Xy_arr = xgb.DMatrix(X, y, weight=w) from_arr = xgb.train(parameters, Xy_arr) np.testing.assert_allclose(from_arr.predict(Xy_it), from_it.predict(Xy_arr)) y -= y.min() y += 0.01 Xy = xgb.QuantileDMatrix(X, y, weight=w) with pytest.raises(ValueError, match=r"Only.*hist.*"): parameters = { "tree_method": "approx", "max_bin": 256, "objective": "reg:gamma", } xgb.train(parameters, Xy) def run_ref_dmatrix(self, rng: Any, device: str, enable_cat: bool) -> None: n_samples, n_features = 2048, 17 if enable_cat: X, y = make_categorical( n_samples, n_features, n_categories=13, onehot=False ) if device == "cuda": import cudf X = cudf.from_pandas(X) y = cudf.from_pandas(y) else: X = rng.normal(loc=0, scale=3, size=n_samples * n_features).reshape( n_samples, n_features ) y = rng.normal(0, 3, size=n_samples) # Use ref Xy = xgb.QuantileDMatrix(X, y, enable_categorical=enable_cat) Xy_valid: xgb.DMatrix = xgb.QuantileDMatrix( X, y, ref=Xy, enable_categorical=enable_cat ) qdm_results: Dict[str, Dict[str, List[float]]] = {} xgb.train( {"tree_method": "hist", "device": device}, Xy, evals=[(Xy, "Train"), (Xy_valid, "valid")], evals_result=qdm_results, ) np.testing.assert_allclose( qdm_results["Train"]["rmse"], qdm_results["valid"]["rmse"] ) # No ref Xy_valid = xgb.DMatrix(X, y, enable_categorical=enable_cat) qdm_results = {} xgb.train( {"tree_method": "hist", "device": device}, Xy, evals=[(Xy, "Train"), (Xy_valid, "valid")], evals_result=qdm_results, ) np.testing.assert_allclose( qdm_results["Train"]["rmse"], qdm_results["valid"]["rmse"] ) # Different number of features Xy = xgb.QuantileDMatrix(X, y, enable_categorical=enable_cat) dXy = xgb.DMatrix(X, y, enable_categorical=enable_cat) n_samples, n_features = 256, 15 X = rng.normal(loc=0, scale=3, size=n_samples * n_features).reshape( n_samples, n_features ) y = rng.normal(0, 3, size=n_samples) with pytest.raises(ValueError, match=r".*features\."): xgb.QuantileDMatrix(X, y, ref=Xy, enable_categorical=enable_cat) # Compare training results n_samples, n_features = 256, 17 if enable_cat: X, y = make_categorical(n_samples, n_features, 13, onehot=False) if device == "cuda": import cudf X = cudf.from_pandas(X) y = cudf.from_pandas(y) else: X = rng.normal(loc=0, scale=3, size=n_samples * n_features).reshape( n_samples, n_features ) y = rng.normal(0, 3, size=n_samples) Xy_valid = xgb.QuantileDMatrix(X, y, ref=Xy, enable_categorical=enable_cat) # use DMatrix as ref Xy_valid_d = xgb.QuantileDMatrix(X, y, ref=dXy, enable_categorical=enable_cat) dXy_valid = xgb.DMatrix(X, y, enable_categorical=enable_cat) qdm_results = {} xgb.train( {"tree_method": "hist", "device": device}, Xy, evals=[(Xy, "Train"), (Xy_valid, "valid")], evals_result=qdm_results, ) dm_results: Dict[str, Dict[str, List[float]]] = {} xgb.train( {"tree_method": "hist", "device": device}, dXy, evals=[(dXy, "Train"), (dXy_valid, "valid"), (Xy_valid_d, "dvalid")], evals_result=dm_results, ) np.testing.assert_allclose( dm_results["Train"]["rmse"], qdm_results["Train"]["rmse"] ) np.testing.assert_allclose( dm_results["valid"]["rmse"], qdm_results["valid"]["rmse"] ) np.testing.assert_allclose( dm_results["dvalid"]["rmse"], qdm_results["valid"]["rmse"] ) Xy_valid = xgb.QuantileDMatrix(X, y, enable_categorical=enable_cat) with pytest.raises(ValueError, match="should be used as a reference"): xgb.train( {"device": device}, dXy, evals=[(dXy, "Train"), (Xy_valid, "Valid")] ) def test_ref_quantile_cut(self) -> None: check_ref_quantile_cut("cpu") @pytest.mark.parametrize("enable_cat", [True, False]) def test_ref_dmatrix(self, enable_cat: bool) -> None: rng = np.random.RandomState(1994) self.run_ref_dmatrix(rng, "cpu", enable_cat) @pytest.mark.parametrize("sparsity", [0.0, 0.5]) def test_predict(self, sparsity: float) -> None: n_samples, n_features = 256, 4 X, y = make_categorical( n_samples, n_features, n_categories=13, onehot=False, sparsity=sparsity ) Xy = xgb.DMatrix(X, y, enable_categorical=True) booster = xgb.train({"tree_method": "hist"}, Xy) Xy = xgb.DMatrix(X, y, enable_categorical=True) a = booster.predict(Xy) qXy = xgb.QuantileDMatrix(X, y, enable_categorical=True) b = booster.predict(qXy) np.testing.assert_allclose(a, b) def test_ltr(self) -> None: X, y, qid, w = make_ltr(100, 3, 3, 5) Xy_qdm = xgb.QuantileDMatrix(X, y, qid=qid, weight=w) Xy = xgb.DMatrix(X, y, qid=qid, weight=w) xgb.train({"tree_method": "hist", "objective": "rank:ndcg"}, Xy) from_qdm = xgb.QuantileDMatrix(X, weight=w, ref=Xy_qdm) from_dm = xgb.QuantileDMatrix(X, weight=w, ref=Xy) assert predictor_equal(from_qdm, from_dm) def test_check_inf(self) -> None: rng = np.random.default_rng(1994) check_inf(rng) # we don't test empty Quantile DMatrix in single node construction. @given( strategies.integers(1, 1000), strategies.integers(1, 100), strategies.fractions(0, 0.99), ) @settings(deadline=None, print_blob=True) def test_to_csr(self, n_samples: int, n_features: int, sparsity: float) -> None: csr, y = make_sparse_regression(n_samples, n_features, sparsity, False) csr = csr.astype(np.float32) qdm = xgb.QuantileDMatrix(data=csr, label=y) ret = qdm.get_data() np.testing.assert_equal(csr.indptr, ret.indptr) np.testing.assert_equal(csr.indices, ret.indices) booster = xgb.train({"tree_method": "hist"}, dtrain=qdm) np.testing.assert_allclose( booster.predict(qdm), booster.predict(xgb.DMatrix(qdm.get_data())) ) def test_dtypes(self) -> None: """Checks for both np array and pd DataFrame.""" n_samples = 128 n_features = 16 for orig, x in np_dtypes(n_samples, n_features): m0 = xgb.QuantileDMatrix(orig) m1 = xgb.QuantileDMatrix(x) assert predictor_equal(m0, m1) # unsupported types for dtype in [ np.bytes_, np.complex64, np.complex128, ]: X: np.ndarray = np.array(orig, dtype=dtype) with pytest.raises(ValueError): xgb.QuantileDMatrix(X) def test_changed_max_bin(self) -> None: n_samples = 128 n_features = 16 csr, y = make_sparse_regression(n_samples, n_features, 0.5, False) Xy = xgb.QuantileDMatrix(csr, y, max_bin=9) booster = xgb.train({"max_bin": 9}, Xy, num_boost_round=2) Xy = xgb.QuantileDMatrix(csr, y, max_bin=11) with pytest.raises(ValueError, match="consistent"): xgb.train({}, Xy, num_boost_round=2, xgb_model=booster) def test_mixed_sparsity(self) -> None: run_mixed_sparsity("cpu") def test_sparse_predict(self) -> None: X, y = make_sparse_regression(512, 16, sparsity=0.9, as_dense=False) Xy: xgb.DMatrix = xgb.QuantileDMatrix(X, y) booster = xgb.train({}, Xy, num_boost_round=8) p0 = booster.predict(Xy) Xy = xgb.DMatrix(X, y) p1 = booster.predict(Xy) np.testing.assert_allclose(p0, p1) X, y = make_categorical(128, 16, 5, onehot=False, sparsity=0.9) Xy = xgb.QuantileDMatrix(X, y, enable_categorical=True) booster = xgb.train({}, Xy, num_boost_round=8) p0 = booster.predict(Xy) Xy = xgb.DMatrix(X, y, enable_categorical=True) p1 = booster.predict(Xy) np.testing.assert_allclose(p0, p1) def test_cv_error(self) -> None: X, y = make_sparse_regression(8, 2, sparsity=0.2, as_dense=False) Xy = xgb.QuantileDMatrix(X, y) with pytest.raises(ValueError, match=""): cv = xgb.cv({}, Xy, 10, nfold=10, early_stopping_rounds=10)