Repository: mad-lab-fau/tpcp (language: python)

Issue:
Optuna Optimizer



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bc059227..8b3bcb13 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- A new class to wrap the optimization framework [Optuna](https://optuna.readthedocs.io/en/stable/).
+  `CustomOptunaOptimize` can be used to create custom wrapper classes for various Optuna optimizations, that play 
+  nicely with `tpcp` and can be nested within tpcp operations. (https://github.com/mad-lab-fau/tpcp/pull/27)
+- A new example for the `CustomOptunaOptimize` wrapper that explains how to create complex custom optimizers using
+  `Optuna` and the new Scorer callbacks (see below) (https://github.com/mad-lab-fau/tpcp/pull/27)
 - `Scorer` now supports an optional callback function, which will be called after each datapoint is scored.
   (https://github.com/mad-lab-fau/tpcp/pull/29)
 - Pipelines, Optimize objects, and `Scorer` are now `Generic`. This improves typing (in particular with VsCode), but 
diff --git a/README.md b/README.md
index 7d9e94c2..ce9fdffa 100644
--- a/README.md
+++ b/README.md
@@ -71,12 +71,21 @@ We are using [poetry](https://python-poetry.org/) to manage dependencies and
 
 To set up the dev environment run the following commands:
 
-```
+```bash
 git clone https://github.com/mad-lab-fau/tpcp
 cd tpcp
 poetry install  -E torch # This might take a while
 ```
 
+To set up the dev environment *including* the required dependencies for using `tpcp` together with `optuna` 
+run the following commands: 
+```bash
+git clone https://github.com/mad-lab-fau/tpcp
+cd tpcp
+poetry install -E optuna  # This might take a while
+```
+
+
 Afterwards you can start to develop and change things.
 If you want to run tests, format your code, build the docs, ..., you can run one of the following `poethepoet` commands
 
@@ -92,7 +101,7 @@ CONFIGURED TASKS
 
 by calling
 
-```
+```bash
 poetry run poe <command name>
 ````
 
diff --git a/docs/conf.py b/docs/conf.py
index 5f39d333..fd01913e 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -147,6 +147,7 @@ def substitute(matchobj):
     "matplotlib": ("https://matplotlib.org/", None),
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None),
     "sklearn": ("https://scikit-learn.org/stable/", None),
+    "optuna": ("https://optuna.readthedocs.io/en/stable", None)
 }
 
 user_agent = "Mozilla/5.0 (X11; Linux x86_64; rv:25.0) Gecko/20100101 Firefox/25.0"
diff --git a/docs/modules/optimize.rst b/docs/modules/optimize.rst
index 85b7a410..533e38a2 100644
--- a/docs/modules/optimize.rst
+++ b/docs/modules/optimize.rst
@@ -18,4 +18,19 @@ Classes
     GridSearch
     GridSearchCV
     DummyOptimize
-    BaseOptimize
\ No newline at end of file
+    BaseOptimize
+
+Optuna Interface
+----------------
+
+.. automodule:: tpcp.optimize.optuna
+    :no-members:
+    :no-inherited-members:
+
+.. currentmodule:: tpcp.optimize
+
+.. autosummary::
+   :toctree: generated/optimize
+   :template: class_with_private.rst
+
+    optuna.CustomOptunaOptimize
\ No newline at end of file
diff --git a/examples/parameter_optimization/_04_custom_optuna_optimizer.py b/examples/parameter_optimization/_04_custom_optuna_optimizer.py
new file mode 100644
index 00000000..3c026776
--- /dev/null
+++ b/examples/parameter_optimization/_04_custom_optuna_optimizer.py
@@ -0,0 +1,472 @@
+r"""
+.. _custom_optuna_optimizer:
+
+Custom Optuna Optimizer
+=======================
+
+.. warning:: This example shows more advanced features of `tpcp` when using
+             `Optuna <https://optuna.readthedocs.io/en/stable/index.html>`_ for hyperparameter optimization.
+             To make this example understandable, you should make yourself familiar with Optuna first and understand
+             how it works, before trying to go through this example.
+
+The most popular method of (hyper-)parameter optimization is GridSearch (or GridSearchCV for optimizable pipelines).
+These methods perform an exhaustive search of the parameter space by simply testing every option.
+Considering that training and testing an algorithm can be very costly, exhaustive gridsearch takes a long time and is
+sometimes not feasible at all due to the required computational load.
+
+For these cases various alternatives exist like :class:`~sklearn.model_selection.RandomizedSearchCV`,
+:class:`~sklearn.model_selection.HalvingGridSearchCV`, or advanced blackbox optimizer like
+:class:`~optuna.samplers.TPESampler`.
+`tpcp` does not implement all of these methods explicitly, as it would simply be too much work.
+However, we try to make it relatively simple to bring such methods into the `tpcp` ecosystem by providing an interface
+for `Optuna <https://optuna.readthedocs.io/en/stable/index.html>`_.
+Optuna is a state-of-the-art hyperparameter optimization framework that allows to implement any of the methods
+mentioned above (and more), and allows to easily create custom samplers and pruners.
+
+However, Optuna uses a (very elegant) functional interface that does not play well with the sklearn-inspired
+interface of `tpcp`.
+Therefore, we provide the :class:`~tpcp.optimize.optuna.CustomOptunaOptimize` class which you can subclass to create
+your own Optuna based optimizer.
+
+.. note:: There is no need to create a custom subclass if you *only* want to run the hyperparameter optimization and
+          *not* nest the optimization into other `tpcp` methods like :func:`~tpcp.validate.cross_validate`.
+          For these cases, you can simply use Optuna with its default interface and just call the respective `tpcp`
+          methods in the objective function.
+
+In this example, we are going to create an optimized gridsearch using custom pruning that terminates trials early
+if we already realise at the first couple of participants that the parameter combination will not work well.
+
+Keep in mind that this example should merely demonstrate the possibility to integrate Optuna with `tpcp`.
+You are very much encouraged to read through the
+`Optuna documentation <https://optuna.readthedocs.io/en/stable/tutorial/index.html>`_ and create your
+own project-specific optimizers.
+"""
+
+# %%
+# The Prerequisites
+# -----------------
+# First, we need a dataset and a pipeline we want to optimize.
+# For this example we are using the `QRSDetector` pipeline (the non-trainable version) and the `ECGExampleData` dataset.
+# Check out the other examples to learn more about them.
+# We will simply copy the code over and create an instance of both objects to be used later.
+#
+# .. note:: We make pretty extensive use of Python's optional typing features (in particular generics) in this example.
+#           This can be a little overwhelming, and you might not need that in your implementation.
+#           So whenever, you see `TpcpClass[SomeClassName]` and you don't understand what it means, you can safely
+#           ignore it.
+#           But just for your understanding, if you see for example `Pipeline[ECGExampleData]` you should mentally
+#           read it as "A pipeline that requires a :class:`~tpcp.Dataset` of type `ECGExampleData` internally.
+#           Whenever you encounter a variable ending with a `T` (e.g. `PipelineT`), these are TypeVar types to type
+#           generics.
+#           You should read that as "Some subclass of :class:`~tpcp.Pipeline`, but we don't know which yet".
+#
+from pathlib import Path
+from typing import Any, Callable, Optional, Sequence, Tuple, Union
+
+import pandas as pd
+
+from examples.algorithms.algorithms_qrs_detection_final import QRSDetector
+from examples.datasets.datasets_final_ecg import ECGExampleData
+from tpcp import Parameter, Pipeline, cf
+
+try:
+    HERE = Path(__file__).parent
+except NameError:
+    HERE = Path(".").resolve()
+data_path = HERE.parent.parent / "example_data/ecg_mit_bih_arrhythmia/data"
+
+# The dataset
+example_data = ECGExampleData(data_path)
+
+
+class MyPipeline(Pipeline[ECGExampleData]):
+    algorithm: Parameter[QRSDetector]
+
+    r_peak_positions_: pd.Series
+
+    def __init__(self, algorithm: QRSDetector = cf(QRSDetector())):
+        self.algorithm = algorithm
+
+    def run(self, datapoint: ECGExampleData):
+        # Note: We need to clone the algorithm instance to make sure we don't leak any data between runs.
+        algo = self.algorithm.clone()
+        algo.detect(datapoint.data["ecg"], datapoint.sampling_rate_hz)
+
+        self.r_peak_positions_ = algo.r_peak_positions_
+        return self
+
+
+# The pipeline
+pipe = MyPipeline()
+
+# %%
+# What We Want To Do
+# ------------------
+# In the `GridSearch Example <grid_search>`__, we already performed a gridsearch using the
+# tpcp-:class:`~tpcp.optimize.GridSearch` class.
+# Here, we want to do something similar, but improve the gridsearch in two key aspects:
+#
+# 1. Instead of doing an exhaustive gridsearch, we use one of Optuna's advanced samplers
+# 2. When we encounter a parameter combination that doesn't work, we want to stop testing as early as possible to not
+#    waste any time on bad parameter combinations.
+#
+# We will start by implementing the first aspect and will then make some modifications to enable the second.
+#
+# The first thing we need for any gridsearch is a score function that tells us how good our parameter combination works.
+# In the `GridSearch Example <grid_search>`__ we used a scorer that returns *accuracy*, *precision* and *f1-score*.
+# We will use basically the same function here, but only return the f1-score, as this is the parameter we want to
+# optimize.
+# We could still calculate and return multiple other scores, but this would complicate the implementation of our
+# Optimizer and hence, is kept as exercise for the reader ;) .
+
+from examples.algorithms.algorithms_qrs_detection_final import match_events_with_reference
+
+
+def f1_score(pipeline: MyPipeline, datapoint: ECGExampleData) -> float:
+    # We use the `safe_run` wrapper instead of just run. This is always a good idea.
+    # We don't need to clone the pipeline here, as GridSearch will already clone the pipeline internally and `run`
+    # will clone it again.
+    pipeline = pipeline.safe_run(datapoint)
+    tolerance_s = 0.02  # We just use 20 ms for this example
+    matches_events, _ = match_events_with_reference(
+        pipeline.r_peak_positions_.to_numpy(),
+        datapoint.r_peak_positions_.to_numpy(),
+        tolerance=tolerance_s * datapoint.sampling_rate_hz,
+    )
+    n_tp = len(matches_events)
+    f1_score = (2 * n_tp) / (len(pipeline.r_peak_positions_) + len(datapoint.r_peak_positions_))
+    return f1_score
+
+
+from optuna import Study, Trial
+
+# %%
+# The Custom Optimizer
+# --------------------
+# Optimizers in `tpcp` are nothing magical – they are simply algorithms that take a pipeline as input parameter and
+# have an action method called `optimize` that takes in a dataset and then optimizes some parameters of the passed
+# pipeline using this data.
+#
+# The :class:`~tpcp.optimize.optuna.CustomOptunaOptimize` class already implements most of that for us and simply
+# requires us to implement our objective function for the optimization like we would need for Optuna anyway.
+#
+# Here, we define the objective function within the `create_objective` method of our custom optimizer and return it.
+# The objective we define here is slightly different from the pure Optuna objective function, as it also takes a
+# :class:`~tpcp.Pipeline` and a :class:`~tpcp.Dataset` as input in addition to the trial-object.
+#
+# The content of our objective function is very similar to our score function, but we do not expect just a single
+# datapoint, but an entire dataset. Also, we need to handle getting and applying our parameters within the objective
+# function.
+#
+# Because we define the function nested within another method, we have access to all class parameters.
+# Hence, if we want to add certain configurations to our objective, we can add parameters to the
+# :class:`~tpcp.optimize.Optimizer` itself and then access it in the objective function.
+#
+# For the :class:`~tpcp.optimize.Optimizer` we want to build we primarily need two custom pieces of configuration:
+#
+# 1. The score function we want to use. We want to make that configurable and not hard-code "f1-score" into our
+#    optimizer.
+# 2. The search space for the parameter search. In Optuna the search space is defined by calls to methods on a
+#    :class:`optuna.Trial` object. Therefore, we take in a callable that gets the trial object passed and returns
+#    the selected parameters. You will see how this works later on.
+#
+# With these two pieces of configuration in place our objective needs to simply do four things:
+# 1. First, we need to call the search space function to get the parameters.
+# 2. Then, apply these parameters to our pipeline.
+# 3. Afterwards, we need to calculate how good the pipeline with the new parameters works for each of the datapoints
+#    within our test dataset.
+# 4. Finally, we return the aggregated score.
+#
+# To avoid writing our own for-loop (for now) for the third step, we use :class:`~tpcp.validate.Scorer` with our
+# custom score function.
+# :class:`~tpcp.validate.Scorer` handles looping and aggregating results over multiple datapoints.
+#
+# With that, our implementation looks as follows:
+from tpcp.optimize.optuna import CustomOptunaOptimize
+from tpcp.types import Dataset_, Pipeline_
+from tpcp.validate import Scorer
+
+
+class OptunaSearch(CustomOptunaOptimize[Pipeline_, Dataset_]):
+    create_search_space: Parameter[Callable[[Trial], None]]
+    score_function: Parameter[Callable[[Pipeline_, Dataset_], float]]
+
+    def __init__(
+        self,
+        pipeline: Pipeline_,
+        study: Study,
+        create_search_space: Callable[[Trial], None],
+        score_function: Callable[[Pipeline_, Dataset_], float],
+        *,
+        n_trials: Optional[int] = None,
+        timeout: Optional[float] = None,
+        return_optimized: bool = True,
+    ) -> None:
+        self.create_search_space = create_search_space
+        self.score_function = score_function
+        super().__init__(pipeline, study, n_trials=n_trials, timeout=timeout, return_optimized=return_optimized)
+
+    def create_objective(self) -> Callable[[Trial, Pipeline_, Dataset_], Union[float, Sequence[float]]]:
+        # Here we define our objective function
+
+        def objective(trial: Trial, pipeline: Pipeline_, dataset: Dataset_) -> float:
+            # First we need to select parameters for the current trial
+            self.create_search_space(trial)
+            # Then we apply these parameters to the pipeline
+            pipeline = pipeline.set_params(**trial.params)
+
+            # We wrap the score function with a scorer to avoid writing our own for-loop to aggregate the results.
+            scorer = Scorer(self.score_function)
+
+            # In the end, we calculate the results per datapoint.
+            # Note that we could expose the `error_score` parameter on an optimizer level.
+            # But let's keep it simple for now.
+            average_score, single_scores = scorer(pipeline, dataset, error_score="raise")
+
+            # As a bonus, we use the custom params option of optuna to store the individual scores per datapoint and the
+            # respective data labels
+            trial.set_user_attr("single_scores", single_scores)
+            trial.set_user_attr("data_labels", dataset.groups)
+
+            return average_score
+
+        return objective
+
+
+# %%
+# Running the optimization
+# ------------------------
+# To run the optimization, we need to create a new Optuna study, a custom sampler and the function that defines our
+# search space:
+from optuna import create_study, samplers
+
+# We use a simple RandomSampler, but every optuna sampler will work
+sampler = samplers.RandomSampler(seed=42)
+# We use a simple in-memory study with the direction "maximize", as we want to optimize for the highest f1-score
+study = create_study(direction="maximize", sampler=sampler)
+
+# %%
+# The search space function requires a little more explanation:
+# In Optuna, we can use the `suggest_...` methods on a trial to get a new value within a given range.
+# This uses our sampler in the background to suggest a new value that makes sense based on the trials that are already
+# completed.
+# The selected parameters are stored in the trial object so that we can access them after the function was called.
+#
+# We use the names of the parameters we want to modify in our pipeline (using the `__` for nested values).
+# This makes applying the parameters to the pipeline later on easy.
+
+
+def create_search_space(trial: Trial):
+    trial.suggest_float("algorithm__min_r_peak_height_over_baseline", 0.1, 2, step=0.1)
+    trial.suggest_float("algorithm__high_pass_filter_cutoff_hz", 0.1, 2, step=0.1)
+
+
+# %%
+# Finally, we are ready to run the pipeline.
+# We create a new instance and set the stopping criteria (in this case 10 random trials).
+# Then we can use the familiar :class:`~tpcp.optimize.Optimize` interface to run everything.
+opti = OptunaSearch(
+    pipe,
+    study,
+    create_search_space=create_search_space,
+    score_function=f1_score,
+    n_trials=10,
+)
+
+opti = opti.optimize(example_data)
+print(
+    f"The best performance was achieved with the parameters {opti.best_params_} and an f1-score of {opti.best_score_}."
+)
+
+# %%
+# We can use `opti.search_results_` to get a full overview over all results.
+# These parameters and parameter names are slightly modified compared to the normal Optuna output, to make it similar
+# to the output of :class:`~tpcp.optimize.GridSearch`.
+pd.DataFrame(opti.search_results_)
+
+# %%
+# If you need even more insides, you can access the study object directly.
+opti.study_
+
+# %%
+# And like with all Optimizers, we can access the `optimized_pipeline_` and call `run` directly on the Optimizer.
+
+opti.optimized_pipeline_
+# %%
+out_pipe = opti.run(example_data[0])
+out_pipe.r_peak_positions_
+
+# %%
+# With this we created a simple random search optimizer (or grid search, or whatever sampler we want to use) using
+# Optuna.
+# By using :class:`~tpcp.optimize.optuna.CustomOptunaOptimize` we get compatibility with the `tpcp` optimizer interface.
+# This means we could throw this optimizer into :func:`~tpcp.validate.cross_validate` and things would just work.
+#
+# A step further: Custom Pruning
+# ------------------------------
+# Simply to demonstrate the power of having full access to all Optuna features, we will implement a custom pruner
+# that stops testing a trial when one datapoint scores below a certain threshold.
+# The idea is that when we iterate through the datapoints and process them one by one, and find one datapoint where
+# the performance is really bad, we already know that this will not be our best choice/a choice we want to use.
+# Hence, there is no need to compute scores for the remaining datapoints.
+#
+# In Optuna we can implement this using a custom Pruner and the *callback* feature of the
+# :class:`~tpcp.validate.Scorer` class.
+# The Pruner will be called everytime we report a new result and will tell us if we should stop evaluating the trial.
+#
+# .. note:: This a unusual usage of pruning. Usually, pruning is used to stop after a certain number of training
+#           epochs of an ML classifier and not to stop half-way through evaluating your dataset.
+#           But it works and is practical.
+#
+# To create our custom pruner we need a new class, sub-classing :class:`~optuna.pruners.BasePruner`.
+# Then we implement a `prune` method that simply checks if the current intermediate value is below a certain threshold.
+# If yes we return True, telling optuna that the trial can be pruned.
+
+from optuna.pruners import BasePruner
+from optuna.study.study import Study
+from optuna.trial import FrozenTrial
+
+
+class MinDatapointPerformancePruner(BasePruner):
+    def __init__(self, min_performance: float):
+        self.min_performance = min_performance
+
+    def prune(self, _: Study, trial: FrozenTrial) -> bool:
+        step = trial.last_step
+
+        if step is not None:
+            score = trial.intermediate_values[step]
+            if score < self.min_performance:
+                return True
+        return False
+
+
+# %%
+# Afterwards, we need to modify our optimizer to work with the pruner.
+# We need to report each calculated value from each datapoint to Optuna as soon as it was calculated and not wait
+# until we ran through the entire dataset.
+# We can do that by passing a callback function to the `Scorer`.
+# This callback will be called after each datapoint is evaluated and allows us to access the most recent score.
+#
+# We define this callback within the objective function to have access to the trial object of the outer scope.
+# Using the `trial` object, we can report the most recent score to Optuna using `trial.report`.
+# This will call the pruner and allows us to check afterwards, if the trial should be pruned.
+# We then write some debug information and end the trial by raising a :class:`~optuna.TrialPruned` exception.
+
+from optuna import TrialPruned
+
+
+class OptunaSearchEarlyStopping(CustomOptunaOptimize[Pipeline_, Dataset_]):
+    create_search_space: Parameter[Callable[[Trial], None]]
+    score_function: Parameter[Callable[[Pipeline_, Dataset_], float]]
+
+    def __init__(
+        self,
+        pipeline: Pipeline_,
+        study: Study,
+        create_search_space: Callable[[Trial], None],
+        score_function: Callable[[Pipeline_, Dataset_], float],
+        *,
+        n_trials: Optional[int] = None,
+        timeout: Optional[float] = None,
+        return_optimized: bool = True,
+    ) -> None:
+        self.create_search_space = create_search_space
+        self.score_function = score_function
+        super().__init__(pipeline, study, n_trials=n_trials, timeout=timeout, return_optimized=return_optimized)
+
+    def create_objective(self) -> Callable[[Trial, Pipeline_, Dataset_], Union[float, Sequence[float]]]:
+        def objective(trial: Trial, pipeline: Pipeline_, dataset: Dataset_) -> float:
+            # First, we need to select parameters for the current trial
+            self.create_search_space(trial)
+            # Then, we apply these parameters to the pipeline
+            pipeline = pipeline.set_params(**trial.params)
+
+            def single_score_callback(*, step: int, dataset: Dataset_, scores: Tuple[float, ...], **_: Any):
+                # We need to report the new score value.
+                # This will call the pruner internally and then tell us if we should stop
+                trial.report(float(scores[step]), step)
+                if trial.should_prune():
+                    # Apparently, our last value was bad, and we should abort.
+                    # However, before we do so, we will save the scores so far as debug information
+                    trial.set_user_attr("single_scores", scores)
+                    trial.set_user_attr("data_labels", dataset[: step + 1].groups)
+                    # And, finally, we abort the trial
+                    raise TrialPruned(
+                        f"Pruned at datapoint {step} ({dataset[step].groups[0]}) with value " f"{scores[step]}."
+                    )
+
+            # We wrap the score function with a Scorer object to avoid writing our own for-loop to aggregate the
+            # results. We pass our callback and `trial` which is passed as a generic kwarg to scorer and hence can be
+            # accessed from within our callback.
+            scorer = Scorer(self.score_function, single_score_callback=single_score_callback)
+
+            # Calculate the results per datapoint.
+            # Note that we could expose the `error_score` parameter on an optimizer level.
+            # But let's keep it simple for now.
+            average_score, single_scores = scorer(pipeline, dataset, error_score="raise")
+
+            # As a bonus, we use the custom params option of Optuna to store the individual scores per datapoint and the
+            # respective data labels.
+            trial.set_user_attr("single_scores", single_scores)
+            trial.set_user_attr("data_labels", dataset.groups)
+
+            return average_score
+
+        return objective
+
+
+# %%
+# Running the new Optimizer stays the same (we even reuse the search space).
+# We only need to add an instance of our pruner to the study.
+sampler = samplers.RandomSampler(seed=42)
+study = create_study(direction="maximize", sampler=sampler, pruner=MinDatapointPerformancePruner(0.3))
+
+opti_early_stop = OptunaSearchEarlyStopping(
+    pipe,
+    study,
+    create_search_space=create_search_space,
+    score_function=f1_score,
+    n_trials=10,
+)
+
+opti_early_stop.optimize(example_data)
+
+# %%
+# And then we can inspect the output.
+# Compared to our previous run, we can see that many trials report NaN as score and "PRUNED" in the `state` column.
+# For each of these values we saved some time.
+# For the other trials, we get the same results as earlier.
+pd.DataFrame(opti_early_stop.search_results_)
+
+# %%
+# Summary
+# -------
+# The tpcp <-> Optuna interface is a little bit more low-level than many other tpcp features.
+# Therefore, here is a short summary of the steps you need:
+#
+# 1. Create a custom optimizer than inherits from `CustomOptunaOptimize`
+# 2. Overwrite the `create_objective` method so that it returns a Callable.
+# 3. The returned callable should expect a :class:`~optuna.Trial`, a :class:`~tpcp.Pipeline`,
+#    and a :class:`~tpcp.Dataset` object as input.
+#    Otherwise, it is identical to the objective function you would write in "plain" Optuna, and hence, should only
+#    return a single cost value for the optimization.
+# 4. If your objective function requires parameter, add them as class attributes via the init.
+# 5. (optional) If you want to report additional values from your optimization, you can do that via the
+#    `set_user_attr` parameter of the :class:`~optuna.Trial` object.
+# 6. (optional) Early stopping and other Pruners can be implemented identical to Optuna.
+#    Using the callback option of :class:`~tpcp.validate.Scorer` you can even hook into the datapoint iteration to
+#    trigger early stopping.
+
+
+# %%
+# Next steps
+# ----------
+# Building a custom optimizer is a little more involved than just using `GridSearch`. However, it allows great
+# flexibility with relatively small overhead compared to a pure implementation in Optuna.
+#
+# In this example we created an objective function that only makes sense for pipelines that don't have an internal
+# optimization.
+# However, instead of just a simple search, you could also create a cross-validation-based search by using
+# :func:`~tpcp.validate.cross_validate` within your objective to split the passed data into multiple train-test sets
+# and optimize hyperparameters similar to :class:`~tpcp.optimize.GridSearchCV`.
diff --git a/pyproject.toml b/pyproject.toml
index 2ee4cc2f..130dda32 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,9 +21,11 @@ pandas = "^1"
 tqdm = "^4.62.3"
 typing-extensions = ">=4.1.1"
 torch = {version = ">=1.6.0", optional = true}
+optuna = {version = "^2.10.0", optional = true}
 
 [tool.poetry.extras]
 torch = ["torch"]
+optuna = ["optuna"]
 
 [tool.poetry.dev-dependencies]
 pytest = "^6.2.1"
diff --git a/tpcp/optimize/optuna.py b/tpcp/optimize/optuna.py
new file mode 100644
index 00000000..d6eea7ae
--- /dev/null
+++ b/tpcp/optimize/optuna.py
@@ -0,0 +1,336 @@
+"""Implementation of methods and classes to wrap the optimization Framework `Optuna`."""
+import warnings
+
+try:
+    import optuna  # noqa: unused-import
+except ImportError as e:
+    raise ImportError(
+        "To use the tpcp Optuna interface, you first need to install optuna (`pip install optuna`)"
+    ) from e
+
+from typing import Any, Callable, Dict, List, Optional, Sequence, TypeVar, Union
+
+import numpy as np
+from optuna import Study, Trial
+from optuna.structs import FrozenTrial
+from optuna.study.study import ObjectiveFuncType
+from typing_extensions import Self
+
+from tpcp import OptimizablePipeline, clone
+from tpcp._dataset import Dataset_
+from tpcp._optimize import BaseOptimize
+from tpcp._pipeline import Pipeline_
+from tpcp.optimize import Optimize
+
+__all__ = ["CustomOptunaOptimize"]
+
+warnings.warn(
+    "The Optuna interface in tpcp is still experimental and we are testing if the workflow makes "
+    "sense even for larger projects. "
+    "This means, the interface for `CustomOptunaOptimize` will likely change in the future.",
+    UserWarning,
+)
+
+CustomOptunaOptimize_ = TypeVar("CustomOptunaOptimize_", bound="CustomOptunaOptimize")
+
+
+class CustomOptunaOptimize(BaseOptimize[Pipeline_, Dataset_]):
+    """Base class for custom Optuna optimizer.
+
+    This provides a relatively simple tpcp compatible interface to Optuna.
+    You basically need to subclass this class and implement the `create_objective` method to return the objective
+    function you want to optimize.
+    The only difference to a normal objective function in Optuna is, that your objective here should expect a
+    pipeline and a dataset object as second and third argument (see Example).
+    If there are parameters you want to make customizable (e.g. which metric to optimize for), expose them in the
+    `__init__` of your subclass.
+
+    Depending on your usecase, your custom optimizers can be single use with a bunch of "hard-coded" logic, or you can
+    try to make them more general, by exposing certain configurability.
+
+    Parameters
+    ----------
+    pipeline
+        A tpcp pipeline with some hyper-parameters that should be optimized.
+        This can either be a normal pipeline or an optimizable-pipeline.
+        This fully depends on your implementation of the `create_objective` method.
+    study
+        The optuna :class:`~optuna.Study` that should be used for optimization.
+    n_trials
+        The number of trials.
+        If this argument is set to :obj:`None`, there is no limitation on the number of trials.
+        In this case you should use :obj:`timeout` instead.
+        Because optuna is called internally by this wrapper, you can not setup a study without limits and end it
+        using CTRL+C (as suggested by the Optuna docs).
+        In this case the entire execution flow would be stopped.
+    timeout
+        Stop study after the given number of second(s).
+        If this argument is set to :obj:`None`, the study is executed without time limitation.
+        In this case you should use :obj:`n_trials` to limit the execution.
+    return_optimized
+        If True, a pipeline object with the overall best parameters is created and re-optimized using all provided data
+        as input.
+        The optimized pipeline object is stored as `optimized_pipeline_`.
+        How the "re-optimization" works depends on the type of pipeline provided.
+        If it is a simple pipeline, no specific re-optimization will be perfomed and `optimized_pipeline_` will simply
+        be an instance of the pipeline with the best parameters indentified in the search.
+        When `pipeline` is a subclass of `OptimizablePipeline`, we attempt to call `pipeline.self_optimize` with the
+        entire dataset provided to the `optimize` method.
+        The result of this self-optimization will be set as `optimized_pipeline`.
+        If this behaviour is undesired, you can overwrite the `return_optimized_pipeline` method in subclass.s
+    callbacks
+        List of callback functions that are invoked at the end of each trial.
+        Each function must accept two parameters with the following types in this order:
+        :class:`~optuna.study.Study` and :class:`~optuna.FrozenTrial`.
+    show_progress_bar
+        Flag to show progress bars or not.
+    gc_after_trial
+        Run the garbage collerctor after each trial.
+        Check the optuna documentation for more detail
+
+    Other Parameters
+    ----------------
+    dataset
+        The dataset instance passed to the optimize method
+
+    Attributes
+    ----------
+    search_results_
+        A dictionary containing all relevant results of the parameter search.
+        The format of this dictionary is designed to be directly passed into the `pd.DataFrame` constructor.
+        Each column then represents the result for one set of parameters.
+
+        The dictionary contains the following entries:
+
+        score
+            The value of the score for a specific trial.
+            If a trial was pruned this value is nan.
+        param_{parameter_name}
+            The value of a respective parameter.
+        params
+            A dictionary representing all parameters.
+        state
+            Whether the trial was completed, pruned, or any type of error occured
+        user_attrs
+        datetime_start
+            When the trial was started
+        datetime_complete
+            When the trial endend
+        duration
+            The duration of the trial
+        user_attrs_...
+            User attributes set within the objective function
+        system_attrs_...
+            System attributes set internally by optuna (usually empty)
+
+        If you need access to further parameter, inspect `self.study_` directly.
+    optimized_pipeline_
+        An instance of the input pipeline with the best parameter set.
+        This is only available if `return_optimized` is not False.
+    best_params_
+        The parameter combination identified in the study
+    best_score_
+        The score achieved in the best trial
+    best_trial_
+        The trial object that resulted in the best results
+    study_
+        The study object itself.
+        This should usually be identical to `self.study`.
+
+    Examples
+    --------
+    >>> from tpcp.validate import Scorer
+    >>> from optuna import create_study
+    >>> from optuna import samplers
+    >>>
+    >>> class MyOptunaOptimizer(CustomOptunaOptimize):
+    ...     def create_objective(self):
+    ...         def objective(trial: Trial, pipeline: Pipeline, dataset: Dataset):
+    ...             trial.suggest_float("my_pipeline_para", 0, 3)
+    ...             mean_score = Scorer(lambda dp: pipeline.score(dp))
+    ...             return mean_score
+    ...         return objective
+    >>>
+    >>> study = create_study(sampler=samplers.RandomSampler())
+    >>> opti = MyOptunaOptimizer(pipeline=MyPipeline(), study=study, n_trials=10)
+    >>> opti = opti.optimize(MyDataset())
+
+    Notes
+    -----
+    As this wrapper attempts to fully encapsule all Optuna calls to make it possible to be run seamlessly in a
+    cross-validation (or similar), you can not start multiple optuna optimizations at the same time which is the
+    preffered way of multi-processing for optuna.
+    In result, you are limited to single-process operations.
+    If you want to get "hacky" you can try the approach suggested
+    `here <https://github.com/optuna/optuna/issues/2862>`__ to create a study that uses joblib for internal
+    multiprocessing.
+
+    """
+
+    pipeline: Pipeline_
+    study: Study
+
+    return_optimized: bool
+
+    # Optuna Parameters that are directly forwarded to study.optimize
+    n_trials: Optional[int]
+    timeout: Optional[float]
+    callbacks: Optional[List[Callable[[Study, FrozenTrial], None]]]
+    gc_after_trial: bool
+    show_progress_bar: bool
+
+    optimized_pipeline_: Pipeline_
+    study_: Study
+
+    def __init__(
+        self,
+        pipeline: Pipeline_,
+        study: Study,
+        *,
+        n_trials: Optional[int] = None,
+        timeout: Optional[float] = None,
+        callbacks: Optional[List[Callable[[Study, FrozenTrial], None]]] = None,
+        gc_after_trial: bool = False,
+        show_progress_bar: bool = False,
+        return_optimized: bool = True,
+    ) -> None:  # noqa: super-init-not-called
+        self.pipeline = pipeline
+        self.study = study
+
+        self.n_trials = n_trials
+        self.timeout = timeout
+        self.callbacks = callbacks
+        self.gc_after_trial = gc_after_trial
+        self.show_progress_bar = show_progress_bar
+
+        self.return_optimized = return_optimized
+
+    @property
+    def search_results_(self) -> Dict[str, Sequence[Any]]:
+        """Detailed results of the study.
+
+        This basically contains the same information as `self.study_.trials_dataframe()`, with some small modifications:
+
+        - columns starting with "params_" are renamed to "param_"
+        - a new column called "params" containing all parameters as dict is added
+        - "value" is renamed to score"
+        - the score of pruned trials is set to `np.nan`
+
+        These changes are made to make the output comparable to the output of :class:`~tpcp.optimize.GridSearch` and
+        :class:`~tpcp.optimize.GridSearchCV`.
+        """
+
+        def rename_param_columns(name: str):
+            param_start = "params_"
+            if name.startswith(param_start):
+                return "param_" + name[len(param_start) :]
+            return name
+
+        base_df = (
+            self.study_.trials_dataframe()
+            .drop("number", axis=1)
+            .rename(columns={"value": "score"})
+            .rename(columns=rename_param_columns)
+        )
+        base_df["params"] = [t.params for t in self.study_.trials]
+
+        # If a trial is pruned we set the score to nan.
+        # This is clearer than showing the last step value.
+        # If this is required people should report them as user values in their objective function
+        base_df.loc[base_df["state"] == "PRUNED", "score"] = np.nan
+
+        return base_df.to_dict(orient="list")
+
+    @property
+    def best_params_(self) -> Dict[str, Any]:
+        """Parameters of the best trial in the :class:`~optuna.study.Study`."""
+        return self.study_.best_params
+
+    @property
+    def best_score_(self) -> float:
+        """Best score reached in the study."""
+        return self.study_.best_value
+
+    @property
+    def best_trial_(self) -> FrozenTrial:
+        """Best trial in the :class:`~optuna.study.Study`."""
+        return self.study_.best_trial
+
+    def optimize(self, dataset: Dataset_, **_: Any) -> Self:
+        """Optimize the objective over the dataset and find the best parameter combination.
+
+        This method calls `self.create_objective` to obtain the objective function that should be optimized.
+
+        Parameters
+        ----------
+        dataset
+            The dataset used for optimization.
+
+        """
+        if self.timeout is None and self.n_trials is None:
+            raise ValueError(
+                "You need to set either `timeout` or `n_trials` to a proper value."
+                "Otherwise the optimization will not stop und run until infinity."
+            )
+
+        self.dataset = dataset
+
+        objective = self._create_objective(self.pipeline, dataset=dataset)
+        self.study_ = self._call_optimize(self.study, objective)
+
+        if self.return_optimized:
+            self.optimized_pipeline_ = self.return_optimized_pipeline(clone(self.pipeline), dataset, self.study_)
+        return self
+
+    def _call_optimize(self, study: Study, objective: ObjectiveFuncType) -> Study:
+        """Call the optuna study.
+
+        This is a separate method to make it easy to modify how the study is called.
+        """
+        study.optimize(
+            objective,
+            n_trials=self.n_trials,
+            timeout=self.timeout,
+            callbacks=self.callbacks,
+            gc_after_trial=self.gc_after_trial,
+            show_progress_bar=self.show_progress_bar,
+        )
+        return study
+
+    def create_objective(self) -> Callable[[Trial, Pipeline_, Dataset_], Union[float, Sequence[float]]]:
+        """Return the objective function that should be optimized.
+
+        This method should be implemented by a child class and return an objective function that is compatible with
+        Optuna.
+        However, compared to a normal Optuna objective function, the function should expect a pipeline and a dataset
+        object as additional inputs to the optimization Trial object.
+        """
+        raise NotImplementedError()
+
+    def _create_objective(self, pipeline: Pipeline_, dataset: Dataset_) -> ObjectiveFuncType:
+        inner_objective = self.create_objective()
+
+        def objective(trial: Trial):
+            inner_pipe = clone(pipeline)
+            return inner_objective(trial, inner_pipe, dataset)
+
+        return objective
+
+    def return_optimized_pipeline(  # noqa: no-self-use
+        self, pipeline: Pipeline_, dataset: Dataset_, study: Study
+    ) -> Pipeline_:
+        """Return the pipeline with the best parameters of a study.
+
+        This either just returns the pipeline with the best parameters set, or if the pipeline is a subclass of
+        `OptimizablePipeline` it attempts a re-optimization of the pipeline using the provided dataset.
+
+        This functionality is a sensible default, but it is expected to overwrite this method in custom subclasses,
+        if specific behaviour is needed.
+
+        Don't call this function on its own! It is only expected to be called internally by optimize.
+        """
+        # Pipeline that will be passed here is already cloned, so no need to clone again.
+        pipeline_with_best_params = pipeline.set_params(**study.best_params)
+        if isinstance(pipeline_with_best_params, OptimizablePipeline):
+            return Optimize(pipeline_with_best_params).optimize(dataset).optimized_pipeline_
+        return pipeline_with_best_params


Test patch:
diff --git a/.github/workflows/test-and-lint.yml b/.github/workflows/test-and-lint.yml
index 96148ee7..b0e756a1 100644
--- a/.github/workflows/test-and-lint.yml
+++ b/.github/workflows/test-and-lint.yml
@@ -25,7 +25,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install poetry
         poetry config virtualenvs.in-project true
-        poetry install -E torch
+        poetry install -E torch -E optuna
     - name: Testing
       run: |
         poetry run poe test
diff --git a/tests/mixins/test_algorithm_mixin.py b/tests/mixins/test_algorithm_mixin.py
index 4dc6085f..d9fbb1b9 100644
--- a/tests/mixins/test_algorithm_mixin.py
+++ b/tests/mixins/test_algorithm_mixin.py
@@ -11,6 +11,7 @@
 
 class TestAlgorithmMixin:
     algorithm_class = None
+    _ignored_names = tuple()
     __test__ = False
 
     @pytest.fixture()
@@ -37,6 +38,8 @@ def test_all_parameters_documented(self):
         documented_names = set(p.name for p in docs["Parameters"])
         actual_names = set(get_param_names(self.algorithm_class))
 
+        actual_names -= set(self._ignored_names)
+
         assert documented_names == actual_names
 
     def test_all_attributes_documented(self, after_action_instance):
diff --git a/tests/test_examples/test_all_examples.py b/tests/test_examples/test_all_examples.py
index 5bacd6cb..e9e93f9c 100644
--- a/tests/test_examples/test_all_examples.py
+++ b/tests/test_examples/test_all_examples.py
@@ -1,5 +1,6 @@
 # This is needed to avoid plots to open
 import matplotlib
+import pandas as pd
 from numpy.testing import assert_almost_equal, assert_array_equal
 
 matplotlib.use("Agg")
@@ -58,3 +59,21 @@ def test_cross_validate():
     from examples.validation._01_cross_validation import results
 
     assert_almost_equal(results["test_f1_score"], [0.9770585, 0.7108303, 0.9250665])
+
+
+def test_optuna():
+    from examples.parameter_optimization._04_custom_optuna_optimizer import opti, opti_early_stop
+
+    assert opti.best_params_ == {
+        "algorithm__min_r_peak_height_over_baseline": 0.4,
+        "algorithm__high_pass_filter_cutoff_hz": 0.4,
+    }
+    assert opti.best_score_ == 0.858757056619628
+
+    # Check number of pruned trials
+    assert pd.DataFrame(opti_early_stop.search_results_)["score"].isna().sum() == 6
+    assert opti_early_stop.best_params_ == {
+        "algorithm__min_r_peak_height_over_baseline": 0.4,
+        "algorithm__high_pass_filter_cutoff_hz": 0.4,
+    }
+    assert opti_early_stop.best_score_ == 0.858757056619628
diff --git a/tests/test_pipelines/test_optimize.py b/tests/test_pipelines/test_optimize.py
index a8b612b2..5910d890 100644
--- a/tests/test_pipelines/test_optimize.py
+++ b/tests/test_pipelines/test_optimize.py
@@ -467,7 +467,7 @@ def test_pure_parameters(self, pure_paras, call_count):
 
     def test_pure_parameters_cache(self):
         """Test that pure parameter cache is deleted after run."""
-        # We just run our test twice. If the cache is not delted, the second run should fail.
+        # We just run our test twice. If the cache is not deleted, the second run should fail.
         self.test_pure_parameters(True, 4)
         self.test_pure_parameters(True, 4)
 
diff --git a/tests/test_pipelines/test_optuna_optimize.py b/tests/test_pipelines/test_optuna_optimize.py
new file mode 100644
index 00000000..b5882c02
--- /dev/null
+++ b/tests/test_pipelines/test_optuna_optimize.py
@@ -0,0 +1,200 @@
+from typing import Callable, List, Optional, Sequence, Union
+from unittest.mock import Mock, patch
+
+import pytest
+from optuna import Study, Trial, create_study
+from optuna.samplers import GridSampler, RandomSampler
+from optuna.structs import FrozenTrial
+
+from tests.mixins.test_algorithm_mixin import TestAlgorithmMixin
+from tests.test_pipelines.conftest import DummyDataset, DummyOptimizablePipeline, DummyPipeline, dummy_single_score_func
+from tpcp import make_optimize_safe
+from tpcp._dataset import Dataset_
+from tpcp._pipeline import OptimizablePipeline, Pipeline_
+from tpcp.optimize.optuna import CustomOptunaOptimize
+from tpcp.validate import Scorer
+
+
+class DummyOptunaOptimizer(CustomOptunaOptimize[Pipeline_, Dataset_]):
+    def __init__(
+        self,
+        pipeline: Pipeline_,
+        study: Study,
+        scoring: Callable,
+        create_search_space: Callable,
+        *,
+        n_trials: Optional[int] = None,
+        timeout: Optional[float] = None,
+        callbacks: Optional[List[Callable[[Study, FrozenTrial], None]]] = None,
+        gc_after_trial: bool = False,
+        show_progress_bar: bool = False,
+        return_optimized: bool = True,
+        mock_objective=None,
+    ) -> None:
+        self.scoring = scoring
+        self.create_search_space = create_search_space
+        self.mock_objective = mock_objective
+        super().__init__(
+            pipeline,
+            study,
+            n_trials=n_trials,
+            timeout=timeout,
+            callbacks=callbacks,
+            gc_after_trial=gc_after_trial,
+            show_progress_bar=show_progress_bar,
+            return_optimized=return_optimized,
+        )
+
+    def create_objective(self) -> Callable[[Trial, Pipeline_, Dataset_], Union[float, Sequence[float]]]:
+        def objective(trial: Trial, pipeline: Pipeline_, dataset: Dataset_) -> float:
+            self.create_search_space(trial)
+            pipeline = pipeline.set_params(**trial.params)
+
+            scorer = Scorer(self.scoring)
+
+            average_score, single_scores = scorer(pipeline, dataset, error_score="raise")
+            trial.set_user_attr("single_scores", single_scores)
+            return average_score
+
+        if self.mock_objective:
+            return self.mock_objective
+        return objective
+
+
+def dummy_search_space(trial: Trial):
+    trial.suggest_categorical("para_1", [1])
+
+
+class TestMetaFunctionalityGridSearch(TestAlgorithmMixin):
+    __test__ = True
+    algorithm_class = DummyOptunaOptimizer
+    _ignored_names = ("create_search_space", "scoring", "mock_objective")
+
+    @pytest.fixture()
+    def after_action_instance(self) -> DummyOptunaOptimizer:
+
+        study = create_study(sampler=RandomSampler(42))
+
+        gs = DummyOptunaOptimizer(
+            DummyOptimizablePipeline(),
+            study,
+            scoring=dummy_single_score_func,
+            create_search_space=dummy_search_space,
+            n_trials=1,
+        )
+        gs.optimize(DummyDataset())
+        return gs
+
+    def test_empty_init(self):
+        pytest.skip()
+
+
+class TestCustomOptunaOptimize:
+    def test_invalid_study_stop(self):
+        with pytest.raises(ValueError):
+            DummyOptunaOptimizer(
+                DummyOptimizablePipeline(),
+                study=create_study(sampler=RandomSampler(42)),
+                scoring=dummy_single_score_func,
+                create_search_space=dummy_search_space,
+                n_trials=None,  # These should not be both None
+                timeout=None,
+            ).optimize(DummyDataset())
+
+    def test_objective_called(self):
+        mock_objective = Mock(return_value=3)
+
+        n_trials = 5
+        dataset = DummyDataset()
+        pipe = DummyOptimizablePipeline()
+
+        DummyOptunaOptimizer(
+            pipe,
+            study=create_study(sampler=RandomSampler(42)),
+            scoring=dummy_single_score_func,
+            create_search_space=dummy_search_space,
+            n_trials=n_trials,
+            timeout=None,
+            mock_objective=mock_objective,
+        ).optimize(dataset)
+
+        assert mock_objective.call_count == n_trials
+        assert mock_objective.call_args[0][2] is dataset
+        # Should be a clone of the pipe, but not the pipe
+        assert mock_objective.call_args[0][1].get_params() == pipe.get_params()
+        assert mock_objective.call_args[0][1] is not pipe
+
+        assert isinstance(mock_objective.call_args[0][0], Trial)
+
+    @pytest.mark.parametrize("return_optimize", (True, False))
+    def test_return_optimized(self, return_optimize):
+        mock_objective = Mock(return_value=3)
+
+        opti = DummyOptunaOptimizer(
+            DummyOptimizablePipeline(),
+            study=create_study(sampler=RandomSampler(42)),
+            scoring=dummy_single_score_func,
+            create_search_space=dummy_search_space,
+            n_trials=1,
+            timeout=None,
+            mock_objective=mock_objective,
+            return_optimized=return_optimize,
+        ).optimize(DummyDataset())
+
+        assert hasattr(opti, "optimized_pipeline_") == return_optimize
+
+    def test_return_optimized_calls_optimize(self):
+        mock_objective = Mock(return_value=3)
+        optimized_pipe = DummyOptimizablePipeline()
+        dataset = DummyDataset()
+
+        with patch.object(DummyOptimizablePipeline, "self_optimize", return_value=optimized_pipe) as mock:
+            mock.__name__ = "self_optimize"
+            DummyOptimizablePipeline.self_optimize = make_optimize_safe(DummyOptimizablePipeline.self_optimize)
+
+            DummyOptunaOptimizer(
+                DummyOptimizablePipeline(),
+                study=create_study(sampler=RandomSampler(42)),
+                scoring=dummy_single_score_func,
+                create_search_space=dummy_search_space,
+                n_trials=1,
+                timeout=None,
+                mock_objective=mock_objective,
+                return_optimized=True,
+            ).optimize(dataset)
+
+        assert mock.call_args[0][1] is dataset
+
+    @pytest.mark.parametrize("pipe", (DummyOptimizablePipeline(), DummyPipeline()))
+    def test_correct_paras_selected(self, pipe):
+        # Should select 1, as it has the highest score
+        scores = {0: 1, 1: 2, 2: 0}
+
+        def create_search_space(trial):
+            trial.suggest_categorical("para_1", list(scores.keys()))
+
+        def scoring(pipe, _):
+            return scores[pipe.para_1]
+
+        opti = DummyOptunaOptimizer(
+            pipe,
+            study=create_study(sampler=GridSampler({"para_1": list(scores.keys())}), direction="maximize"),
+            scoring=scoring,
+            create_search_space=create_search_space,
+            n_trials=3,
+            timeout=None,
+            return_optimized=True,
+        ).optimize(DummyDataset())
+
+        assert opti.best_params_ == {"para_1": 1}
+        assert opti.best_score_ == 2
+        assert opti.optimized_pipeline_.para_1 == 1
+
+        r = opti.search_results_
+
+        assert set(r["param_para_1"]) == set(scores.keys())
+        assert set(tuple(p.items()) for p in r["params"]) == set((("para_1", k),) for k in scores.keys())
+
+        if isinstance(pipe, OptimizablePipeline):
+            # That is expected when self_optimize was called correctly.
+            assert opti.optimized_pipeline_.optimized == opti.optimized_pipeline_.para_2