chore: raise warning for unregistered namespaces

Author: abelaba

src/optimagic/parameters/tree_registry.py

@@ -1,5 +1,6 @@
 """Wrapper around optree to tailor it to optimagic."""
 
+import warnings
 from functools import partial
 from itertools import product
 
@@ -23,13 +24,15 @@
 def tree_flatten(tree, is_leaf=None, namespace=""):
     """Flatten a pytree."""
     _register_namespaces()
+    _check_namespace(namespace)
 
     return _with_insertion_order(namespace, optree.tree_flatten, tree, is_leaf=is_leaf)
 
 
 def tree_just_flatten(tree, is_leaf=None, namespace=""):
     """Get the leaves of a pytree."""
     _register_namespaces()
+    _check_namespace(namespace)
 
     return _with_insertion_order(namespace, optree.tree_leaves, tree, is_leaf=is_leaf)
 
@@ -39,6 +42,7 @@ def tree_unflatten(treedef, leaves, namespace=""):
     _register_namespaces()
 
     if not isinstance(treedef, PyTreeSpec):
+        _check_namespace(namespace)
         treedef = _with_insertion_order(namespace, optree.tree_structure, treedef)
 
     # optree.tree_unflatten doesn't need to be wrapped with _with_insertion_order
@@ -53,12 +57,15 @@ def tree_map(func, tree, is_leaf=None, namespace=""):
     require to be wrapped with _with_insertion_order.
     """
     _register_namespaces()
+    _check_namespace(namespace)
+
     return optree.tree_map(func, tree, is_leaf=is_leaf, namespace=namespace)
 
 
 def leaf_names(tree, is_leaf=None, namespace="", separator="_"):
     """Get the path names for tree leaves."""
     _register_namespaces()
+    _check_namespace(namespace)
 
     paths, _, _ = _with_insertion_order(
         namespace, optree.tree_flatten_with_path, tree, is_leaf=is_leaf
@@ -128,6 +135,16 @@ def _get_equality_checkers():
     return equality_checkers
 
 
+def _check_namespace(namespace: str) -> None:
+    """Checks if the namespace is a registered and raise a warning."""
+    if namespace and namespace not in OPTREE_NAMESPACES:
+        warnings.warn(
+            f"Namespace '{namespace}' is not registered. "
+            f"Registered namespaces are: {','.join(OPTREE_NAMESPACES)}. "
+            "Pytree method is being parsed with the default optree namespace."
+        )
+
+
 def _register_namespaces() -> None:
     """Register pytree flatten/unflatten methods for each namespace.
 

tests/optimagic/parameters/test_tree_registry.py

@@ -132,10 +132,20 @@ def test_tree_methods_with_optimagic_namespace(namespace, bounds_df):
     assert_frame_equal(tree, expected)
 
 
-def test_tree_flatten_with_unregisted_namespace(value_df):
+@pytest.mark.parametrize(
+    "tree_method",
+    [tree_flatten, tree_just_flatten, leaf_names, tree_map, tree_flatten],
+)
+def test_tree_methods_raise_warning_with_unregisted_namespace(tree_method, value_df):
     """If namespace is not registered optree method fallbacks to default behaviour."""
-    leaves, _ = tree_flatten(value_df, namespace="unregistered_namespace")
-    assert leaves == [value_df]
+    unregistered_namespace = "unregistered_namespace"
+    with pytest.warns(match="is not registered."):
+        if tree_method == tree_map:
+            _ = tree_map(lambda x: x, value_df, namespace=unregistered_namespace)
+        elif tree_method == tree_unflatten:
+            _ = tree_method(value_df, [], namespace=unregistered_namespace)
+        else:
+            _ = tree_method(value_df, namespace=unregistered_namespace)
 
 
 def test_tree_flatten_and_unflatten_with_None():