fix: add default namespace for dict insertion ordering

Author: abelaba

src/optimagic/parameters/tree_registry.py

@@ -9,7 +9,7 @@
 import pandas as pd
 from optree.pytree import PyTreeSpec
 
-from optimagic.typing import OPTREE_NAMESPACES
+from optimagic.typing import DEFAULT_NAMESPACE, OPTREE_NAMESPACES
 
 try:
     import jax.numpy as jnp  # type: ignore[import-not-found]
@@ -21,72 +21,61 @@
 _are_namespaces_registered = False
 
 
-def tree_flatten(tree, is_leaf=None, namespace=""):
+def tree_flatten(tree, is_leaf=None, namespace=DEFAULT_NAMESPACE):
     """Flatten a pytree."""
     _register_namespaces()
     _check_namespace(namespace)
+    with optree.dict_insertion_ordered(True, namespace=namespace):
+        return optree.tree_flatten(tree, is_leaf=is_leaf, namespace=namespace)
 
-    return _with_insertion_order(namespace, optree.tree_flatten, tree, is_leaf=is_leaf)
 
-
-def tree_just_flatten(tree, is_leaf=None, namespace=""):
+def tree_just_flatten(tree, is_leaf=None, namespace=DEFAULT_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)
+    with optree.dict_insertion_ordered(True, namespace=namespace):
+        return optree.tree_leaves(tree, is_leaf=is_leaf, namespace=namespace)
 
 
-def tree_unflatten(treedef, leaves, namespace=""):
+def tree_unflatten(treedef, leaves, namespace=DEFAULT_NAMESPACE):
     """Reconstruct a pytree from the tree definition and the leaves."""
     _register_namespaces()
 
     if not isinstance(treedef, PyTreeSpec):
         _check_namespace(namespace)
-        treedef = _with_insertion_order(namespace, optree.tree_structure, treedef)
+        with optree.dict_insertion_ordered(True, namespace=namespace):
+            treedef = optree.tree_structure(treedef, namespace=namespace)
 
-    # optree.tree_unflatten doesn't need to be wrapped with _with_insertion_order
-    # because it keeps the insertion order for dictionaries.
+    # Doesn't need to be wrapped with dict_insertion_ordered
+    # because it keeps the insertion order for dictionaries by default.
     return optree.tree_unflatten(treedef, leaves)
 
 
-def tree_map(func, tree, is_leaf=None, namespace=""):
-    """Map an input function over pytree args to produce a new pytree.
-
-    optree.tree_map always respects insertion order for dictionaries and doesn't
-    require to be wrapped with _with_insertion_order.
-    """
+def tree_map(func, tree, is_leaf=None, namespace=DEFAULT_NAMESPACE):
+    """Map an input function over pytree args to produce a new pytree."""
     _register_namespaces()
     _check_namespace(namespace)
 
+    # Doesn't need to be wrapped with dict_insertion_ordered
+    # because it keeps the insertion order for dictionaries by default.
     return optree.tree_map(func, tree, is_leaf=is_leaf, namespace=namespace)
 
 
-def leaf_names(tree, is_leaf=None, namespace="", separator="_"):
+def leaf_names(tree, is_leaf=None, namespace=DEFAULT_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
-    )
+    with optree.dict_insertion_ordered(True, namespace=namespace):
+        paths, _, _ = optree.tree_flatten_with_path(
+            tree, is_leaf=is_leaf, namespace=namespace
+        )
     return [separator.join(str(p) for p in path) for path in paths]
 
 
-def _with_insertion_order(namespace, optree_func, *args, **kwargs):
-    """Call an optree function, preserving dict key order within a namespace.
-
-    By default, optree sorts dictionary keys. When a namespace is provided,
-    this wrapper enables dict_insertion_ordered mode so that the original
-    key order is preserved in the output.
-    """
-    if namespace:
-        with optree.dict_insertion_ordered(True, namespace=namespace):
-            return optree_func(*args, namespace=namespace, **kwargs)
-    return optree_func(*args, namespace=namespace, **kwargs)
-
-
-def tree_equal(tree, other, is_leaf=None, namespace="", equality_checkers=None):
+def tree_equal(
+    tree, other, is_leaf=None, namespace=DEFAULT_NAMESPACE, equality_checkers=None
+):
     """Check the equality between two trees."""
     equality_checkers = (
         _get_equality_checkers()
@@ -136,8 +125,8 @@ def _get_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:
+    """Checks if the namespace is registered and raise a warning."""
+    if namespace != DEFAULT_NAMESPACE and namespace not in OPTREE_NAMESPACES:
         warnings.warn(
             f"Namespace '{namespace}' is not registered. "
             f"Registered namespaces are: {','.join(OPTREE_NAMESPACES)}. "

src/optimagic/typing.py

@@ -174,12 +174,12 @@ class MultiStartIterationHistory(TupleLikeAccess):
     exploration: IterationHistory | None = None
 
 
+DEFAULT_NAMESPACE = "optimagic_namespace"
 OPTREE_NAMESPACES = (
     "value",
     "lower_bound",
     "upper_bound",
     "soft_lower_bound",
     "soft_upper_bound",
 )
-
 VALUE_NAMESPACE = OPTREE_NAMESPACES[0]

tests/optimagic/parameters/test_tree_registry.py

@@ -158,41 +158,38 @@ def test_tree_flatten_and_unflatten_with_None():
     assert tree == [None]
 
 
-def test_dict_insertion_ordering_is_respected_for_registered_namespaces():
+@pytest.mark.parametrize("namespace", OPTREE_NAMESPACES)
+def test_dict_insertion_ordering_is_respected_for_registered_namespaces(namespace):
     params = {"b": [1, 4], "a": [8, 9]}
-    leaves, _ = tree_flatten(params, namespace=VALUE_NAMESPACE)
+    leaves, _ = tree_flatten(params, namespace=namespace)
     assert leaves == [1, 4, 8, 9]
-    leaves2 = tree_just_flatten(params, namespace=VALUE_NAMESPACE)
-    assert leaves2 == [1, 4, 8, 9]
-    names = leaf_names(params, namespace=VALUE_NAMESPACE)
-    assert names == ["b_0", "b_1", "a_0", "a_1"]
 
+    tree = tree_unflatten(params, [1, 4, 8, 9], namespace=namespace)
+    assert list(tree.items()) == [("b", [1, 4]), ("a", [8, 9])]
 
-def test_dict_ordering_default_behaviour_is_by_name():
-    params = {"b": [1, 4], "a": [8, 9]}
-    leaves, _ = tree_flatten(params)
-    assert leaves == [8, 9, 1, 4]
+    leaves2 = tree_just_flatten(params, namespace=namespace)
+    assert leaves2 == [1, 4, 8, 9]
 
-    leaves2 = tree_just_flatten(params)
-    assert leaves2 == [8, 9, 1, 4]
+    tree = tree_map(lambda x: x, params, namespace=namespace)
+    assert list(tree.items()) == [("b", [1, 4]), ("a", [8, 9])]
 
-    names = leaf_names(params)
-    assert names == ["a_0", "a_1", "b_0", "b_1"]
+    names = leaf_names(params, namespace=namespace)
+    assert names == ["b_0", "b_1", "a_0", "a_1"]
 
 
-def test_unflatten_respects_insertion_order():
+def test_dict_insertion_ordering_is_respected_for_default_namespace():
     params = {"b": [1, 4], "a": [8, 9]}
-    leaves, treespec = tree_flatten(params)
-    tree = tree_unflatten(treespec, leaves)
+    leaves, _ = tree_flatten(params)
+    assert leaves == [1, 4, 8, 9]
+
+    tree = tree_unflatten(params, [1, 4, 8, 9])
     assert list(tree.items()) == [("b", [1, 4]), ("a", [8, 9])]
-    leaves2, treespec2 = tree_flatten(params, namespace=VALUE_NAMESPACE)
-    tree2 = tree_unflatten(treespec2, leaves2)
-    assert list(tree2.items()) == [("b", [1, 4]), ("a", [8, 9])]
 
+    leaves2 = tree_just_flatten(params)
+    assert leaves2 == [1, 4, 8, 9]
 
-def test_map_always_respects_insertion_order():
-    params = {"b": [1, 4], "a": [8, 9]}
     tree = tree_map(lambda x: x, params)
     assert list(tree.items()) == [("b", [1, 4]), ("a", [8, 9])]
-    tree2 = tree_map(lambda x: x, params, namespace=VALUE_NAMESPACE)
-    assert list(tree2.items()) == [("b", [1, 4]), ("a", [8, 9])]
+
+    names = leaf_names(params)
+    assert names == ["b_0", "b_1", "a_0", "a_1"]