chore: rename tree_just_flatten to tree_leaves

Author: abelaba

src/estimagic/bootstrap.py

@@ -15,7 +15,7 @@
 from optimagic.parameters.tree_registry import (
     leaf_names,
     tree_flatten,
-    tree_just_flatten,
+    tree_leaves,
     tree_unflatten,
 )
 from optimagic.typing import VALUE_NAMESPACE
@@ -108,8 +108,7 @@ def bootstrap(
     # ==================================================================================
 
     flat_outcomes = [
-        tree_just_flatten(_outcome, namespace=VALUE_NAMESPACE)
-        for _outcome in all_outcomes
+        tree_leaves(_outcome, namespace=VALUE_NAMESPACE) for _outcome in all_outcomes
     ]
     internal_outcomes = np.array(flat_outcomes)
 

src/estimagic/estimate_msm.py

@@ -52,7 +52,7 @@
 from optimagic.parameters.space_conversion import InternalParams
 from optimagic.parameters.tree_registry import (
     leaf_names,
-    tree_just_flatten,
+    tree_leaves,
 )
 from optimagic.shared.check_option_dicts import (
     check_optimization_options,
@@ -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_leaves(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_leaves(empirical_moments, namespace=VALUE_NAMESPACE)
 
     _simulate_moments = _partial_kwargs(simulate_moments, simulate_moments_kwargs)
     _jacobian = _partial_kwargs(jacobian, jacobian_kwargs)
@@ -455,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, namespace=namespace))
+        simulated_flat = np.array(tree_leaves(simulated, namespace=namespace))
 
     deviations = simulated_flat - flat_empirical_moments
     residuals = deviations @ chol_weights

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 tree_just_flatten
+from optimagic.parameters.tree_registry import tree_leaves
 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_leaves(raw, namespace=VALUE_NAMESPACE)
         )  # 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 (
-    tree_just_flatten,
+    tree_leaves,
     tree_unflatten,
 )
 from optimagic.typing import VALUE_NAMESPACE
@@ -150,7 +150,7 @@ def calculate_estimation_summary(
     # ==================================================================================
 
     flat_data = {
-        key: tree_just_flatten(val, namespace=VALUE_NAMESPACE)
+        key: tree_leaves(val, namespace=VALUE_NAMESPACE)
         for key, val in summary_data.items()
     }
 
@@ -171,8 +171,8 @@ def calculate_estimation_summary(
     # create tree with values corresponding to indices of df
     indices = tree_unflatten(summary_data["value"], names, namespace=VALUE_NAMESPACE)
 
-    estimates_flat = tree_just_flatten(summary_data["value"])
-    indices_flat = tree_just_flatten(indices)
+    estimates_flat = tree_leaves(summary_data["value"])
+    indices_flat = tree_leaves(indices)
 
     # use index chunks in indices_flat to access the corresponding sub data frame of df,
     # and use the index information stored in estimates_flat to form the correct (multi)
@@ -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_leaves(estimates, namespace=VALUE_NAMESPACE))
 
     free_estimates = FreeParams(
         values=external_flat[mask],

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 tree_just_flatten
+from optimagic.parameters.tree_registry import tree_leaves
 from optimagic.typing import VALUE_NAMESPACE
 
 
@@ -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_leaves(_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_leaves(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

@@ -22,10 +22,9 @@
 from optimagic.parameters.bounds import Bounds, get_internal_bounds, pre_process_bounds
 from optimagic.parameters.tree_registry import (
     tree_flatten,
-    tree_just_flatten,
+    tree_leaves,
     tree_unflatten,
 )
-from optimagic.parameters.tree_registry import tree_just_flatten as tree_leaves
 from optimagic.typing import VALUE_NAMESPACE, BatchEvaluatorLiteral, PyTree
 
 
@@ -226,18 +225,14 @@ def first_derivative(
 
         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_leaves(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)
-            )
+            min_steps = np.array(tree_leaves(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)
-            )
+            step_size = np.array(tree_leaves(step_size, namespace=VALUE_NAMESPACE))
     else:
         x = params.astype(np.float64)
 
@@ -544,18 +539,14 @@ def second_derivative(
 
         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_leaves(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)
-            )
+            min_steps = np.array(tree_leaves(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)
-            )
+            step_size = np.array(tree_leaves(step_size, namespace=VALUE_NAMESPACE))
     else:
         x = params.astype(np.float64)
 

src/optimagic/examples/criterion_functions.py

@@ -17,7 +17,7 @@
 )
 from optimagic.parameters.block_trees import matrix_to_block_tree
 from optimagic.parameters.tree_registry import (
-    tree_just_flatten,
+    tree_leaves,
     tree_unflatten,
 )
 from optimagic.typing import VALUE_NAMESPACE, PyTree
@@ -214,9 +214,7 @@ 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=VALUE_NAMESPACE), dtype=np.float64
-        )
+        x = np.array(tree_leaves(params, namespace=VALUE_NAMESPACE), dtype=np.float64)
     return x
 
 

src/optimagic/optimization/fun_value.py

@@ -7,7 +7,7 @@
 from numpy.typing import NDArray
 
 from optimagic.exceptions import InvalidFunctionError
-from optimagic.parameters.tree_registry import tree_just_flatten
+from optimagic.parameters.tree_registry import tree_leaves
 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_leaves(value, namespace=VALUE_NAMESPACE)
 
     flat_arr = np.asarray(flat, dtype=np.float64)
     return flat_arr

src/optimagic/optimization/history.py

@@ -9,7 +9,7 @@
 
 from optimagic.parameters.tree_registry import (
     leaf_names,
-    tree_just_flatten,
+    tree_leaves,
 )
 from optimagic.timing import CostModel
 from optimagic.typing import VALUE_NAMESPACE, Direction, EvalTask, PyTree
@@ -400,7 +400,7 @@ def _get_flat_params(params: list[PyTree]) -> list[list[float]]:
     if fast_path:
         flatten = lambda x: x.tolist()
     else:
-        flatten = partial(tree_just_flatten, namespace=VALUE_NAMESPACE)
+        flatten = partial(tree_leaves, namespace=VALUE_NAMESPACE)
 
     return [flatten(p) for p in params]
 

src/optimagic/parameters/block_trees.py

@@ -5,9 +5,9 @@
 
 from optimagic.parameters.tree_registry import (
     tree_flatten,
+    tree_leaves,
     tree_unflatten,
 )
-from optimagic.parameters.tree_registry import tree_just_flatten as tree_leaves
 from optimagic.typing import VALUE_NAMESPACE