Repository: OpenSourceEconomics/estimagic

Issue:
Improve Benchmarking 
### Problems:

1. The optimization histories for benchmarks are from a log database. The intention behind this was to get thread safety for optimizers that parallelize. However, when running benchmarks it is almost always a good idea not to parallelize on the level of function evaluations but on the benchmark problem level. 
2. The processing of optimization histories for the benchmark assumes that parameters are a DataFrame. It also does all calculations in pandas which might not be ideal. It would currently not run with general pytree parameters.

### Solution

1. Add an option to `maximize` and `minimize` that allows to collect the optimization histories (i.e. function evaluations and parameters) in memory. The mechanism can be similar to how the cache works, i.e. a list is partialled into `internal_criterion_and_derivative_template` and in each iteration, the parameters and function value are appended to that list. In contrast to the cache, this behavior is switched on by the user and not determined from the `algo_info`. If it is on, the histories are added to the optimization result. The following things should be available:
    - the scalar criterion value 
    - the raw criterion output 
    - the external params
2. Rewrite the processing of the optimization histories for general pytrees. Add a fast path in case that parameters are numpy arrays. Use mainly numpy, not pandas for the calculations.
3. Adjust convergence_plot 
4. Adjust profile_plot


Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) Report
> Merging [#333](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) (fe13d38) into [main](https://codecov.io/gh/OpenSourceEconomics/estimagic/commit/b97c488b1444c8784a4c65cced3dd45778866f34?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) (b97c488) will **decrease** coverage by `0.29%`.
> The diff coverage is `52.72%`.

```diff
@@            Coverage Diff             @@
##             main     #333      +/-   ##
==========================================
- Coverage   90.72%   90.42%   -0.30%     
==========================================
  Files         173      173              
  Lines       12825    12862      +37     
==========================================
- Hits        11635    11630       -5     
- Misses       1190     1232      +42     
```


| [Impacted Files](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) | Coverage Δ | |
|---|---|---|
| [...imagic/optimization/internal\_criterion\_template.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-c3JjL2VzdGltYWdpYy9vcHRpbWl6YXRpb24vaW50ZXJuYWxfY3JpdGVyaW9uX3RlbXBsYXRlLnB5) | `81.95% <ø> (ø)` | |
| [src/estimagic/benchmarking/run\_benchmark.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-c3JjL2VzdGltYWdpYy9iZW5jaG1hcmtpbmcvcnVuX2JlbmNobWFyay5weQ==) | `66.26% <52.72%> (-29.39%)` | :arrow_down: |
| [src/estimagic/logging/read\_log.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-c3JjL2VzdGltYWdpYy9sb2dnaW5nL3JlYWRfbG9nLnB5) | `69.56% <0.00%> (-18.85%)` | :arrow_down: |
| [...magic/optimization/subsolvers/\_trsbox\_quadratic.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-c3JjL2VzdGltYWdpYy9vcHRpbWl6YXRpb24vc3Vic29sdmVycy9fdHJzYm94X3F1YWRyYXRpYy5weQ==) | `84.08% <0.00%> (-0.82%)` | :arrow_down: |
| [src/estimagic/logging/database\_utilities.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-c3JjL2VzdGltYWdpYy9sb2dnaW5nL2RhdGFiYXNlX3V0aWxpdGllcy5weQ==) | `86.36% <0.00%> (-0.51%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics). Last update [b97c488...fe13d38](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/333?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics).


Patch:
diff --git a/src/estimagic/benchmarking/get_benchmark_problems.py b/src/estimagic/benchmarking/get_benchmark_problems.py
index c10c0e6..bdb5f3e 100644
--- a/src/estimagic/benchmarking/get_benchmark_problems.py
+++ b/src/estimagic/benchmarking/get_benchmark_problems.py
@@ -2,7 +2,6 @@ import warnings
 from functools import partial
 
 import numpy as np
-import pandas as pd
 from estimagic.benchmarking.cartis_roberts import CARTIS_ROBERTS_PROBLEMS
 from estimagic.benchmarking.more_wild import MORE_WILD_PROBLEMS
 from estimagic.benchmarking.noise_distributions import NOISE_DISTRIBUTIONS
@@ -54,7 +53,7 @@ def get_benchmark_problems(
             {"name": {"inputs": {...}, "solution": {...}, "info": {...}}}
             where "inputs" are keyword arguments for ``minimize`` such as the criterion
             function and start parameters. "solution" contains the entries "params" and
-            "value" and "info" might  contain information about the test problem.
+            "value" and "info" might contain information about the test problem.
 
     """
     raw_problems = _get_raw_problems(name)
@@ -127,9 +126,7 @@ def _create_problem_inputs(specification, additive_options, multiplicative_optio
     )
     _x = specification["start_x"]
 
-    _params = pd.DataFrame(_x.reshape(-1, 1), columns=["value"])
-
-    inputs = {"criterion": _criterion, "params": _params}
+    inputs = {"criterion": _criterion, "params": _x}
     return inputs
 
 
@@ -138,7 +135,7 @@ def _create_problem_solution(specification):
     if _solution_x is None:
         _solution_x = specification["start_x"] * np.nan
 
-    _params = pd.DataFrame(_solution_x.reshape(-1, 1), columns=["value"])
+    _params = _solution_x
     _value = specification["solution_criterion"]
 
     solution = {
@@ -151,8 +148,7 @@ def _create_problem_solution(specification):
 def _internal_criterion_template(
     params, criterion, additive_options, multiplicative_options
 ):
-    x = params["value"].to_numpy()
-    critval = criterion(x)
+    critval = criterion(params)
 
     noise = _get_combined_noise(
         critval,
@@ -169,6 +165,7 @@ def _internal_criterion_template(
         }
     else:
         out = noisy_critval
+
     return out
 
 
diff --git a/src/estimagic/benchmarking/process_benchmark_results.py b/src/estimagic/benchmarking/process_benchmark_results.py
index 4869527..e642cf1 100644
--- a/src/estimagic/benchmarking/process_benchmark_results.py
+++ b/src/estimagic/benchmarking/process_benchmark_results.py
@@ -1,5 +1,7 @@
 import numpy as np
 import pandas as pd
+from estimagic.parameters.tree_registry import get_registry
+from pybaum import tree_just_flatten
 
 
 def create_convergence_histories(
@@ -40,15 +42,16 @@ def create_convergence_histories(
             - parameter_distance_normalized
             - monotone_parameter_distance
             - monotone_parameter_distance_normalized
-
     """
     # get solution values for each problem
+    registry = get_registry(extended=True)
+    x_opt = {
+        name: tree_just_flatten(prob["solution"]["params"], registry=registry)
+        for name, prob in problems.items()
+    }
     f_opt = pd.Series(
         {name: prob["solution"]["value"] for name, prob in problems.items()}
     )
-    x_opt = {
-        name: prob["solution"]["params"]["value"] for name, prob in problems.items()
-    }
 
     # build df from results
     time_sr = _get_history_as_stacked_sr_from_results(results, "time_history")
@@ -136,6 +139,7 @@ def _get_history_of_the_parameter_distance(results, x_opt):
     x_dist_history = {}
     for (prob, algo), res in results.items():
         param_history = res["params_history"]
+
         x_dist_history[(prob, algo)] = pd.Series(
             np.linalg.norm(param_history - x_opt[prob], axis=1)
         )
diff --git a/src/estimagic/benchmarking/run_benchmark.py b/src/estimagic/benchmarking/run_benchmark.py
index f72587b..06d0b56 100644
--- a/src/estimagic/benchmarking/run_benchmark.py
+++ b/src/estimagic/benchmarking/run_benchmark.py
@@ -5,25 +5,22 @@ TO-DO:
     - finish medium scale problems from https://arxiv.org/pdf/1710.11005.pdf, Page 34.
     - add scalar problems from https://github.com/AxelThevenot
 - Add option for deterministic noise or wiggle.
-
 """
-from pathlib import Path
-
 import numpy as np
+import pandas as pd
 from estimagic import batch_evaluators
-from estimagic.logging.read_log import read_optimization_histories
 from estimagic.optimization.optimize import minimize
+from estimagic.parameters.tree_registry import get_registry
+from pybaum import tree_just_flatten
 
 
 def run_benchmark(
     problems,
     optimize_options,
-    logging_directory,
     *,
     batch_evaluator="joblib",
     n_cores=1,
     error_handling="continue",
-    fast_logging=True,
     seed=None,
 ):
     """Run problems with different optimize options.
@@ -42,48 +39,25 @@ def run_benchmark(
             Alternatively, the values can just be an algorithm which is then benchmarked
             at default settings.
         batch_evaluator (str or callable): See :ref:`batch_evaluators`.
-        logging_directory (pathlib.Path): Directory in which the log databases are
-            saved.
         n_cores (int): Number of optimizations that is run in parallel. Note that in
             addition to that an optimizer might parallelize.
         error_handling (str): One of "raise", "continue".
-        fast_logging (bool): Whether the slightly unsafe but much faster database
-            configuration is chosen.
 
     Returns:
         dict: Nested Dictionary with information on the benchmark run. The outer keys
             are tuples where the first entry is the name of the problem and the second
             the name of the optimize options. The values are dicts with the entries:
-            "runtime", "params_history", "criterion_history", "solution"
-
+            "params_history", "criterion_history", "time_history" and "solution".
     """
     np.random.seed(seed)
-    logging_directory = Path(logging_directory)
-    logging_directory.mkdir(parents=True, exist_ok=True)
 
     if isinstance(batch_evaluator, str):
         batch_evaluator = getattr(
             batch_evaluators, f"{batch_evaluator}_batch_evaluator"
         )
-
     opt_options = _process_optimize_options(optimize_options)
 
-    log_options = {"fast_logging": fast_logging, "if_table_exists": "replace"}
-
-    kwargs_list = []
-    names = []
-    for prob_name, problem in problems.items():
-        for option_name, options in opt_options.items():
-            kwargs = {
-                **options,
-                **problem["inputs"],
-                "logging": logging_directory / f"{prob_name}_{option_name}.db",
-                "log_options": log_options,
-            }
-            kwargs_list.append(kwargs)
-            names.append((prob_name, option_name))
-
-    log_paths = [kwargs["logging"] for kwargs in kwargs_list]
+    kwargs_list, names = _get_kwargs_list_and_names(problems, opt_options)
 
     raw_results = batch_evaluator(
         func=minimize,
@@ -93,20 +67,7 @@ def run_benchmark(
         unpack_symbol="**",
     )
 
-    results = {}
-    for name, result, log_path in zip(names, raw_results, log_paths):
-        histories = read_optimization_histories(log_path)
-        stop = histories["metadata"]["timestamps"].max()
-        start = histories["metadata"]["timestamps"].min()
-        runtime = (stop - start).total_seconds()
-
-        results[name] = {
-            "params_history": histories["params"],
-            "criterion_history": histories["values"],
-            "time_history": histories["metadata"]["timestamps"] - start,
-            "solution": result,
-            "runtime": runtime,
-        }
+    results = _get_results(names, raw_results, kwargs_list)
 
     return results
 
@@ -127,11 +88,63 @@ def _process_optimize_options(raw_options):
         if not isinstance(option, dict):
             option = {"algorithm": option}
 
-        if "log_options" in option:
-            raise ValueError(
-                "Log options cannot be specified as part of optimize_options. Logging "
-                "behavior is configured by the run_benchmark function."
-            )
         out_options[name] = option
 
     return out_options
+
+
+def _get_kwargs_list_and_names(problems, opt_options):
+    kwargs_list = []
+    names = []
+
+    for prob_name, problem in problems.items():
+        for option_name, options in opt_options.items():
+            kwargs = {**options, **problem["inputs"]}
+            kwargs_list.append(kwargs)
+            names.append((prob_name, option_name))
+
+    return kwargs_list, names
+
+
+def _get_results(names, raw_results, kwargs_list):
+    registry = get_registry(extended=True)
+    results = {}
+
+    for name, result, inputs in zip(names, raw_results, kwargs_list):
+
+        if isinstance(result, dict):
+            params_history = pd.DataFrame(
+                [
+                    tree_just_flatten(hist["params"], registry=registry)
+                    for hist in result["history"]
+                ]
+            )
+            criterion_history = pd.Series(
+                [hist["scalar_criterion"] for hist in result["history"]]
+            )
+
+            timestamps = np.array([hist["timestamp"] for hist in result["history"]])
+            start = timestamps[0]
+            time_history = pd.Series(timestamps - start)
+        elif isinstance(result, str):
+            _criterion = inputs["criterion"]
+
+            params_history = pd.DataFrame(
+                tree_just_flatten(inputs["params"], registry=registry)
+            ).T
+            criterion_history = pd.Series(_criterion(inputs["params"])["value"])
+
+            time_history = pd.Series([np.inf])
+        else:
+            raise ValueError(
+                "'result' object is expected to be of type 'dict' or 'str'."
+            )
+
+        results[name] = {
+            "params_history": params_history,
+            "criterion_history": criterion_history,
+            "time_history": time_history,
+            "solution": result,
+        }
+
+    return results
diff --git a/src/estimagic/optimization/internal_criterion_template.py b/src/estimagic/optimization/internal_criterion_template.py
index 2bf6406..3a02f1a 100644
--- a/src/estimagic/optimization/internal_criterion_template.py
+++ b/src/estimagic/optimization/internal_criterion_template.py
@@ -1,4 +1,5 @@
 import datetime
+import time
 import warnings
 
 from estimagic.differentiation.derivatives import first_derivative
@@ -67,7 +68,7 @@ def internal_criterion_and_derivative_template(
         numdiff_options (dict): Keyword arguments for the calculation of numerical
             derivatives. See :ref:`first_derivative` for details. Note that the default
             method is changed to "forward" for speed reasons.
-        logging (bool): Wether logging is used.
+        logging (bool): Whether logging is used.
         db_kwargs (dict): Dictionary with entries "database", "path" and "fast_logging".
         error_handling (str): Either "raise" or "continue". Note that "continue" does
             not absolutely guarantee that no error is raised but we try to handle as
@@ -248,8 +249,10 @@ def internal_criterion_and_derivative_template(
             "params": current_params,
             "criterion": new_criterion,
             "scalar_criterion": scalar_critval,
+            "timestamp": time.perf_counter(),
         }
         history_container.append(hist_entry)
+
     return res
 
 
@@ -310,7 +313,6 @@ def _log_new_evaluations(
     Note: There are some seemingly unnecessary type conversions because sqlalchemy
     can fail silently when called with numpy dtypes instead of the equivalent python
     types.
-
     """
     data = {
         "params": external_x,
diff --git a/src/estimagic/visualization/profile_plot.py b/src/estimagic/visualization/profile_plot.py
index 1e56456..61a895f 100644
--- a/src/estimagic/visualization/profile_plot.py
+++ b/src/estimagic/visualization/profile_plot.py
@@ -1,5 +1,3 @@
-import warnings
-
 import numpy as np
 import pandas as pd
 import plotly.express as px
@@ -67,13 +65,17 @@ def profile_plot(
 
     Returns:
         plotly.Figure
-
     """
-
     if stopping_criterion is None:
         raise ValueError(
             "You must specify a stopping criterion for the performance plot. "
         )
+    if runtime_measure not in ["walltime", "n_evaluations"]:
+        raise ValueError(
+            "Only 'walltime' or 'n_evaluations' are allowed as "
+            f"runtime_measure. You specified {runtime_measure}."
+        )
+
     df, converged_info = create_convergence_histories(
         problems=problems,
         results=results,
@@ -90,20 +92,8 @@ def profile_plot(
 
     if normalize_runtime:
         solution_times = solution_times.divide(solution_times.min(axis=1), axis=0)
-        # set again to inf because no inf Timedeltas were allowed.
         solution_times[~converged_info] = np.inf
-    else:
-        if (
-            runtime_measure == "walltime"
-            and (solution_times == pd.Timedelta(weeks=1000)).any().any()
-        ):
-            warnings.warn(
-                "Some algorithms did not converge. Their walltime has been "
-                "set to a very high value instead of infinity because Timedeltas do not"
-                "support infinite values."
-            )
-
-    # create performance profiles
+
     alphas = _determine_alpha_grid(solution_times)
     for_each_alpha = pd.concat(
         {alpha: solution_times <= alpha for alpha in alphas},
@@ -111,12 +101,8 @@ def profile_plot(
     )
     performance_profiles = for_each_alpha.groupby("alpha").mean().stack().reset_index()
 
-    # Build plot
-
     fig = px.line(performance_profiles, x="alpha", y=0, color="algorithm")
-    # dropped some styling parameters
 
-    # Plot Styling
     xlabels = {
         ("n_evaluations", True): "Multiple of Minimal Number of Function Evaluations\n"
         "Needed to Solve the Problem",
@@ -138,7 +124,7 @@ def profile_plot(
         template=template,
     )
 
-    fig.add_hline(y=1)  # dropped styling
+    fig.add_hline(y=1)
     return fig
 
 
@@ -164,25 +150,13 @@ def _create_solution_times(df, runtime_measure, converged_info):
     solution_times = df.groupby(["problem", "algorithm"])[runtime_measure].max()
     solution_times = solution_times.unstack()
 
-    # inf not allowed for timedeltas so put something very large
-    if runtime_measure == "walltime":
-        inf_value = pd.Timedelta(weeks=1000)
-    elif runtime_measure == "n_evaluations":
-        inf_value = np.inf
-    else:
-        raise ValueError(
-            "Only 'walltime' or 'n_evaluations' are allowed as "
-            f"runtime_measure. You specified {runtime_measure}."
-        )
-
-    solution_times[~converged_info] = inf_value
+    solution_times[~converged_info] = np.inf
     return solution_times
 
 
 def _determine_alpha_grid(solution_times):
     switch_points = _find_switch_points(solution_times=solution_times)
 
-    # add point to the right
     point_to_right = switch_points[-1] * 1.05
     extended_switch_points = np.append(switch_points, point_to_right)
     mid_points = (extended_switch_points[:-1] + extended_switch_points[1:]) / 2


Test patch:
diff --git a/tests/benchmarking/test_run_benchmark.py b/tests/benchmarking/test_run_benchmark.py
index fa4c5c5..965513e 100644
--- a/tests/benchmarking/test_run_benchmark.py
+++ b/tests/benchmarking/test_run_benchmark.py
@@ -1,8 +1,20 @@
+from functools import partial
+
+import numpy as np
+import pandas as pd
+import pytest
 from estimagic import get_benchmark_problems
+from estimagic.benchmarking.more_wild import linear_full_rank
+from estimagic.benchmarking.more_wild import linear_full_rank_solution_x
+from estimagic.benchmarking.more_wild import rosenbrock
 from estimagic.benchmarking.run_benchmark import run_benchmark
 
+# ======================================================================================
+# Internal benchmark suite
+# ======================================================================================
+
 
-def test_run_benchmark_dict_options(tmpdir):
+def test_run_benchmark_dict_options():
     all_problems = get_benchmark_problems("more_wild")
     first_two_names = list(all_problems)[:2]
     first_two = {name: all_problems[name] for name in first_two_names}
@@ -15,12 +27,9 @@ def test_run_benchmark_dict_options(tmpdir):
         },
     }
 
-    logging_directory = tmpdir / "benchmark_logs"
-
-    res = run_benchmark(
+    result = run_benchmark(
         problems=first_two,
         optimize_options=optimize_options,
-        logging_directory=logging_directory,
     )
 
     expected_keys = {
@@ -29,23 +38,18 @@ def test_run_benchmark_dict_options(tmpdir):
         ("linear_full_rank_good_start", "tuned_lbfgsb"),
         ("linear_full_rank_bad_start", "tuned_lbfgsb"),
     }
-
-    assert set(res) == expected_keys
+    assert set(result) == expected_keys
 
 
-def test_run_benchmark_list_options(tmpdir):
+def test_run_benchmark_list_options():
     all_problems = get_benchmark_problems("example")
     first_two_names = list(all_problems)[:2]
     first_two = {name: all_problems[name] for name in first_two_names}
-
     optimize_options = ["scipy_lbfgsb", "scipy_neldermead"]
 
-    logging_directory = tmpdir / "benchmark_logs"
-
-    res = run_benchmark(
+    result = run_benchmark(
         problems=first_two,
         optimize_options=optimize_options,
-        logging_directory=logging_directory,
     )
 
     expected_keys = {
@@ -54,5 +58,148 @@ def test_run_benchmark_list_options(tmpdir):
         ("linear_full_rank_good_start", "scipy_neldermead"),
         ("rosenbrock_good_start", "scipy_neldermead"),
     }
+    assert set(result) == expected_keys
+
+
+def test_run_benchmark_failing():
+    all_problems = get_benchmark_problems("more_wild")
+    failing_name = "jennrich_sampson"
+    failing = {failing_name: all_problems[failing_name]}
+
+    optimize_options = ["scipy_lbfgsb"]
+
+    with pytest.warns():
+        result = run_benchmark(problems=failing, optimize_options=optimize_options)
+
+    key = (failing_name, "scipy_lbfgsb")
+    assert isinstance(result[key]["solution"], str)
+
+
+# ======================================================================================
+# Custom benchmark functions
+# ======================================================================================
+
+
+def get_benchmark_problems_custom(raw_problems, internal_criterion):
+    problems = {}
+    for name, specification in raw_problems.items():
+        inputs = _create_problem_inputs_custom(specification, internal_criterion)
+
+        problems[name] = {
+            "inputs": inputs,
+            "solution": _create_problem_solution_custom(specification),
+            "info": specification.get("info", {}),
+        }
+
+    return problems
+
+
+def _create_problem_inputs_custom(specification, internal_criterion_func):
+    _criterion = partial(
+        internal_criterion_func,
+        criterion=specification["criterion"],
+    )
+
+    _params = specification["start_x"]
+    inputs = {"criterion": _criterion, "params": _params}
+
+    return inputs
+
+
+def _create_problem_solution_custom(specification):
+    _solution_x = specification.get("solution_x")
+    if _solution_x is None:
+        _solution_x = specification["start_x"] * np.nan
+
+    _value = specification["solution_criterion"]
 
-    assert set(res) == expected_keys
+    solution = {
+        "params": _solution_x,
+        "value": _value,
+    }
+    return solution
+
+
+def _internal_criterion_pandas(params, criterion):
+    x = params["value"].to_numpy()
+    critval = criterion(x)
+
+    out = {
+        "root_contributions": critval,
+        "value": critval @ critval,
+    }
+
+    return out
+
+
+def _internal_criterion_dict(params, criterion):
+    if "b" in params:
+        x = np.array([params["a"]] + params["b"].flatten().tolist())
+    else:
+        x = params["a"]
+    critval = criterion(x)
+
+    out = {
+        "root_contributions": critval,
+        "value": critval @ critval,
+    }
+
+    return out
+
+
+problems_pandas_input = {
+    "linear_full_rank_good_start": {
+        "criterion": partial(linear_full_rank, dim_out=45),
+        "start_x": pd.DataFrame(np.ones(9), columns=["value"]),
+        "solution_x": pd.DataFrame(linear_full_rank_solution_x, columns=["value"]),
+        "start_criterion": 72,
+        "solution_criterion": 36,
+    },
+    "rosenbrock_good_start": {
+        "criterion": rosenbrock,
+        "start_x": pd.DataFrame([-1.2, 1], columns=["value"]),
+        "solution_x": pd.DataFrame(np.ones(2), columns=["value"]),
+        "start_criterion": 24.2,
+        "solution_criterion": 0,
+    },
+}
+
+
+prolems_dict_input = {
+    "linear_full_rank_good_start": {
+        "criterion": partial(linear_full_rank, dim_out=45),
+        "start_x": {"a": 1, "b": np.ones((2, 2, 2))},
+        "solution_x": {
+            "a": linear_full_rank_solution_x[0],
+            "b": linear_full_rank_solution_x[1:].reshape(2, 2, 2),
+        },
+        "start_criterion": 72,
+        "solution_criterion": 36,
+    },
+    "rosenbrock_good_start": {
+        "criterion": rosenbrock,
+        "start_x": {"a": np.array([-1.2, 1])},
+        "solution_x": {"a": np.ones(2)},
+        "start_criterion": 24.2,
+        "solution_criterion": 0,
+    },
+}
+
+TEST_CASES = [
+    (problems_pandas_input, _internal_criterion_pandas),
+    (prolems_dict_input, _internal_criterion_dict),
+]
+
+
+@pytest.mark.parametrize("raw_problems, internal_criterion_func", TEST_CASES)
+def test_custom_benchmarks(raw_problems, internal_criterion_func):
+    problems = get_benchmark_problems_custom(raw_problems, internal_criterion_func)
+
+    optimize_options = ["scipy_lbfgsb"]
+    result = run_benchmark(problems=problems, optimize_options=optimize_options)
+
+    expected_keys = {
+        ("linear_full_rank_good_start", "scipy_lbfgsb"),
+        ("rosenbrock_good_start", "scipy_lbfgsb"),
+    }
+    assert set(result) == expected_keys
diff --git a/tests/visualization/test_convergence_plot.py b/tests/visualization/test_convergence_plot.py
index 3b4cc8a..1311e81 100644
--- a/tests/visualization/test_convergence_plot.py
+++ b/tests/visualization/test_convergence_plot.py
@@ -35,7 +35,6 @@ def test_convergence_plot_options(options, grid):
         problems,
         optimizers,
         n_cores=1,  # must be 1 for the test to work
-        logging_directory="logging",
     )
 
     convergence_plot(
diff --git a/tests/visualization/test_profile_plot.py b/tests/visualization/test_profile_plot.py
index 63b5100..7a60200 100644
--- a/tests/visualization/test_profile_plot.py
+++ b/tests/visualization/test_profile_plot.py
@@ -73,16 +73,16 @@ def test_create_solution_times_walltime():
     df = pd.DataFrame(
         columns=["problem", "algorithm", "n_evaluations", "walltime"],
         data=[
-            ["prob1", "algo1", 0, pd.Timedelta(seconds=0)],
-            ["prob1", "algo1", 1, pd.Timedelta(seconds=1)],
+            ["prob1", "algo1", 0, 0],
+            ["prob1", "algo1", 1, 1],
             #
-            ["prob1", "algo2", 2, pd.Timedelta(seconds=2)],
-            ["prob1", "algo2", 3, pd.Timedelta(seconds=3)],
+            ["prob1", "algo2", 2, 2],
+            ["prob1", "algo2", 3, 3],
             #
-            ["prob2", "algo1", 5, pd.Timedelta(seconds=5)],
+            ["prob2", "algo1", 5, 5],
             #
-            ["prob2", "algo2", 0, pd.Timedelta(seconds=0)],
-            ["prob2", "algo2", 1, pd.Timedelta(seconds=1)],
+            ["prob2", "algo2", 0, 0],
+            ["prob2", "algo2", 1, 1],
         ],
     )
     info = pd.DataFrame(
@@ -94,8 +94,8 @@ def test_create_solution_times_walltime():
     )
     expected = pd.DataFrame(
         {
-            "algo1": [pd.Timedelta(seconds=1), pd.Timedelta(seconds=5)],
-            "algo2": [pd.Timedelta(seconds=3), pd.Timedelta(weeks=1000)],
+            "algo1": [1, 5],
+            "algo2": [3, np.inf],
         },
         index=pd.Index(["prob1", "prob2"], name="problem"),
     )
@@ -130,7 +130,6 @@ def test_profile_plot_options(options):
         problems,
         optimizers,
         n_cores=1,  # must be 1 for the test to work
-        logging_directory="logging",
     )
 
     profile_plot(problems=problems, results=results, **options)