chore: move namespaces list to typing.py

Author: abelaba

src/estimagic/bootstrap.py

@@ -13,12 +13,12 @@
 from optimagic.batch_evaluators import joblib_batch_evaluator
 from optimagic.parameters.block_trees import matrix_to_block_tree
 from optimagic.parameters.tree_registry import (
-    extended,
     leaf_names,
     tree_flatten,
     tree_just_flatten,
     tree_unflatten,
 )
+from optimagic.typing import value_namespace
 from optimagic.utilities import get_rng
 
 
@@ -108,7 +108,8 @@ def bootstrap(
     # ==================================================================================
 
     flat_outcomes = [
-        tree_just_flatten(_outcome, namespace=extended) for _outcome in all_outcomes
+        tree_just_flatten(_outcome, namespace=value_namespace)
+        for _outcome in all_outcomes
     ]
     internal_outcomes = np.array(flat_outcomes)
 
@@ -166,10 +167,10 @@ def outcomes(self):
             List[Any]: The boostrap outcomes as a list of pytrees.
 
         """
-        _, treedef = tree_flatten(self._base_outcome, namespace=extended)
+        _, treedef = tree_flatten(self._base_outcome, namespace=value_namespace)
 
         outcomes = [
-            tree_unflatten(treedef, out, namespace=extended)
+            tree_unflatten(treedef, out, namespace=value_namespace)
             for out in self._internal_outcomes
         ]
         return outcomes
@@ -185,9 +186,9 @@ def se(self):
         cov = self._internal_cov
         se = np.sqrt(np.diagonal(cov))
 
-        _, treedef = tree_flatten(self._base_outcome, namespace=extended)
+        _, treedef = tree_flatten(self._base_outcome, namespace=value_namespace)
 
-        se = tree_unflatten(treedef, se, namespace=extended)
+        se = tree_unflatten(treedef, se, namespace=value_namespace)
         return se
 
     def cov(self, return_type="pytree"):
@@ -208,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=extended))
+            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)
@@ -236,15 +237,15 @@ def ci(self, ci_method="percentile", ci_level=0.95):
 
         """
         base_outcome_flat, treedef = tree_flatten(
-            self._base_outcome, namespace=extended
+            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=extended)
-        upper = tree_unflatten(treedef, upper_flat, namespace=extended)
+        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):
@@ -273,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=extended)
+        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/estimate_msm.py

@@ -51,13 +51,13 @@
 from optimagic.parameters.conversion import Converter, get_converter
 from optimagic.parameters.space_conversion import InternalParams
 from optimagic.parameters.tree_registry import (
-    extended,
     leaf_names,
     tree_just_flatten,
 )
 from optimagic.shared.check_option_dicts import (
     check_optimization_options,
 )
+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=extended))
+            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=extended)
+    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=extended,
+            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=extended)
+            col_names = leaf_names(self._empirical_moments, namespace=value_namespace)
             out = pd.DataFrame(
                 data=raw,
                 index=row_names,

src/estimagic/msm_weighting.py

@@ -6,7 +6,8 @@
 
 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 extended, tree_just_flatten
+from optimagic.parameters.tree_registry import tree_just_flatten
+from optimagic.typing import value_namespace
 from optimagic.utilities import robust_inverse
 
 
@@ -54,7 +55,7 @@ def get_moments_cov(
     def func(data, **kwargs):
         raw = calculate_moments(data, **kwargs)
         out = pd.Series(
-            tree_just_flatten(raw, namespace=extended)
+            tree_just_flatten(raw, namespace=value_namespace)
         )  # xxxx won't be necessary soon!
         return out
 

src/estimagic/shared_covs.py

@@ -6,10 +6,10 @@
 
 from optimagic.parameters.block_trees import matrix_to_block_tree
 from optimagic.parameters.tree_registry import (
-    extended,
     tree_just_flatten,
     tree_unflatten,
 )
+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=extended)
+        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=extended)
+    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=extended))
+    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=extended)
+    pytree = tree_unflatten(params, flat, namespace=value_namespace)
     return pytree
 
 

src/optimagic/benchmarking/run_benchmark.py

@@ -13,7 +13,8 @@
 from optimagic import batch_evaluators
 from optimagic.algorithms import AVAILABLE_ALGORITHMS
 from optimagic.optimization.optimize import minimize
-from optimagic.parameters.tree_registry import extended, tree_just_flatten
+from optimagic.parameters.tree_registry import tree_just_flatten
+from optimagic.typing import value_namespace
 
 
 def run_benchmark(
@@ -189,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=extended)]
+        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=extended) 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/differentiation/derivatives.py

@@ -21,13 +21,12 @@
 from optimagic.parameters.block_trees import hessian_to_block_tree, matrix_to_block_tree
 from optimagic.parameters.bounds import Bounds, get_internal_bounds, pre_process_bounds
 from optimagic.parameters.tree_registry import (
-    extended,
     tree_flatten,
     tree_just_flatten,
     tree_unflatten,
 )
 from optimagic.parameters.tree_registry import tree_just_flatten as tree_leaves
-from optimagic.typing import BatchEvaluatorLiteral, PyTree
+from optimagic.typing import BatchEvaluatorLiteral, PyTree, value_namespace
 
 
 @dataclass(frozen=True)
@@ -222,19 +221,23 @@ def first_derivative(
     is_fast_path = _is_1d_array(params)
 
     if not is_fast_path:
-        x, params_treedef = tree_flatten(params, namespace=extended)
+        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=extended)
+                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=extended))
+            min_steps = np.array(
+                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=extended))
+            step_size = np.array(
+                tree_just_flatten(step_size, namespace=value_namespace)
+            )
     else:
         x = params.astype(np.float64)
 
@@ -288,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, extended)
+            _unflatten_if_not_nan(p, params_treedef, value_namespace)
             for p in evaluation_points
         ]
 
@@ -327,14 +330,14 @@ def first_derivative(
     elif vector_out:
         f0 = f0_tree.astype(float)
     else:
-        f0 = tree_leaves(f0_tree, namespace=extended)
+        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=extended,
+        namespace=value_namespace,
         is_scalar_out=scalar_out,
         is_vector_out=vector_out,
     )
@@ -536,19 +539,23 @@ def second_derivative(
     is_fast_path = _is_1d_array(params)
 
     if not is_fast_path:
-        x, params_treedef = tree_flatten(params, namespace=extended)
+        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=extended)
+                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=extended))
+            min_steps = np.array(
+                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=extended))
+            step_size = np.array(
+                tree_just_flatten(step_size, namespace=value_namespace)
+            )
     else:
         x = params.astype(np.float64)
 
@@ -624,7 +631,8 @@ 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, extended) for p in points
+                _unflatten_if_not_nan(p, params_treedef, value_namespace)
+                for p in points
             ]
             for step_type, points in evaluation_points.items()
         }
@@ -663,13 +671,13 @@ def second_derivative(
     func_value = f0
 
     f0_tree = unpacker(f0)
-    f0 = tree_leaves(f0_tree, namespace=extended)
+    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=extended
+            raw_evals=evals, unpacker=unpacker, namespace=value_namespace
         )
         for step_type, evals in raw_evals.items()
     }

src/optimagic/examples/criterion_functions.py

@@ -17,11 +17,10 @@
 )
 from optimagic.parameters.block_trees import matrix_to_block_tree
 from optimagic.parameters.tree_registry import (
-    extended,
     tree_just_flatten,
     tree_unflatten,
 )
-from optimagic.typing import PyTree
+from optimagic.typing import PyTree, value_namespace
 
 
 @mark.scalar
@@ -215,10 +214,12 @@ def _get_x(params: PyTree) -> NDArray[np.float64]:
     if isinstance(params, np.ndarray) and params.ndim == 1:
         x = params.astype(float)
     else:
-        x = np.array(tree_just_flatten(params, namespace=extended), dtype=np.float64)
+        x = np.array(
+            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=extended)
+    out = tree_unflatten(params, flat.tolist(), namespace=value_namespace)
     return out