diff --git a/tests/conftest.py b/tests/conftest.py index 449ab24..862131b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -75,8 +75,9 @@ def mindset_data(): return load_mindset_data() -@pytest.fixture(scope="function") -def twins_data(rng): +@pytest.fixture(scope="session") +def twins_data(): + rng = np.random.default_rng(_SEED) ( chosen_df, outcome_column, @@ -94,28 +95,30 @@ def twins_data(rng): ) -@pytest.fixture(scope="module") +@pytest.fixture(scope="session") def n_numericals(): return 25 -@pytest.fixture(scope="module") +@pytest.fixture(scope="session") def n_categoricals(): return 5 -@pytest.fixture(scope="module") +@pytest.fixture(scope="session") def sample_size(): return 100_000 -@pytest.fixture(scope="function") -def numerical_covariates(sample_size, n_numericals, rng): +@pytest.fixture(scope="session") +def numerical_covariates(sample_size, n_numericals): + rng = np.random.default_rng(_SEED) return generate_covariates(sample_size, n_numericals, format="numpy", rng=rng) -@pytest.fixture(scope="function") -def mixed_covariates(sample_size, n_numericals, n_categoricals, rng): +@pytest.fixture(scope="session") +def mixed_covariates(sample_size, n_numericals, n_categoricals): + rng = np.random.default_rng(_SEED) return generate_covariates( sample_size, n_numericals + n_categoricals, @@ -125,52 +128,72 @@ def mixed_covariates(sample_size, n_numericals, n_categoricals, rng): ) -@pytest.fixture(scope="function") +@pytest.fixture(scope="session") def numerical_experiment_dataset_continuous_outcome_binary_treatment_linear_te( - numerical_covariates, rng + sample_size, n_numericals ): - covariates, _, _ = numerical_covariates + rng = np.random.default_rng(_SEED) + covariates, _, _ = generate_covariates( + sample_size, n_numericals, format="numpy", rng=rng + ) return _generate_rct_experiment_data(covariates, False, rng, 0.3, None) -@pytest.fixture(scope="function") +@pytest.fixture(scope="session") def numerical_experiment_dataset_binary_outcome_binary_treatment_linear_te( - numerical_covariates, rng + sample_size, n_numericals ): - covariates, _, _ = numerical_covariates + rng = np.random.default_rng(_SEED) + covariates, _, _ = generate_covariates( + sample_size, n_numericals, format="numpy", rng=rng + ) return _generate_rct_experiment_data(covariates, True, rng, 0.3, None) -@pytest.fixture(scope="function") +@pytest.fixture(scope="session") def mixed_experiment_dataset_continuous_outcome_binary_treatment_linear_te( - mixed_covariates, rng + sample_size, n_numericals, n_categoricals ): - covariates, _, _ = mixed_covariates + rng = np.random.default_rng(_SEED) + covariates, _, _ = generate_covariates( + sample_size, + n_numericals + n_categoricals, + n_categoricals=n_categoricals, + format="pandas", + rng=rng, + ) return _generate_rct_experiment_data(covariates, False, rng, 0.3, None) -@pytest.fixture(scope="function") +@pytest.fixture(scope="session") def numerical_experiment_dataset_continuous_outcome_multi_treatment_linear_te( - numerical_covariates, rng + sample_size, n_numericals ): - covariates, _, _ = numerical_covariates + rng = np.random.default_rng(_SEED) + covariates, _, _ = generate_covariates( + sample_size, n_numericals, format="numpy", rng=rng + ) return _generate_rct_experiment_data( covariates, False, rng, [0.2, 0.1, 0.3, 0.15, 0.25], None ) -@pytest.fixture(scope="function") +@pytest.fixture(scope="session") def numerical_experiment_dataset_continuous_outcome_multi_treatment_constant_te( - numerical_covariates, rng + sample_size, n_numericals ): - covariates, _, _ = numerical_covariates + rng = np.random.default_rng(_SEED) + covariates, _, _ = generate_covariates( + sample_size, n_numericals, format="numpy", rng=rng + ) return _generate_rct_experiment_data( covariates, False, rng, [0.2, 0.1, 0.3, 0.15, 0.25], np.array([-2, 5, 0, 3]) ) -@pytest.fixture -def dummy_dataset(rng): +@pytest.fixture(scope="session") +def dummy_dataset(): + rng = np.random.default_rng(_SEED) sample_size = 100 n_features = 10 X = rng.standard_normal((sample_size, n_features)) @@ -179,8 +202,9 @@ def dummy_dataset(rng): return X, y, w -@pytest.fixture(scope="function") -def feature_importance_dataset(rng): +@pytest.fixture(scope="session") +def feature_importance_dataset(): + rng = np.random.default_rng(_SEED) n_samples = 10000 x0 = rng.normal(10, 1, n_samples) x1 = rng.normal(2, 1, n_samples) diff --git a/tests/test_learner.py b/tests/test_learner.py index e76018e..e40a167 100644 --- a/tests/test_learner.py +++ b/tests/test_learner.py @@ -312,9 +312,8 @@ def test_learner_twins(metalearner, reference_value, twins_data, rng): @pytest.mark.parametrize("n_classes", [2, 5, 10]) @pytest.mark.parametrize("n_variants", [2, 5]) @pytest.mark.parametrize("is_classification", [True, False]) -def test_learner_evaluate( - metalearner, is_classification, rng, sample_size, n_classes, n_variants -): +def test_learner_evaluate(metalearner, is_classification, rng, n_classes, n_variants): + sample_size = 1000 factory = metalearner_factory(metalearner) if n_variants > 2 and not factory._supports_multi_treatment(): pytest.skip() @@ -617,8 +616,9 @@ def test_conditional_average_outcomes_smoke( @pytest.mark.parametrize("n_classes", [5, 10]) @pytest.mark.parametrize("n_variants", [2, 5]) def test_conditional_average_outcomes_smoke_multi_class( - metalearner_prefix, rng, sample_size, n_classes, n_variants + metalearner_prefix, rng, n_classes, n_variants ): + sample_size = 1000 factory = metalearner_factory(metalearner_prefix) X = rng.standard_normal((sample_size, 10)) @@ -648,8 +648,9 @@ def test_conditional_average_outcomes_smoke_multi_class( @pytest.mark.parametrize("n_variants", [2, 5]) @pytest.mark.parametrize("is_classification", [True, False]) def test_predict_smoke( - metalearner_prefix, is_classification, rng, sample_size, n_classes, n_variants + metalearner_prefix, is_classification, rng, n_classes, n_variants ): + sample_size = 1000 factory = metalearner_factory(metalearner_prefix) if n_variants > 2 and not factory._supports_multi_treatment(): pytest.skip() diff --git a/tests/test_metalearner.py b/tests/test_metalearner.py index 9de1af3..2e89913 100644 --- a/tests/test_metalearner.py +++ b/tests/test_metalearner.py @@ -152,7 +152,7 @@ def test_metalearner_init( @pytest.mark.parametrize( "implementation", - [_TestMetaLearner, TLearner, SLearner, XLearner, RLearner, DRLearner], + [TLearner, SLearner, XLearner, RLearner, DRLearner], ) def test_metalearner_categorical( mixed_experiment_dataset_continuous_outcome_binary_treatment_linear_te, @@ -198,7 +198,7 @@ def test_metalearner_categorical( @pytest.mark.parametrize( "implementation", - [_TestMetaLearner, TLearner, SLearner, XLearner, RLearner, DRLearner], + [TLearner, SLearner, XLearner, RLearner, DRLearner], ) def test_metalearner_missing_data_smoke( mixed_experiment_dataset_continuous_outcome_binary_treatment_linear_te, @@ -227,7 +227,7 @@ def test_metalearner_missing_data_smoke( @pytest.mark.parametrize( "implementation", - [_TestMetaLearner, TLearner, SLearner, XLearner, RLearner, DRLearner], + [TLearner, SLearner, XLearner, RLearner, DRLearner], ) def test_metalearner_missing_data_error( numerical_experiment_dataset_continuous_outcome_binary_treatment_linear_te, @@ -258,7 +258,7 @@ def test_metalearner_missing_data_error( @pytest.mark.parametrize( "implementation", - [_TestMetaLearner, TLearner, SLearner, XLearner, RLearner, DRLearner], + [TLearner, SLearner, XLearner, RLearner, DRLearner], ) def test_metalearner_format_consistent( numerical_experiment_dataset_continuous_outcome_binary_treatment_linear_te, @@ -345,7 +345,7 @@ def test_n_folds(n_folds): @pytest.mark.parametrize( "implementation", - [_TestMetaLearner, TLearner, SLearner, XLearner, RLearner, DRLearner], + [TLearner, SLearner, XLearner, RLearner, DRLearner], ) def test_metalearner_model_names(implementation): set1 = set(implementation.nuisance_model_specifications().keys()) @@ -702,7 +702,6 @@ def test_fit_params_rlearner_error(dummy_dataset): @pytest.mark.parametrize( "implementation, needs_estimates", [ - (_TestMetaLearner, True), (TLearner, True), (SLearner, True), (XLearner, True),