chore: replace tree_equal method with optree impl

Author: abelaba

src/optimagic/parameters/tree_registry.py

@@ -185,3 +185,35 @@ def _unflatten_df_optree(aux_data, leaves):
     lambda aux_data, leaves: np.array(leaves).reshape(aux_data),
     namespace=extended_namespace,
 )
+
+EQUALITY_CHECKERS = {}
+EQUALITY_CHECKERS[np.ndarray] = lambda a, b: bool((a == b).all())
+EQUALITY_CHECKERS[pd.Series] = lambda a, b: a.equals(b)
+EQUALITY_CHECKERS[pd.DataFrame] = lambda a, b: a.equals(b)
+
+
+def tree_equal(tree, other, is_leaf=None, registry=None, equality_checkers=None):
+    equality_checkers = (
+        EQUALITY_CHECKERS
+        if equality_checkers is None
+        else {**EQUALITY_CHECKERS, **equality_checkers}
+    )
+
+    first_flat, first_treespec = tree_flatten(tree, is_leaf=is_leaf, registry=registry)
+    second_flat, second_treespec = tree_flatten(
+        other, is_leaf=is_leaf, registry=registry
+    )
+
+    first_names = leaf_names(tree, is_leaf=is_leaf, registry=registry)
+    second_names = leaf_names(tree, is_leaf=is_leaf, registry=registry)
+
+    equal = first_names == second_names and first_treespec == second_treespec
+
+    if equal:
+        for first, second in zip(first_flat, second_flat, strict=True):
+            check_func = equality_checkers.get(type(first), lambda a, b: a == b)
+            equal = equal and check_func(first, second)
+            if not equal:
+                break
+
+    return equal

tests/estimagic/test_shared.py

@@ -4,7 +4,6 @@
 import pandas as pd
 import pytest
 from numpy.testing import assert_array_almost_equal as aaae
-from pybaum import tree_equal
 
 from estimagic.shared_covs import (
     _to_numpy,
@@ -15,7 +14,7 @@
     transform_free_cov_to_cov,
     transform_free_values_to_params_tree,
 )
-from optimagic.parameters.tree_registry import get_registry, leaf_names
+from optimagic.parameters.tree_registry import get_registry, leaf_names, tree_equal
 from optimagic.utilities import get_rng
 
 

tests/optimagic/differentiation/test_compare_derivatives_with_jax.py

@@ -7,10 +7,10 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_array_almost_equal as aaae
-from pybaum import tree_equal
 
 from optimagic.config import IS_JAX_INSTALLED
 from optimagic.differentiation.derivatives import first_derivative, second_derivative
+from optimagic.parameters.tree_registry import tree_equal
 
 if not IS_JAX_INSTALLED:
     pytestmark = pytest.mark.skip(reason="jax is not installed.")

tests/optimagic/logging/test_logger.py

@@ -3,7 +3,6 @@
 import numpy as np
 import pandas as pd
 import pytest
-from pybaum import tree_equal
 
 from optimagic.logging.logger import (
     LogOptions,
@@ -13,7 +12,11 @@
     SQLiteLogReader,
 )
 from optimagic.optimization.optimize import minimize
-from optimagic.parameters.tree_registry import get_registry, tree_just_flatten
+from optimagic.parameters.tree_registry import (
+    get_registry,
+    tree_equal,
+    tree_just_flatten,
+)
 from optimagic.typing import Direction
 
 

tests/optimagic/parameters/test_block_trees.py

@@ -2,7 +2,6 @@
 import pandas as pd
 import pytest
 from numpy.testing import assert_array_equal
-from pybaum import tree_equal
 
 from optimagic import second_derivative
 from optimagic.parameters.block_trees import (
@@ -11,7 +10,7 @@
     hessian_to_block_tree,
     matrix_to_block_tree,
 )
-from optimagic.parameters.tree_registry import get_registry
+from optimagic.parameters.tree_registry import get_registry, tree_equal
 from optimagic.parameters.tree_registry import tree_just_flatten as tree_leaves