chore: use namespaces for passing data_col value for dataframes

Author: abelaba

src/optimagic/parameters/bounds.py

@@ -11,7 +11,6 @@
 from optimagic.parameters.tree_registry import (
     extended,
     leaf_names,
-    set_data_col_df_attribute,
     tree_map,
 )
 from optimagic.parameters.tree_registry import (
@@ -181,9 +180,7 @@ def _update_bounds_and_flatten(
         np.ndarray: The updated and flattened bounds.
 
     """
-    flat_nan_tree = tree_leaves(
-        set_data_col_df_attribute(nan_tree, data_col=kind), namespace=extended
-    )
+    flat_nan_tree = tree_leaves(nan_tree, namespace=kind)
     if bounds is not None:
         flat_bounds = tree_leaves(bounds, namespace=extended)
 

src/optimagic/parameters/tree_registry.py

@@ -1,38 +1,34 @@
 """Wrapper around pybaum get_registry to tailor it to optimagic."""
 
 import itertools
+from functools import partial
 from itertools import product
 
 import numpy as np
 import optree
 import pandas as pd
 from optree.pytree import PyTreeSpec
 
-from optimagic.typing import extended_namespace
+extended = "value"
+namespaces = [
+    extended,
+    "lower_bound",
+    "upper_bound",
+    "soft_lower_bound",
+    "soft_upper_bound",
+]
 
-
-def _get_df_names(df):
-    index_strings = list(df.index.map(_index_element_to_string))
-    if "value" in df:
-        out = index_strings
-    else:
-        out = ["_".join([loc, col]) for loc, col in product(index_strings, df.columns)]
-
-    return out
-
-
-def _index_element_to_string(element):
-    if isinstance(element, (tuple, list)):
-        as_strings = [str(entry) for entry in element]
-        res_string = "_".join(as_strings)
-    else:
-        res_string = str(element)
-
-    return res_string
+EQUALITY_CHECKERS = {}
+EQUALITY_CHECKERS[np.ndarray.__name__] = lambda a, b: bool((a == b).all())
+EQUALITY_CHECKERS[pd.Series.__name__] = lambda a, b: a.equals(b)
+EQUALITY_CHECKERS[pd.DataFrame.__name__] = lambda a, b: a.equals(b)
 
 
 def tree_flatten(tree, is_leaf=None, namespace=""):
-    with optree.dict_insertion_ordered(True, namespace=extended_namespace):
+    if namespace:
+        with optree.dict_insertion_ordered(True, namespace=namespace):
+            return optree.tree_flatten(tree, is_leaf=is_leaf, namespace=namespace)
+    else:
         return optree.tree_flatten(tree, is_leaf=is_leaf, namespace=namespace)
 
 
@@ -41,9 +37,6 @@ def tree_just_flatten(tree, is_leaf=None, namespace=""):
     return leaves
 
 
-extended = extended_namespace
-
-
 def tree_unflatten(treedef, leaves, is_leaf=None, namespace=""):
     if not isinstance(treedef, PyTreeSpec):
         _, treedef = tree_flatten(treedef, is_leaf=is_leaf, namespace=namespace)
@@ -60,14 +53,35 @@ def leaf_names(tree, is_leaf=None, namespace="", separator="_"):
     return [separator.join(str(p) for p in path) for path in paths]
 
 
-def set_data_col_df_attribute(tree, data_col):
-    def set_attr(node):
-        if isinstance(node, pd.DataFrame):
-            node = node.copy()
-            node.attrs["data_col"] = data_col
-        return node
+def tree_equal(tree, other, is_leaf=None, namespace="", 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, namespace=namespace
+    )
+    second_flat, second_treespec = tree_flatten(
+        other, is_leaf=is_leaf, namespace=namespace
+    )
 
-    return tree_map(set_attr, tree)
+    first_names = leaf_names(tree, is_leaf=is_leaf, namespace=namespace)
+    second_names = leaf_names(other, is_leaf=is_leaf, namespace=namespace)
+
+    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).__name__, lambda a, b: a == b
+            )
+            equal = equal and check_func(first, second)
+            if not equal:
+                break
+
+    return equal
 
 
 def _array_element_names(arr):
@@ -76,8 +90,27 @@ def _array_element_names(arr):
     return names
 
 
-def _flatten_df_optree(df):
-    data_col = df.attrs.get("data_col", "value")
+def _get_df_names(df):
+    index_strings = list(df.index.map(_index_element_to_string))
+    if "value" in df:
+        out = index_strings
+    else:
+        out = ["_".join([loc, col]) for loc, col in product(index_strings, df.columns)]
+
+    return out
+
+
+def _index_element_to_string(element):
+    if isinstance(element, (tuple, list)):
+        as_strings = [str(entry) for entry in element]
+        res_string = "_".join(as_strings)
+    else:
+        res_string = str(element)
+
+    return res_string
+
+
+def _flatten_df(df, data_col):
     is_value_df = "value" in df
     if is_value_df:
         flat = df.get(data_col, default=np.full(len(df), np.nan)).tolist()
@@ -91,8 +124,7 @@ def _flatten_df_optree(df):
     return flat, aux_data, _get_df_names(df)
 
 
-def _unflatten_df_optree(aux_data, leaves):
-    data_col = aux_data["df"].attrs.get("data_col", "value")
+def _unflatten_df(aux_data, leaves, data_col):
     if aux_data["is_value_df"]:
         out = aux_data["df"].assign(**{data_col: leaves})
     else:
@@ -104,61 +136,44 @@ def _unflatten_df_optree(aux_data, leaves):
     return out
 
 
-optree.register_pytree_node(
-    pd.DataFrame,
-    _flatten_df_optree,
-    _unflatten_df_optree,
-    namespace=extended_namespace,
-)
-
-optree.register_pytree_node(
-    pd.Series,
-    lambda sr: (
-        sr.tolist(),
-        {"index": sr.index, "name": sr.name},
-        list(sr.index.map(_index_element_to_string)),
-    ),
-    lambda aux_data, leaves: pd.Series(leaves, **aux_data),
-    namespace=extended_namespace,
-)
-
-optree.register_pytree_node(
-    np.ndarray,
-    lambda arr: (arr.flatten().tolist(), arr.shape, _array_element_names(arr)),
-    lambda aux_data, leaves: np.array(leaves).reshape(aux_data),
-    namespace=extended_namespace,
-)
+def _flatten_series(series: pd.Series):
+    return (
+        series.tolist(),
+        {"index": series.index, "name": series.name},
+        list(series.index.map(_index_element_to_string)),
+    )
 
-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 _unflatten_series(aux_data, leaves):
+    return pd.Series(leaves, **aux_data)
 
-def tree_equal(tree, other, is_leaf=None, namespace="", 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, namespace=namespace
-    )
-    second_flat, second_treespec = tree_flatten(
-        other, is_leaf=is_leaf, namespace=namespace
-    )
+def _flatten_ndarray(arr: np.ndarray):
+    return arr.flatten().tolist(), arr.shape, _array_element_names(arr)
 
-    first_names = leaf_names(tree, is_leaf=is_leaf, namespace=namespace)
-    second_names = leaf_names(tree, is_leaf=is_leaf, namespace=namespace)
 
-    equal = first_names == second_names and first_treespec == second_treespec
+def _unflatten_ndarray(aux_data, leaves):
+    return np.array(leaves).reshape(aux_data)
 
-    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
+for namespace in namespaces:
+    optree.register_pytree_node(
+        pd.DataFrame,
+        partial(_flatten_df, data_col=namespace),
+        partial(_unflatten_df, data_col=namespace),
+        namespace=namespace,
+    )
+
+    optree.register_pytree_node(
+        pd.Series,
+        _flatten_series,
+        _unflatten_series,
+        namespace=namespace,
+    )
+
+    optree.register_pytree_node(
+        np.ndarray,
+        _flatten_ndarray,
+        _unflatten_ndarray,
+        namespace=namespace,
+    )

src/optimagic/typing.py

@@ -22,7 +22,6 @@
 Scalar = Any
 
 T = TypeVar("T")
-extended_namespace = "extended_namespace"
 
 
 class AggregationLevel(Enum):

tests/optimagic/parameters/test_tree_registry.py

@@ -6,7 +6,6 @@
 from optimagic.parameters.tree_registry import (
     extended,
     leaf_names,
-    set_data_col_df_attribute,
     tree_flatten,
     tree_unflatten,
 )
@@ -61,28 +60,3 @@ def test_unflatten_partially_numeric_df(other_df):
 def test_leaf_names_partially_numeric_df(other_df):
     names = leaf_names(other_df, namespace=extended)
     assert names == ["alpha_b", "alpha_c", "beta_b", "beta_c", "gamma_b", "gamma_c"]
-
-
-def test_set_data_col_attribute_assigns_attribute(value_df):
-    df = set_data_col_df_attribute(value_df, data_col="attr")
-    assert df.attrs.get("data_col") == "attr"
-    assert value_df.attrs.get("data_col") is None
-
-
-def test_set_data_col_attribute_unflattened_tree_has_attribute(value_df):
-    df = set_data_col_df_attribute(value_df, data_col="attr")
-    tree, treedef = tree_flatten(df, namespace=extended)
-    df = tree_unflatten(treedef, tree)
-    assert df.attrs.get("data_col") == "attr"
-
-
-def test_set_data_col_attribute_returns_nan(value_df):
-    df = set_data_col_df_attribute(value_df, data_col="attr")
-    tree, treedef = tree_flatten(df, namespace=extended)
-    assert all(np.isnan(value) for value in tree)
-
-
-def test_set_data_col_attribute_returs_column_values(value_df):
-    df = set_data_col_df_attribute(value_df, data_col="a")
-    tree, treedef = tree_flatten(df, namespace=extended)
-    assert tree == [0, 2, 4]