chore: remove get_registry method and use namespace arugment

Author: abelaba

src/estimagic/bootstrap.py

@@ -13,7 +13,7 @@
 from optimagic.batch_evaluators import joblib_batch_evaluator
 from optimagic.parameters.block_trees import matrix_to_block_tree
 from optimagic.parameters.tree_registry import (
-    get_registry,
+    extended,
     leaf_names,
     tree_flatten,
     tree_just_flatten,
@@ -107,9 +107,8 @@ def bootstrap(
     # Process results
     # ==================================================================================
 
-    registry = get_registry(extended=True)
     flat_outcomes = [
-        tree_just_flatten(_outcome, registry=registry) for _outcome in all_outcomes
+        tree_just_flatten(_outcome, namespace=extended) for _outcome in all_outcomes
     ]
     internal_outcomes = np.array(flat_outcomes)
 
@@ -167,11 +166,10 @@ def outcomes(self):
             List[Any]: The boostrap outcomes as a list of pytrees.
 
         """
-        registry = get_registry(extended=True)
-        _, treedef = tree_flatten(self._base_outcome, registry=registry)
+        _, treedef = tree_flatten(self._base_outcome, namespace=extended)
 
         outcomes = [
-            tree_unflatten(treedef, out, registry=registry)
+            tree_unflatten(treedef, out, namespace=extended)
             for out in self._internal_outcomes
         ]
         return outcomes
@@ -187,10 +185,9 @@ def se(self):
         cov = self._internal_cov
         se = np.sqrt(np.diagonal(cov))
 
-        registry = get_registry(extended=True)
-        _, treedef = tree_flatten(self._base_outcome, registry=registry)
+        _, treedef = tree_flatten(self._base_outcome, namespace=extended)
 
-        se = tree_unflatten(treedef, se, registry=registry)
+        se = tree_unflatten(treedef, se, namespace=extended)
         return se
 
     def cov(self, return_type="pytree"):
@@ -211,8 +208,7 @@ def cov(self, return_type="pytree"):
         cov = self._internal_cov
 
         if return_type == "dataframe":
-            registry = get_registry(extended=True)
-            names = np.array(leaf_names(self._base_outcome, registry=registry))
+            names = np.array(leaf_names(self._base_outcome, namespace=extended))
             cov = pd.DataFrame(cov, columns=names, index=names)
         elif return_type == "pytree":
             cov = matrix_to_block_tree(cov, self._base_outcome, self._base_outcome)
@@ -239,15 +235,16 @@ def ci(self, ci_method="percentile", ci_level=0.95):
                 bounds of confidence intervals.
 
         """
-        registry = get_registry(extended=True)
-        base_outcome_flat, treedef = tree_flatten(self._base_outcome, registry=registry)
+        base_outcome_flat, treedef = tree_flatten(
+            self._base_outcome, namespace=extended
+        )
 
         lower_flat, upper_flat = calculate_ci(
             base_outcome_flat, self._internal_outcomes, ci_method, ci_level
         )
 
-        lower = tree_unflatten(treedef, lower_flat, registry=registry)
-        upper = tree_unflatten(treedef, upper_flat, registry=registry)
+        lower = tree_unflatten(treedef, lower_flat, namespace=extended)
+        upper = tree_unflatten(treedef, upper_flat, namespace=extended)
         return lower, upper
 
     def p_values(self):
@@ -276,8 +273,7 @@ def summary(self, ci_method="percentile", ci_level=0.95):
                 Soon this will be a pytree.
 
         """
-        registry = get_registry(extended=True)
-        names = leaf_names(self.base_outcome, registry=registry)
+        names = leaf_names(self.base_outcome, namespace=extended)
         summary_data = _calulcate_summary_data_bootstrap(
             self, ci_method=ci_method, ci_level=ci_level
         )

src/estimagic/estimate_msm.py

@@ -51,7 +51,7 @@
 from optimagic.parameters.conversion import Converter, get_converter
 from optimagic.parameters.space_conversion import InternalParams
 from optimagic.parameters.tree_registry import (
-    get_registry,
+    extended,
     leaf_names,
     tree_just_flatten,
 )
@@ -321,8 +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"]
-            registry = get_registry(extended=True)
-            out = np.array(tree_just_flatten(sim_mom, registry=registry))
+            out = np.array(tree_just_flatten(sim_mom, namespace=extended))
             return out
 
         int_jac = first_derivative(
@@ -421,8 +420,7 @@ def get_msm_optimization_functions(
 
     chol_weights = np.linalg.cholesky(flat_weights)
 
-    registry = get_registry(extended=True)
-    flat_emp_mom = tree_just_flatten(empirical_moments, registry=registry)
+    flat_emp_mom = tree_just_flatten(empirical_moments, namespace=extended)
 
     _simulate_moments = _partial_kwargs(simulate_moments, simulate_moments_kwargs)
     _jacobian = _partial_kwargs(jacobian, jacobian_kwargs)
@@ -433,7 +431,7 @@ def get_msm_optimization_functions(
             simulate_moments=_simulate_moments,
             flat_empirical_moments=flat_emp_mom,
             chol_weights=chol_weights,
-            registry=registry,
+            namespace=extended,
         )
     )
 
@@ -448,7 +446,7 @@ def get_msm_optimization_functions(
 
 
 def _msm_criterion(
-    params, simulate_moments, flat_empirical_moments, chol_weights, registry
+    params, simulate_moments, flat_empirical_moments, chol_weights, namespace
 ):
     """Calculate msm criterion given parameters and building blocks."""
     simulated = simulate_moments(params)
@@ -457,7 +455,7 @@ def _msm_criterion(
     if isinstance(simulated, np.ndarray) and simulated.ndim == 1:
         simulated_flat = simulated
     else:
-        simulated_flat = np.array(tree_just_flatten(simulated, registry=registry))
+        simulated_flat = np.array(tree_just_flatten(simulated, namespace=namespace))
 
     deviations = simulated_flat - flat_empirical_moments
     residuals = deviations @ chol_weights
@@ -978,9 +976,8 @@ def sensitivity(
                 inner_tree=self._empirical_moments,
             )
         elif return_type == "dataframe":
-            registry = get_registry(extended=True)
             row_names = self._internal_estimates.names
-            col_names = leaf_names(self._empirical_moments, registry=registry)
+            col_names = leaf_names(self._empirical_moments, namespace=extended)
             out = pd.DataFrame(
                 data=raw,
                 index=row_names,

src/estimagic/msm_weighting.py

@@ -6,7 +6,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 get_registry, tree_just_flatten
+from optimagic.parameters.tree_registry import extended, tree_just_flatten
 from optimagic.utilities import robust_inverse
 
 
@@ -50,13 +50,11 @@ def get_moments_cov(
 
     first_eval = calculate_moments(data, **moment_kwargs)
 
-    registry = get_registry(extended=True)
-
     @functools.wraps(calculate_moments)
     def func(data, **kwargs):
         raw = calculate_moments(data, **kwargs)
         out = pd.Series(
-            tree_just_flatten(raw, registry=registry)
+            tree_just_flatten(raw, namespace=extended)
         )  # xxxx won't be necessary soon!
         return out
 

src/estimagic/shared_covs.py

@@ -6,7 +6,7 @@
 
 from optimagic.parameters.block_trees import matrix_to_block_tree
 from optimagic.parameters.tree_registry import (
-    get_registry,
+    extended,
     tree_just_flatten,
     tree_unflatten,
 )
@@ -149,9 +149,8 @@ def calculate_estimation_summary(
     # Flatten summary and construct data frame for flat estimates
     # ==================================================================================
 
-    registry = get_registry(extended=True)
     flat_data = {
-        key: tree_just_flatten(val, registry=registry)
+        key: tree_just_flatten(val, namespace=extended)
         for key, val in summary_data.items()
     }
 
@@ -170,7 +169,7 @@ def calculate_estimation_summary(
     # ==================================================================================
 
     # create tree with values corresponding to indices of df
-    indices = tree_unflatten(summary_data["value"], names, registry=registry)
+    indices = tree_unflatten(summary_data["value"], names, namespace=extended)
 
     estimates_flat = tree_just_flatten(summary_data["value"])
     indices_flat = tree_just_flatten(indices)
@@ -319,8 +318,7 @@ def calculate_free_estimates(estimates, internal_estimates):
     mask = internal_estimates.free_mask
     names = internal_estimates.names
 
-    registry = get_registry(extended=True)
-    external_flat = np.array(tree_just_flatten(estimates, registry=registry))
+    external_flat = np.array(tree_just_flatten(estimates, namespace=extended))
 
     free_estimates = FreeParams(
         values=external_flat[mask],
@@ -354,8 +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
-    registry = get_registry(extended=True)
-    pytree = tree_unflatten(params, flat, registry=registry)
+    pytree = tree_unflatten(params, flat, namespace=extended)
     return pytree
 
 

src/optimagic/benchmarking/run_benchmark.py

@@ -13,7 +13,7 @@
 from optimagic import batch_evaluators
 from optimagic.algorithms import AVAILABLE_ALGORITHMS
 from optimagic.optimization.optimize import minimize
-from optimagic.parameters.tree_registry import get_registry, tree_just_flatten
+from optimagic.parameters.tree_registry import extended, tree_just_flatten
 
 
 def run_benchmark(
@@ -179,7 +179,6 @@ def _process_one_result(optimize_result, problem):
         dict: Processed result.
 
     """
-    _registry = get_registry(extended=True)
     _criterion = problem["noise_free_fun"]
     _start_x = problem["inputs"]["params"]
     _start_crit_value = _criterion(_start_x)
@@ -190,15 +189,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, registry=_registry)]
+        params_history_flat = [tree_just_flatten(_start_x, namespace=extended)]
         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, registry=_registry) for p in params_history
+            tree_just_flatten(p, namespace=extended) for p in params_history
         ]
         if _is_noisy:
             criterion_history = np.array([_criterion(p) for p in params_history])