chore: add tests for tree methods and rearrange tree-registry methods

Author: abelaba

pixi.lock

@@ -18,7 +18,7 @@
     tree_just_flatten,
     tree_unflatten,
 )
-from optimagic.typing import value_namespace
+from optimagic.typing import VALUE_NAMESPACE
 from optimagic.utilities import get_rng
 
 
@@ -108,7 +108,7 @@ def bootstrap(
     # ==================================================================================
 
     flat_outcomes = [
-        tree_just_flatten(_outcome, namespace=value_namespace)
+        tree_just_flatten(_outcome, namespace=VALUE_NAMESPACE)
         for _outcome in all_outcomes
     ]
     internal_outcomes = np.array(flat_outcomes)
@@ -167,10 +167,10 @@ def outcomes(self):
             List[Any]: The boostrap outcomes as a list of pytrees.
 
         """
-        _, treedef = tree_flatten(self._base_outcome, namespace=value_namespace)
+        _, treedef = tree_flatten(self._base_outcome, namespace=VALUE_NAMESPACE)
 
         outcomes = [
-            tree_unflatten(treedef, out, namespace=value_namespace)
+            tree_unflatten(treedef, out, namespace=VALUE_NAMESPACE)
             for out in self._internal_outcomes
         ]
         return outcomes
@@ -186,9 +186,9 @@ def se(self):
         cov = self._internal_cov
         se = np.sqrt(np.diagonal(cov))
 
-        _, treedef = tree_flatten(self._base_outcome, namespace=value_namespace)
+        _, treedef = tree_flatten(self._base_outcome, namespace=VALUE_NAMESPACE)
 
-        se = tree_unflatten(treedef, se, namespace=value_namespace)
+        se = tree_unflatten(treedef, se, namespace=VALUE_NAMESPACE)
         return se
 
     def cov(self, return_type="pytree"):
@@ -209,7 +209,7 @@ def cov(self, return_type="pytree"):
         cov = self._internal_cov
 
         if return_type == "dataframe":
-            names = np.array(leaf_names(self._base_outcome, namespace=value_namespace))
+            names = np.array(leaf_names(self._base_outcome, namespace=VALUE_NAMESPACE))
             cov = pd.DataFrame(cov, columns=names, index=names)
         elif return_type == "pytree":
             cov = matrix_to_block_tree(cov, self._base_outcome, self._base_outcome)
@@ -237,15 +237,15 @@ def ci(self, ci_method="percentile", ci_level=0.95):
 
         """
         base_outcome_flat, treedef = tree_flatten(
-            self._base_outcome, namespace=value_namespace
+            self._base_outcome, namespace=VALUE_NAMESPACE
         )
 
         lower_flat, upper_flat = calculate_ci(
             base_outcome_flat, self._internal_outcomes, ci_method, ci_level
         )
 
-        lower = tree_unflatten(treedef, lower_flat, namespace=value_namespace)
-        upper = tree_unflatten(treedef, upper_flat, namespace=value_namespace)
+        lower = tree_unflatten(treedef, lower_flat, namespace=VALUE_NAMESPACE)
+        upper = tree_unflatten(treedef, upper_flat, namespace=VALUE_NAMESPACE)
         return lower, upper
 
     def p_values(self):
@@ -274,7 +274,7 @@ def summary(self, ci_method="percentile", ci_level=0.95):
                 Soon this will be a pytree.
 
         """
-        names = leaf_names(self.base_outcome, namespace=value_namespace)
+        names = leaf_names(self.base_outcome, namespace=VALUE_NAMESPACE)
         summary_data = _calulcate_summary_data_bootstrap(
             self, ci_method=ci_method, ci_level=ci_level
         )

src/estimagic/bootstrap.py

@@ -57,7 +57,7 @@
 from optimagic.shared.check_option_dicts import (
     check_optimization_options,
 )
-from optimagic.typing import value_namespace
+from optimagic.typing import VALUE_NAMESPACE
 from optimagic.utilities import get_rng, to_pickle
 
 
@@ -321,7 +321,7 @@ def func(x):
             sim_mom = simulate_moments(params, **simulate_moments_kwargs)
             if isinstance(sim_mom, dict) and "simulated_moments" in sim_mom:
                 sim_mom = sim_mom["simulated_moments"]
-            out = np.array(tree_just_flatten(sim_mom, namespace=value_namespace))
+            out = np.array(tree_just_flatten(sim_mom, namespace=VALUE_NAMESPACE))
             return out
 
         int_jac = first_derivative(
@@ -420,7 +420,7 @@ def get_msm_optimization_functions(
 
     chol_weights = np.linalg.cholesky(flat_weights)
 
-    flat_emp_mom = tree_just_flatten(empirical_moments, namespace=value_namespace)
+    flat_emp_mom = tree_just_flatten(empirical_moments, namespace=VALUE_NAMESPACE)
 
     _simulate_moments = _partial_kwargs(simulate_moments, simulate_moments_kwargs)
     _jacobian = _partial_kwargs(jacobian, jacobian_kwargs)
@@ -431,7 +431,7 @@ def get_msm_optimization_functions(
             simulate_moments=_simulate_moments,
             flat_empirical_moments=flat_emp_mom,
             chol_weights=chol_weights,
-            namespace=value_namespace,
+            namespace=VALUE_NAMESPACE,
         )
     )
 
@@ -977,7 +977,7 @@ def sensitivity(
             )
         elif return_type == "dataframe":
             row_names = self._internal_estimates.names
-            col_names = leaf_names(self._empirical_moments, namespace=value_namespace)
+            col_names = leaf_names(self._empirical_moments, namespace=VALUE_NAMESPACE)
             out = pd.DataFrame(
                 data=raw,
                 index=row_names,

src/estimagic/estimate_msm.py

@@ -7,7 +7,7 @@
 from estimagic.bootstrap import bootstrap
 from optimagic.parameters.block_trees import block_tree_to_matrix, matrix_to_block_tree
 from optimagic.parameters.tree_registry import tree_just_flatten
-from optimagic.typing import value_namespace
+from optimagic.typing import VALUE_NAMESPACE
 from optimagic.utilities import robust_inverse
 
 
@@ -55,7 +55,7 @@ def get_moments_cov(
     def func(data, **kwargs):
         raw = calculate_moments(data, **kwargs)
         out = pd.Series(
-            tree_just_flatten(raw, namespace=value_namespace)
+            tree_just_flatten(raw, namespace=VALUE_NAMESPACE)
         )  # xxxx won't be necessary soon!
         return out
 

src/estimagic/msm_weighting.py

@@ -9,7 +9,7 @@
     tree_just_flatten,
     tree_unflatten,
 )
-from optimagic.typing import value_namespace
+from optimagic.typing import VALUE_NAMESPACE
 
 
 def transform_covariance(
@@ -150,7 +150,7 @@ def calculate_estimation_summary(
     # ==================================================================================
 
     flat_data = {
-        key: tree_just_flatten(val, namespace=value_namespace)
+        key: tree_just_flatten(val, namespace=VALUE_NAMESPACE)
         for key, val in summary_data.items()
     }
 
@@ -169,7 +169,7 @@ def calculate_estimation_summary(
     # ==================================================================================
 
     # create tree with values corresponding to indices of df
-    indices = tree_unflatten(summary_data["value"], names, namespace=value_namespace)
+    indices = tree_unflatten(summary_data["value"], names, namespace=VALUE_NAMESPACE)
 
     estimates_flat = tree_just_flatten(summary_data["value"])
     indices_flat = tree_just_flatten(indices)
@@ -318,7 +318,7 @@ def calculate_free_estimates(estimates, internal_estimates):
     mask = internal_estimates.free_mask
     names = internal_estimates.names
 
-    external_flat = np.array(tree_just_flatten(estimates, namespace=value_namespace))
+    external_flat = np.array(tree_just_flatten(estimates, namespace=VALUE_NAMESPACE))
 
     free_estimates = FreeParams(
         values=external_flat[mask],
@@ -352,7 +352,7 @@ def transform_free_values_to_params_tree(values, free_params, params):
     mask = free_params.free_mask
     flat = np.full(len(mask), np.nan)
     flat[np.ix_(mask)] = values
-    pytree = tree_unflatten(params, flat, namespace=value_namespace)
+    pytree = tree_unflatten(params, flat, namespace=VALUE_NAMESPACE)
     return pytree
 
 

src/estimagic/shared_covs.py

@@ -14,7 +14,7 @@
 from optimagic.algorithms import AVAILABLE_ALGORITHMS
 from optimagic.optimization.optimize import minimize
 from optimagic.parameters.tree_registry import tree_just_flatten
-from optimagic.typing import value_namespace
+from optimagic.typing import VALUE_NAMESPACE
 
 
 def run_benchmark(
@@ -190,15 +190,15 @@ def _process_one_result(optimize_result, problem):
 
     # This will happen if the optimization raised an error
     if isinstance(optimize_result, str):
-        params_history_flat = [tree_just_flatten(_start_x, namespace=value_namespace)]
+        params_history_flat = [tree_just_flatten(_start_x, namespace=VALUE_NAMESPACE)]
         criterion_history = [_start_crit_value]
         time_history = [np.inf]
         batches_history = [0]
     else:
         history = optimize_result.history
         params_history = history.params
         params_history_flat = [
-            tree_just_flatten(p, namespace=value_namespace) for p in params_history
+            tree_just_flatten(p, namespace=VALUE_NAMESPACE) for p in params_history
         ]
         if _is_noisy:
             criterion_history = np.array([_criterion(p) for p in params_history])

src/optimagic/benchmarking/run_benchmark.py

@@ -26,7 +26,7 @@
     tree_unflatten,
 )
 from optimagic.parameters.tree_registry import tree_just_flatten as tree_leaves
-from optimagic.typing import BatchEvaluatorLiteral, PyTree, value_namespace
+from optimagic.typing import VALUE_NAMESPACE, BatchEvaluatorLiteral, PyTree
 
 
 @dataclass(frozen=True)
@@ -221,22 +221,22 @@ def first_derivative(
     is_fast_path = _is_1d_array(params)
 
     if not is_fast_path:
-        x, params_treedef = tree_flatten(params, namespace=value_namespace)
+        x, params_treedef = tree_flatten(params, namespace=VALUE_NAMESPACE)
         x = np.array(x, dtype=np.float64)
 
         if scaling_factor is not None and not np.isscalar(scaling_factor):
             scaling_factor = np.array(
-                tree_just_flatten(scaling_factor, namespace=value_namespace)
+                tree_just_flatten(scaling_factor, namespace=VALUE_NAMESPACE)
             )
 
         if min_steps is not None and not np.isscalar(min_steps):
             min_steps = np.array(
-                tree_just_flatten(min_steps, namespace=value_namespace)
+                tree_just_flatten(min_steps, namespace=VALUE_NAMESPACE)
             )
 
         if step_size is not None and not np.isscalar(step_size):
             step_size = np.array(
-                tree_just_flatten(step_size, namespace=value_namespace)
+                tree_just_flatten(step_size, namespace=VALUE_NAMESPACE)
             )
     else:
         x = params.astype(np.float64)
@@ -291,7 +291,7 @@ def first_derivative(
     if not is_fast_path:
         evaluation_points = [
             # entries are either a numpy.ndarray or np.nan
-            _unflatten_if_not_nan(p, params_treedef, value_namespace)
+            _unflatten_if_not_nan(p, params_treedef, VALUE_NAMESPACE)
             for p in evaluation_points
         ]
 
@@ -330,14 +330,14 @@ def first_derivative(
     elif vector_out:
         f0 = f0_tree.astype(float)
     else:
-        f0 = tree_leaves(f0_tree, namespace=value_namespace)
+        f0 = tree_leaves(f0_tree, namespace=VALUE_NAMESPACE)
         f0 = np.array(f0, dtype=np.float64)
 
     # convert the raw evaluations to numpy arrays
     raw_evals_arr = _convert_evals_to_numpy(
         raw_evals=raw_evals,
         unpacker=unpacker,
-        namespace=value_namespace,
+        namespace=VALUE_NAMESPACE,
         is_scalar_out=scalar_out,
         is_vector_out=vector_out,
     )
@@ -539,22 +539,22 @@ def second_derivative(
     is_fast_path = _is_1d_array(params)
 
     if not is_fast_path:
-        x, params_treedef = tree_flatten(params, namespace=value_namespace)
+        x, params_treedef = tree_flatten(params, namespace=VALUE_NAMESPACE)
         x = np.array(x, dtype=np.float64)
 
         if scaling_factor is not None and not np.isscalar(scaling_factor):
             scaling_factor = np.array(
-                tree_just_flatten(scaling_factor, namespace=value_namespace)
+                tree_just_flatten(scaling_factor, namespace=VALUE_NAMESPACE)
             )
 
         if min_steps is not None and not np.isscalar(min_steps):
             min_steps = np.array(
-                tree_just_flatten(min_steps, namespace=value_namespace)
+                tree_just_flatten(min_steps, namespace=VALUE_NAMESPACE)
             )
 
         if step_size is not None and not np.isscalar(step_size):
             step_size = np.array(
-                tree_just_flatten(step_size, namespace=value_namespace)
+                tree_just_flatten(step_size, namespace=VALUE_NAMESPACE)
             )
     else:
         x = params.astype(np.float64)
@@ -631,7 +631,7 @@ def second_derivative(
         evaluation_points = {
             # entries are either a numpy.ndarray or np.nan, we unflatten only
             step_type: [
-                _unflatten_if_not_nan(p, params_treedef, value_namespace)
+                _unflatten_if_not_nan(p, params_treedef, VALUE_NAMESPACE)
                 for p in points
             ]
             for step_type, points in evaluation_points.items()
@@ -671,13 +671,13 @@ def second_derivative(
     func_value = f0
 
     f0_tree = unpacker(f0)
-    f0 = tree_leaves(f0_tree, namespace=value_namespace)
+    f0 = tree_leaves(f0_tree, namespace=VALUE_NAMESPACE)
     f0 = np.array(f0, dtype=np.float64)
 
     # convert the raw evaluations to numpy arrays
     raw_evals = {
         step_type: _convert_evals_to_numpy(
-            raw_evals=evals, unpacker=unpacker, namespace=value_namespace
+            raw_evals=evals, unpacker=unpacker, namespace=VALUE_NAMESPACE
         )
         for step_type, evals in raw_evals.items()
     }

src/optimagic/differentiation/derivatives.py

@@ -20,7 +20,7 @@
     tree_just_flatten,
     tree_unflatten,
 )
-from optimagic.typing import PyTree, value_namespace
+from optimagic.typing import VALUE_NAMESPACE, PyTree
 
 
 @mark.scalar
@@ -215,11 +215,11 @@ def _get_x(params: PyTree) -> NDArray[np.float64]:
         x = params.astype(float)
     else:
         x = np.array(
-            tree_just_flatten(params, namespace=value_namespace), dtype=np.float64
+            tree_just_flatten(params, namespace=VALUE_NAMESPACE), dtype=np.float64
         )
     return x
 
 
 def _unflatten_gradient(flat: NDArray[np.float64], params: PyTree) -> PyTree:
-    out = tree_unflatten(params, flat.tolist(), namespace=value_namespace)
+    out = tree_unflatten(params, flat.tolist(), namespace=VALUE_NAMESPACE)
     return out

src/optimagic/examples/criterion_functions.py

@@ -8,7 +8,7 @@
 
 from optimagic.exceptions import InvalidFunctionError
 from optimagic.parameters.tree_registry import tree_just_flatten
-from optimagic.typing import AggregationLevel, PyTree, Scalar, value_namespace
+from optimagic.typing import VALUE_NAMESPACE, AggregationLevel, PyTree, Scalar
 from optimagic.utilities import isscalar
 
 
@@ -123,7 +123,7 @@ def _get_flat_value(value: PyTree) -> NDArray[np.float64]:
     elif isinstance(value, np.ndarray):
         flat = value.flatten()
     else:
-        flat = tree_just_flatten(value, namespace=value_namespace)
+        flat = tree_just_flatten(value, namespace=VALUE_NAMESPACE)
 
     flat_arr = np.asarray(flat, dtype=np.float64)
     return flat_arr