diff --git a/src/dreamer/trees/__init__.py b/src/dreamer/trees/__init__.py index e1ac6acf..eba2ff3d 100644 --- a/src/dreamer/trees/__init__.py +++ b/src/dreamer/trees/__init__.py @@ -29,6 +29,15 @@ def create_tree(tree_type: str, **kwargs: Any) -> SurprisalTree: Raises: ValueError: If tree_type is not recognized """ + # `surprisal.py` calls this factory with a uniform `k=settings.DREAM.SURPRISAL.TREE_K` kwarg for every tree type, + # but `k` is only meaningful for the kNN-based trees (kdtree, balltree, graph). + # The other 4 use different tunables and raise TypeError if `k` is passed. + # Drop it here so the factory accepts a uniform kwargs dict. + + trees_without_k = {"rptree", "covertree", "lsh", "prototype"} + if tree_type in trees_without_k: + kwargs.pop("k", None) + if tree_type == "rptree": return RPTree(**kwargs) elif tree_type == "kdtree": diff --git a/tests/dreamer/test_trees.py b/tests/dreamer/test_trees.py new file mode 100644 index 00000000..5c09a241 --- /dev/null +++ b/tests/dreamer/test_trees.py @@ -0,0 +1,51 @@ +import pytest + +from src.dreamer.trees import ( + CoverTree, + GraphSurprisal, + LSHSurprisal, + PrototypeSurprisal, + RPTree, + SklearnTreeWrapper, + SurprisalTree, + create_tree, +) + +ALL_TREE_TYPES = [ + "kdtree", + "balltree", + "rptree", + "covertree", + "lsh", + "graph", + "prototype", +] + +EXPECTED_CLASS = { + "kdtree": SklearnTreeWrapper, + "balltree": SklearnTreeWrapper, + "rptree": RPTree, + "covertree": CoverTree, + "lsh": LSHSurprisal, + "graph": GraphSurprisal, + "prototype": PrototypeSurprisal, +} + + +@pytest.mark.parametrize("tree_type", ALL_TREE_TYPES) +def test_create_tree_accepts_uniform_k_kwarg(tree_type: str): + tree = create_tree(tree_type=tree_type, k=5) + assert isinstance(tree, SurprisalTree) + assert isinstance(tree, EXPECTED_CLASS[tree_type]) + + +@pytest.mark.parametrize("tree_type", ALL_TREE_TYPES) +def test_create_tree_without_k(tree_type: str): + """The factory should also work when no ``k`` is supplied.""" + tree = create_tree(tree_type=tree_type) + assert isinstance(tree, EXPECTED_CLASS[tree_type]) + + +def test_create_tree_unknown_type_raises(): + with pytest.raises(ValueError, match="Unknown tree type"): + create_tree(tree_type="not_a_tree")