Repository: OpenSourceEconomics/estimagic

Issue:
Allow pytrees in bootstrap
## Problem

- Parameters are assumed to be DataFrames
- Statistics are assumed to be Series

## Solution

- Allow both to be general pytrees

Discussion:

Patch:
diff --git a/src/estimagic/estimation/estimate_ml.py b/src/estimagic/estimation/estimate_ml.py
index a45fa08..7a8ab7f 100644
--- a/src/estimagic/estimation/estimate_ml.py
+++ b/src/estimagic/estimation/estimate_ml.py
@@ -17,6 +17,7 @@ from estimagic.inference.shared import calculate_ci
 from estimagic.inference.shared import calculate_inference_quantities
 from estimagic.inference.shared import calculate_p_values
 from estimagic.inference.shared import check_is_optimized_and_derivative_case
+from estimagic.inference.shared import convert_flat_params_to_pytree
 from estimagic.inference.shared import get_derivative_case
 from estimagic.inference.shared import transform_covariance
 from estimagic.optimization.optimize import maximize
@@ -502,8 +503,7 @@ class LikelihoodResult:
 
         Returns:
             Any: A pytree with the same structure as params containing standard errors
-            for the parameter estimates.
-
+                for the parameter estimates.
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -512,10 +512,8 @@ class LikelihoodResult:
             seed=seed,
         )
 
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = np.sqrt(np.diagonal(free_cov))
-
-        out = self._converter.params_from_internal(helper)
+        free_se = np.sqrt(np.diagonal(free_cov))
+        out = convert_flat_params_to_pytree(free_se, self._flat_params, self._converter)
 
         return out
 
@@ -556,7 +554,6 @@ class LikelihoodResult:
         Returns:
             Any: The covariance matrix of the estimated parameters as block-pytree or
                 numpy array.
-
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -611,10 +608,8 @@ class LikelihoodResult:
             seed (int): Seed for the random number generator. Only used if there are
                 transforming constraints.
 
-
         Returns:
             Any: The estimation summary as pytree of DataFrames.
-
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -668,7 +663,6 @@ class LikelihoodResult:
                 confidence intervals.
             Any: Pytree with the same structure as params containing upper bounds of
                 confidence intervals.
-
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -682,13 +676,12 @@ class LikelihoodResult:
 
         free_lower, free_upper = calculate_ci(free_values, free_se, ci_level)
 
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = free_lower
-        lower = self._converter.params_from_internal(helper)
-
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = free_upper
-        upper = self._converter.params_from_internal(helper)
+        lower = convert_flat_params_to_pytree(
+            free_lower, self._flat_params, self._converter
+        )
+        upper = convert_flat_params_to_pytree(
+            free_upper, self._flat_params, self._converter
+        )
 
         return lower, upper
 
@@ -699,7 +692,7 @@ class LikelihoodResult:
         bounds_handling="clip",
         seed=None,
     ):
-        """Calculate p_values.
+        """Calculate p-values.
 
         Args:
             method (str): One of "jacobian", "hessian", "robust", "cluster_robust",
@@ -724,11 +717,8 @@ class LikelihoodResult:
                 transforming constraints.
 
         Returns:
-            Any: Pytree with the same structure as params containing lower bounds of
-                confidence intervals.
-            Any: Pytree with the same structure as params containing upper bounds of
-                confidence intervals.
-
+            Any: Pytree with the same structure as params containing p-values.
+            Any: Pytree with the same structure as params containing p-values.
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -739,12 +729,11 @@ class LikelihoodResult:
 
         free_values = self._flat_params.values[self._flat_params.free_mask]
         free_se = np.sqrt(np.diagonal(free_cov))
-
         free_p_values = calculate_p_values(free_values, free_se)
 
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = free_p_values
-        out = self._converter.params_from_internal(helper)
+        out = convert_flat_params_to_pytree(
+            free_p_values, self._flat_params, self._converter
+        )
 
         return out
 
@@ -753,6 +742,5 @@ class LikelihoodResult:
 
         Args:
             path (str, pathlib.Path): A str or pathlib.path ending in .pkl or .pickle.
-
         """
         to_pickle(self, path=path)
diff --git a/src/estimagic/estimation/estimate_msm.py b/src/estimagic/estimation/estimate_msm.py
index 883b937..0025ab0 100644
--- a/src/estimagic/estimation/estimate_msm.py
+++ b/src/estimagic/estimation/estimate_msm.py
@@ -17,6 +17,7 @@ from estimagic.inference.shared import calculate_ci
 from estimagic.inference.shared import calculate_inference_quantities
 from estimagic.inference.shared import calculate_p_values
 from estimagic.inference.shared import check_is_optimized_and_derivative_case
+from estimagic.inference.shared import convert_flat_params_to_pytree
 from estimagic.inference.shared import get_derivative_case
 from estimagic.inference.shared import transform_covariance
 from estimagic.optimization.optimize import minimize
@@ -432,7 +433,6 @@ def _partial_kwargs(func, kwargs):
 
     In contrast to normal partial this works if kwargs in None. If func is not a
     callable it simply returns None.
-
     """
     if isinstance(func, Callable):
         if kwargs not in (None, {}):
@@ -550,8 +550,7 @@ class MomentsResult:
 
         Returns:
             Any: A pytree with the same structure as params containing standard errors
-            for the parameter estimates.
-
+                for the parameter estimates.
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -560,10 +559,8 @@ class MomentsResult:
             seed=seed,
         )
 
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = np.sqrt(np.diagonal(free_cov))
-
-        out = self._converter.params_from_internal(helper)
+        free_se = np.sqrt(np.diagonal(free_cov))
+        out = convert_flat_params_to_pytree(free_se, self._flat_params, self._converter)
 
         return out
 
@@ -605,7 +602,6 @@ class MomentsResult:
         Returns:
             Any: The covariance matrix of the estimated parameters as block-pytree or
                 numpy array.
-
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -613,6 +609,7 @@ class MomentsResult:
             bounds_handling=bounds_handling,
             seed=seed,
         )
+
         if return_type == "array":
             out = free_cov
         elif return_type == "dataframe":
@@ -662,7 +659,6 @@ class MomentsResult:
 
         Returns:
             Any: The estimation summary as pytree of DataFrames.
-
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -717,7 +713,6 @@ class MomentsResult:
                 confidence intervals.
             Any: Pytree with the same structure as params containing upper bounds of
                 confidence intervals.
-
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -728,16 +723,14 @@ class MomentsResult:
 
         free_values = self._flat_params.values[self._flat_params.free_mask]
         free_se = np.sqrt(np.diagonal(free_cov))
-
         free_lower, free_upper = calculate_ci(free_values, free_se, ci_level)
 
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = free_lower
-        lower = self._converter.params_from_internal(helper)
-
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = free_upper
-        upper = self._converter.params_from_internal(helper)
+        lower = convert_flat_params_to_pytree(
+            free_lower, self._flat_params, self._converter
+        )
+        upper = convert_flat_params_to_pytree(
+            free_upper, self._flat_params, self._converter
+        )
 
         return lower, upper
 
@@ -748,7 +741,7 @@ class MomentsResult:
         bounds_handling="clip",
         seed=None,
     ):
-        """Calculate confidence intervals.
+        """Calculate p-values.
 
         Args:
             method (str): One of "robust", "optimal". Despite the name, "optimal" is
@@ -771,11 +764,8 @@ class MomentsResult:
                 transforming constraints.
 
         Returns:
-            Any: Pytree with the same structure as params containing lower bounds of
-                confidence intervals.
-            Any: Pytree with the same structure as params containing upper bounds of
-                confidence intervals.
-
+            Any: Pytree with the same structure as params containing p-values.
+            Any: Pytree with the same structure as params containing p-values.
         """
         free_cov = self._get_free_cov(
             method=method,
@@ -786,12 +776,11 @@ class MomentsResult:
 
         free_values = self._flat_params.values[self._flat_params.free_mask]
         free_se = np.sqrt(np.diagonal(free_cov))
-
         free_p_values = calculate_p_values(free_values, free_se)
 
-        helper = np.full(len(self._flat_params.values), np.nan)
-        helper[self._flat_params.free_mask] = free_p_values
-        out = self._converter.params_from_internal(helper)
+        out = convert_flat_params_to_pytree(
+            free_p_values, self._flat_params, self._converter
+        )
 
         return out
 
@@ -860,7 +849,6 @@ class MomentsResult:
             Any: The sensitivity measure as a pytree, numpy array or DataFrame.
                 In 2d formats, the sensitivity measures have one row per estimated
                 parameter and one column per moment.
-
         """
         if self._has_constraints:
             raise NotImplementedError(
@@ -956,6 +944,5 @@ class MomentsResult:
 
         Args:
             path (str, pathlib.Path): A str or pathlib.path ending in .pkl or .pickle.
-
         """
         to_pickle(self, path=path)
diff --git a/src/estimagic/estimation/msm_weighting.py b/src/estimagic/estimation/msm_weighting.py
index 515088c..12c9f44 100644
--- a/src/estimagic/estimation/msm_weighting.py
+++ b/src/estimagic/estimation/msm_weighting.py
@@ -58,7 +58,7 @@ def get_moments_cov(
         )  # xxxx won't be necessary soon!
         return out
 
-    cov_arr = bootstrap(data=data, outcome=func, outcome_kwargs=moment_kwargs)["cov"]
+    cov_arr = bootstrap(data=data, outcome=func, outcome_kwargs=moment_kwargs).cov()
 
     if isinstance(cov_arr, pd.DataFrame):
         cov_arr = cov_arr.to_numpy()  # xxxx won't be necessary soon
diff --git a/src/estimagic/inference/bootstrap.py b/src/estimagic/inference/bootstrap.py
index 0eb91e1..f712cad 100644
--- a/src/estimagic/inference/bootstrap.py
+++ b/src/estimagic/inference/bootstrap.py
@@ -1,8 +1,19 @@
+import functools
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
 import pandas as pd
 from estimagic.batch_evaluators import joblib_batch_evaluator
 from estimagic.inference.bootstrap_ci import compute_ci
 from estimagic.inference.bootstrap_helpers import check_inputs
 from estimagic.inference.bootstrap_outcomes import get_bootstrap_outcomes
+from estimagic.parameters.block_trees import matrix_to_block_tree
+from estimagic.parameters.tree_registry import get_registry
+from pybaum import leaf_names
+from pybaum import tree_flatten
+from pybaum import tree_just_flatten
+from pybaum import tree_unflatten
 
 
 def bootstrap(
@@ -12,8 +23,6 @@ def bootstrap(
     outcome_kwargs=None,
     n_draws=1_000,
     cluster_by=None,
-    ci_method="percentile",
-    alpha=0.05,
     seed=None,
     n_cores=1,
     error_handling="continue",
@@ -23,15 +32,13 @@ def bootstrap(
     for statistic of interest in given original sample.
 
     Args:
-        data (pandas.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-            Needs to return a pandas Series.
+        data (pandas.DataFrame): Original dataset.
+        outcome (callable): Function of the data calculating statistic of interest.
+            Returns a general pytree (e.g. pandas Series, dict, numpy array, etc.).
         outcome_kwargs (dict): Additional keyword arguments for outcome.
         n_draws (int): number of bootstrap samples to draw.
-        cluster_by (str): column name of variable to cluster by or None.
-        ci_method (str): method of choice for confidence interval computation.
-        alpha (float): significance level of choice.
-        seeds (numpy.array): array of seeds for bootstrap samples, default is none.
+        cluster_by (str): Column name of variable to cluster by or None.
+        seeds (numpy.array): Array of seeds for bootstrap samples, default is none.
         n_cores (int): number of jobs for parallelization.
         error_handling (str): One of "continue", "raise". Default "continue" which means
             that bootstrap estimates are only calculated for those samples where no
@@ -41,17 +48,17 @@ def bootstrap(
             as the estimagic batch_evaluators. See :ref:`batch_evaluators`.
 
     Returns:
-        results (pandas.DataFrame): DataFrame where k'th row contains mean estimate,
-        standard error, and confidence interval of k'th parameter.
-
+        BootstrapResult: A BootstrapResult object storing information on summary
+            statistics, the covariance matrix, and estimated boostrap outcomes.
     """
+    check_inputs(data, cluster_by)
 
-    check_inputs(data, cluster_by, ci_method, alpha)
+    if outcome_kwargs is not None:
+        outcome = functools.partial(outcome, **outcome_kwargs)
 
-    estimates = get_bootstrap_outcomes(
+    bootstrap_outcomes = get_bootstrap_outcomes(
         data=data,
         outcome=outcome,
-        outcome_kwargs=outcome_kwargs,
         cluster_by=cluster_by,
         seed=seed,
         n_draws=n_draws,
@@ -60,46 +67,199 @@ def bootstrap(
         batch_evaluator=batch_evaluator,
     )
 
-    out = bootstrap_from_outcomes(
-        data, outcome, estimates, ci_method=ci_method, alpha=alpha, n_cores=n_cores
-    )
+    base_outcome = outcome(data)
+    out = bootstrap_from_outcomes(base_outcome, bootstrap_outcomes)
 
     return out
 
 
-def bootstrap_from_outcomes(
-    data, outcome, bootstrap_outcomes, *, ci_method="percentile", alpha=0.05, n_cores=1
-):
-    """Set up results table containing mean, standard deviation and confidence interval
-    for each estimated parameter.
+def bootstrap_from_outcomes(base_outcome, bootstrap_outcomes):
+    """Create BootstrapResults object.
 
     Args:
-        data (pandas.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-            Needs to return a pandas Series.
-        bootstrap_outcomes (pandas.DataFrame): DataFrame of bootstrap_outcomes in the
-            bootstrap samples.
-        ci_method (str): method of choice for confidence interval computation.
-        n_cores (int): number of jobs for parallelization.
-        alpha (float): significance level of choice.
+        base_outcome (pytree): Pytree of base outcomes, i.e. the outcome
+            statistic(s) evaluated on the original data set.
+        bootstrap_outcomes (list): List of pytrees of estimated
+            bootstrap outcomes.
 
     Returns:
-        results (pandas.DataFrame): table of results.
-
+        BootstrapResult: A BootstrapResult object storing information on summary
+            statistics, the covariance matrix, and the estimated boostrap outcomes.
     """
+    if isinstance(bootstrap_outcomes, list):
+        registry = get_registry(extended=True)
 
-    check_inputs(data=data, ci_method=ci_method, alpha=alpha)
+        flat_outcomes = [
+            tree_just_flatten(est, registry=registry) for est in bootstrap_outcomes
+        ]
+        internal_outcomes = np.array(flat_outcomes)
+    else:
+        raise TypeError(
+            "bootstrap_outcomes must be a list of pytrees, "
+            f"not {type(bootstrap_outcomes)}."
+        )
 
-    summary = pd.DataFrame(bootstrap_outcomes.mean(axis=0), columns=["mean"])
+    out = BootstrapResult(
+        _base_outcome=base_outcome,
+        _internal_outcomes=internal_outcomes,
+    )
 
-    summary["std"] = bootstrap_outcomes.std(axis=0)
+    return out
 
-    cis = compute_ci(data, outcome, bootstrap_outcomes, ci_method, alpha, n_cores)
-    summary["lower_ci"] = cis["lower_ci"]
-    summary["upper_ci"] = cis["upper_ci"]
 
-    cov = bootstrap_outcomes.cov()
+@dataclass
+class BootstrapResult:
+    _base_outcome: Any
+    _internal_outcomes: np.ndarray
 
-    out = {"summary": summary, "cov": cov, "outcomes": bootstrap_outcomes}
+    @property
+    def _outcomes(self):
+        return self.outcomes()
 
-    return out
+    @property
+    def _se(self):
+        return self.se()
+
+    @property
+    def _cov(self):
+        return self.cov()
+
+    @property
+    def _ci(self):
+        return self.ci()
+
+    @property
+    def _p_values(self):
+        return self.p_values()
+
+    @property
+    def _summary(self):
+        return self.summary()
+
+    @property
+    def base_outcome(self):
+        """Returns the base outcome statistic(s).
+
+        Returns:
+            pytree: Pytree of base outcomes, i.e. the outcome statistic(s) evaluated
+                on the original data set.
+        """
+        return self._base_outcome
+
+    def outcomes(self):
+        """Returns the estimated bootstrap outcomes.
+
+        Returns:
+            Any: The boostrap outcomes as a list of pytrees.
+        """
+        registry = get_registry(extended=True)
+        _, treedef = tree_flatten(self._base_outcome, registry=registry)
+
+        outcomes = [
+            tree_unflatten(treedef, out, registry=registry)
+            for out in self._internal_outcomes
+        ]
+
+        return outcomes
+
+    def se(self):
+        """Calculate standard errors.
+
+        Returns:
+            list: A list of pytrees containing standard errors for the bootstrapped
+                statistic.
+        """
+        cov = np.cov(self._internal_outcomes, rowvar=False)
+        se = np.sqrt(np.diagonal(cov))
+
+        registry = get_registry(extended=True)
+        _, treedef = tree_flatten(self._base_outcome, registry=registry)
+
+        out = tree_unflatten(treedef, se, registry=registry)
+
+        return out
+
+    def cov(self, return_type="pytree"):
+        """Calculate the variance-covariance matrix of the estimated parameters.
+
+        Args:
+            return_type (str): One of "pytree", "array" or "dataframe". Default pytree.
+                If "array", a 2d numpy array with the covariance is returned. If
+                "dataframe", a pandas DataFrame with parameter names in the
+                index and columns are returned.
+                The default is "pytree".
+
+        Returns:
+            Any: The covariance matrix of the estimated parameters as a block-pytree,
+                numpy array, or pandas DataFrame.
+        """
+        cov = np.cov(self._internal_outcomes, rowvar=False)
+
+        if return_type == "array":
+            out = cov
+        elif return_type == "dataframe":
+            registry = get_registry(extended=True)
+            leafnames = leaf_names(self._base_outcome, registry=registry)
+            free_index = np.array(leafnames)
+            out = pd.DataFrame(data=cov, columns=free_index, index=free_index)
+        elif return_type == "pytree":
+            out = matrix_to_block_tree(cov, self._base_outcome, self._base_outcome)
+        else:
+            raise TypeError(
+                "return_type must be one of pytree, array, or dataframe, "
+                f"not {return_type}."
+            )
+
+        return out
+
+    def ci(self, ci_method="percentile", ci_level=0.95):
+        """Calculate confidence intervals.
+
+        Args:
+            ci_method (str): Method of choice for confidence interval computation.
+                The default is "percentile".
+            ci_level (float): Confidence level for the calculation of confidence
+                intervals. The default is 0.95.
+
+        Returns:
+            Any: Pytree with the same structure as base_outcome containing lower
+                bounds of confidence intervals.
+            Any: Pytree with the same structure as base_outcome containing upper
+                bounds of confidence intervals.
+        """
+        registry = get_registry(extended=True)
+        base_outcome_flat, treedef = tree_flatten(self._base_outcome, registry=registry)
+
+        lower_flat, upper_flat = compute_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)
+
+        return lower, upper
+
+    def p_values(self):
+        """Calculate p-values.
+
+        Returns:
+            Any: A pytree with the same structure as base_outcome containing p-values
+                for the parameter estimates.
+        """
+        raise NotImplementedError("Bootstrapped p-values are not implemented yet.")
+
+    def summary(self, ci_method="percentile", ci_level=0.95):
+        """Create a summary of bootstrap results.
+
+        Args:
+            ci_method (str): Method of choice for confidence interval computation.
+                The default is "percentile".
+            ci_level (float): Confidence level for the calculation of confidence
+                intervals. The default is 0.95.
+
+        Returns:
+            pd.DataFrame: The estimation summary as a DataFrame containing information
+                on the mean, standard errors, as well as the confidence intervals.
+                Soon this will be a pytree.
+        """
+        raise NotImplementedError("summary is not implemented yet.")
diff --git a/src/estimagic/inference/bootstrap_ci.py b/src/estimagic/inference/bootstrap_ci.py
index 6531a19..a8a4f4d 100644
--- a/src/estimagic/inference/bootstrap_ci.py
+++ b/src/estimagic/inference/bootstrap_ci.py
@@ -1,285 +1,209 @@
 import numpy as np
-import pandas as pd
 from estimagic.inference.bootstrap_helpers import check_inputs
-from joblib import delayed
-from joblib import Parallel
 from scipy.stats import norm
 
 
-def compute_ci(data, outcome, estimates, ci_method="percentile", alpha=0.05, n_cores=1):
-    """Compute confidence interval of bootstrap estimates. Parts of the code of the
-    subfunctions of this function are taken from Daniel Saxton's resample library, as
-    found on https://github.com/dsaxton/resample/ .
+def compute_ci(
+    base_outcome_flat,
+    estimates,
+    ci_method="percentile",
+    ci_level=0.95,
+):
+    """Compute confidence interval of bootstrap estimates.
+
+    Parts of the code of the subfunctions of this function are taken from
+    Daniel Saxton's resample library, as found on
+    https://github.com/dsaxton/resample/
 
 
     Args:
-        data (pandas.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-            Needs to return a pandas Series.
-        estimates (pandas.DataFrame): DataFrame of estimates in the bootstrap samples.
-        ci_method (str): method of choice for confidence interval computation.
-        alpha (float): significance level of choice.
-        n_cores (int): number of jobs for parallelization.
+        base_outcome_flat (list): List of flat base outcomes, i.e. the outcome
+            statistic(s) evaluated on the original data set.
+        estimates (np.ndarray): Array of estimates computed on the bootstrapped
+            samples.
+        ci_method (str): Method of choice for confidence interval computation.
+        ci_level (float): Confidence level for the calculation of confidence
+            intervals. The default is 0.95.
 
     Returns:
-        cis (pandas.DataFrame): DataFrame where k'th row contains CI for k'th parameter.
-
+        np.ndarray: 1d array of the lower confidence interval, where the k'th entry
+            contains the lower confidence interval for the k'th parameter.
+        np.ndarray: 1d array of the upper confidence interval, where the k'th entry
+            contains the upper confidence interval for the k'th parameter.
     """
+    check_inputs(ci_method=ci_method, ci_level=ci_level)
 
-    check_inputs(data=data, alpha=alpha, ci_method=ci_method)
+    alpha = 1 - ci_level
 
-    funcname = "_ci_" + ci_method
+    if ci_method == "percentile":
+        cis = _ci_percentile(estimates, alpha)
+    elif ci_method == "bc":
+        cis = _ci_bc(estimates, base_outcome_flat, alpha)
+    elif ci_method == "t":
+        cis = _ci_t(estimates, base_outcome_flat, alpha)
+    elif ci_method == "basic":
+        cis = _ci_basic(estimates, base_outcome_flat, alpha)
+    elif ci_method == "normal":
+        cis = _ci_normal(estimates, base_outcome_flat, alpha)
 
-    cis = globals()[funcname](data, outcome, estimates, alpha, n_cores)
+    return cis[:, 0], cis[:, 1]
 
-    return pd.DataFrame(
-        cis, index=estimates.columns.tolist(), columns=["lower_ci", "upper_ci"]
-    )
 
-
-def _ci_percentile(data, outcome, estimates, alpha, n_cores):
+def _ci_percentile(estimates, alpha):
     """Compute percentile type confidence interval of bootstrap estimates.
 
     Args:
-        data (pd.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-        estimates (data.Frame): DataFrame of estimates in the bootstrap samples.
-        alpha (float): significance level of choice.
-        n_cores (int): number of jobs for parallelization.
+        estimates (np.ndarray): Array of estimates computed on the bootstrapped
+            samples.
+        alpha (float): Statistical significance level of choice.
 
     Returns:
-        cis (np.array): array where k'th row contains CI for k'th parameter.
-
+        cis (np.ndarray): 2d array where k'th row contains the upper and lower CI
+            for k'th parameter.
     """
-
     num_params = estimates.shape[1]
-    boot_est = estimates.values
     cis = np.zeros((num_params, 2))
 
     for k in range(num_params):
 
-        q = _eqf(boot_est[:, k])
-        cis[k, :] = np.array([q(alpha / 2), q(1 - alpha / 2)])
+        q = _eqf(estimates[:, k])
+        cis[k, :] = q(alpha / 2), q(1 - alpha / 2)
 
     return cis
 
 
-def _ci_bca(data, outcome, estimates, alpha, n_cores):
-    """Compute bca type confidence interval of bootstrap estimates.
-
-    Args:
-        data (pd.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-        estimates (data.Frame): DataFrame of estimates in the bootstrap samples.
-        alpha (float): significance level of choice.
-        n_cores (int): number of jobs for parallelization.
-
-    Returns:
-        cis (np.array): array where k'th row contains CI for k'th parameter.
-
-    """
-
-    num_params = estimates.shape[1]
-    boot_est = estimates.values
-    cis = np.zeros((num_params, 2))
-
-    theta = outcome(data)
-
-    jack_est = _jackknife(data, outcome, n_cores)
-    jack_mean = np.mean(jack_est, axis=0)
-
-    for k in range(num_params):
-
-        q = _eqf(boot_est[:, k])
-        params = boot_est[:, k]
-
-        # bias correction
-        z_naught = norm.ppf(np.mean(params <= theta[k]))
-        z_low = norm.ppf(alpha)
-        z_high = norm.ppf(1 - alpha)
-
-        # accelaration
-        acc = np.sum((jack_mean[k] - jack_est[k]) ** 3) / (
-            6 * np.sum((jack_mean[k] - jack_est[k]) ** 2) ** (3 / 2)
-        )
-
-        p1 = norm.cdf(z_naught + (z_naught + z_low) / (1 - acc * (z_naught + z_low)))
-        p2 = norm.cdf(z_naught + (z_naught + z_high) / (1 - acc * (z_naught + z_high)))
-
-        cis[k, :] = np.array([q(p1), q(p2)])
-
-    return cis
-
-
-def _ci_bc(data, outcome, estimates, alpha, n_cores):
+def _ci_bc(estimates, base_outcome, alpha):
     """Compute bc type confidence interval of bootstrap estimates.
 
     Args:
-        data (pd.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-        estimates (data.Frame): DataFrame of estimates in the bootstrap samples.
-        alpha (float): significance level of choice.
-        n_cores (int): number of jobs for parallelization.
+        estimates (np.ndarray): Array of estimates computed on the bootstrapped
+            samples.
+        base_outcome (list): List of flat base outcomes, i.e. the outcome
+            statistics evaluated on the original data set.
+        alpha (float): Statistical significance level of choice.
 
     Returns:
-        cis (np.array): array where k'th row contains CI for k'th parameter.
-
+        cis (np.ndarray): 2d array where k'th row contains the upper and lower CI
+            for k'th parameter.
     """
-
     num_params = estimates.shape[1]
-    boot_est = estimates.values
     cis = np.zeros((num_params, 2))
 
-    theta = outcome(data)
-
     for k in range(num_params):
 
-        q = _eqf(boot_est[:, k])
-        params = boot_est[:, k]
+        q = _eqf(estimates[:, k])
+        params = estimates[:, k]
 
-        # bias correction
-        z_naught = norm.ppf(np.mean(params <= theta[k]))
+        # Bias correction
+        z_naught = norm.ppf(np.mean(params <= base_outcome[k]))
         z_low = norm.ppf(alpha)
         z_high = norm.ppf(1 - alpha)
 
         p1 = norm.cdf(z_naught + (z_naught + z_low))
         p2 = norm.cdf(z_naught + (z_naught + z_high))
 
-        cis[k, :] = np.array([q(p1), q(p2)])
+        cis[k, :] = q(p1), q(p2)
 
     return cis
 
 
-def _ci_t(data, outcome, estimates, alpha, n_cores):
+def _ci_t(estimates, base_outcome, alpha):
     """Compute studentized confidence interval of bootstrap estimates.
 
     Args:
-        data (pd.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-        estimates (data.Frame): DataFrame of estimates in the bootstrap samples.
-        alpha (float): significance level of choice.
-        n_cores (int): number of jobs for parallelization.
+        estimates (np.ndarray): Array of estimates computed on the bootstrapped
+            samples.
+        base_outcome (list): List of flat base outcomes, i.e. the outcome
+            statistics evaluated on the original data set.
+        alpha (float): Statistical significance level of choice.
 
     Returns:
-        cis (np.array): array where k'th row contains CI for k'th parameter.
-
+        cis (np.ndarray): 2d array where k'th row contains the upper and lower CI
+            for k'th parameter.
     """
-
     num_params = estimates.shape[1]
-    boot_est = estimates.values
     cis = np.zeros((num_params, 2))
-    theta = outcome(data)
 
     for k in range(num_params):
 
-        params = boot_est[:, k]
+        params = estimates[:, k]
 
         theta_std = np.std(params)
 
-        tq = _eqf((params - theta[k]) / theta_std)
+        tq = _eqf((params - base_outcome[k]) / theta_std)
         t1 = tq(1 - alpha / 2)
         t2 = tq(alpha / 2)
 
-        cis[k, :] = np.array([theta[k] - theta_std * t1, theta[k] - theta_std * t2])
+        cis[k, :] = base_outcome[k] - theta_std * t1, base_outcome[k] - theta_std * t2
 
     return cis
 
 
-def _ci_normal(data, outcome, estimates, alpha, n_cores):
+def _ci_normal(estimates, base_outcome, alpha):
     """Compute approximate normal confidence interval of bootstrap estimates.
 
     Args:
-        data (pd.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-        estimates (data.Frame): DataFrame of estimates in the bootstrap samples.
-        alpha (float): significance level of choice.
-        n_cores (int): number of jobs for parallelization.
+        estimates (np.ndarray): Array of estimates computed on the bootstrapped
+            samples.
+        base_outcome (list): List of flat base outcomes, i.e. the outcome
+            statistics evaluated on the original data set.
+        alpha (float): Statistical significance level of choice.
 
     Returns:
-        cis (np.array): array where k'th row contains CI for k'th parameter.
-
+        cis (np.ndarray): 2d array where k'th row contains the upper and lower CI
+            for k'th parameter.
     """
-
     num_params = estimates.shape[1]
-    boot_est = estimates.values
     cis = np.zeros((num_params, 2))
-    theta = outcome(data)
 
     for k in range(num_params):
 
-        params = boot_est[:, k]
+        params = estimates[:, k]
         theta_std = np.std(params)
         t = norm.ppf(alpha / 2)
 
-        cis[k, :] = np.array([theta[k] + theta_std * t, theta[k] - theta_std * t])
+        cis[k, :] = base_outcome[k] + theta_std * t, base_outcome[k] - theta_std * t
 
     return cis
 
 
-def _ci_basic(data, outcome, estimates, alpha, n_cores):
+def _ci_basic(estimates, base_outcome, alpha):
     """Compute basic bootstrap confidence interval of bootstrap estimates.
 
     Args:
-        data (pd.DataFrame): original dataset.
-        outcome (callable): function of the data calculating statistic of interest.
-        estimates (data.Frame): DataFrame of estimates in the bootstrap samples.
-        alpha (float): significance level of choice.
-        n_cores (int): number of jobs for parallelization.
+        estimates (np.ndarray): Array of estimates computed on the bootstrapped
+            samples.
+        base_outcome (list): List of flat base outcomes, i.e. the outcome
+            statistics evaluated on the original data set.
+        alpha (float): Statistical significance level of choice.
 
     Returns:
-        cis (np.array): array where k'th row contains CI for k'th parameter.
-
+        cis (np.ndarray): 2d array where k'th row contains the upper and lower CI
+            for k'th parameter.
     """
-
     num_params = estimates.shape[1]
-    boot_est = estimates.values
     cis = np.zeros((num_params, 2))
-    theta = outcome(data)
 
     for k in range(num_params):
 
-        q = _eqf(boot_est[:, k])
+        q = _eqf(estimates[:, k])
 
-        cis[k, :] = np.array(
-            [2 * theta[k] - q(1 - alpha / 2), 2 * theta[k] - q(alpha / 2)]
+        cis[k, :] = (
+            2 * base_outcome[k] - q(1 - alpha / 2),
+            2 * base_outcome[k] - q(alpha / 2),
         )
 
     return cis
 
 
-def _jackknife(data, outcome, n_cores=1):
-    """Calculate leave-one-out estimator.
-
-    Args:
-        data (pd.DataFrame): original dataset.
-        n_cores (int): number of jobs for parallelization.
-
-    Returns:
-        jk_estimates (pd.DataFrame): DataFrame of estimated parameters.
-
-    """
-
-    n = len(data)
-
-    def loop(i):
-
-        df = data.drop(index=i)
-        return outcome(df)
-
-    jk_estimates = Parallel(n_jobs=n_cores)(delayed(loop)(i) for i in range(n))
-
-    return np.array(jk_estimates)
-
-
 def _eqf(sample):
     """Return empirical quantile function of the given sample.
 
     Args:
-        sample (pd.DataFrame): sample to base quantile function on.
+        sample (np.ndarray): Sample to base quantile function on.
 
     Returns:
-        f (callable): quantile function for given sample.
-
+        f (callable): Quantile function for given sample.
     """
 
     def f(x):
diff --git a/src/estimagic/inference/bootstrap_helpers.py b/src/estimagic/inference/bootstrap_helpers.py
index fc77f5c..06a20c8 100644
--- a/src/estimagic/inference/bootstrap_helpers.py
+++ b/src/estimagic/inference/bootstrap_helpers.py
@@ -1,19 +1,19 @@
 import pandas as pd
 
 
-def check_inputs(data, cluster_by=None, ci_method="percentile", alpha=0.05):
+def check_inputs(data=None, cluster_by=None, ci_method="percentile", ci_level=0.95):
     """Check validity of inputs.
+
     Args:
         data (pd.DataFrame): original dataset.
         cluster_by (str): column name of variable to cluster by.
         ci_method (str): method of choice for confidence interval computation.
         alpha (float): significance level of choice.
-
     """
 
-    ci_method_list = ["percentile", "bca", "bc", "t", "normal", "basic"]
+    ci_method_list = ["percentile", "bc", "t", "normal", "basic"]
 
-    if not isinstance(data, pd.DataFrame):
+    if not isinstance(data, pd.DataFrame) and data is not None:
         raise ValueError("Input 'data' must be DataFrame.")
 
     elif (cluster_by is not None) and (cluster_by not in data.columns.tolist()):
@@ -24,9 +24,9 @@ def check_inputs(data, cluster_by=None, ci_method="percentile", alpha=0.05):
     elif ci_method not in ci_method_list:
         raise ValueError(
             "ci_method must be 'percentile', 'bc',"
-            " 'bca', 't', 'basic' or 'normal', '{method}'"
+            " 't', 'basic' or 'normal', '{method}'"
             " was supplied".format(method=ci_method)
         )
 
-    elif alpha > 1 or alpha < 0:
-        raise ValueError("Input 'alpha' must be in [0,1].")
+    elif ci_level > 1 or ci_level < 0:
+        raise ValueError("Input 'ci_level' must be in [0,1].")
diff --git a/src/estimagic/inference/bootstrap_outcomes.py b/src/estimagic/inference/bootstrap_outcomes.py
index 04e4bd9..5fce516 100644
--- a/src/estimagic/inference/bootstrap_outcomes.py
+++ b/src/estimagic/inference/bootstrap_outcomes.py
@@ -1,6 +1,3 @@
-from functools import partial
-
-import pandas as pd
 from estimagic.batch_evaluators import process_batch_evaluator
 from estimagic.inference.bootstrap_helpers import check_inputs
 from estimagic.inference.bootstrap_samples import get_bootstrap_indices
@@ -9,7 +6,6 @@ from estimagic.inference.bootstrap_samples import get_bootstrap_indices
 def get_bootstrap_outcomes(
     data,
     outcome,
-    outcome_kwargs=None,
     cluster_by=None,
     seed=None,
     n_draws=1000,
@@ -22,7 +18,7 @@ def get_bootstrap_outcomes(
     Args:
         data (pandas.DataFrame): original dataset.
         outcome (callable): function of the dataset calculating statistic of interest.
-            Needs to return array-like object or pd.Series.
+            Returns a general pytree (e.g. pandas Series, dict, numpy array, etc.).
         cluster_by (str): column name of the variable to cluster by.
         seed (int): Random seed.
         n_draws (int): number of draws, only relevant if seeds is None.
@@ -35,16 +31,10 @@ def get_bootstrap_outcomes(
             as the estimagic batch_evaluators. See :ref:`batch_evaluators`.
 
     Returns:
-        estimates (pandas.DataFrame): Outcomes for different bootstrap samples. The
-            columns are the index of the result of ``outcome``.
-
+        estimates (list):  List of pytrees of estimated bootstrap outcomes.
     """
-    batch_evaluator = process_batch_evaluator(batch_evaluator)
-
     check_inputs(data=data, cluster_by=cluster_by)
-
-    if outcome_kwargs is not None:
-        outcome = partial(outcome, **outcome_kwargs)
+    batch_evaluator = process_batch_evaluator(batch_evaluator)
 
     indices = get_bootstrap_indices(
         data=data,
@@ -73,7 +63,6 @@ def _get_bootstrap_outcomes_from_indices(
     error_handling,
     batch_evaluator,
 ):
-
     arguments = [{"data": data, "indices": ind, "outcome": outcome} for ind in indices]
 
     raw_estimates = batch_evaluator(
@@ -87,9 +76,7 @@ def _get_bootstrap_outcomes_from_indices(
     estimates = [est for est in raw_estimates if not isinstance(est, str)]
     tracebacks = [est for est in raw_estimates if isinstance(est, str)]
 
-    if estimates:
-        estimates_df = pd.concat(estimates, axis=1).T
-    else:
+    if not estimates:
         msg = (
             "Calculating of all bootstrap outcomes failed. The tracebacks of the "
             "raised Exceptions are reproduced below:"
@@ -104,7 +91,7 @@ def _get_bootstrap_outcomes_from_indices(
             "anything but diagnostic purposes. Check warnings for more information. "
         )
 
-    return estimates_df
+    return estimates
 
 
 def _take_indices_and_calculate_outcome(indices, data, outcome):
diff --git a/src/estimagic/inference/shared.py b/src/estimagic/inference/shared.py
index 130a4eb..c9a337d 100644
--- a/src/estimagic/inference/shared.py
+++ b/src/estimagic/inference/shared.py
@@ -37,7 +37,6 @@ def transform_covariance(
             parameters. If parameters were fixed (explicitly or by other constraints),
             the index is a subset of params.index. The columns are the same as the
             index.
-
     """
     if isinstance(internal_cov, pd.DataFrame):
         internal_cov = internal_cov.to_numpy()
@@ -293,3 +292,11 @@ def calculate_p_values(flat_values, flat_standard_error):
     tvalues = flat_values / np.clip(flat_standard_error, 1e-300, np.inf)
     pvalues = 2 * scipy.stats.norm.sf(np.abs(tvalues))
     return pvalues
+
+
+def convert_flat_params_to_pytree(free_statistic, flat_params, converter):
+    helper = np.full(len(flat_params.values), np.nan)
+    helper[flat_params.free_mask] = free_statistic
+    out = converter.params_from_internal(helper)
+
+    return out


Test patch:
diff --git a/tests/inference/test_bootstrap.py b/tests/inference/test_bootstrap.py
index ddc0cd5..be3d9de 100644
--- a/tests/inference/test_bootstrap.py
+++ b/tests/inference/test_bootstrap.py
@@ -1,8 +1,10 @@
 import numpy as np
 import pandas as pd
 import pytest
+from estimagic.inference.bootstrap import bootstrap
 from estimagic.inference.bootstrap import bootstrap_from_outcomes
 from pandas.testing import assert_frame_equal as afe
+from pandas.testing import assert_series_equal as ase
 
 
 @pytest.fixture
@@ -13,8 +15,12 @@ def setup():
         np.array([[1, 10], [2, 7], [3, 6], [4, 5]]), columns=["x1", "x2"]
     )
 
-    x = np.array([[2.0, 8.0], [2.0, 8.0], [2.5, 7.0], [3.0, 6.0], [3.25, 5.75]])
-    out["estimates"] = pd.DataFrame(x, columns=["x1", "x2"])
+    y = np.array([[2.0, 8.0], [2.0, 8.0], [2.5, 7.0], [3.0, 6.0], [3.25, 5.75]])
+    out["estimates_arr"] = y
+    out["estimates_df"] = pd.DataFrame(y, columns=["x1", "x2"])
+    out["estimates_dict"] = {"x1": [2, 2, 2.5, 3, 3.25], "x2": [8, 8, 7, 6, 5.75]}
+    out["estimates_pytree"] = [pd.Series(row, index=["x1", "x2"]) for row in y]
+    out["estimates_pytree_x1"] = [pd.Series(row, index=["x1"]) for row in y[:, 0]]
 
     return out
 
@@ -24,27 +30,159 @@ def expected():
     out = {}
 
     z = np.array([[2.55, 0.5701, 2, 3.225], [6.95, 1.0665, 5.775, 8]])
+    cov = np.array([[0.325, -0.60625], [-0.60625, 1.1375]])
+    p_values = np.array([0.0, 0.0])
+    lower_ci = np.array([2, 5.775])
+    upper_ci = np.array([3.225, 8])
 
-    out["results"] = pd.DataFrame(
+    out["summary"] = pd.DataFrame(
         z, columns=["mean", "std", "lower_ci", "upper_ci"], index=["x1", "x2"]
     )
+    out["lower_ci"] = pd.Series(lower_ci, index=["x1", "x2"])
+    out["upper_ci"] = pd.Series(upper_ci, index=["x1", "x2"])
+    out["lower_ci_x1"] = pd.Series(lower_ci[0], index=["x1"])
+    out["upper_ci_x1"] = pd.Series(upper_ci[0], index=["x1"])
+    out["cov"] = pd.DataFrame(cov, columns=["x1", "x2"], index=["x1", "x2"])
+    out["se"] = pd.Series(np.sqrt(np.diagonal(cov)), index=["x1", "x2"])
+    out["p_values"] = pd.Series(p_values, index=["x1", "x2"])
+    out["p_value_x1"] = pd.Series(p_values[0], index=["x1"])
 
     return out
 
 
-def g(data):
-    return data.mean(axis=0)
+def _outcome_func(data, shift=0):
+    """Compute column means.
 
+    Args:
+        data (pd.Series or pd.DataFrame): The data set.
+        shift (float): Scalar that is added to the column means.
+
+    Returns:
+        pd.Series: Series where the k-th row corresponds to the mean
+            of the k-th column of the input data.
+    """
+    # Return pd.Series when .mean() is applied to a Series
+    # Only applying .mean() to a pd.Series would yield a float
+    return pd.DataFrame(data).mean(axis=0) + shift
 
-def test_bootstrap_from_outcomes(setup, expected):
 
-    results = bootstrap_from_outcomes(
+@pytest.mark.parametrize("shift", [0, 10, -10])
+def test_bootstrap_with_outcome_kwargs(shift, setup):
+    result = bootstrap(
         data=setup["df"],
-        outcome=g,
-        bootstrap_outcomes=setup["estimates"],
-        ci_method="percentile",
-    )["summary"]
-
-    # use rounding to adjust precision because there is no other way of handling this
-    # such that it is compatible across all supported pandas versions.
-    afe(results.round(2), expected["results"].round(2))
+        outcome=_outcome_func,
+        seed=123,
+        outcome_kwargs={"shift": shift},
+    )
+
+    expected = pd.Series([2.5, 7.0], index=["x1", "x2"])
+    ase(result.base_outcome, expected + shift)
+
+
+def test_bootstrap_from_outcomes(setup, expected):
+
+    result = bootstrap_from_outcomes(
+        base_outcome=_outcome_func(setup["df"]),
+        bootstrap_outcomes=setup["estimates_pytree"],
+    )
+
+    outcomes = result.outcomes()
+    lower_ci, upper_ci = result.ci()
+    covariance = result.cov()
+    standard_errors = result.se()
+
+    with pytest.raises(NotImplementedError) as error:
+        assert result._p_values()
+    assert str(error.value) == "Bootstrapped p-values are not implemented yet."
+
+    with pytest.raises(NotImplementedError) as error:
+        assert result._summary()
+    assert str(error.value) == "summary is not implemented yet."
+
+    # Use rounding to adjust precision and ensure reproducibility accross all
+    # supported pandas versions.
+    for i in range(len(outcomes)):
+        ase(outcomes[i], setup["estimates_pytree"][i])
+
+    ase(lower_ci.round(2), expected["lower_ci"].round(2))
+    ase(upper_ci.round(2), expected["upper_ci"].round(2))
+    afe(covariance.round(2), expected["cov"].round(2))
+    ase(standard_errors.round(2), expected["se"].round(2))
+
+
+def test_bootstrap_from_outcomes_private_methods(setup, expected):
+
+    result = bootstrap_from_outcomes(
+        base_outcome=_outcome_func(setup["df"]),
+        bootstrap_outcomes=setup["estimates_pytree"],
+    )
+
+    outcomes = result._outcomes
+    lower_ci, upper_ci = result._ci
+    covariance = result._cov
+    standard_errors = result._se
+
+    with pytest.raises(NotImplementedError) as error:
+        assert result._p_values()
+    assert str(error.value) == "Bootstrapped p-values are not implemented yet."
+
+    with pytest.raises(NotImplementedError) as error:
+        assert result._summary()
+    assert str(error.value) == "summary is not implemented yet."
+
+    for i in range(len(outcomes)):
+        ase(outcomes[i], setup["estimates_pytree"][i])
+
+    ase(lower_ci.round(2), expected["lower_ci"].round(2))
+    ase(upper_ci.round(2), expected["upper_ci"].round(2))
+    afe(covariance.round(2), expected["cov"].round(2))
+    ase(standard_errors.round(2), expected["se"].round(2))
+
+
+def test_bootstrap_from_outcomes_single_outcome(setup, expected):
+
+    result = bootstrap_from_outcomes(
+        base_outcome=_outcome_func(setup["df"]["x1"]),
+        bootstrap_outcomes=setup["estimates_pytree_x1"],
+    )
+
+    outcomes = result.outcomes()
+    lower_ci, upper_ci = result.ci()
+
+    for i in range(len(outcomes)):
+        ase(outcomes[i], setup["estimates_pytree_x1"][i])
+
+    ase(lower_ci.round(2), expected["lower_ci_x1"].round(2))
+    ase(upper_ci.round(2), expected["upper_ci_x1"].round(2))
+
+
+@pytest.mark.parametrize("input_type", ["arr", "df", "dict"])
+def test_wrong_input_types(input_type, setup):
+    with pytest.raises(TypeError):
+        assert bootstrap_from_outcomes(
+            base_outcome=_outcome_func(setup["df"]),
+            bootstrap_outcomes=setup["estimates_" + input_type],
+        )
+
+
+@pytest.mark.parametrize("return_type", ["array", "dataframe", "pytree"])
+def test_cov_correct_return_type(return_type, setup):
+    result = bootstrap_from_outcomes(
+        base_outcome=_outcome_func(setup["df"]),
+        bootstrap_outcomes=setup["estimates_pytree"],
+    )
+    _ = result.cov(return_type=return_type)
+
+
+def test_cov_wrong_return_type(setup):
+    result = bootstrap_from_outcomes(
+        base_outcome=_outcome_func(setup["df"]),
+        bootstrap_outcomes=setup["estimates_pytree"],
+    )
+
+    expected_msg = "return_type must be one of pytree, array, or dataframe, not dict."
+
+    with pytest.raises(TypeError) as error:
+        assert result.cov(return_type="dict")
+
+    assert str(error.value) == expected_msg
diff --git a/tests/inference/test_bootstrap_ci.py b/tests/inference/test_bootstrap_ci.py
index 6f57362..acbb1ef 100644
--- a/tests/inference/test_bootstrap_ci.py
+++ b/tests/inference/test_bootstrap_ci.py
@@ -1,10 +1,13 @@
+import itertools
+
 import numpy as np
 import pandas as pd
 import pytest
-from estimagic.inference.bootstrap_ci import _jackknife
 from estimagic.inference.bootstrap_ci import compute_ci
 from estimagic.inference.bootstrap_helpers import check_inputs
+from estimagic.parameters.tree_registry import get_registry
 from numpy.testing import assert_array_almost_equal as aaae
+from pybaum import tree_just_flatten
 
 
 @pytest.fixture
@@ -14,9 +17,9 @@ def setup():
     out["df"] = pd.DataFrame(
         np.array([[1, 10], [2, 7], [3, 6], [4, 5]]), columns=["x1", "x2"]
     )
-
-    x = np.array([[2.0, 8.0], [2.0, 8.0], [2.5, 7.0], [3.0, 6.0], [3.25, 5.75]])
-    out["estimates"] = pd.DataFrame(x, columns=["x1", "x2"])
+    out["estimates"] = np.array(
+        [[2.0, 8.0], [2.0, 8.0], [2.5, 7.0], [3.0, 6.0], [3.25, 5.75]]
+    )
 
     return out
 
@@ -38,83 +41,51 @@ def expected():
 
     out["bc_ci"] = np.array([[2, 3.2342835077057543], [5.877526959881923, 8]])
 
-    out["bca_ci"] = np.array([[2, 3.2058815797003826], [5.9404612332752915, 8]])
-
     out["t_ci"] = np.array([[1.775, 3], [6.0, 8.225]])
 
-    out["jk_estimates"] = np.array(
-        [
-            [3, 6],
-            [2.6666666666666665, 7],
-            [2.3333333333333335, 7.333333333333333],
-            [2, 7.666666666666667],
-        ]
-    )
-
     return out
 
 
-def g(data):
+def _outcome_fun_series(data):
     return data.mean(axis=0)
 
 
-def g_arr(data):
-    return np.array(data.mean(axis=0))
-
-
-def h(data):
-    return 2 * g(data)
-
-
-def h_arr(data):
-    return 2 * g_arr(data)
-
-
-def test_percentile_ci(setup, expected):
-    percentile_ci = compute_ci(
-        setup["df"], g, setup["estimates"], ci_method="percentile"
-    )
-    aaae(percentile_ci, expected["percentile_ci"])
-
-
-def test_normal_ci(setup, expected):
-    normal_ci = compute_ci(setup["df"], g, setup["estimates"], ci_method="normal")
-    aaae(normal_ci, expected["normal_ci"])
+def _outcome_func_dict(data):
+    return data.mean(axis=0).to_dict()
 
 
-def test_basic_ci(setup, expected):
-    basic_ci = compute_ci(setup["df"], g, setup["estimates"], ci_method="basic")
-    aaae(basic_ci, expected["basic_ci"])
-
+def _outcome_func_arr(data):
+    return np.array(data.mean(axis=0))
 
-def test_bc_ci(setup, expected):
-    bc_ci = compute_ci(setup["df"], g, setup["estimates"], ci_method="bc")
-    aaae(bc_ci, expected["bc_ci"])
 
+TEST_CASES = itertools.product(
+    [_outcome_fun_series, _outcome_func_dict, _outcome_func_arr],
+    ["percentile", "normal", "basic", "bc", "t"],
+)
 
-def test_bca_ci(setup, expected):
-    bca_ci = compute_ci(setup["df"], g, setup["estimates"], ci_method="bca")
-    aaae(bca_ci, expected["bca_ci"])
 
+@pytest.mark.parametrize("outcome, method", TEST_CASES)
+def test_ci(outcome, method, setup, expected):
+    registry = get_registry(extended=True)
 
-def test_t_ci(setup, expected):
-    t_ci = compute_ci(setup["df"], g, setup["estimates"], ci_method="t")
-    aaae(t_ci, expected["t_ci"])
+    def outcome_flat(data):
+        return tree_just_flatten(outcome(data), registry=registry)
 
+    base_outcome = outcome_flat(setup["df"])
+    lower, upper = compute_ci(base_outcome, setup["estimates"], ci_method=method)
 
-def test_jackknife(setup, expected):
-    jk_estimates = _jackknife(setup["df"], g)
-    aaae(jk_estimates, expected["jk_estimates"])
+    aaae(lower, expected[method + "_ci"][:, 0])
+    aaae(upper, expected[method + "_ci"][:, 1])
 
 
-def test_check_inputs_data(setup, expected):
+def test_check_inputs_data():
     data = "this is not a data frame"
     with pytest.raises(ValueError) as excinfo:
         check_inputs(data=data)
     assert "Input 'data' must be DataFrame." == str(excinfo.value)
 
 
-def test_check_inputs_cluster_by(setup, expected):
+def test_check_inputs_cluster_by(setup):
     cluster_by = "this is not a column name of df"
     with pytest.raises(ValueError) as excinfo:
         check_inputs(data=setup["df"], cluster_by=cluster_by)
@@ -123,20 +94,20 @@ def test_check_inputs_cluster_by(setup, expected):
     )
 
 
-def test_check_inputs_ci_method(setup, expected):
+def test_check_inputs_ci_method(setup):
     ci_method = 4
     with pytest.raises(ValueError) as excinfo:
         check_inputs(data=setup["df"], ci_method=ci_method)
     expected_msg = (
         "ci_method must be 'percentile', 'bc',"
-        f" 'bca', 't', 'basic' or 'normal', '{ci_method}'"
+        f" 't', 'basic' or 'normal', '{ci_method}'"
         f" was supplied"
     )
     assert str(excinfo.value) == expected_msg
 
 
-def test_check_inputs_alpha(setup, expected):
-    alpha = 666
+def test_check_inputs_ci_level(setup):
+    ci_level = 666
     with pytest.raises(ValueError) as excinfo:
-        check_inputs(data=setup["df"], alpha=alpha)
-    assert "Input 'alpha' must be in [0,1]." == str(excinfo.value)
+        check_inputs(data=setup["df"], ci_level=ci_level)
+    assert "Input 'ci_level' must be in [0,1]." == str(excinfo.value)
diff --git a/tests/inference/test_bootstrap_outcomes.py b/tests/inference/test_bootstrap_outcomes.py
index f802c1e..aa1da3e 100644
--- a/tests/inference/test_bootstrap_outcomes.py
+++ b/tests/inference/test_bootstrap_outcomes.py
@@ -4,11 +4,9 @@ import numpy as np
 import pandas as pd
 import pytest
 from estimagic.batch_evaluators import joblib_batch_evaluator
-from estimagic.inference.bootstrap_outcomes import (
-    _get_bootstrap_outcomes_from_indices,
-)
+from estimagic.inference.bootstrap_outcomes import _get_bootstrap_outcomes_from_indices
 from estimagic.inference.bootstrap_outcomes import get_bootstrap_outcomes
-from pandas.testing import assert_frame_equal as afe
+from numpy.testing import assert_array_almost_equal as aaae
 
 
 @pytest.fixture
@@ -17,10 +15,34 @@ def data():
     return df
 
 
-def test_get_bootstrap_estimates_runs(data):
+def _mean_return_series(data):
+    out = np.mean(data, axis=0)
+    return out
+
+
+def _mean_return_dict(data):
+    out = np.mean(data, axis=0)
+    return out.to_dict()
+
+
+def _mean_return_array(data):
+    out = np.mean(data, axis=0).to_numpy()
+    return out
+
+
+@pytest.mark.parametrize(
+    "outcome",
+    [
+        (functools.partial(np.mean, axis=0)),
+        (_mean_return_series),
+        (_mean_return_dict),
+        (_mean_return_array),
+    ],
+)
+def test_get_bootstrap_estimates_runs(outcome, data):
     get_bootstrap_outcomes(
         data=data,
-        outcome=functools.partial(np.mean, axis=0),
+        outcome=outcome,
         n_draws=5,
     )
 
@@ -35,8 +57,8 @@ def test_bootstrap_estimates_from_indices_without_errors(data):
         batch_evaluator=joblib_batch_evaluator,
     )
 
-    expected = pd.DataFrame([[3.0, 6.0], [2, 8]], columns=["x1", "x2"])
-    afe(calculated, expected)
+    expected = [[3.0, 6.0], [2, 8]]
+    aaae(calculated, expected)
 
 
 def test_get_bootstrap_estimates_with_error_and_raise(data):
@@ -69,7 +91,7 @@ def test_get_bootstrap_estimates_with_some_errors_and_continue(data):
         return data.mean()
 
     with pytest.warns(UserWarning):
-        res = get_bootstrap_outcomes(
+        res_flat = get_bootstrap_outcomes(
             data=data,
             outcome=_raise_assertion_error_sometimes,
             n_draws=100,
@@ -77,4 +99,4 @@ def test_get_bootstrap_estimates_with_some_errors_and_continue(data):
             seed=123,
         )
 
-    assert 30 <= len(res) <= 70
+    assert 30 <= len(res_flat) <= 70