From 251392f47d6eb3fde1595dc075c5d3c78bc0c42a Mon Sep 17 00:00:00 2001 From: Ryan Hausen Date: Wed, 21 Aug 2024 16:49:59 -0400 Subject: [PATCH] added separate test for nanmean_f, reset the order of pytest fixtures --- treeple/stats/tests/test_forest.py | 2 +- treeple/stats/tests/test_utils.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/treeple/stats/tests/test_forest.py b/treeple/stats/tests/test_forest.py index 06cbeeca..e32070e5 100644 --- a/treeple/stats/tests/test_forest.py +++ b/treeple/stats/tests/test_forest.py @@ -238,7 +238,7 @@ def test_comight_repeated_feature_sets(seed): @pytest.mark.parametrize( ("use_bottleneck", "use_sparse"), - itertools.product([False, True], [False, True]), + itertools.product([True, False], [True, False]), ) def test_build_coleman_forest(use_bottleneck: bool, use_sparse: bool): """Simple test for building a Coleman forest. diff --git a/treeple/stats/tests/test_utils.py b/treeple/stats/tests/test_utils.py index d3b16908..7c664dc1 100644 --- a/treeple/stats/tests/test_utils.py +++ b/treeple/stats/tests/test_utils.py @@ -67,3 +67,24 @@ def test_non_nan_samples(use_bottleneck: bool): expected = np.array([0, 2]) actual = utils._non_nan_samples(posterior_array) np.testing.assert_array_equal(expected, actual) + + +@pytest.mark.parametrize("use_bottleneck", [True, False]) +def test_nanmean_f(use_bottleneck: bool): + if use_bottleneck and utils.DISABLE_BN_ENV_VAR in os.environ: + del os.environ[utils.DISABLE_BN_ENV_VAR] + importlib.reload(utils) + else: + os.environ[utils.DISABLE_BN_ENV_VAR] = "1" + importlib.reload(utils) + + posterior_array = np.array( + [ + [1, 2, np.nan], + [3, 4, np.nan], + ] + ) + + expected = np.array([1.5, 3.5]) + actual = utils.nanmean_f(posterior_array, axis=1) + np.testing.assert_array_equal(expected, actual)