Repository: cog-imperial/entmoot (language: python)

Issue:
Add tests in the Pyomo version
- [x]  pyomo with gurobi yields same results as gurobipy
- [ ]  pyomo with CBC solver should find good points on benchmark problems
- [ ] Compare BO-Loops
- [x] Test small example with known minimum
- [ ] Test, if Pyomo user can run ENTMOOT without having gurobipy installed and vice versa


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
new file mode 100644
index 0000000..01c8eaf
--- /dev/null
+++ b/.github/workflows/python-package.yml
@@ -0,0 +1,46 @@
+# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Tests
+
+on:
+  push:
+    branches: [ "master", "entmoot-v2"]
+  pull_request:
+    branches: [ "master" ]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11"]
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v3
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install flake8 pytest
+        pip install -r requirements-dev.txt
+    - shell: bash
+      id: write-license
+      env:
+        WLSLICENSE: ${{ secrets.WLSLICENSE   }}
+      run: |
+        echo "$WLSLICENSE" > $PWD/gurobi.lic
+        echo "grb_license_file=$PWD/gurobi.lic" >> $GITHUB_OUTPUT
+    - name: Test with pytest
+      env:
+        CICD_ACTIVE: ${{ secrets.CICD_ACTIVE }}
+        GRB_LICENSE_FILE: ${{ steps.write-license.outputs.grb_license_file }}
+      run: |
+        pytest --cov
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v3
diff --git a/.gitignore b/.gitignore
index ab56d4c..a2603c8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,10 +4,14 @@ __pycache__
 *.json
 *.gz
 *.pkl
-test.py
-.DS_Store
-env
 *.log
 *.sh
 .idea
 .pytest_cache
+.DS_Store
+env
+/docs/notebooks/.ipynb_checkpoints/
+docs/_build
+.coverage
+Makefile
+make.bat
\ No newline at end of file
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
new file mode 100644
index 0000000..6843574
--- /dev/null
+++ b/.readthedocs.yaml
@@ -0,0 +1,20 @@
+# .readthedocs.yaml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+version: 2
+
+# Set the version of Python and other tools you might need
+build:
+  os: ubuntu-20.04
+  tools:
+    python: "3.9"
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+   configuration: docs/conf.py
+
+# Optionally declare the Python requirements required to build your docs
+python:
+   install:
+   - requirements: requirements-dev.txt
\ No newline at end of file
diff --git a/README.md b/README.md
index 2782295..77e03aa 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,21 @@
+[![Tests](https://github.com/cog-imperial/entmoot/actions/workflows/python-package.yml/badge.svg?branch=entmoot-v2)](https://github.com/cog-imperial/entmoot/actions/workflows/python-package.yml)
+[![codecov](https://codecov.io/gh/cog-imperial/entmoot/branch/entmoot-v2/graph/badge.svg)](https://codecov.io/gh/cog-imperial/entmoot)
+![Python versions](https://img.shields.io/badge/python-3.7%20%7C%203.8%20%7C%203.9%20%7C%203.10%20%7C%203.11-blue)
+[![Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) 
+[![Read the Docs](https://readthedocs.org/projects/entmoot/badge/?version=latest)](https://entmoot.readthedocs.io/en/latest/)
 
 <img src="media/entmoot_logo.png" width="400">
 
-`ENTMOOT` (**EN**semble **T**ree **MO**del **O**ptimization **T**ool) is a framework to handle tree-based surrogate models in Bayesian optimization applications. Gradient-boosted
-tree models from `lightgbm` are combined with a distance-based uncertainty 
+`ENTMOOT` (**EN**semble **T**ree **MO**del **O**ptimization **T**ool) is a framework to perform Bayesian Optimization using tree-based surrogate models. Gradient-boosted tree models from `lightgbm` are combined with a distance-based uncertainty 
 measure in a deterministic global optimization framework to optimize black-box functions. More
 details on the method here: https://arxiv.org/abs/2003.04774.
 
+## Documentation
+
+The docs can be found here: https://entmoot.readthedocs.io/
+
+## How to reference ENTMOOT
+
 When using any `ENTMOOT` for any publications please reference this software package as:
 ```
 @article{thebelt2021entmoot,
@@ -19,215 +29,10 @@ When using any `ENTMOOT` for any publications please reference this software pac
 }
 ```
 
-The "ask-tell" optimizer interface for an intuitive incorporation of own 
-black-box functions was included from 
-[scikit-optimize](https://github.com/scikit-optimize/scikit-optimize/) 
-which is available under the BSD 3-Clause License.
-
-## Installation - Linux & Mac OS
-### Requirements
-* python >= 3.7
-* numpy >= 1.18.4
-* scikit-learn >= 0.21.3
-* pyyaml >= 5.3.1
-* lightgbm >= 2.3.1
-
-On Mac, you will also need to install libomp: `brew install libomp`.
-
-### Optional
-* gurobi >= 9.0.1
-
-
-### Installing ENTMOOT
-Install all required packages by running the command
-```
-pip install -r requirements.txt
-```
-To install `ENTMOOT`, run the following command
-```
-pip install git+https://github.com/cog-imperial/entmoot
-```
-
-### Uninstalling ENTMOOT
-The ENTMOOT package can be uninstalled by running
-```
-pip uninstall entmoot
-```
-
-## Installation - Windows
-The installation on Windows-based systems is similar with some exceptions for
-individual commands. Make sure that `python` and `git` is installed and 
-accessible in the shell you are using to install `ENTMOOT`. Follow the 
-instructions by executing all commands in `cmd.exe`.
-
-### Installing ENTMOOT
-Install all required packages by running the command
-```
-pip install -r requirements.txt
-```
-To install `ENTMOOT`, run the following command
-```
-pip install git+https://github.com/cog-imperial/entmoot
-```
-
-### Uninstalling ENTMOOT
-The ENTMOOT package can be uninstalled by running
-```
-pip uninstall entmoot
-```
-
-## Installing Gurobi
-To use the `acq_optimizer= 'global'` setting in `ENTMOOT`, the solver 
-software [Gurobi v9.0](https://www.gurobi.com/resource/overview-of-gurobi-9-0/) 
-or newer is required. Gurobi is a commercial mathematical optimization solver and 
-free of charge for academic research. It is available on Linux, Windows and 
-Mac OS. 
-
-Please follow the instructions to obtain a [free academic license](https://www.gurobi.com/academia/academic-program-and-licenses/). Once Gurobi 
-is installed on your system, follow the steps to setup the Python interface [gurobipy](https://www.gurobi.com/documentation/9.0/quickstart_mac/the_grb_python_interface_f.html) in your virtual environment created for `ENTMOOT`.
-
-## Example: Single Objective Optimization with Constraints
-For a simple example we use Rosenbrock as our black-box function. We import the
-`entmoot_minimize` function which automates the optimization procedure. The 
-black-box function is given as `func` to the solver. For more details on 
-individual settings we refer to the source code which explains all parameters
-in detail.
-
-An example script that minimizes the Rosenbrock function using `ENTMOOT` is
-given in the following.
-```
-from entmoot.optimizer.optimizer import Optimizer
-from entmoot.benchmarks import Rosenbrock
-
-func = Rosenbrock()
-
-opt = Optimizer(func.get_bounds(10),
-                base_estimator="ENTING",
-                n_initial_points=20,
-                initial_point_generator="random",
-                acq_func="LCB",
-                acq_optimizer="sampling",
-                random_state=100,
-                model_queue_size=None,
-                base_estimator_kwargs={
-                    "lgbm_params": {"min_child_samples": 1}
-                },
-                verbose=True,
-                )
-
-# run optimizer for 20 iterations
-res = opt.run(func, n_iter=20)
-print(f"-> best solution found {res.fun}")
-```
-
-The following example shows how own constraints can be enforced on the input 
-variables. This only works when Gurobi is selected as the solution strategy,
- i. e. `acq_optimizer="global"`. Constraints are formulated according to 
- [documentation](https://www.gurobi.com/documentation/9.0/refman/py_model_addconstr.html).
-```
-from entmoot.optimizer.optimizer import Optimizer
-from entmoot.benchmarks import SimpleCat
-
-func = SimpleCat()
-
-# initialize the search space manually
-from entmoot.space.space import Space
-space = Space(func.get_bounds())
-
-# get the core of the gurobi model from helper function 'get_core_gurobi_model'
-from entmoot.optimizer.gurobi_utils import get_core_gurobi_model
-core_model = get_core_gurobi_model(space)
-
-# ordering of variable indices is dependent on space definition
-
-# cont_var_dict contains all continuous variables
-x0 = core_model._cont_var_dict[0]
-x1 = core_model._cont_var_dict[1]
-
-# cat_var_dict contains all categorical variables
-x2 = core_model._cat_var_dict[2]
-
-# define constraints accordingly
-core_model.addConstr(x0 + x1 >= 2)
-core_model.addConstr(x1 == 1)
-core_model.update()
-
-opt = Optimizer(func.get_bounds(),
-                base_estimator="ENTING",
-                n_initial_points=0,
-                initial_point_generator="random",
-                acq_func="LCB",
-                acq_optimizer="global",
-                random_state=100,
-                acq_func_kwargs=None,
-                acq_optimizer_kwargs={
-                    "add_model_core": core_model
-                },
-                model_queue_size=None,
-                base_estimator_kwargs={
-                    "lgbm_params": {"min_child_samples": 1}
-                },
-                verbose=True,
-                )
-
-# add initial points that are feasible
-for x in [(1.1, 1.0, 'mult6'), (1.2, 1.0, 'mult6'), (1.3, 1.0, 'pow2')]:
-    opt.tell(x, func(x))
-
-# run optimizer for 20 iterations
-res = opt.run(func, n_iter=20)
-print(f"-> best solution found {res.fun}")
-```
-
-## Example: Multiple Objective Optimization with Constraints
-
-`ENTMOOT` also supports multi-objective optimization according to:
-
-```
-@article{thebelt2022multi,
-  title={Multi-objective constrained optimization for energy applications via tree ensembles},
-  author={Thebelt, Alexander and Tsay, Calvin and Lee, Robert M and Sudermann-Merx, Nathan and Walz, David and Tranter, Tom and Misener, Ruth},
-  journal={Applied Energy},
-  volume={306},
-  pages={118061},
-  year={2022},
-  publisher={Elsevier}
-}
-```
-
-An example that derives Pareto-optimal points of the 
-[Fonzeca Freming](https://en.wikipedia.org/wiki/Test_functions_for_optimization) is given in the
-following:
-
-```
-from entmoot.benchmarks import FonzecaFleming
-from entmoot.optimizer import Optimizer
-
-# initialize multi-objective test function
-funcMulti = FonzecaFleming()
-
-# define optimizer object and specify num_obj=2
-opt = Optimizer(funcMulti.get_bounds(),
-                num_obj=2,
-                n_initial_points=10,
-                random_state=100)
-
-# main BO loop that derives pareto-optimal points
-for _ in range(50):
-    next_x = opt.ask()
-    next_y = funcMulti(next_x)
-    opt.tell(next_x,next_y)
-```
-
-Using multi-objective functionality in `ENTMOOT` requires the specification of `num_obj` which 
-informs the solver about the number of objectives that we optimize for. `ENTMOOT` minimizes
-objectives which requires the modification of maximization problems, i.e. minimizing the 
-negative objective.
 
 ## Authors
 * **[Alexander Thebelt](https://optimisation.doc.ic.ac.uk/person/alexander-thebelt/)** ([ThebTron](https://github.com/ThebTron)) - Imperial College London
 * **[Nathan Sudermann-Merx](https://www.mannheim.dhbw.de/profile/sudermann-merx)** ([spiralulam](https://github.com/spiralulam)) - Cooperative State University Mannheim
-* **[David Walz](https://www.linkedin.com/in/walzds/?originalSubdomain=de)** ([DavidWalz](https://github.com/DavidWalz)) - BASF SE
 ## License
 The ENTMOOT package is released under the BSD 3-Clause License. Please refer to the [LICENSE](https://github.com/cog-imperial/entmoot/blob/master/LICENSE) file for details.
 
diff --git a/docs/apidoc/entmoot.enting.rst b/docs/apidoc/entmoot.enting.rst
new file mode 100644
index 0000000..6ab233d
--- /dev/null
+++ b/docs/apidoc/entmoot.enting.rst
@@ -0,0 +1,7 @@
+Enting
+==========================
+
+.. autoclass:: entmoot.models.enting.Enting
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/apidoc/entmoot.rst b/docs/apidoc/entmoot.rst
new file mode 100644
index 0000000..3b19532
--- /dev/null
+++ b/docs/apidoc/entmoot.rst
@@ -0,0 +1,7 @@
+ENTMOOT
+=======
+
+.. toctree::
+   :maxdepth: 2
+
+   entmoot.enting
diff --git a/docs/conf.py b/docs/conf.py
new file mode 100644
index 0000000..dbb5370
--- /dev/null
+++ b/docs/conf.py
@@ -0,0 +1,78 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+import os
+import sys
+import inspect
+
+__location__ = os.path.join(
+    os.getcwd(), os.path.dirname(inspect.getfile(inspect.currentframe()))
+)
+
+sys.path.insert(0, os.path.join(__location__, "../../entmoot"))
+sys.path.insert(0, os.path.join(__location__, "../entmoot"))
+sys.path.insert(0, os.path.join(__location__, ".."))
+
+# -- Run sphinx-apidoc -------------------------------------------------------
+# This hack is necessary since RTD does not issue `sphinx-apidoc` before running
+# `sphinx-build -b html . _build/html`. See Issue:
+# https://github.com/rtfd/readthedocs.org/issues/1139
+# DON'T FORGET: Check the box "Install your project inside a virtualenv using
+# setup.py install" in the RTD Advanced Settings.
+# Additionally it helps us to avoid running apidoc manually
+
+try:  # for Sphinx >= 1.7
+    from sphinx.ext import apidoc
+except ImportError:
+    from sphinx import apidoc
+
+# -- Project information -----------------------------------------------------
+
+project = "ENTMOOT"
+copyright = "2022, Alexander Thebelt, Nathan Sudermann-Merx"
+author = "Alexander Thebelt, Nathan Sudermann-Merx"
+
+
+# -- General configuration ---------------------------------------------------
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = ["sphinx.ext.autodoc", "nbsphinx"]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "**.ipynb_checkpoints"]
+
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = "sphinx_rtd_theme"
+
+# Path to the entmoot logo
+html_logo = "../media/entmoot_logo.png"
+
+html_theme_options = {"logo_only": True, "style_nav_header_background": "white"}
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = []
+
+nbsphinx_execute = "never"
diff --git a/docs/index.rst b/docs/index.rst
new file mode 100644
index 0000000..34f9f4e
--- /dev/null
+++ b/docs/index.rst
@@ -0,0 +1,70 @@
+Welcome to ENTMOOT's documentation!
+===================================
+
+**ENTMOOT** (**EN**\semble **T**\ree **MO**\del **O**\ptimization **T**\ool) is a framework to perform Bayesian
+Optimization using tree-based surrogate models. Gradient-boosted tree models from `LightGBM <https://lightgbm.readthedocs.io>`__
+are combined with a distance-based uncertainty measure in a deterministic global optimization framework to optimize
+black-box functions. More details on the method can be found here: https://arxiv.org/abs/2003.04774.
+
+Appetizer
+---------
+This small example illustrates how to use ENTMOOT.
+
+.. code-block:: python
+
+    from entmoot import Enting, ProblemConfig, PyomoOptimizer
+    import numpy as np
+    import random
+
+    # This is the function you want to minimize.
+    def my_func(x: float) -> float:
+        # randomly disturbed function f(x) = x^2 + 1 + eps
+        return x**2 + 1 + random.uniform(-0.2, 0.2)
+
+    # Define a one-dimensional minimization problem with one real variable bounded by -2 and 3
+    problem_config = ProblemConfig()
+    problem_config.add_feature("real", (-2, 3))
+    problem_config.add_min_objective()
+
+    # Create training data for the tree model
+    X_train = np.reshape(np.linspace(-2, 3, 10), (-1, 1))
+    y_train = np.reshape([my_func(x) for x in X_train], (-1, 1))
+
+    # Define Bayesian optimization parameters uding an l1-distance based uncertainty measure and
+    # penalizing the distance from well-known areas, i.e. exploitation (instead of exploration)
+    params_bo = {"unc_params": {"dist_metric": "l1", "acq_sense": "penalty"}}
+
+    # Define an Enting object which holds information about the problem as well as the parameters...
+    enting = Enting(problem_config, params=params_bo)
+    # ... and train the underlying tree model.
+    enting.fit(X_train, y_train)
+
+    # Create an PyomoOptimizer object that solves the optimization problem using the solver "GLPK"
+    params_pyo = {"solver_name": "glpk"}
+    opt_pyo = PyomoOptimizer(problem_config, params=params_pyo)
+    res = opt_pyo.solve(enting)
+
+    # Inspect the result. The optimal point should be close to zero.
+    print(f"Optimal point: {res.opt_point[0]}")
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Contents:
+
+   Installation Guide <installation>
+   Examples <notebooks>
+   API Reference <apidoc/entmoot>
+
+Authors
+~~~~~~~
+* `Alexander Thebelt <https://optimisation.doc.ic.ac.uk/person/alexander-thebelt>`_ (`ThebTron <https://github.com/ThebTron>`_) - Imperial College London
+* `Nathan Sudermann-Merx <https://www.mannheim.dhbw.de/profile/sudermann-merx>`_ (`spiralulam <https://github.com/spiralulam>`_) - Cooperative State University Mannheim
+
+License
+~~~~~~~
+The ENTMOOT package is released under the BSD 3-Clause License. Please refer to the
+`LICENSE <https://github.com/cog-imperial/entmoot/blob/master/LICENSE>`_ file for details.
+
+Acknowledgements
+~~~~~~~~~~~~~~~~
+The support of BASF SE, Lugwigshafen am Rhein is gratefully acknowledged.
\ No newline at end of file
diff --git a/docs/installation.rst b/docs/installation.rst
new file mode 100644
index 0000000..1c27373
--- /dev/null
+++ b/docs/installation.rst
@@ -0,0 +1,45 @@
+How to install ENTMOOT
+======================
+ENTMOOT supports the Python (CPython) versions 3.7 - 3.11.
+
+Using PIP
+-------------
+
+You can install ENTMOOT with
+
+::
+
+   pip install entmoot
+
+**On Mac**, you also need to install libomp:
+
+::
+
+    brew install libomp
+
+In addition, you need to install **one of the following** Python packages.
+
+* gurobipy >= 10.0.0
+* pyomo >= 6.4.4
+
+The package gurobipy belongs to the commercial solver `Gurobi <https://www.gurobi.com/>`__ whereas pyomo refers to the
+open source framework `Pyomo <http://www.pyomo.org/>`__ which supports several commercial and noncommercial solvers.
+
+
+Gurobi or Pyomo?
+-----------------
+You can work with ENTMOOT using Gurobi or Pyomo.
+
+* `Gurobi <https://www.gurobi.com/>`__ is a commercial mathematical optimization solver and very fast. Not only
+  in solving the corresponding optimization problem but also in building it. Using Gurobi requires a valid license. Note
+  that Gurobi offers `free academic licenses <https://www.gurobi.com/academia/academic-program-and-licenses/>`__.
+* `Pyomo <http://www.pyomo.org/>`__ is an open-source optimization modeling language with APIs to several commercial and
+  noncommercial solvers. Note that noncommercial solvers are generally (much) slower than commercial ones.
+
+If you decide to work with Pyomo, you have to install additional solvers that can solve mixed-integer optimization
+problems. The most common choices are the commercial solvers `Gurobi <https://www.gurobi.com/>`__,
+`CPLEX <https://www.ibm.com/de-de/analytics/cplex-optimizer>`__ and
+`Xpress <https://www.fico.com/en/products/fico-xpress-optimization>`__ as well as the noncommercial solvers
+`GLPK <https://www.gnu.org/software/glpk/>`__,
+`CBC <https://github.com/coin-or/Cbc/>`__  and `SCIP <https://www.scipopt.org
+/>`__.
\ No newline at end of file
diff --git a/docs/notebooks.rst b/docs/notebooks.rst
new file mode 100644
index 0000000..93adeb4
--- /dev/null
+++ b/docs/notebooks.rst
@@ -0,0 +1,11 @@
+Examples
+========
+These Jupyter Notebooks provide some examples on how to work with ENTMOOT. Their source code can be found
+`here <https://github.com/cog-imperial/entmoot/tree/entmoot-v2/docs/notebooks>`_.
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+   notebooks/multi_obj_pareto
+   notebooks/multi_obj_with_constraints
diff --git a/docs/notebooks/multi_obj_pareto.ipynb b/docs/notebooks/multi_obj_pareto.ipynb
new file mode 100644
index 0000000..07996a4
--- /dev/null
+++ b/docs/notebooks/multi_obj_pareto.ipynb
@@ -0,0 +1,123 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "6f01675f-cd1a-4712-9ca4-3cc1dbbd9056",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Multiple Objective Optimization"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "567ea5b7-f871-4ce3-8970-6f82728de9bf",
+   "metadata": {
+    "jupyter": {
+     "outputs_hidden": true
+    },
+    "tags": []
+   },
+   "source": [
+    "An example that computes a Pareto-optimal point of a benchmark function using Chebyshev weights is given in the\n",
+    "following:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cae5fecd-f30c-4a5c-966e-c14008a1af8a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from entmoot import Enting, ProblemConfig, GurobiOptimizer, PyomoOptimizer\n",
+    "from entmoot.benchmarks import (\n",
+    "    build_multi_obj_categorical_problem,\n",
+    "    eval_multi_obj_cat_testfunc,\n",
+    ")\n",
+    "\n",
+    "# define problem\n",
+    "problem_config = ProblemConfig(rnd_seed=73)\n",
+    "# number of objectives\n",
+    "number_objectives = 2\n",
+    "build_multi_obj_categorical_problem(problem_config, n_obj=number_objectives)\n",
+    "\n",
+    "# sample data\n",
+    "rnd_sample = problem_config.get_rnd_sample_list(num_samples=20)\n",
+    "testfunc_evals = eval_multi_obj_cat_testfunc(rnd_sample, n_obj=number_objectives)\n",
+    "\n",
+    "      \n",
+    "params = {\"unc_params\": {\"dist_metric\": \"l1\", \"acq_sense\": \"exploration\"}}\n",
+    "enting = Enting(problem_config, params=params)\n",
+    "# fit tree ensemble\n",
+    "enting.fit(rnd_sample, testfunc_evals)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7b88056d-e265-4097-9711-6bdcec723b44",
+   "metadata": {},
+   "source": [
+    "### Gurobi Version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "311c88fa-f757-44f3-8ae5-555f715fc1b4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Build GurobiOptimizer object and solve optimization problem\n",
+    "params_gurobi = {\"MIPGap\": 0}\n",
+    "opt_gur = GurobiOptimizer(problem_config, params=params_gurobi)\n",
+    "res_gur = opt_gur.solve(enting)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ab64ab28-9578-43ea-8943-821637681acf",
+   "metadata": {},
+   "source": [
+    "### Pyomo Version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cbf5ea4c-fe6b-4f29-8b35-ae88cfdb2379",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Build PyomoOptimizer object with GLPK as solver and solve optimization problem\n",
+    "params_pyo = {\"solver_name\": \"glpk\"}\n",
+    "opt_pyo = PyomoOptimizer(problem_config, params=params_pyo)\n",
+    "res_pyo = opt_pyo.solve(enting)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/docs/notebooks/multi_obj_with_constraints.ipynb b/docs/notebooks/multi_obj_with_constraints.ipynb
new file mode 100644
index 0000000..653e16f
--- /dev/null
+++ b/docs/notebooks/multi_obj_with_constraints.ipynb
@@ -0,0 +1,206 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "6f01675f-cd1a-4712-9ca4-3cc1dbbd9056",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Bring your own constraints!\n",
+    "The following example shows how your own constraints can be enforced on the input\n",
+    "variables."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c8da843-670c-4f07-bd66-471ec19d3601",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from entmoot import Enting, ProblemConfig, GurobiOptimizer, PyomoOptimizer\n",
+    "from entmoot.benchmarks import (\n",
+    "    build_multi_obj_categorical_problem,\n",
+    "    eval_multi_obj_cat_testfunc,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "20d1e71a-2eff-49d1-8d45-412ebdc5412a",
+   "metadata": {},
+   "source": [
+    "This part is pretty standard."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "311c88fa-f757-44f3-8ae5-555f715fc1b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define problem\n",
+    "problem_config = ProblemConfig(rnd_seed=73)\n",
+    "# number of objectives\n",
+    "number_objectives = 2\n",
+    "build_multi_obj_categorical_problem(problem_config, n_obj=number_objectives)\n",
+    "\n",
+    "# sample data\n",
+    "rnd_sample = problem_config.get_rnd_sample_list(num_samples=20)\n",
+    "testfunc_evals = eval_multi_obj_cat_testfunc(rnd_sample, n_obj=number_objectives)\n",
+    "\n",
+    "params = {\"unc_params\": {\"dist_metric\": \"l1\", \"acq_sense\": \"exploration\"}}\n",
+    "enting = Enting(problem_config, params=params)\n",
+    "# fit tree ensemble\n",
+    "enting.fit(rnd_sample, testfunc_evals)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "818791f2-9784-4b33-acf7-fba8a4747a93",
+   "metadata": {},
+   "source": [
+    "How to add constraints depends on wether you are using Gurobi or Pyomo"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "943e15bb-c253-4d2f-8937-4708cd090fec",
+   "metadata": {},
+   "source": [
+    "### Gurobi Version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0dc70de-f14f-42a4-9202-0e4777d33bec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# get optimization model\n",
+    "model_gur = problem_config.get_gurobi_model_core()\n",
+    "# extract decision variables\n",
+    "x = model_gur._all_feat[3]\n",
+    "y = model_gur._all_feat[4]\n",
+    "z = model_gur._all_feat[5]\n",
+    "# add constraint that all variables should coincide\n",
+    "model_gur.addConstr(x == y)\n",
+    "model_gur.addConstr(y == z)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ae8f39bd-224d-4bd7-9b6b-5bbc6a01d355",
+   "metadata": {},
+   "source": [
+    "It is important to update the Gurobi model. Otherwise the constraints will not be added to the model!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "61a8e860-92ee-4535-ad73-9b2acf132ad8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_gur.update()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0d9a901-c7b4-46d9-a932-f1b2292476cb",
+   "metadata": {},
+   "source": [
+    "Now you can run the optimization and verify that the variable values indeed equal each other."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4ae2f026-dded-4376-866b-6ca34adc5bf9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Build GurobiOptimizer object and solve optimization problem\n",
+    "params_gurobi = {\"MIPGap\": 0}\n",
+    "opt_gur = GurobiOptimizer(problem_config, params=params_gurobi)\n",
+    "\n",
+    "res_gur = opt_gur.solve(enting, model_core=model_gur)\n",
+    "x_opt, y_opt, z_opt = res_gur.opt_point[3:]\n",
+    "\n",
+    "assert round(x_opt, 5) == round(y_opt, 5) and round(y_opt, 5) == round(z_opt, 5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "24b3d335-d601-41c1-ad46-d80949c0cfcc",
+   "metadata": {},
+   "source": [
+    "### Pyomo Version"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1644fd93-d413-4aaa-9d91-2467f2c59900",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Pyomo version\n",
+    "import pyomo.environ as pyo\n",
+    "\n",
+    "model_pyo = problem_config.get_pyomo_model_core()\n",
+    "# extract decision variables\n",
+    "x = model_pyo._all_feat[3]\n",
+    "y = model_pyo._all_feat[4]\n",
+    "z = model_pyo._all_feat[5]\n",
+    "# add constraint that all variables should coincide\n",
+    "model_pyo.xy_equal_constr = pyo.Constraint(expr=x == y)\n",
+    "model_pyo.yz_equal_constr = pyo.Constraint(expr=y == z)\n",
+    "\n",
+    "# Build GurobiOptimizer object and solve optimization problem\n",
+    "params_pyomo = {\"solver_name\": \"gurobi\"}\n",
+    "opt_pyo = PyomoOptimizer(problem_config, params=params_pyomo)\n",
+    "\n",
+    "res_pyo = opt_pyo.solve(enting, model_core=model_pyo)\n",
+    "x_opt, y_opt, z_opt = res_pyo.opt_point[3:]\n",
+    "\n",
+    "assert round(x_opt, 5) == round(y_opt, 5) and round(y_opt, 5) == round(z_opt, 5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c9f1cad4-8d2b-4bc7-bc86-de06d4a141c3",
+   "metadata": {},
+   "source": [
+    "Note that no model update is required in the Pyomo version in contrast to the Gurobi variant"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/entmoot/__init__.py b/entmoot/__init__.py
index 008e2e1..c443348 100644
--- a/entmoot/__init__.py
+++ b/entmoot/__init__.py
@@ -1,8 +1,4 @@
-from .optimizer import Optimizer
-from .optimizer import entmoot_minimize
-from .space import Space
-from .benchmarks import *
-
-__all__ = (
-    "Optimizer"
-)
\ No newline at end of file
+from entmoot.problem_config import ProblemConfig
+from entmoot.models.enting import Enting
+from entmoot.optimizers.gurobi_opt import GurobiOptimizer
+from entmoot.optimizers.pyomo_opt import PyomoOptimizer
diff --git a/entmoot/__version__.py b/entmoot/__version__.py
index a04c123..75ae201 100644
--- a/entmoot/__version__.py
+++ b/entmoot/__version__.py
@@ -1,4 +1,4 @@
-__version__      = '1.0.4'
-__author__       = 'Alexander Thebelt'
-__author_email__ = 'alexander.thebelt18@imperial.ac.uk'
-__license__      = 'BSD 3-Clause License'
+__version__ = "1.0.4"
+__author__ = "Alexander Thebelt"
+__author_email__ = "alexander.thebelt18@imperial.ac.uk"
+__license__ = "BSD 3-Clause License"
diff --git a/entmoot/acquisition.py b/entmoot/acquisition.py
deleted file mode 100644
index 4131d14..0000000
--- a/entmoot/acquisition.py
+++ /dev/null
@@ -1,115 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-NOTE: Changes were made to the scikit-optimize source code included here. 
-For the most recent version of scikit-optimize we refer to:
-https://github.com/scikit-optimize/scikit-optimize/
-
-Copyright (c) 2019-2020 Alexander Thebelt.
-"""
-
-import numpy as np
-import warnings
-import sys
-
-
-def _gaussian_acquisition(X, model,
-                          num_obj: int =1,
-                          acq_func: str = "LCB",
-                          acq_func_kwargs: dict = None):
-    """
-    Wrapper so that the output of this function can be
-    directly passed to a minimizer.
-    """
-    # not available for multi-obj prediction
-    assert num_obj == 1, f"acquisition predict is not available for 'num_obj > 1'"
-
-    # Check inputs
-    X = np.asarray(X)
-    if X.ndim != 2:
-        raise ValueError("X is {}-dimensional, however,"
-                         " it must be 2-dimensional.".format(X.ndim))
-
-    if acq_func_kwargs is None:
-        acq_func_kwargs = {}
-
-    kappa = acq_func_kwargs.get("kappa", 1.96)
-
-    # Evaluate acquisition function
-    if acq_func == "LCB":
-        acq_vals = gaussian_lcb(X, model, kappa)
-    else:
-        raise ValueError("Acquisition function not implemented.")
-
-    return acq_vals
-
-
-def gaussian_lcb(X, model, kappa: float = 1.96):
-    """
-    Use the lower confidence bound to estimate the acquisition
-    values.
-
-    The trade-off between exploitation and exploration is left to
-    be controlled by the user through the parameter ``kappa``.
-
-    Parameters
-    ----------
-    X : array-like, shape (n_samples, n_features)
-        Values where the acquisition function should be computed.
-
-    model : sklearn estimator that implements predict with ``return_std``
-        The fit estimator that approximates the function through the
-        method ``predict``.
-        It should have a ``return_std`` parameter that returns the standard
-        deviation.
-
-    kappa : float, default 1.96 or 'inf'
-        Controls how much of the variance in the predicted values should be
-        taken into account. If set to be very high, then we are favouring
-        exploration over exploitation and vice versa.
-        If set to 'inf', the acquisition function will only use the variance
-        which is useful in a pure exploration setting.
-        Useless if ``method`` is not set to "LCB".
-
-    Returns
-    -------
-    values : array-like, shape (X.shape[0],)
-        Acquisition function values computed at X.
-
-    """
-    # Compute posterior
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-
-        mu, std = model.predict(X, return_std=True)
-
-        if kappa == "inf":
-            return -std
-
-        return mu - kappa * std
diff --git a/entmoot/benchmarks.py b/entmoot/benchmarks.py
index 6d511ab..53d8469 100644
--- a/entmoot/benchmarks.py
+++ b/entmoot/benchmarks.py
@@ -1,154 +1,72 @@
-"""Different benchmark problems for testing"""
-
 import numpy as np
-from entmoot.utils import is_2Dlistlike
-from abc import ABC, abstractmethod
-
-
-class BenchmarkFunction(ABC):
-
-    def __init__(self, func_config: dict = None):
-        if func_config is None:
-            func_config = {}
-        self.name = 'benchmark_function'
-        self.func_config = func_config
-        self.y_opt = 0.0
-
-    def get_bounds(self, n_dim=2):
-        pass
-
-    def get_X_opt(self, n_dim=2):
-        pass
-
-    def __call__(self, X):
-        # check if multiple points are given
-        if not is_2Dlistlike(X):
-            X = [X]
-
-        res = []
-
-        for x in X:
-            res.append(self._eval_point(x))
-        
-        if len(res) == 1:
-            res = res[0]
-
-        return res
-    
-    @abstractmethod
-    def _eval_point(self,x):
-        pass
-
-
-class Rosenbrock(BenchmarkFunction):
-
-    def __init__(self, func_config: dict = None):
-        if func_config is None:
-            func_config = {}
-        self.name = 'rosenbrock'
-        self.func_config = func_config
-
-    def get_bounds(self, n_dim: int = 2):
-        return [(-2.048,2.048) for _ in range(n_dim)]
-
-    def _eval_point(self, X):
-        X = np.asarray_chkfinite(X)
-        X0 = X[:-1]
-        X1 = X[1:]
-
-        add1 = sum((1.0 - X0)**2.0)
-        add2 = 100.0 * sum((X1 - X0**2.0)**2.0)
-        return add1 + add2
-
-
-class SimpleCat(BenchmarkFunction):
-
-    def __init__(self, func_config: dict = None):
-        if func_config is None:
-            func_config = {}
-        from entmoot.space.space import Categorical
-        self.cat_dims = [
-            Categorical(['mult6','pow2'])
-        ]
-        self.name = 'benchmark_function'
-        self.func_config = func_config
-
-    def get_bounds(self, n_dim: int = 2):
-        temp_bounds = [(-2.0, 2.0) for _ in range(n_dim)]
-        temp_bounds.extend(self.cat_dims)
-        return temp_bounds
-
-    def _eval_point(self, X):
-        cat = X[-1]
-        X = np.asarray_chkfinite(X[:-1])
-        X0 = X[:-1]
-        X1 = X[1:]
-
-        add1 = X0[0]
-        add2 = X1[0]
-        
-        if cat == 'mult6':
-            return 6*(add1 + add2)
-        elif cat == 'pow2':
-            return (add1 + add2)**2
+from numpy.typing import ArrayLike
+from typing import Iterable
+from entmoot import ProblemConfig
+
+
+def build_multi_obj_categorical_problem(problem_config: ProblemConfig, n_obj: int = 2):
+    """
+    Builds a small test example which is frequently used by the tests.
+    :param problem_config: ProblemConfig object where features and objectives are added
+    :param n_obj: Number of objectives
+    :return: None, the problem definition happens "inplace"
+    """
+    problem_config.add_feature("categorical", ("blue", "orange", "gray"))
+    problem_config.add_feature("integer", (5, 6))
+    problem_config.add_feature("binary")
+    problem_config.add_feature("real", (5.0, 6.0))
+    problem_config.add_feature("real", (4.6, 6.0))
+    problem_config.add_feature("real", (5.0, 6.0))
+
+    for _ in range(n_obj):
+        problem_config.add_min_objective()
+
+
+def eval_multi_obj_cat_testfunc(X: ArrayLike, n_obj: int = 2) -> np.ndarray:
+    """
+    Benchmark function with at least four input variables and one or two outputs
+    :param X: Usually a numpy array or a list of tuples. Each row (or tuple) consists of at least four entries with the
+    following properties:
+    - The first one is a categorical variable with the possible values "blue", "orange" and "gray".
+    - The second one is an integer variable.
+    - The third one is a binary variable.
+    - The fourth one and all additional variables are real numbers.
+    :param n_obj: number of objectives (one or two)
+    :return: objective value(s) for each data point
+    """
+
+    # without the dtype=object paramer, each entry of X is converted into a string
+    X = np.array(X, dtype=object)
+
+    def compute_objectives(xi: Iterable):
+        if xi[0] == "blue":
+            return (
+                xi[1] * xi[2] * np.sin(sum(xi[3:]))
+                + xi[1] * (1 - xi[2]) * np.cos(sum(xi[3:])),
+                xi[1] * xi[2] * sum(xi[3:]) - xi[1] * (1 - xi[2]) * sum(xi[3:]),
+            )
+        elif xi[0] == "orange":
+            return (
+                xi[1] * xi[2] * (sum(xi[3:]) / len(xi[3:])) ** 2
+                + xi[1] * (1 - xi[2]) * (sum(xi[3:]) / len(xi[3:])) ** 3,
+                np.sqrt(
+                    abs(xi[1] * xi[2] * sum(xi[3:]) + xi[1] * (1 - xi[2]) * sum(xi[3:]))
+                ),
+            )
+        elif xi[0] == "gray":
+            return (xi[2] * xi[3] ** xi[1], -(1 - xi[2]) * xi[3] ** xi[1])
         else:
-            raise ValueError("Please pick a category from '['mult2','pow2']' for X[-1].")
-
-
-class FonzecaFleming(BenchmarkFunction):
-    # test function from: https://en.wikipedia.org/wiki/Test_functions_for_optimization
-    # code based on: https://github.com/scwolof/HBMOO/blob/master/test_problems.py
-
-    def __init__ (self, n: int = 2):
-        self.name = 'FonzecaFleming'
-        self.n = n
-
-    @property
-    def invn(self):
-        return 1. / np.sqrt(self.n)
-
-    @property
-    def r (self):
-        return np.array([[0., 1.],[0., 1.]])
-
-    def get_bounds(self):
-        return np.array([[-4., 4.]]*self.n)
-
-    def f1 (self, x):
-        return self.f(x, 1)
-
-    def f2 (self, x):
-        return self.f(x, 2)
-
-    def f (self, x, i):
-        x = np.asarray(x)
-        xx = (x + (2*i-3)*self.invn)**2
-        sx = np.sum(xx, axis=None if x.ndim == 1 else 1)
-        return 1. - np.exp(-sx)
-
-    def _eval_point(self, X):
-        return self.f1(X), self.f2(X)
-
-
-class RosenbrockMulti(BenchmarkFunction):
-
-    def __init__(self, func_config: dict = None):
-        if func_config is None:
-            func_config = {}
-        self.name = 'rosenbrock_multi'
-        self.func_config = func_config
-
-    def get_bounds(self, n_dim: int = 2):
-        return [(-2.048,2.048) for _ in range(n_dim)]
-
-    def get_X_opt(self, n_dim: int = 2):
-        return [[ 1.0 for _ in range(n_dim) ]]
-
-    def _eval_point(self, X):
-        X = np.asarray_chkfinite(X)
-        X0 = X[:-1]
-        X1 = X[1:]
-
-        add1 = sum( (1.0 - X0)**2.0 )
-        add2 = 100.0 * sum( (X1 - X0**2.0)**2.0 )
-        return add1 + add2, add2*5/add1
+            raise IOError(
+                f"You provided the illegal value {xi[0]} for the categorical variable. Allowed values are "
+                f"'blue', 'orange' and 'gray'"
+            )
+
+    if n_obj == 2:
+        return np.array([compute_objectives(xi) for xi in X])
+    elif n_obj == 1:
+        return np.array([[sum(compute_objectives(xi)) / 2] for xi in X])
+    else:
+        raise IOError(
+            f"You provided the illegal value {n_obj} for the number of objectives. "
+            f"Allowed values are 1 and 2"
+        )
diff --git a/entmoot/learning/__init__.py b/entmoot/learning/__init__.py
deleted file mode 100644
index d3992a0..0000000
--- a/entmoot/learning/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from .tree_model import EntingRegressor
-
-__all__ = ("EntingRegressor")
\ No newline at end of file
diff --git a/entmoot/learning/distance_based_std.py b/entmoot/learning/distance_based_std.py
deleted file mode 100644
index 5aafb45..0000000
--- a/entmoot/learning/distance_based_std.py
+++ /dev/null
@@ -1,1252 +0,0 @@
-from sklearn.preprocessing import StandardScaler
-from abc import ABC, abstractmethod
-from entmoot.space import Space
-from typing import Optional
-
-import numpy as np
-
-
-class DistanceMetric(ABC):
-    """Computes distances and defines the optimization model for both 
-    exploration and penalty.
-
-    Parameters
-    ----------
-    -
-
-    Attributes
-    ----------
-    -
-    """
-
-    @abstractmethod
-    def get_distance(self, x_left: np.array, x_right: np.array):
-        """Compute the distance between `x_left` and `x_right` per row of x_left.
-
-        Parameters
-        ----------
-        x_left : np.array, shape (n_rows, n_dims)
-            Each row of `n_rows` is a reference point. Each dimension of `n_dim` 
-            is the numerical value of a continuous variable.
-
-        x_right : np.array, shape (n_dims,)
-            Each dimension is the numerical value of a continuous variable.
-
-        Returns
-        -------
-        dist : np.array, shape(n_rows,)
-            Distance between `x_left` and `x_right`. If multiple rows are given 
-            for `x_left`, a 1-dimensional array, i.e. `n_rows` > 1, is returned.
-        """
-
-        pass
-
-    def get_max_space_scaled_dist(self, x_shifts, x_scalers, model):
-        # computes maximum distance in search space
-        lb = np.asarray(
-            [model._cont_var_dict[i].lb for i in model._cont_var_dict]
-        )
-        ub = np.asarray(
-            [model._cont_var_dict[i].ub for i in model._cont_var_dict]
-        )
-
-        lb_std = np.divide(lb - x_shifts, x_scalers)
-        ub_std = np.divide(ub - x_shifts, x_scalers)
-
-        max_dist = self.get_distance(
-            lb_std,
-            ub_std
-        )
-        return max_dist
-
-class SquaredEuclidean(DistanceMetric):
-    """Computes distances and defines the optimization model for both 
-    exploration and penalty. The distance metric used is the squared euclidean
-    distance.
-
-    Parameters
-    ----------
-    -
-
-    Attributes
-    ----------
-    -
-    """
-
-    @staticmethod
-    def get_distance(x_left: np.array, x_right: np.array):
-        """Compute the distance between `x_left` and `x_right` per row of x_left.
-        Here the squared euclidean distance is used.
-
-        Parameters
-        ----------
-        x_left : np.array, shape (n_rows, n_dims)
-            Each row of `n_rows` is a reference point. Each dimension of `n_dim` 
-            is the numerical value of a continuous variable.
-
-        x_right : np.array, shape (n_dims,)
-            Each dimension is the numerical value of a continuous variable.
-
-        Returns
-        -------
-        dist : np.array, shape(n_rows,)
-            Distance between `x_left` and `x_right`. If multiple rows are given 
-            for `x_left`, a 1-dimensional array, i.e. `n_rows` > 1, is returned.
-        """
-
-        if x_left.ndim == 1:
-            dist = np.sum((x_left - x_right)**2)
-        else:
-            dist = np.sum((x_left - x_right)**2, axis=1)
-        return dist
-
-    def add_exploration_to_gurobi_model(self,
-        ref_points: np.array, x_shifts: np.array, x_scalers: np.array, dim_scaler: int, distance_bound: float, model,
-                                        add_rhs=None):
-        """Adds exploration constraints to a gurobi optimization model.
-        Incentivizes solutions far away from reference points.
-
-        Parameters
-        ----------
-        ref_points : np.array, shape (n_rows, n_dims)
-            Each row of `n_rows` is a reference point. Each dimension of `n_dim` 
-            is the numerical value of a continuous variable.
-
-        x_shifts : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is shifted.
-
-        x_scalers : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is scaled.
-
-        dim_scaler : int
-            Value by which the sum of uncertainties is scaled, i.e. either 1 or n_dims
-
-        distance_bound : float
-            Defines the maximum value that the exploration term can take.
-
-        Returns
-        -------
-        -
-        """
-
-        from gurobipy import quicksum
-
-        n_ref_points = len(ref_points)
-
-        # variable alpha captures distance measure
-        alpha_bound = distance_bound
-        model._alpha = \
-            model.addVar(
-                lb=0, 
-                ub=alpha_bound, 
-                name="alpha", 
-                vtype='C'
-            )
-
-        def distance_ref_point_i(model, xi_ref, x_shifts: np.array, x_scalers: np.array, dim_scaler: int, add_rhs=None):
-            # function returns constraints capturing the standardized
-            # exploration distance
-            c_x = model._cont_var_dict
-            alpha = model._alpha
-
-            diff_to_ref_point_i = quicksum(
-                ( (xi_ref[j] - (c_x[key]-x_shifts[j]) / x_scalers[j]) * \
-                    (xi_ref[j] - (c_x[key]-x_shifts[j]) / x_scalers[j]) )
-                for j,key in enumerate(model._cont_var_dict)
-            )
-
-            if add_rhs is None:
-                return dim_scaler*alpha <= diff_to_ref_point_i
-            else:
-                return dim_scaler*alpha <= diff_to_ref_point_i + add_rhs
-
-        # add exploration distances as quadratic constraints to the model
-        for i in range(n_ref_points):
-            if add_rhs:
-                temp_rhs = add_rhs[i]
-            else:
-                temp_rhs = None
-
-            model.addQConstr(
-                distance_ref_point_i(
-                    model, ref_points[i], x_shifts, x_scalers, dim_scaler, temp_rhs
-                ),
-                name=f"std_const_{i}"
-            )
-        model.update()
-
-    def add_penalty_to_gurobi_model(self,
-        ref_points: np.array, x_shifts: np.array, x_scalers: np.array, dim_scaler: int, model, add_rhs=None):
-        """Adds penalty constraints to a gurobi optimization model.
-        Incentivizes solutions close to reference points.
-
-        Parameters
-        ----------
-        ref_points : np.array, shape (n_rows, n_dims)
-            Each row of n_rows is a reference point. Each dimension of n_dim is 
-            the numerical value of a continuous variable.
-
-        x_shifts : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is shifted.
-
-        x_scalers : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is scaled.
-
-        dim_scaler : int
-            Value by which the sum of uncertainties is scaled, i.e. either 1 or n_dims
-
-        distance_bound : float
-            Defines the maximum value that the exploration term can take.
-
-        Returns
-        -------
-        -
-        """
-
-        from gurobipy import GRB, quicksum
-
-        n_ref_points = len(ref_points)
-
-        # big m is required to formulate the constraints
-        model._big_m = \
-            self.get_max_space_scaled_dist(ref_points, x_shifts, x_scalers, model)
-        
-        # binary variables b_ref correspond to active cluster centers
-        model._b_ref = \
-            model.addVars(
-                range(n_ref_points),
-                name="b_ref", 
-                vtype=GRB.BINARY
-            )
-
-        # variable alpha captures distance measure
-        model._alpha = \
-            model.addVar(
-                ub=GRB.INFINITY,
-                lb=0.0,
-                name="alpha", 
-                vtype='C'
-            )
-
-        def distance_ref_point_k(model, xk_ref, k_ref, x_shifts, x_scalers, dim_scaler, add_rhs=None):
-            # function returns constraints capturing the standardized
-            # penalty distance
-            c_x = model._cont_var_dict
-            b_ref = model._b_ref
-            alpha = model._alpha
-            n_features = len(xk_ref)
-
-            diff_to_ref_point_k = quicksum(
-                ( (xk_ref[j]  - (c_x[key]-x_shifts[j]) / x_scalers[j]) * \
-                    (xk_ref[j]  - (c_x[key]-x_shifts[j]) / x_scalers[j]) )
-                for j,key in enumerate(model._cont_var_dict)
-            )
-            big_m_term = model._big_m*(1-b_ref[k_ref])
-            if add_rhs is None:
-                return diff_to_ref_point_k <= dim_scaler*(alpha + big_m_term)
-            else:
-                return diff_to_ref_point_k + add_rhs[k_ref] <= dim_scaler*(alpha + big_m_term)
-
-        # add penalty distances as quadratic constraints to the model
-        for k in range(n_ref_points):
-            model.addQConstr(
-                distance_ref_point_k(
-                    model, 
-                    ref_points[k], 
-                    k, 
-                    x_shifts,
-                    x_scalers,
-                    dim_scaler,
-                    add_rhs
-                ),
-                name=f"std_const_{k}"
-            )
-
-        def sum_ref_point_vars(n_ref_points: int, model):
-            return quicksum(model._b_ref[k] \
-                for k in range(n_ref_points))== 1
-
-        # add additional sum constraints forcing only one ref_point to
-        # be active
-        model.addConstr(
-            sum_ref_point_vars(n_ref_points, model),
-            name="std_ref_sum"
-        )
-
-
-class Manhattan(DistanceMetric):
-    """Computes distances and defines the optimization model for both 
-    exploration and penalty. The distance metric used is the manhattan
-    distance.
-
-    Parameters
-    ----------
-    -
-
-    Attributes
-    ----------
-    -
-    """
-
-    @staticmethod
-    def get_distance(x_left: np.array, x_right: np.array):
-        """Compute the distance between `x_left` and `x_right` per row of 
-        `x_left`. Here the manhattan distance is used.
-
-        Parameters
-        ----------
-        x_left : np.array, shape (n_rows, n_dims)
-            Each row of `n_rows` is a reference point. Each dimension of n_dim is 
-            the numerical value of a continuous variable.
-
-        x_right : np.array, shape (n_dims,)
-            Each dimension is the numerical value of a continuous variable.
-
-        Returns
-        -------
-        dist : np.array, shape(n_rows,)
-            Distance between `x_left` and `x_right`. If multiple rows are given 
-            for `x_left`, a 1-dimensional array, i.e. `n_rows` > 1, is returned.
-        """
-
-        if x_left.ndim == 1:
-            dist = np.sum( np.abs(x_left - x_right) )
-        else:
-            dist = np.sum( np.abs(x_left - x_right), axis=1)
-        return dist
-
-    def add_exploration_to_gurobi_model(self,
-        ref_points: np.array, x_shifts: np.array, x_scalers: np.array, dim_scaler: int, distance_bound: float, model,
-                                        add_rhs=None):
-        """Adds exploration constraints to a gurobi optimization model.
-        Incentivizes solutions far away from reference points.
-
-        Parameters
-        ----------
-        ref_points : np.array, shape (n_rows, n_dims)
-            Each row of `n_rows` is a reference point. Each dimension of `n_dim` 
-            is the numerical value of a continuous variable.
-
-        x_shifts : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is shifted.
-
-        x_scalers : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is scaled.
-
-        dim_scaler : int
-            Value by which the sum of uncertainties is scaled, i.e. either 1 or n_dims
-
-        distance_bound : float
-            Defines the maximum value that the exploration term can take.
-
-        Returns
-        -------
-        -
-        """
-
-        from gurobipy import GRB, quicksum
-
-        n_ref_points = len(ref_points)
-        cont_feat_ids = model._cont_var_dict
-
-        # two sets of variables are used to capture positive and negative 
-        # parts of manhattan distance
-        model._c_x_aux_pos = \
-            model.addVars(range(n_ref_points), cont_feat_ids, 
-                        name="c_x_aux_pos", vtype='C')
-
-        model._c_x_aux_neg = \
-            model.addVars(range(n_ref_points), cont_feat_ids, 
-                        name="c_x_aux_neg", vtype='C')
-        
-        # variable alpha captures distance measure
-        alpha_bound = distance_bound
-        model._alpha = \
-            model.addVar(lb=0, 
-                        ub=alpha_bound, 
-                        name="alpha", vtype='C')
-
-        def distance_ref_point_i_for_feat_j(
-            model, xi_ref, i_ref, feat_j, i_var, x_shifts, x_scalers, add_rhs=None):
-            # function returns constraints capturing the standardized
-            # exploration distance
-            c_x = model._cont_var_dict
-
-            diff_to_ref_point_i = \
-                ( xi_ref[feat_j] - (c_x[i_var]-x_shifts[feat_j]) / \
-                    x_scalers[feat_j] )
-            return diff_to_ref_point_i == model._c_x_aux_pos[i_ref, i_var] - \
-                model._c_x_aux_neg[i_ref, i_var]
-
-        for i_ref in range(n_ref_points):
-            for feat_j,key in enumerate(cont_feat_ids):
-                # add constraints to capture distances in variables
-                # _c_x_aux_pos and _c_x_aux_neg
-                model.addConstr(
-                    distance_ref_point_i_for_feat_j(model, 
-                        ref_points[i_ref], i_ref, feat_j, key, x_shifts, x_scalers),
-                    name=f"std_const_feat_{key}_{i_ref}"
-                )
-                # add sos constraints that allow only one of the +/- vars, 
-                # i.e. _c_x_aux_pos / _c_x_aux_neg to be active
-                model.addSOS(GRB.SOS_TYPE1, 
-                    [
-                        model._c_x_aux_pos[i_ref, key], 
-                        model._c_x_aux_neg[i_ref, key]
-                    ]
-                )
-
-            # add exploration distances as linear constraints to the model
-            if add_rhs is None:
-                model.addConstr(
-                    dim_scaler*model._alpha <= quicksum(
-                        (model._c_x_aux_pos[i_ref, j] + model._c_x_aux_neg[i_ref, j])
-                        for j in cont_feat_ids
-                    ),
-                    name=f"alpha_sum"
-                )
-            else:
-                model.addConstr(
-                    dim_scaler*model._alpha <= quicksum(
-                        (model._c_x_aux_pos[i_ref, j] + model._c_x_aux_neg[i_ref, j])
-                        for j in cont_feat_ids
-                    ) + add_rhs[i_ref],
-                    name=f"alpha_sum"
-                )            
-        model.update()
-
-    def add_penalty_to_gurobi_model(self,
-        ref_points: np.array, x_shifts: np.array, x_scalers: np.array, dim_scaler: int, model, add_rhs=None):
-        """Adds penalty constraints to a gurobi optimization model.
-        Incentivizes solutions close to reference points.
-
-        Parameters
-        ----------
-        ref_points : np.array, shape (n_rows, n_dims)
-            Each row of `n_rows` is a reference point. Each dimension of `n_dim` 
-            is the numerical value of a continuous variable.
-
-        x_shifts : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is shifted.
-
-        x_scalers : np.array, shape (n_dims,)
-            Each dimension is the value by which the continuous variable is scaled.
-
-        dim_scaler : int
-            Value by which the sum of uncertainties is scaled, i.e. either 1 or n_dims
-
-        distance_bound : float
-            Defines the maximum value that the exploration term can take.
-
-        Returns
-        -------
-        -
-        """
-
-        from gurobipy import GRB, quicksum
-
-        n_ref_points = len(ref_points)
-        cont_feat_ids = model._cont_var_dict
-
-        # big m is required to formulate the constraints
-        model._big_m = \
-            self.get_max_space_scaled_dist(ref_points, x_shifts, x_scalers, model)
-        
-        # two sets of variables are used to capture positive and negative 
-        # parts of manhattan distance
-        model._c_x_aux_pos = \
-            model.addVars(range(n_ref_points), cont_feat_ids, 
-                        name="c_x_aux_pos", vtype='C')
-
-        model._c_x_aux_neg = \
-            model.addVars(range(n_ref_points), cont_feat_ids, 
-                        name="c_x_aux_neg", vtype='C')
-        
-        # binary variables b_ref correspond to active cluster centers
-        model._b_ref = \
-            model.addVars(n_ref_points,
-                        name="b_ref", vtype=GRB.BINARY)
-
-        # variable alpha captures distance measure
-        model._alpha = \
-            model.addVar(ub=GRB.INFINITY,
-                        lb=0.0,
-                        name="alpha", vtype='C')
-
-        def distance_ref_point_i_for_feat_j(
-            model, xi_ref, i_ref, feat_j, i_var, x_shifts, x_scalers, add_rhs=None):
-            # function returns constraints capturing the standardized
-            # exploration distance
-            c_x = model._cont_var_dict
-
-            diff_to_ref_point_i = \
-                ( xi_ref[feat_j] - (c_x[i_var]-x_shifts[feat_j]) / \
-                    x_scalers[feat_j] )
-            return diff_to_ref_point_i == model._c_x_aux_pos[i_ref, i_var] - \
-                model._c_x_aux_neg[i_ref, i_var]
-
-        for i_ref in range(n_ref_points):
-            for feat_j,key in enumerate(cont_feat_ids):
-                # add constraints to capture distances in variables
-                # _c_x_aux_pos and _c_x_aux_neg
-                model.addConstr(
-                    distance_ref_point_i_for_feat_j(
-                        model, ref_points[i_ref], i_ref, feat_j, key, x_shifts, x_scalers),
-                    name=f"std_const_feat_{key}_{i_ref}"
-                )
-                # add sos constraints that allow only one of the +/- vars, 
-                # i.e. _c_x_aux_pos / _c_x_aux_neg to be active
-                model.addSOS(GRB.SOS_TYPE1, 
-                    [
-                        model._c_x_aux_pos[i_ref, key], 
-                        model._c_x_aux_neg[i_ref, key]
-                    ]
-                )
-
-        # add penalty distances as linear constraints to the model
-            if add_rhs is None:
-                model.addConstr(
-                    quicksum(
-                        (model._c_x_aux_pos[i_ref, j] + model._c_x_aux_neg[i_ref, j])
-                        for j in cont_feat_ids
-                    ) <= dim_scaler*(model._alpha + model._big_m*(1-model._b_ref[i_ref])),
-                    name=f"alpha_sum"
-                )
-            else:
-                model.addConstr(
-                    quicksum(
-                        (model._c_x_aux_pos[i_ref, j] + model._c_x_aux_neg[i_ref, j])
-                        for j in cont_feat_ids
-                    )  + add_rhs[i_ref] <= \
-                        dim_scaler*(model._alpha + model._big_m*(1-model._b_ref[i_ref])),
-                    name=f"alpha_sum"
-                )    
-        model.update()
-
-        def sum_ref_point_vars(n_ref_points: int, model):
-            return quicksum(model._b_ref[k] for k in range(n_ref_points))== 1
-
-        # add additional sum constraints forcing only one ref_point to
-        # be active
-        model.addConstr(
-            sum_ref_point_vars(n_ref_points, model),
-            name="std_ref_sum"
-        )
-
-class NonSimilarityMetric:
-    def __init__(self):
-        pass
-
-    def get_non_similarity(self,x_left, x_right):
-        dsim = self.get_similarity(x_left, x_right)
-        n_cat = x_left.shape[1]
-        
-        dsim = (-1) * dsim + n_cat*1
-        return dsim
-
-    def update(self, Xi_cat, cat_dims):
-        pass
-
-    def get_gurobi_model_rhs(self, cat_idx, ref_points, model):
-        if cat_idx:
-            constr_list = []
-            for X in ref_points:
-                temp_constr = 0
-
-                for cat_id in range(len(X)):
-                    for cat in range(len(self.cat_mat[cat_id])):
-                        temp_cat_mat = self.cat_mat[cat_id][cat,X[cat_id]]
-                        temp_constr += \
-                            (1 - temp_cat_mat)*model._cat_var_dict[cat_idx[cat_id]][cat]
-
-                constr_list.append(temp_constr)
-            return constr_list
-        else:
-            return None
-
-
-class Overlap(NonSimilarityMetric):
-    def __init__(self):
-        pass
-
-    def update(self, cat_dims):
-
-        # generate matrix rules
-        cat_mat_rule = \
-            lambda i, j, cat_id: 1 if i==j else 0
-
-        self.cat_mat = []
-
-        for cat_id, cat in enumerate(cat_dims):
-            temp_cat = cat.transform(cat.categories)
-            n_cat = len(temp_cat)
-
-            # populate matrix
-            temp_mat = np.fromfunction(
-                np.vectorize(cat_mat_rule,), (n_cat,n_cat), dtype=int, cat_id=cat_id
-            )
-
-            self.cat_mat.append(temp_mat)
-
-    def get_similarity(self, x_left, x_right):
-        dist_cat = np.sum(x_left == x_right, axis=1)
-        return dist_cat
-
-
-class Goodall4(NonSimilarityMetric):
-    def __init__(self):
-        pass
-
-    def update(self, Xi_cat, cat_dims):
-
-        def get_pk2(cat_rows, cat):
-            count_cat = np.sum(cat_rows == cat)
-            n_rows = len(cat_rows)
-            return ( count_cat*(count_cat-1) ) / ( n_rows*(n_rows-1) )
-
-        Xi_cat = np.asarray(Xi_cat)
-
-        # generate matrix rules
-        cat_mat_rule = \
-            lambda i,j,cat_id: get_pk2( Xi_cat[:,cat_id],i ) \
-                if i == j else 0
-
-        self.cat_mat = []
-
-        for cat_id,cat in enumerate(cat_dims):
-            temp_cat = cat.transform(cat.categories)
-            n_cat = len(temp_cat)
-
-            # populate matrix
-            temp_mat = np.fromfunction(
-                np.vectorize(cat_mat_rule,), (n_cat,n_cat), dtype=int, cat_id=cat_id
-            )
-
-            self.cat_mat.append(temp_mat)
-
-    def get_similarity(self, x_left, x_right):
-        # create zero vector that contains similarities for all data points
-        sim_vec = np.zeros(x_left.shape[0])
-
-        # iterate through cat features and populate sim_mat entries
-        for cat_id in range(x_left.shape[1]):
-
-            # compute individual similarities
-            sim_vec += \
-                self.cat_mat[cat_id][int(x_right[cat_id]),x_left[:,cat_id]]
-        return sim_vec
-
-class OF(NonSimilarityMetric):
-    def __init__(self):
-        pass
-
-    def update(self, Xi_cat, cat_dims):
-
-        def get_of_frac(cat_rows, cat_left, cat_right):
-            from math import log
-            count_cat_left = np.sum(cat_rows == cat_left)
-            count_cat_right = np.sum(cat_rows == cat_right)
-            n_rows = len(cat_rows)
-            denom = \
-                1 + log( n_rows/count_cat_left )*log( n_rows/count_cat_right )
-            return 1 / denom
-
-        Xi_cat = np.asarray(Xi_cat)
-
-        # generate matrix rules
-        cat_mat_rule = \
-            lambda i,j,cat_id: get_of_frac(Xi_cat[:,cat_id], i, j) \
-                if i is not j else 1
-
-        self.cat_mat = []
-
-        for cat_id,cat in enumerate(cat_dims):
-            temp_cat = cat.transform(cat.categories)
-            n_cat = len(temp_cat)
-
-            # populate matrix
-            temp_mat = np.fromfunction(
-                np.vectorize(cat_mat_rule,), (n_cat,n_cat), dtype=int, cat_id=cat_id
-            )
-
-            self.cat_mat.append(temp_mat)
-
-    def get_similarity(self, x_left, x_right):
-        # create zero vector that contains similarities for all data points
-        sim_vec = np.zeros(x_left.shape[0])
-
-        # iterate through cat features and populate sim_mat entries
-        for cat_id in range(x_left.shape[1]):
-
-            # compute individual similarities
-            sim_vec += \
-                self.cat_mat[cat_id][int(x_right[cat_id]), x_left[:, cat_id]]
-        return sim_vec
-
-
-class DistanceBasedStd(ABC):
-    """
-    Define a distance-based standard estimator.
-
-    A `DistanceBasedStd` object is used to quantify model uncertainty based 
-    on distance to reference points, e.g. data points. The underlying assumption 
-    is that base estimator predictions are good close to training data.
-
-    Use this class as a template if you want to develop your own distance-based
-    measure.
-
-    :params space: Space, search space object that contains all vars
-    :params unc_scaling: str, type of scaling used for uncertainty
-        "standard": uses sample mean as shift and sample std as scaling
-        "normalize": normalizes vars by lower and upper bounds provided
-    :params dist_metric: str,
-        type of distance metric used for non-categorical vars, i.e.
-        currently supported ["squared_euclidean", "manhattan"]
-    :params cat_dist_metric: str,
-        type of dist. metric used for uncertainty of categorical vars, i.e. check:
-            {S. Boriah, V. Chandola, V. Kumar,
-            Similarity Measures for Categorical Data: A Comparative Evaluation,
-            2008, SIAM}
-        currently available options in ["overlap", "goodall4", "of"]
-    """
-    def __init__(self,
-        space: Space,
-        unc_scaling: str = "standard",
-        dist_metric: str = 'squared_euclidean',
-        cat_dist_metric: str = 'overlap'):
-
-        self.std_type = 'distance'
-
-        self.unc_scaling = unc_scaling
-
-        # set distance metric for cont variables
-        if dist_metric == 'squared_euclidean':
-            from entmoot.learning.distance_based_std import SquaredEuclidean
-            self.dist_metric = SquaredEuclidean()
-
-        elif dist_metric == 'manhattan':
-            from entmoot.learning.distance_based_std import Manhattan
-            self.dist_metric = Manhattan()
-
-        self.space = space
-
-        # set similarity metric for cat variables
-        if cat_dist_metric == "overlap":
-            from entmoot.learning.distance_based_std import Overlap
-            self.cat_sim_metric = Overlap()
-            
-        elif cat_dist_metric == "goodall4":
-            from entmoot.learning.distance_based_std import Goodall4
-            self.cat_sim_metric = Goodall4()
-
-        elif cat_dist_metric == "of":
-            from entmoot.learning.distance_based_std import OF
-            self.cat_sim_metric = OF()
-
-    def get_x_cont_vals(self, xi):
-        return [x for idx,x in enumerate(xi) \
-            if idx not in self.cat_idx]
-
-    def get_x_cat_vals(self, xi):
-        return [x for idx,x in enumerate(xi) \
-            if idx in self.cat_idx]
-
-
-    def set_params(self, **kwargs):
-        """Sets parameters related to distance-based standard estimator.
-
-        Parameters
-        ----------
-        kwargs : dict
-            Additional arguments to be passed to the standard estimator
-
-        Returns
-        -------
-        -
-        """
-        pass
-
-    def update(self, Xi: list, yi: float, cat_column=None):
-        """Update available data points which is usually done after every
-        iteration.
-
-        Parameters
-        ----------
-        Xi : list
-            Points at which objective has been evaluated.
-        yi : scalar
-            Values of objective at corresponding points in `Xi`.
-
-        Returns
-        -------
-        -
-        """
-        if cat_column is None:
-            cat_column = []
-
-        self.cat_idx = cat_column
-        
-        # collect cont and cat variable values
-        if self.cat_idx:
-            self.Xi_cont = []
-            self.Xi_cat = []
-            
-            for xi in Xi:
-
-                # check if all 
-
-                # collect continuous vars
-                # check if all vars are categorical
-                if len(self.cat_idx) != len(xi):
-                    self.Xi_cont.append( [x for idx,x in enumerate(xi) \
-                        if idx not in self.cat_idx]
-                    )
-
-                # collect categorical vars
-                self.Xi_cat.append( [x for idx,x in enumerate(xi) \
-                    if idx in self.cat_idx]
-                )
-        else:
-            self.Xi_cont = Xi
-            self.Xi_cat = []
-
-        # update data set attributes
-        self.Xi = Xi
-        self.yi = yi
-
-        self.n_features = Xi.shape[1]
-
-        if list(self.Xi_cont):
-            # standardize points for comparison
-            if self.unc_scaling == "standard":
-                standard_scaler = StandardScaler()
-                projected_features = standard_scaler.fit_transform(self.Xi_cont)
-                self.x_shifts = standard_scaler.mean_
-                self.x_scalers = standard_scaler.scale_
-                self.dim_scaler = 1
-
-                # standardize dataset
-                self.Xi_cont_scaled = self.scale_with_Xi(self.Xi_cont)
-
-            # normalize points for comparison
-            elif self.unc_scaling == "normalize":
-                self.x_shifts = np.asarray(
-                    [dim.low
-                     for idx, dim in enumerate(self.space.dimensions)
-                     if idx not in self.cat_idx]
-                )
-
-                self.x_scalers = np.asarray(
-                    [dim.high - dim.low
-                     for idx, dim in enumerate(self.space.dimensions)
-                     if idx not in self.cat_idx]
-                )
-                self.dim_scaler = len(self.space.dimensions)
-
-                # standardize dataset
-                self.Xi_cont_scaled = self.scale_with_Xi(self.Xi_cont)
-
-        # update cat_sim_metric with new data
-        cat_dims = self.get_x_cat_vals(self.space.dimensions)
-        self.cat_sim_metric.update(self.Xi_cat, cat_dims)
-
-    def scale_with_Xi(self, X: np.array):
-        """Standardize given input `X` based on attribute `Xi`.
-
-        Parameters
-        ----------
-        X : numpy array, shape (n_rows, n_dims)
-            Each row of n_rows is a point in `X`.
-            Points which are standardized based on `Xi`
-
-        Returns
-        -------
-        x_standard : numpy array, shape (n_rows, n_dims)
-            Standardized array of `X`
-        """
-        x_standard = np.divide(X - self.x_shifts, self.x_scalers)
-        return x_standard
-
-    def get_closest_point_distance(self, X: np.array):
-        """Get distance to point of attribute `ref_points` which is closest to 
-        point given as parameter `X`.
-
-        Parameters
-        ----------
-        X : numpy array, shape (n_dims,)
-            Point to which the distance of closest reference point is 
-            computed
-
-        Returns
-        -------
-        dist : numpy array, shape (1, )
-            Returns distance to closest `ref_point`.
-        """
-        if list(self.Xi_cont):
-            ref_points_cont = np.asarray(self.ref_points_cont)
-
-            X_cont = self.get_x_cont_vals(X)
-
-            # standardize cont variables
-            x_cont_scaled = np.asarray(self.scale_with_Xi(X_cont))
-
-            # compute all distances for cont variables
-            dist_cont = \
-                self.dist_metric.get_distance(ref_points_cont, x_cont_scaled)
-
-        # compute all distances for cat variables
-        if self.Xi_cat:
-            X_cat = X[self.cat_idx]
-            dist_cat = \
-                self.cat_sim_metric.get_non_similarity(
-                    self.ref_points_cat, X_cat)
-
-            if list(self.Xi_cont):
-                dist_cont += dist_cat
-            else:
-                dist_cont = dist_cat
-
-        return np.min(dist_cont)
-
-    def predict(self, X: np.array):
-        """Predict standard estimate at location `X`.
-
-        Parameters
-        ----------
-        X : numpy array, shape (n_rows, n_dims)
-            Points at which the standard estimator is evaluated.
-
-        Returns
-        -------
-        dist : numpy array, shape (n_rows,)
-            Returns distances to closest `ref_point` for every point per row
-            in `n_rows`.
-        """
-        dist = np.empty([X.shape[0],])
-
-        for row_res,Xi in enumerate(X):
-            ref_distance = self.get_closest_point_distance(Xi)
-            dist[row_res] = ref_distance
-
-        if self.unc_scaling == "standard":
-            return dist
-        elif self.unc_scaling == "normalize":
-            return dist / len(self.space.dimensions)
-
-    @abstractmethod
-    def add_to_gurobi_model(self, model):
-        """Add standard estimator to gurobi model. Model details are 
-        implemented in child class.
-
-        Parameters
-        ----------
-        model : gurobipy.Model,
-            Model to which the standard estimator is added.
-        """
-        pass
-
-    @abstractmethod
-    def get_gurobi_obj(self, model):
-        """Get contribution of standard estimator to gurobi model objective
-        function.
-
-        Parameters
-        ----------
-        model : gurobipy.Model,
-            Model to which the standard estimator is added.
-
-        s
-        ----------
-        alpha : gurobipy.Var,
-            Model variable that takes the value of the uncertainty measure.
-        """
-        pass
-
-
-class DistanceBasedExploration(DistanceBasedStd):
-    """Defines a child class based on `DistanceBasedStd`. Exploration 
-    refers to how the distance measure contributes to the acquisition
-    function. Exploration refers to incentivizing distance to reference points
-    leading to a negative contribution of the distance measure to the objective
-    function.
-
-
-    :params space: Space, search space object that contains all vars
-    :params unc_scaling: str, type of scaling used for uncertainty
-        "standard": uses sample mean as shift and sample std as scaling
-        "normalize": normalizes vars by lower and upper bounds provided
-    :params dist_metric: str,
-        type of distance metric used for non-categorical vars, i.e.
-        currently supported ["squared_euclidean", "manhattan"]
-    :params cat_dist_metric: str,
-        type of dist. metric used for uncertainty of categorical vars, i.e. check:
-            {S. Boriah, V. Chandola, V. Kumar,
-            Similarity Measures for Categorical Data: A Comparative Evaluation,
-            2008, SIAM}
-        currently available options in ["overlap", "goodall4", "of"]
-    :params zeta: Optional[float],
-        coefficient determining how the distance measure is bounded,
-        i.e. bound = abs(zeta * np.var(yi))
-    """
-    def __init__(self,
-        space: Space,
-        unc_scaling: str = "standard",
-        dist_metric: str = 'squared_euclidean',
-        cat_dist_metric: str = 'overlap',
-        zeta: Optional[float] = 0.5):
-
-        super().__init__(space,
-                         unc_scaling=unc_scaling,
-                         dist_metric=dist_metric,
-                         cat_dist_metric=cat_dist_metric)
-        self.zeta = zeta
-
-    def set_params(self, **kwargs):
-        """Sets parameters related to distance-based standard estimator.
-
-        Parameters
-        ----------
-        kwargs : dict
-            Additional arguments to be passed to the standard estimator
-
-        Returns
-        -------
-        -
-        """
-        zeta = kwargs.get("zeta", 0.5)
-        self.zeta = zeta
-
-    def update(self, Xi: list, yi: float, cat_column=None):
-        """Update available data points which is usually done after every
-        iteration.
-
-        Parameters
-        ----------
-        Xi : list
-            Points at which objective has been evaluated.
-        yi : scalar
-            Values of objective at corresponding points in `Xi`.
-
-        Returns
-        -------
-        -
-        """
-        if cat_column is None:
-            cat_column = []
-
-        super().update(Xi, yi, cat_column=cat_column)
-
-        self.ref_points_unscaled = self.Xi_cont
-
-        if list(self.Xi_cont):
-            self.ref_points_cont = self.Xi_cont_scaled
-        self.ref_points_cat = np.asarray(self.Xi_cat, dtype=int)
-
-        # compute upper bound of uncertainty
-        if self.unc_scaling == "standard":
-            y_var = np.var(yi)
-            self.distance_bound = abs(self.zeta*y_var)
-        elif self.unc_scaling == "normalize":
-            self.distance_bound = 1.0
-
-    def predict(self, X):
-        """Predict standard estimate at location `X`. By default `dist` is
-        bounded by attribute `distance_bound`.
-
-        Parameters
-        ----------
-        X : numpy array, shape (n_rows, n_dims)
-            Points at which the standard estimator is evaluated.
-
-        Returns
-        -------
-        dist : numpy array, shape (n_rows,)
-            Returns distances to closest `ref_point` for every point per row
-            in `n_rows`.
-        """
-        dist = super().predict(X)
-
-        # prediction has max out at `distance_bound`
-        dist[dist > self.distance_bound] = self.distance_bound
-
-        if self.unc_scaling == "standard":
-            return dist
-        elif self.unc_scaling == "normalize":
-            return dist / len(self.space.dimensions)
-
-    def get_gurobi_obj(self, model):
-        """Get contribution of standard estimator to gurobi model objective
-        function.
-
-        Parameters
-        ----------
-        model : gurobipy.Model,
-            Model to which the standard estimator is added.
-
-        Returns
-        -------
-        alpha : gurobipy.Var,
-            Model variable that takes the value of the uncertainty measure.
-        """
-        # negative contributation of alpha requires non-convex flag in gurobi.
-        model.Params.NonConvex=2
-        return -model._alpha
-
-    def add_to_gurobi_model(self, model):
-        """Add standard estimator to gurobi model. Model details are 
-        implemented in child class.
-
-        Parameters
-        ----------
-        model : gurobipy.Model,
-            Model to which the standard estimator is added.
-        """
-        cat_rhs = self.cat_sim_metric.get_gurobi_model_rhs(
-            self.cat_idx, self.ref_points_cat, model
-        )
-
-        self.dist_metric.add_exploration_to_gurobi_model(
-            self.ref_points_cont,
-            self.x_shifts,
-            self.x_scalers,
-            self.dim_scaler,
-            self.distance_bound, 
-            model,
-            add_rhs=cat_rhs
-        )
-
-class DistanceBasedPenalty(DistanceBasedStd):
-    """Defines a child class based on `DistanceBasedStd`. Penalty
-    refers to how the distance measure contributes to the acquisition
-    function. Penalty refers to penalizing distance to reference points
-    leading to a positive contribution of the distance measure to the objective
-    function.
-
-
-    :params space: Space, search space object that contains all vars
-    :params unc_scaling: str, type of scaling used for uncertainty
-        "standard": uses sample mean as shift and sample std as scaling
-        "normalize": normalizes vars by lower and upper bounds provided
-    :params dist_metric: str,
-        type of distance metric used for non-categorical vars, i.e.
-        currently supported ["squared_euclidean", "manhattan"]
-    :params cat_dist_metric: str,
-        type of dist. metric used for uncertainty of categorical vars, i.e. check:
-            {S. Boriah, V. Chandola, V. Kumar,
-            Similarity Measures for Categorical Data: A Comparative Evaluation,
-            2008, SIAM}
-        currently available options in ["overlap", "goodall4", "of"]
-    """
-
-    def __init__(self,
-        space: Space,
-        unc_scaling: str = "standard",
-        dist_metric: str = 'squared_euclidean',
-        cat_dist_metric: str = 'overlap'):
-
-        super().__init__(space,
-                         unc_scaling=unc_scaling,
-                         dist_metric=dist_metric,
-                         cat_dist_metric=cat_dist_metric)
-
-    def set_params(self, **kwargs):
-        """Sets parameters related to distance-based standard estimator.
-
-        Parameters
-        ----------
-        kwargs : dict
-            Additional arguments to be passed to the standard estimator
-
-        Returns
-        -------
-        -
-        """
-        pass
-
-    def update(self, Xi: list, yi: float, cat_column=None):
-        """Update available data points which is usually done after every
-        iteration.
-
-        Parameters
-        ----------
-        Xi : list
-            Points at which objective has been evaluated.
-        yi : scalar
-            Values of objective at corresponding points in `Xi`.
-
-        Returns
-        -------
-        -
-        """
-        if cat_column is None:
-            cat_column = []
-
-        super().update(Xi, yi, cat_column=cat_column)
-
-        self.ref_points_unscaled = self.Xi_cont
-
-        #self.ref_points = self.Xi_standard
-
-        if list(self.Xi_cont):
-            self.ref_points_cont = self.Xi_cont_scaled
-        self.ref_points_cat = np.asarray(self.Xi_cat, dtype=int)
-
-
-    def predict(self, X: np.array):
-        """Predict standard estimate at location `X`. Sign of `dist` is negative
-        because it contributes as a penalty.
-
-        Parameters
-        ----------
-        X : numpy array, shape (n_rows, n_dims)
-            Points at which the standard estimator is evaluated.
-
-        Returns
-        -------
-        dist : numpy array, shape (n_rows,)
-            Returns distances to closest `ref_point` for every point per row
-            in `n_rows`.
-        """
-        dist = super().predict(X)
-        if self.unc_scaling == "standard":
-            return -dist
-        elif self.unc_scaling == "normalize":
-            return -dist / len(self.space.dimensions)
-
-    def get_gurobi_obj(self, model):
-        """Get contribution of standard estimator to gurobi model objective
-        function.
-
-        Parameters
-        ----------
-        model : gurobipy.Model,
-            Model to which the standard estimator is added.
-
-        Returns
-        -------
-        alpha : gurobipy.Var,
-            Model variable that takes the value of the uncertainty measure.
-        """
-        return model._alpha
-
-    def add_to_gurobi_model(self, model):
-        """Add standard estimator to gurobi model. Model details are 
-        implemented in child class.
-
-        Parameters
-        ----------
-        model : gurobipy.Model,
-            Model to which the standard estimator is added.
-        """
-        cat_rhs = self.cat_sim_metric.get_gurobi_model_rhs(
-            self.cat_idx, self.ref_points_cat, model
-        )
-
-        self.dist_metric.add_penalty_to_gurobi_model(
-            self.ref_points_cont, 
-            self.x_shifts,
-            self.x_scalers,
-            self.dim_scaler,
-            model,
-            add_rhs=cat_rhs
-        )
\ No newline at end of file
diff --git a/entmoot/learning/gbm_model.py b/entmoot/learning/gbm_model.py
deleted file mode 100644
index 252a4e8..0000000
--- a/entmoot/learning/gbm_model.py
+++ /dev/null
@@ -1,484 +0,0 @@
-import collections as coll
-
-
-class GbmModel:
-    """Define a gbm model.
-
-    A `GbmModel` is a object-oriented class to enocode tree-based models from
-    different libraries. Supported tree training libraries are: LightGBM.
-
-    Use this class as an interface to read in model data from tree model
-    training libraries and translate it to the ENTMOOT native structure.
-
-    Parameters
-    ----------
-    tree_list : list
-        Ordered tree list that obeys ENTMOOT input specifications
-
-    Attributes
-    ----------
-    -
-
-    """
-    def __init__(self, tree_list: list):
-        self.load_ordered_tree_dict(tree_list)
-
-    def _build_tree(self, tree):
-        """Define `GbmNode` based on single tree.
-
-        Parameters
-        ----------
-        trees : list of dicts,
-            Specifies list of node dicts
-        n_trees : number of trees in `tree_list`
-
-        Returns
-        -------
-        gbm_node : GbmNode
-        """
-        node = tree.pop(0)
-        split_var = node['split_var']
-        split_code_pred = node['split_code_pred']
-
-        return GbmNode(
-            split_var=split_var,
-            split_code_pred=split_code_pred,
-            tree=tree
-        )
-
-    def load_ordered_tree_dict(self, tree_list: list):
-        """Define attributes used in `GbmModel`.
-
-        Parameters
-        ----------
-        tree_list : list of trees,
-            Specifies list of trees which contain list of node dicts
-
-        Returns
-        -------
-        -
-        """
-        self.n_trees = len(tree_list)
-        self.trees = [self._build_tree(tree) for tree in tree_list]
-
-    def get_leaf_encodings(self, tree):
-        yield from self.trees[tree].get_leaf_encodings()
-
-    def get_branch_encodings(self, tree):
-        yield from self.trees[tree].get_branch_encodings()
-
-    def get_leaf_weight(self, tree, encoding):
-        return self.trees[tree].get_leaf_weight(encoding)
-
-    def get_leaf_weights(self, tree):
-        return self.trees[tree].get_leaf_weights()
-
-    def get_branch_partition_pair(self, tree, encoding):
-        return self.trees[tree].get_branch_partition_pair(encoding)
-
-    def get_left_leaves(self, tree, encoding):
-        yield from (
-            encoding + s
-            for s in self.trees[tree].get_left_leaves(encoding)
-        )
-
-    def get_right_leaves(self, tree, encoding):
-        yield from (
-            encoding + s
-            for s in self.trees[tree].get_right_leaves(encoding)
-        )
-
-    def get_branch_partition_pairs(self, tree, leaf_encoding):
-        yield from self.trees[tree].get_branch_partition_pairs(leaf_encoding)
-
-    def get_participating_variables(self, tree, leaf):
-        return set(self.trees[tree].get_participating_variables(leaf))
-
-    def get_all_participating_variables(self, tree):
-        return set(self.trees[tree].get_all_participating_variables())
-
-    def get_var_lower(self, tree, leaf, var, lower):
-        return self.trees[tree].get_var_lower(leaf, var, lower)
-
-    def get_var_upper(self, tree, leaf, var, upper):
-        return self.trees[tree].get_var_upper(leaf, var, upper)
-
-    def get_var_interval(self, tree, leaf, var, var_interval):
-        return self.trees[tree].get_var_interval(leaf, var, var_interval)
-
-    def get_var_break_points(self):
-        var_breakpoints = {}
-        for tree in self.trees:
-            for var, breakpoint in tree.get_all_partition_pairs():
-                try:
-                    if isinstance(breakpoint, list):
-                        var_breakpoints[var].append(breakpoint)
-                    else:
-                        var_breakpoints[var].add(breakpoint)
-                except KeyError:
-                    if isinstance(breakpoint, list):
-                        var_breakpoints[var] = [breakpoint]
-                    else:    
-                        var_breakpoints[var] = set([breakpoint])
-
-        for k in var_breakpoints:
-            if isinstance(var_breakpoints[k], set):
-                var_breakpoints[k] = sorted(var_breakpoints[k])
-
-        return var_breakpoints
-
-    def get_leaf_count(self):
-        resulting_counter = sum(
-            (tree.get_leaf_count() for tree in self.trees),
-            coll.Counter()
-        )
-        del resulting_counter['leaf']
-        return resulting_counter
-
-    def get_active_leaves(self, X):
-        all_active_leaves = []
-        for tree in self.trees:
-            active_leaf = []
-            tree._populate_active_leaf_encodings(active_leaf, X)
-            all_active_leaves.append(''.join(active_leaf))
-        return all_active_leaves
-
-    def get_active_area(self, X, cat_idx=[], space=None, volume=False):
-            active_splits = {}
-            for tree in self.trees:
-                tree._populate_active_splits(active_splits, X)
-
-            all_active_splits = {}
-            for idx,dim in enumerate(space.dimensions):
-                if idx not in cat_idx:
-
-                    if idx in active_splits:
-                        # if tree splits on this conti var
-                        all_active_splits[idx] = active_splits[idx]
-                        all_active_splits[idx].insert(0,dim.transformed_bounds[0])
-                        all_active_splits[idx].append(dim.transformed_bounds[1])
-                    else:
-                        # add actual bounds of var if tree doesn't split on var
-                        all_active_splits[idx] = \
-                            [
-                                dim.transformed_bounds[0],
-                                dim.transformed_bounds[1]
-                            ]
-            # sort all splits and extract modified bounds for vars
-            for key in all_active_splits:
-                all_active_splits[key] = \
-                    sorted( list(set(all_active_splits[key])) )[:2]
-
-            # return hypervolume if required
-            if volume:
-                hyper_vol = 1
-                for key in all_active_splits:
-                    hyper_vol *= \
-                        abs(all_active_splits[key][0] - all_active_splits[key][1])
-                return all_active_splits, hyper_vol
-            else:
-                return all_active_splits
-
-
-class GbmType:
-    def _populate_active_splits(self, active_splits, X):
-        if isinstance(self.split_code_pred, list):
-            if X[self.split_var] in self.split_code_pred:
-                self.left._populate_active_splits(active_splits, X)
-            else:
-                self.right._populate_active_splits(active_splits, X)
-        else:
-            if self.split_var != -1:                
-                if not self.split_var in active_splits:
-                    active_splits[self.split_var] = []
-
-                if X[self.split_var] <= self.split_code_pred:
-                    self.left._populate_active_splits(active_splits, X)
-                    active_splits[self.split_var].append(self.split_code_pred)
-                else:
-                    self.right._populate_active_splits(active_splits, X)
-
-    def _populate_active_leaf_encodings(self, active_leaf, X):
-        if isinstance(self.split_code_pred, list):
-            if X[self.split_var] in self.split_code_pred:
-                active_leaf.append('0')
-                self.left._populate_active_leaf_encodings(active_leaf, X)
-            else:
-                active_leaf.append('1')
-                self.right._populate_active_leaf_encodings(active_leaf, X)
-        else:
-            if self.split_var != -1:                
-                if X[self.split_var] <= self.split_code_pred:
-                    active_leaf.append('0')
-                    self.left._populate_active_leaf_encodings(active_leaf, X)
-                else:
-                    active_leaf.append('1')
-                    self.right._populate_active_leaf_encodings(active_leaf, X)
-
-
-class GbmNode(GbmType):
-    """Defines a gbm node which can be a split or leaf.
-
-    Initializing `GbmNode` triggers recursive initialization of all nodes that 
-    appear in a tree
-
-    Parameters
-    ----------
-    split_var : int
-        Index of split variable used
-    split_code_pred : list
-        Value that defines split
-    tree : list
-        List of node dicts that define the tree
-
-    Attributes
-    ----------
-    split_var : int
-        Index of split variable used
-    split_code_pred : list
-        Value that defines split
-    tree : list
-        List of node dicts that define the tree
-    """
-    def __init__(
-        self,
-        split_var,
-        split_code_pred,
-        tree
-    ):
-        self.split_var = split_var
-        self.split_code_pred = split_code_pred
-
-        # check if tree is empty
-        if not tree:
-            print("EmptyTreeModelError: LightGBM was not able to train a "
-            "tree model based on your parameter specifications. This can "
-            "usually be fixed by increasing the number of `n_initial_points` or "
-            "reducing `min_child_samples` via `base_estimator_kwargs` (default "
-            "is 2). Alternatively, change `acq_optimizer='sampling'`.")
-            import sys
-            sys.exit(1)
-
-        # read left node 
-        node = tree.pop(0)
-        split_var = node['split_var']
-        split_code_pred = node['split_code_pred']
-
-        # split_var value of -1 refers to leaf node
-        if split_var == -1:
-            self.left = LeafNode(
-                split_code_pred = split_code_pred
-            )
-        else:
-            self.left = GbmNode(
-                split_var = split_var,
-                split_code_pred = split_code_pred,
-                tree = tree
-            )
-
-        # read right node
-        node = tree.pop(0)
-        split_var = node['split_var']
-        split_code_pred = node['split_code_pred']
-
-        if split_var == -1:
-            self.right = LeafNode(
-                split_code_pred=split_code_pred
-            )
-        else:
-            self.right = GbmNode(
-                split_var=split_var,
-                split_code_pred=split_code_pred,
-                tree=tree
-            )
-
-    def __repr__(self):
-        return ', '.join([
-            str(x)
-            for x in [
-                self.split_var,
-                self.split_code_pred]
-        ])
-
-    def _get_next_node(self, direction):
-        return self.right if int(direction) else self.left
-
-    def get_leaf_encodings(self, current_string=''):
-        yield from self.left.get_leaf_encodings(current_string+'0')
-        yield from self.right.get_leaf_encodings(current_string+'1')
-
-    def get_branch_encodings(self, current_string=''):
-        yield current_string
-        yield from self.left.get_branch_encodings(current_string+'0')
-        yield from self.right.get_branch_encodings(current_string+'1')
-
-    def get_leaf_weight(self, encoding):
-        next_node = self.right if int(encoding[0]) else self.left
-        return next_node.get_leaf_weight(encoding[1:])
-
-    def get_leaf_weights(self):
-        yield from self.left.get_leaf_weights()
-        yield from self.right.get_leaf_weights()
-
-    def get_branch_partition_pair(self, encoding):
-        if not encoding:
-            return self.split_var, self.split_code_pred
-        else:
-            next_node = self._get_next_node(encoding[0])
-            return next_node.get_branch_partition_pair(encoding[1:])
-
-    def get_left_leaves(self, encoding):
-        if encoding:
-            next_node = self._get_next_node(encoding[0])
-            yield from next_node.get_left_leaves(encoding[1:])
-        else:
-            yield from self.left.get_leaf_encodings('0')
-
-    def get_right_leaves(self, encoding):
-        if encoding:
-            next_node = self._get_next_node(encoding[0])
-            yield from next_node.get_right_leaves(encoding[1:])
-        else:
-            yield from self.right.get_leaf_encodings('1')
-
-    def get_branch_partition_pairs(self, encoding):
-        yield (self.split_var, self.split_code_pred)
-        try:
-            next_node = self.right if int(encoding[0]) else self.left
-            next_gen = next_node.get_branch_partition_pairs(encoding[1:])
-        except IndexError:
-            next_gen = []
-        yield from next_gen
-
-    def get_participating_variables(self, encoding):
-        yield self.split_var
-        next_node = self.right if int(encoding[0]) else self.left
-        yield from next_node.get_participating_variables(encoding[1:])
-
-    def get_all_participating_variables(self):
-        yield self.split_var
-        yield from self.left.get_all_participating_variables()
-        yield from self.right.get_all_participating_variables()
-
-    def get_all_partition_pairs(self):
-        yield (self.split_var, self.split_code_pred)
-        yield from self.left.get_all_partition_pairs()
-        yield from self.right.get_all_partition_pairs()
-
-    def get_var_lower(self, encoding, var, lower):
-        if encoding:
-            if self.split_var == var:
-                if int(encoding[0]) == 1:
-                    assert lower <= self.split_code_pred
-                    lower = self.split_code_pred
-            next_node = self.right if int(encoding[0]) else self.left
-            return next_node.get_var_lower(encoding[1:], var, lower)
-        else:
-            return lower
-
-    def get_var_upper(self, encoding, var, upper):
-        if encoding:
-            if self.split_var == var:
-                if int(encoding[0]) == 0:
-                    assert upper >= self.split_code_pred
-                    upper = self.split_code_pred
-            next_node = self.right if int(encoding[0]) else self.left
-            return next_node.get_var_upper(encoding[1:], var, upper)
-        else:
-            return upper
-
-    def get_var_interval(self, encoding, var, var_interval):
-        if int(encoding[0]):
-            next_node = self.right
-            if self.split_var == var:
-                var_interval = (self.split_code_pred, var_interval[1])
-        else:
-            next_node = self.left
-            if self.split_var == var:
-                var_interval = (var_interval[0], self.split_code_pred)
-        return next_node.get_var_interval(encoding[1:], var, var_interval)
-
-    def get_leaf_count(self):
-        left_count = self.left.get_leaf_count()
-        right_count = self.right.get_leaf_count()
-        joint_count = left_count+right_count
-
-        left_key = (self.split_var, self.split_code_pred, 0)
-        right_key = (self.split_var, self.split_code_pred, 1)
-        joint_count[left_key] += left_count['leaf']
-        joint_count[right_key] += right_count['leaf']
-        return joint_count
-
-
-class LeafNode(GbmType):
-    """Defines a child class of `GbmType`. Leaf nodes have `split_var = -1` and 
-    `split_code_pred` as leaf value defined by training."""
-    def __init__(self,
-                 split_code_pred):
-        self.split_var = -1
-        self.split_code_pred = split_code_pred
-
-    def __repr__(self):
-        return ', '.join([
-            str(x)
-            for x in [
-                'Leaf',
-                self.split_code_pred]
-        ])
-
-    def switch_to_maximisation(self):
-        """Changes the sign of tree model prediction by changing signs of 
-        leaf values."""
-        self.split_code_pred = -1*self.split_code_pred
-
-
-    def get_leaf_encodings(self, current_string=''):
-        yield current_string
-
-    def get_branch_encodings(self, current_string=''):
-        yield from []
-
-    def get_leaf_weight(self, encoding):
-        assert not encoding
-        return self.split_code_pred
-
-    def get_leaf_weights(self):
-        yield self.split_code_pred
-
-    def get_branch_partition_pair(self, encoding):
-        raise Exception('Should not get here.')
-
-    def get_left_leaves(self, encoding):
-        raise Exception('Should not get here.')
-
-    def get_right_leaves(self, encoding):
-        raise Exception('Should not get here.')
-
-    def get_branch_partition_pairs(self, encoding):
-        assert not encoding
-        yield from []
-
-    def get_participating_variables(self, encoding):
-        assert not encoding
-        yield from []
-
-    def get_all_participating_variables(self):
-        yield from []
-
-    def get_all_partition_pairs(self):
-        yield from []
-
-    def get_var_lower(self, encoding, lower):
-        assert not encoding
-        return lower
-
-    def get_var_upper(self, encoding, upper):
-        assert not encoding
-        return upper
-
-    def get_var_interval(self, encoding, var_interval):
-        assert not encoding
-        return var_interval
-
-    def get_leaf_count(self):
-        return coll.Counter({'leaf': 1})
\ No newline at end of file
diff --git a/entmoot/learning/lgbm_processing.py b/entmoot/learning/lgbm_processing.py
deleted file mode 100644
index f33402f..0000000
--- a/entmoot/learning/lgbm_processing.py
+++ /dev/null
@@ -1,128 +0,0 @@
-
-def order_tree_model_dict(tree_model_dict: dict, cat_column=None):
-    """
-    Sorts a tree model dict provided by LightGBM from .json file to make it 
-    suitable for use in ENTMOOT.
-
-    Key name `tree_info` is specific to LightGBM and needs to be changed for
-    other tree training libraries that support model dump to .json format.
-
-    Parameters
-    ----------
-    tree_model_dict : dict,
-        Includes all information on trees trained in LightGBM
-
-    Returns
-    -------
-    ordered_tree_list : list,
-        Ordered tree list compatible with `GbmModel`
-    """
-    if cat_column is None:
-        cat_column = []
-    tree_list = tree_model_dict['tree_info']
-    ordered_tree_list = order_tree_list(tree_list, cat_column=cat_column)
-    return ordered_tree_list
-
-
-def order_tree_list(tree_list, cat_column=None):
-    """
-    Sorts a tree list provided by LightGBM from .json file to make it 
-    suitable for use in ENTMOOT.
-
-    Parameters
-    ----------
-    tree_list : list,
-        Unsorted list of trees from LightGBM
-
-    Returns
-    -------
-    ordered_tree_list : list,
-        Ordered tree list compatible with `GbmModel`
-    """
-    if cat_column is None:
-        cat_column = []
-    ordered_tree_list = []
-    for tree in tree_list:
-        ordered_node_list = order_node_list(tree, cat_column=cat_column)
-        if ordered_node_list:
-            ordered_tree_list.append(ordered_node_list)
-    return ordered_tree_list
-
-
-def order_node_list(tree, cat_column=None):
-    """
-    Sorts a list of node dict from a LightGBM instance. Key `tree_structure` is 
-    specific for LightGBM.
-
-    Parameters
-    ----------
-    tree : list,
-        Unsorted list of node dicts
-
-    Returns
-    -------
-    ordered_node_list : list,
-        Ordered list of node dicts compatible with `GbmModel`
-    """
-    if cat_column is None:
-        cat_column = []
-    node = []
-    node.append(tree['tree_structure'])
-    ordered_node_list = []
-    add_next_nodes(ordered_node_list, node, cat_column=cat_column)
-    return ordered_node_list
-
-
-def add_next_nodes(ordered_node_list, node, cat_column=None):
-    """
-    Processes LightGBM node and adds sorted node to `ordered_node_list`.
-
-    Parameters
-    ----------
-    ordered_node_list : list,
-        List of nodes to which sorted node is added
-
-    node : list,
-        List of unsorted LightGBM nodes for one tree
-
-    Returns
-    -------
-    -
-    """
-    if cat_column is None:
-        cat_column = []
-
-    if node:   
-        new_node = {}
-        
-        try:
-            new_node['split_var'] = node[-1]['split_feature']
-            if not new_node['split_var'] in cat_column:
-                # read numerical variables
-                temp_node_val = node[-1]['threshold']
-                new_node['split_code_pred'] = round(temp_node_val,5)
-            else:
-                # read categorical variables
-                cat_set = node[-1]['threshold'].split("||")
-                temp_node_val = [int(cat) for cat in cat_set]
-                new_node['split_code_pred'] = temp_node_val
-
-        except KeyError:
-            new_node['split_var'] = -1
-            temp_node_val = node[-1]['leaf_value']
-            new_node['split_code_pred'] = round(temp_node_val,5)
-
-        ordered_node_list.append(new_node)
-
-        temp_node = node.pop(-1)
-        try:
-            node.append(temp_node['right_child'])
-        except KeyError:
-            pass
-
-        try:
-            node.append(temp_node['left_child'])
-        except KeyError:
-            pass
-
-        add_next_nodes(ordered_node_list, node, cat_column=cat_column)
diff --git a/entmoot/learning/proximit_based_std.py b/entmoot/learning/proximit_based_std.py
deleted file mode 100644
index b11a737..0000000
--- a/entmoot/learning/proximit_based_std.py
+++ /dev/null
@@ -1,110 +0,0 @@
-from .distance_based_std import DistanceMetric
-import numpy as np
-
-class ProximityMetric(DistanceMetric):
-
-    def __init__(self, space):
-        self.space = space
-        self.std_type = 'proximity'
-        self.active_leaves_per_data = []
-        self.obj_label = 0
-
-    @staticmethod
-    def get_distance(x_left, x_right):
-        pass
-
-    def add_to_gurobi_model(self, model):
-
-        model._alpha = \
-            model.addVar(
-                lb=0, 
-                ub=1, 
-                name="alpha", 
-                vtype='C'
-            )
-
-        # iterate over all leaf activations for data points
-        for data_id, data_enc in enumerate(self.active_leaves_per_data):
-            temp_lhs = 0
-            for tree_id, leaf_enc in enumerate(data_enc):
-                temp_lhs += model._z_l[self.obj_label, tree_id, leaf_enc]
-
-            temp_constr = 1-1/len(data_enc)*temp_lhs >= model._alpha
-            model.addConstr(
-                temp_constr,
-                name=f"prox_{self.obj_label}_{data_id}"
-            )
-        model.update()
-
-    def update(self, Xi, gbm_model):
-        # collect active leaves for all datapoints
-
-        self.gbm_model = gbm_model
-        self.active_leaves_per_data = np.asarray(
-            [
-                self.gbm_model.get_active_leaves(x) for x in Xi
-            ]
-        )
-
-    def set_params(self,**kwargs):
-        """Sets parameters related to distance-based standard estimator.
-
-        Parameters
-        ----------
-        kwargs : dict
-            Additional arguments to be passed to the standard estimator
-
-        Returns
-        -------
-        -
-        """
-        zeta = kwargs.get("zeta", 0.5)
-        self.zeta = zeta
-
-    def predict(self, X: np.array):
-        """Predict standard estimate at location `X`.
-
-        Parameters
-        ----------
-        X : numpy array, shape (n_rows, n_dims)
-            Points at which the standard estimator is evaluated.
-
-        Returns
-        -------
-        dist : numpy array, shape (n_rows,)
-            Returns distances to closest `ref_point` for every point per row
-            in `n_rows`.
-        """
-        dist = np.empty([X.shape[0],])
-
-        for row_res,Xi in enumerate(X):
-            ref_dist = self.get_distance(Xi)
-            dist[row_res] = ref_dist
-        return dist
-
-    def get_distance(self, X):
-        active_leaves = np.asarray(self.gbm_model.get_active_leaves(X))
-
-        #compute overlap of active leaves for X and the entire dataset 
-        overlap_mat = self.active_leaves_per_data == active_leaves
-
-        # scaling with number of trees and substracting from 1 makes distance
-        prox_mat = (-np.sum(overlap_mat,1)/len(active_leaves)) + 1
-        return np.min(prox_mat)
-
-    def get_gurobi_obj(self, model, scaled=False):
-        """Get contribution of standard estimator to gurobi model objective
-        function.
-
-        Parameters
-        ----------
-        model : gurobipy.Model,
-            Model to which the standard estimator is added.
-
-        Returns
-        -------
-        alpha : gurobipy.Var,
-            Model variable that takes the value of the uncertainty measure.
-        """
-
-        return -model._alpha
\ No newline at end of file
diff --git a/entmoot/learning/tree_model.py b/entmoot/learning/tree_model.py
deleted file mode 100644
index c34ddfb..0000000
--- a/entmoot/learning/tree_model.py
+++ /dev/null
@@ -1,392 +0,0 @@
-import numpy as np
-from sklearn.base import clone
-import copy
-
-
-class EntingRegressor:
-    """Predict with LightGBM tree model and include model uncertainty 
-    defined by distance-based standard estimator.
-
-
-    Parameters
-    ----------
-    base_estimator : LGBMRegressor instance, EntingRegressor instance or 
-        None (default). If LGBMRegressor instance of None is given: new 
-        EntingRegressor is defined. If EntingRegressor is given: base_estimator
-        and std_estimator are given to new instance.
-    std_estimator : DistanceBasedStd instance,
-        Determines which measure is used to capture uncertainty.
-    random_state : int, RandomState instance, or None (default)
-        Set random state to something other than None for reproducible
-        results.
-    """
-
-    def __init__(self,
-                space,
-                base_estimator,
-                std_estimator,
-                random_state:int = None,
-                cat_idx=None):
-
-        if cat_idx is None:
-            cat_idx = []
-
-        np.random.seed(random_state)
-
-        self.random_state = random_state
-        self.space = space
-        self.base_estimator = base_estimator
-        self.std_estimator = std_estimator
-        self.cat_idx = cat_idx
-        self.num_obj = len(self.base_estimator)
-
-        # check if base_estimator is EntingRegressor
-        if isinstance(base_estimator[0], EntingRegressor):
-            self.base_estimator = base_estimator.base_estimator
-            self.std_estimator = base_estimator.std_estimator
-
-    def fit(self, X, y):
-        """Fit model and standard estimator to observations.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_samples, n_features)
-            Training vectors, where `n_samples` is the number of samples
-            and `n_features` is the number of features.
-
-        y : array-like, shape=(n_samples,)
-            Target values (real numbers in regression)
-
-        Returns
-        -------
-        -
-        """
-        self.regressor_ = []
-        y = np.asarray(y)
-        self._y = y
-
-        for i,est in enumerate(self.base_estimator):
-            # suppress lgbm output
-            est.set_params(
-                random_state=self.random_state,
-                verbose=-1
-            )
-
-            # clone base_estimator (only supported for sklearn estimators)
-            self.regressor_.append(clone(est))
-
-            # update std estimator
-            self.std_estimator.update(X, y, cat_column=self.cat_idx)
-
-            # update tree model regressor
-            import warnings
-
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore")
-
-                if self.num_obj > 1:
-                    if self.cat_idx:
-                        self.regressor_[-1].fit(X, y[:,i], categorical_feature=self.cat_idx)
-                    else:
-                        self.regressor_[-1].fit(X, y[:,i])
-                else:
-                    if self.cat_idx:
-                        self.regressor_[-1].fit(X, y, categorical_feature=self.cat_idx)
-                    else:
-                        self.regressor_[-1].fit(X, y)
-
-    def set_params(self, **params):
-        """Sets parameters related to tree model estimator. All parameter 
-        options for LightGBM are given here: 
-            https://lightgbm.readthedocs.io/en/latest/Parameters.html
-
-        Parameters
-        ----------
-        kwargs : dict
-            Additional arguments to be passed to the tree model
-
-        Returns
-        -------
-        -
-        """
-        if not "min_child_samples" in params:
-            params.update({"min_child_samples":2})
-
-        for i in range(len(self.base_estimator)):
-            self.base_estimator[i].set_params(**params)
-
-    def predict(self, X, return_std=False):
-        """Predict.
-
-        Parameters
-        ----------
-        X : array-like, shape (n_samples, n_features)
-            where `n_samples` is the number of samples
-            and `n_features` is the number of features.
-        return_std: boolean, If `return_std` is set to False, only tree model prediction is returned. If set to True,
-            tree model prediction and predicted standard deviation is returned.
-
-        Returns
-        -------
-        mean or (mean, std): np.array, shape (n_rows, n_dims) 
-            or tuple(np.array, np.array), depending on value of `return_std`.
-        """
-
-        if self.num_obj == 1:
-            mean = self.regressor_[0].predict(X)
-        else:
-            mean = []
-            for tree in self.regressor_:
-                mean.append(tree.predict(X))
-
-        if return_std:
-            std = self.std_estimator.predict(X)
-            return mean, std
-
-        # return the mean
-        return mean
-
-    def get_gbm_model(self):
-        """
-        Returns `GbmModel` instance of the tree model.
-
-        Parameters
-        ----------
-        -
-
-        Returns
-        -------
-        gbm_model : `GbmModel`
-            ENTMOOT native tree model format to formulate optimization model
-        """
-
-        from entmoot.learning.lgbm_processing import order_tree_model_dict
-        from entmoot.learning.gbm_model import GbmModel
-        gbm_model_dict = {}
-        for i,tree in enumerate(self.regressor_):
-            original_tree_model_dict = tree._Booster.dump_model()
-
-            # import json
-            # with open('tree_dict_next.json', 'w') as f:
-            #     json.dump(
-            #         original_tree_model_dict,
-            #         f,
-            #         indent=4,
-            #         sort_keys=False
-            #     )
-
-            ordered_tree_model_dict = \
-                order_tree_model_dict(
-                    original_tree_model_dict,
-                    cat_column=self.cat_idx
-                )
-
-            gbm_model_dict[i] = GbmModel(ordered_tree_model_dict)
-        return gbm_model_dict
-
-    def get_global_next_x(self,
-                          acq_func,
-                          acq_func_kwargs,
-                          add_model_core,
-                          weight,
-                          verbose,
-                          gurobi_env,
-                          gurobi_timelimit,
-                          has_unc=True):
-        from entmoot.optimizer.gurobi_utils import \
-            get_core_gurobi_model, add_gbm_to_gurobi_model, \
-            add_std_to_gurobi_model, add_acq_to_gurobi_model, \
-            set_gurobi_init_to_ref, get_gbm_multi_obj_from_model
-
-        # suppress  output to command window
-        import logging
-        logger = logging.getLogger()
-        logger.setLevel(logging.CRITICAL)
-
-        # start building model
-        gurobi_model = \
-            get_core_gurobi_model(
-                self.space, add_model_core, env=gurobi_env
-            )
-
-        if verbose == 2:
-            print("")
-            print("")
-            print("")
-            print("SOLVER: *** start gurobi solve ***")
-            gurobi_model.Params.LogToConsole = 1
-        else:
-            gurobi_model.Params.LogToConsole = 0
-
-        # convert into gbm_model format
-        # and add to gurobi model
-        gbm_model_dict = self.get_gbm_model()
-        add_gbm_to_gurobi_model(
-            self.space, gbm_model_dict, gurobi_model
-        )
-
-        # add std estimator to gurobi model
-        if has_unc:
-            add_std_to_gurobi_model(self, gurobi_model)
-
-        # collect different objective function contributions
-        from entmoot.optimizer.gurobi_utils import get_gbm_model_multi_obj_mu, get_gbm_model_mu
-
-        if self.num_obj > 1:
-            model_mu = get_gbm_model_multi_obj_mu(gurobi_model, self._y)
-
-            if has_unc:
-                model_unc = self.std_estimator.get_gurobi_obj(gurobi_model)
-            else: model_unc = None
-        else:
-            model_mu = get_gbm_model_mu(gurobi_model, self._y, norm=False)
-            model_unc = self.std_estimator.get_gurobi_obj(gurobi_model)
-
-        # add obj to gurobi model
-        from opti.sampling.simplex import sample
-
-        weight = sample(self.num_obj,1)[0] if weight is None else weight
-
-        add_acq_to_gurobi_model(gurobi_model, model_mu, model_unc,
-                                self,
-                                weights=weight,
-                                num_obj=self.num_obj,
-                                acq_func=acq_func,
-                                acq_func_kwargs=acq_func_kwargs)
-
-        # set initial gurobi model vars to std_est reference points
-        if has_unc:
-            if self.std_estimator.std_type == 'distance':
-                set_gurobi_init_to_ref(self, gurobi_model)
-
-        # set gurobi time limit
-        if gurobi_timelimit is not None:
-            gurobi_model.Params.TimeLimit = gurobi_timelimit
-        gurobi_model.Params.OutputFlag = 1
-
-        gurobi_model.update()
-        gurobi_model.optimize()
-
-        assert gurobi_model.SolCount >= 1, "gurobi couldn't find a feasible " + \
-                                           "solution. Try increasing the timelimit if specified. " + \
-                                           "In case you specify your own 'add_model_core' " + \
-                                           "please check that the model is feasible."
-
-        # store optimality gap of gurobi computation
-        gurobi_mipgap = None
-        if acq_func not in ["HLCB"]:
-            gurobi_mipgap = gurobi_model.mipgap
-
-        # for i in range(2): REMOVEX
-        #     gurobi_model.params.ObjNumber = i
-        #     # Query the o-th objective value
-        #     print(f"{i}:")
-        #     print(gurobi_model.ObjNVal)
-
-        # output next_x
-        next_x = np.empty(len(self.space.dimensions))
-
-        # cont features
-        for i in gurobi_model._cont_var_dict:
-            next_x[i] = gurobi_model._cont_var_dict[i].x
-
-        # cat features
-        for i in gurobi_model._cat_var_dict:
-            cat = \
-                [
-                    key
-                    for key in gurobi_model._cat_var_dict[i]
-                    if int(
-                    round(gurobi_model._cat_var_dict[i][key].x, 1)
-                ) == 1
-                ]
-
-            next_x[i] = cat[0]
-
-        model_mu = get_gbm_multi_obj_from_model(gurobi_model)
-        if has_unc:
-            model_std = gurobi_model._alpha.x
-        else:
-            model_std = None
-
-        return next_x, model_mu, model_std, gurobi_mipgap
-
-    def copy(self):
-        return copy.copy(self)
-
-
-class MisicRegressor(EntingRegressor):
-
-    def __init__(self,
-                space,
-                base_estimator=None,
-                std_estimator=None,
-                random_state=None,
-                cat_idx=None):
-
-        if cat_idx is None:
-            cat_idx = []
-
-        self.random_state = random_state
-        self.space = space
-        self.base_estimator = base_estimator
-        self.std_estimator = std_estimator
-        self.cat_idx = cat_idx
-        self.num_obj = len(self.base_estimator)
-
-        # check if base_estimator is EntingRegressor
-        if isinstance(base_estimator, MisicRegressor):
-            self.base_estimator = base_estimator.base_estimator
-            self.std_estimator = base_estimator.std_estimator
-
-    def fit(self, X, y):
-        """Fit model and standard estimator to observations.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_samples, n_features)
-            Training vectors, where `n_samples` is the number of samples
-            and `n_features` is the number of features.
-
-        y : array-like, shape=(n_samples,)
-            Target values (real numbers in regression)
-
-        Returns
-        -------
-        -
-        """
-        self.regressor_ = []
-        y = np.asarray(y)
-        self._y = y
-
-        for i,est in enumerate(self.base_estimator):
-            # suppress lgbm output
-            est.set_params(
-                random_state=self.random_state,
-                verbose=-1
-            )
-
-            # clone base_estimator (only supported for sklearn estimators)
-            self.regressor_.append(clone(est))
-
-            # update tree model regressor
-            import warnings
-
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore")
-
-                if self.num_obj > 1:
-                    if self.cat_idx:
-                        self.regressor_[-1].fit(X, y[:,i], categorical_feature=self.cat_idx)
-                    else:
-                        self.regressor_[-1].fit(X, y[:,i])
-                else:
-                    if self.cat_idx:
-                        self.regressor_[-1].fit(X, y, categorical_feature=self.cat_idx)
-                    else:
-                        self.regressor_[-1].fit(X, y)
-
-        gbm_model = self.get_gbm_model()[0]
-
-        # update std estimator
-        self.std_estimator.update(X, y, gbm_model, cat_column=self.cat_idx)
diff --git a/entmoot/models/__init__.py b/entmoot/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/entmoot/models/base_model.py b/entmoot/models/base_model.py
new file mode 100644
index 0000000..88e1eba
--- /dev/null
+++ b/entmoot/models/base_model.py
@@ -0,0 +1,15 @@
+class BaseModel:
+    def __init__(self, problem_config, num_obj, params):
+        raise NotImplementedError()
+
+    def predict(self, X):
+        raise NotImplementedError()
+
+    def fit(self, X, y):
+        raise NotImplementedError()
+
+    def add_to_gurobipy_model(self, model_core):
+        raise NotImplementedError()
+
+    def add_to_pyomo_model(self, model_core):
+        raise NotImplementedError()
diff --git a/entmoot/models/enting.py b/entmoot/models/enting.py
new file mode 100644
index 0000000..7612ee9
--- /dev/null
+++ b/entmoot/models/enting.py
@@ -0,0 +1,233 @@
+from entmoot import ProblemConfig
+from entmoot.models.base_model import BaseModel
+from entmoot.models.mean_models.tree_ensemble import TreeEnsemble
+from entmoot.models.uncertainty_models.distance_based_uncertainty import (
+    DistanceBasedUncertainty,
+)
+import numpy as np
+
+
+class Enting(BaseModel):
+    """
+    This class contains a nice living space for your tree model. You can fit your model, use it for predictions and
+    provide information that are needed to build optimization models that incorporate the tree structure of your model.
+
+    Example:
+        .. code-block:: python
+
+            from entmoot import ProblemConfig
+            import numpy as np
+            import random
+
+            # Define a one-dimensional minimization problem with one real variable bounded by -2 and 3
+            problem_config = ProblemConfig()
+            problem_config.add_feature("real", (-2, 3))
+            problem_config.add_min_objective()
+
+            # Create training data using the randomly disturbed function f(x) = x^2 + 1 + eps
+            X_train = np.linspace(-2, 3, 10)
+            y_train = [x**2 + 1 + random.uniform(-0.2, 0.2) for x in X_train]
+
+            # Define enting object and corresponding parameters
+            params = {"unc_params": {"dist_metric": "l1"}}
+            enting = Enting(problem_config, params=params)
+            # Fit tree model
+            enting.fit(X_train, y_train)
+            # Compute the predictions for training data and see that light gbm fitted a step function
+            # with three steps
+            enting.predict(X_train)
+
+            # Define parameters needed during optimization
+            params_pyo = {"solver_name": "gurobi", "solver_options": {"MIPGap": 0}}
+            # Build PyomoOptimizer object. This step will internally call the methods
+            # add_to_pyomo_model() or add_to_gurobipy_model(), resp., depending on the choice
+            # of your optimizer
+            opt_pyo = PyomoOptimizer(problem_config, params=params_pyo)
+            # As expected, the optimal input of the tree model is near the origin (cf. X_opt_pyo)
+            X_opt_pyo, _, _ = opt_pyo.solve(enting)
+    """
+
+    def __init__(self, problem_config: ProblemConfig, params: dict = None):
+
+        if params is None:
+            params = {}
+
+        self._problem_config = problem_config
+
+        # check params values
+        tree_training_params = params.get("tree_train_params", {})
+
+        # initialize mean model
+        self.mean_model = TreeEnsemble(
+            problem_config=problem_config, params=tree_training_params
+        )
+
+        # initialize unc model
+        unc_params = params.get("unc_params", {})
+        self._acq_sense = unc_params.get("acq_sense", "exploration")
+        assert self._acq_sense in (
+            "exploration",
+            "penalty",
+        ), f"Pick 'acq_sense' '{self._acq_sense}' in '('exploration', 'penalty')'."
+
+        self._beta = unc_params.get("beta", 1.96)
+        assert (
+            self._beta >= 0.0
+        ), f"Value for 'beta' is {self._beta} but must be '>= 0.0'."
+
+        if self._acq_sense == "exploration":
+            self._beta = -self._beta
+
+        self.unc_model = DistanceBasedUncertainty(
+            problem_config=problem_config, params=unc_params
+        )
+
+    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
+        """
+        performs the training of you tree model using training data and labels
+        """
+        # encode categorical features
+        X = self._problem_config.encode(X)
+
+        # check dims of X and y
+        if X.ndim == 1:
+            X = np.atleast_2d(X)
+
+        assert X.shape[-1] == len(self._problem_config.feat_list), (
+            "Argument 'X' has wrong dimensions. "
+            f"Expected '(num_samples, {len(self._problem_config.feat_list)})', got '{X.shape}'."
+        )
+
+        if y.ndim == 1:
+            y = np.atleast_2d(y)
+
+        assert (y.shape[-1] == 2 and len(self._problem_config.obj_list) == 1) or (
+            y.shape[-1] == len(self._problem_config.obj_list)
+        ), (
+            "Argument 'y' has wrong dimensions. "
+            f"Expected '(num_samples, {len(self._problem_config.obj_list)})', got '{y.shape}'."
+        )
+
+        self.mean_model.fit(X, y)
+        self.unc_model.fit(X, y)
+
+    def predict(self, X: np.ndarray) -> list:
+        """
+        Computes prediction value of tree model for X.
+        """
+        # encode categorical features
+        X = self._problem_config.encode(X)
+
+        # check dims of X
+        if X.ndim == 1:
+            X = np.atleast_2d(X)
+
+        assert X.shape[-1] == len(self._problem_config.feat_list), (
+            "Argument 'X' has wrong dimensions. "
+            f"Expected '(num_samples, {len(self._problem_config.feat_list)})', got '{X.shape}'."
+        )
+
+        mean_pred = self.mean_model.predict(X).tolist()
+        unc_pred = self.unc_model.predict(X)
+
+        comb_pred = [(mean, unc) for mean, unc in zip(mean_pred, unc_pred)]
+        return comb_pred
+
+    def predict_pareto(self):
+        pass
+
+    def predict_acq(self, X: np.ndarray) -> list:
+        """
+        predicts value of acquisition function (which contains not only the mean value but also the uncertainty)
+        """
+        acq_pred = []
+        comb_pred = self.predict(X)
+        for mean, unc in comb_pred:
+            acq_pred.append(mean + self._beta * unc)
+        return acq_pred
+
+    def add_to_gurobipy_model(self, core_model, weights: tuple = None) -> None:
+        """
+        This method enriches the core model by adding variables and constraints based on information
+        from the tree model.
+        """
+        from gurobipy import GRB
+        from entmoot.utils import sample
+
+        # add uncertainty model part
+        self.unc_model.add_to_gurobipy_model(core_model)
+
+        core_model._mu = core_model.addVar(
+            lb=-GRB.INFINITY, ub=GRB.INFINITY, name=f"mean_obj", vtype="C"
+        )
+
+        if len(self._problem_config.obj_list) == 1:
+            # single objective case
+            self.mean_model.add_to_gurobipy_model(core_model, add_mu_var=True)
+            core_model.addConstr(core_model._mu == core_model._aux_mu[0])
+        else:
+            # multi-objective case
+            self.mean_model.add_to_gurobipy_model(
+                core_model, add_mu_var=True, normalize_mean=True
+            )
+            if weights is not None:
+                moo_weights = weights
+            else:
+                moo_weights = sample(len(self._problem_config.obj_list), 1)[0]
+
+            for idx, obj in enumerate(self._problem_config.obj_list):
+                core_model.addConstr(
+                    core_model._mu >= moo_weights[idx] * core_model._aux_mu[idx],
+                    name=f"weighted_mean_obj_{idx}",
+                )
+
+        core_model.setObjective(core_model._mu + self._beta * core_model._unc)
+        core_model.update()
+
+    def add_to_pyomo_model(self, core_model, weights: tuple = None) -> None:
+        """
+        This method enriches the core model by adding variables and constraints based on information
+        from the tree model.
+        """
+        import pyomo.environ as pyo
+        from entmoot.utils import sample
+
+        # add uncertainty model part
+        self.unc_model.add_to_pyomo_model(core_model)
+
+        core_model._mu = pyo.Var(domain=pyo.Reals)
+
+        if len(self._problem_config.obj_list) == 1:
+            # single objective case
+            self.mean_model.add_to_pyomo_model(core_model, add_mu_var=True)
+            # Get objective name
+            obj_name = self._problem_config.obj_list[0].name
+            core_model.constraint_link_mu_auxmu = pyo.Constraint(
+                expr=core_model._aux_mu[obj_name] == core_model._mu
+            )
+        else:
+            # multi-objective case
+            self.mean_model.add_to_pyomo_model(
+                core_model, add_mu_var=True, normalize_mean=True
+            )
+            if weights is not None:
+                moo_weights = weights
+            else:
+                moo_weights = sample(len(self._problem_config.obj_list), 1)[0]
+
+            objectives_position_name = list(
+                enumerate([obj.name for obj in self._problem_config.obj_list])
+            )
+
+            def constrs_weights_auxmus(model, pos, objname):
+                return model._mu >= moo_weights[pos] * core_model._aux_mu[objname]
+
+            core_model.constr_coupling_mu_auxmu = pyo.Constraint(
+                objectives_position_name,
+                rule=constrs_weights_auxmus,
+            )
+
+        # Define objective function
+        core_model.obj = pyo.Objective(
+            expr=core_model._mu + self._beta * core_model._unc, sense=pyo.minimize
+        )
diff --git a/entmoot/models/leaf_gp.py b/entmoot/models/leaf_gp.py
new file mode 100644
index 0000000..c8d7202
--- /dev/null
+++ b/entmoot/models/leaf_gp.py
@@ -0,0 +1,9 @@
+from entmoot.models.base_model import BaseModel
+
+
+class TreeKernelEntmoot(BaseModel):
+    def add_to_gurobipy_model(core_model, gurobi_env):
+        return NotImplementedError()
+
+    def add_to_pyomo_model(core_model):
+        return NotImplementedError()
diff --git a/entmoot/models/mean_models/__init__.py b/entmoot/models/mean_models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/entmoot/models/mean_models/lgbm_utils.py b/entmoot/models/mean_models/lgbm_utils.py
new file mode 100644
index 0000000..489aefc
--- /dev/null
+++ b/entmoot/models/mean_models/lgbm_utils.py
@@ -0,0 +1,56 @@
+def read_lgbm_tree_model_dict(tree_model_dict, cat_idx):
+    ordered_tree_list = []
+
+    for tree in tree_model_dict["tree_info"]:
+
+        # generate list of nodes in tree
+        root_node = [tree["tree_structure"]]
+        node_list = []
+
+        # populate node_list and add to ordered_tree_list if non-empty
+        add_next_nodes(node_list=node_list, node=root_node, cat_idx=cat_idx)
+
+        if node_list:
+            ordered_tree_list.append(node_list)
+
+    return ordered_tree_list
+
+
+def add_next_nodes(node_list, node, cat_idx):
+    if node:
+        # add new node definition
+        new_node = {}
+
+        try:
+            new_node["split_var"] = node[-1]["split_feature"]
+            if not new_node["split_var"] in cat_idx:
+                # read numerical variables, solver accuracy 10e-5
+                temp_node_val = round(node[-1]["threshold"], 5)
+                new_node["split_code_pred"] = temp_node_val
+            else:
+                # read categorical variables
+                cat_set = node[-1]["threshold"].split("||")
+                temp_node_val = [int(cat) for cat in cat_set]
+                new_node["split_code_pred"] = temp_node_val
+
+        except KeyError:
+            # arrived at leaf node
+            new_node["split_var"] = -1
+            temp_node_val = node[-1]["leaf_value"]
+            new_node["split_code_pred"] = temp_node_val
+
+        node_list.append(new_node)
+
+        # move on to next node in tree
+        temp_node = node.pop(-1)
+        try:
+            node.append(temp_node["right_child"])
+        except KeyError:
+            pass
+
+        try:
+            node.append(temp_node["left_child"])
+        except KeyError:
+            pass
+
+        add_next_nodes(node_list, node, cat_idx=cat_idx)
diff --git a/entmoot/models/mean_models/meta_tree_ensemble.py b/entmoot/models/mean_models/meta_tree_ensemble.py
new file mode 100644
index 0000000..681a579
--- /dev/null
+++ b/entmoot/models/mean_models/meta_tree_ensemble.py
@@ -0,0 +1,207 @@
+class MetaTreeModel:
+    def __init__(self, tree_list: list):
+        self._num_trees = len(tree_list)
+        self.load_ordered_tree_list(tree_list)
+
+    @property
+    def num_trees(self):
+        return self._num_trees
+
+    def _build_tree(self, tree):
+        node = tree.pop(0)
+        split_var = node["split_var"]
+        split_code_pred = node["split_code_pred"]
+
+        return TreeNode(split_var=split_var, split_code_pred=split_code_pred, tree=tree)
+
+    def load_ordered_tree_list(self, tree_list: list):
+        self.trees = [self._build_tree(tree) for tree in tree_list]
+
+    def get_leaf_encodings(self, tree):
+        yield from self.trees[tree].get_leaf_encodings()
+
+    def get_leaf_weight(self, tree, encoding):
+        return self.trees[tree].get_leaf_weight(encoding)
+
+    def get_var_break_points(self):
+        var_breakpoints = {}
+        for tree in self.trees:
+            for var, breakpoint in tree.get_all_partition_pairs():
+                try:
+                    if isinstance(breakpoint, list):
+                        var_breakpoints[var].append(breakpoint)
+                    else:
+                        var_breakpoints[var].add(breakpoint)
+                except KeyError:
+                    if isinstance(breakpoint, list):
+                        var_breakpoints[var] = [breakpoint]
+                    else:
+                        var_breakpoints[var] = set([breakpoint])
+
+        for k in var_breakpoints:
+            if isinstance(var_breakpoints[k], set):
+                var_breakpoints[k] = sorted(var_breakpoints[k])
+
+        return var_breakpoints
+
+    def get_participating_variables(self, tree, leaf):
+        return set(self.trees[tree].get_participating_variables(leaf))
+
+    def get_branch_encodings(self, tree):
+        yield from self.trees[tree].get_branch_encodings()
+
+    def get_branch_partition_pair(self, tree, encoding):
+        return self.trees[tree].get_branch_partition_pair(encoding)
+
+    def get_left_leaves(self, tree, encoding):
+        yield from (encoding + s for s in self.trees[tree].get_left_leaves(encoding))
+
+    def get_right_leaves(self, tree, encoding):
+        yield from (encoding + s for s in self.trees[tree].get_right_leaves(encoding))
+
+    def prune_var_bnds(self, encodings, var_bnds):
+        # prunes var_bnds provided according to which splits are active
+        for tree_id, leaf_enc in encodings:
+            self.trees[tree_id]._prune_var_bnds(0, leaf_enc, var_bnds)
+        return var_bnds
+
+
+class TreeType:
+    pass
+
+
+class TreeNode(TreeType):
+    def __init__(self, split_var, split_code_pred, tree):
+        self.split_var = split_var
+        self.split_code_pred = split_code_pred
+
+        # check if tree is empty
+        assert (
+            tree
+        ), "Given the data and train configuration, no tree ensemble could be trained."
+
+        # read left node
+        node = tree.pop(0)
+        split_var = node["split_var"]
+        split_code_pred = node["split_code_pred"]
+
+        # split_var value of -1 refers to leaf node
+        if split_var == -1:
+            self.left = LeafNode(split_code_pred=split_code_pred)
+        else:
+            self.left = TreeNode(
+                split_var=split_var, split_code_pred=split_code_pred, tree=tree
+            )
+
+        # read right node
+        node = tree.pop(0)
+        split_var = node["split_var"]
+        split_code_pred = node["split_code_pred"]
+
+        if split_var == -1:
+            self.right = LeafNode(split_code_pred=split_code_pred)
+        else:
+            self.right = TreeNode(
+                split_var=split_var, split_code_pred=split_code_pred, tree=tree
+            )
+
+    def __repr__(self):
+        return ", ".join([str(x) for x in [self.split_var, self.split_code_pred]])
+
+    def get_leaf_encodings(self, current_string=""):
+        yield from self.left.get_leaf_encodings(current_string + "0")
+        yield from self.right.get_leaf_encodings(current_string + "1")
+
+    def get_leaf_weight(self, encoding):
+        next_node = self.right if int(encoding[0]) else self.left
+        return next_node.get_leaf_weight(encoding[1:])
+
+    def get_all_partition_pairs(self):
+        yield (self.split_var, self.split_code_pred)
+        yield from self.left.get_all_partition_pairs()
+        yield from self.right.get_all_partition_pairs()
+
+    def get_participating_variables(self, encoding):
+        yield self.split_var
+        next_node = self.right if int(encoding[0]) else self.left
+        yield from next_node.get_participating_variables(encoding[1:])
+
+    def get_branch_encodings(self, current_string=""):
+        yield current_string
+        yield from self.left.get_branch_encodings(current_string + "0")
+        yield from self.right.get_branch_encodings(current_string + "1")
+
+    def _get_next_node(self, direction):
+        return self.right if int(direction) else self.left
+
+    def get_branch_partition_pair(self, encoding):
+        if not encoding:
+            return self.split_var, self.split_code_pred
+        else:
+            next_node = self._get_next_node(encoding[0])
+            return next_node.get_branch_partition_pair(encoding[1:])
+
+    def get_left_leaves(self, encoding):
+        if encoding:
+            next_node = self._get_next_node(encoding[0])
+            yield from next_node.get_left_leaves(encoding[1:])
+        else:
+            yield from self.left.get_leaf_encodings("0")
+
+    def get_right_leaves(self, encoding):
+        if encoding:
+            next_node = self._get_next_node(encoding[0])
+            yield from next_node.get_right_leaves(encoding[1:])
+        else:
+            yield from self.right.get_leaf_encodings("1")
+
+    def _prune_var_bnds(self, curr_depth, leaf_enc, var_bnds):
+        if self.split_var != -1:
+            if isinstance(self.split_code_pred, list):
+                # categorical variable
+                cat_set = set(self.split_code_pred)
+                if leaf_enc[curr_depth] == '0':
+                    var_bnds[self.split_var] = \
+                        set(var_bnds[self.split_var]).intersection(cat_set)
+                    self.left._prune_var_bnds(curr_depth + 1, leaf_enc, var_bnds)
+                else:
+                    var_bnds[self.split_var] = \
+                        set(var_bnds[self.split_var]).difference(cat_set)
+                    self.right._prune_var_bnds(curr_depth + 1, leaf_enc, var_bnds)
+            else:
+                # continuous variable
+                lb, ub = var_bnds[self.split_var]
+                if leaf_enc[curr_depth] == '0':
+                    ub = min(ub,self.split_code_pred)
+                    var_bnds[self.split_var] = (lb, ub)
+                    self.left._prune_var_bnds(curr_depth + 1, leaf_enc, var_bnds)
+                else: # if value is '1'
+                    lb = max(lb,self.split_code_pred)
+                    var_bnds[self.split_var] = (lb, ub)
+                    self.right._prune_var_bnds(curr_depth + 1, leaf_enc, var_bnds)
+
+
+class LeafNode(TreeType):
+    def __init__(self, split_code_pred):
+        self.split_var = -1
+        self.split_code_pred = split_code_pred
+
+    def __repr__(self):
+        return ", ".join([str(x) for x in ["Leaf", self.split_code_pred]])
+
+    def get_leaf_encodings(self, current_string=""):
+        yield current_string
+
+    def get_leaf_weight(self, encoding):
+        assert not encoding
+        return self.split_code_pred
+
+    def get_all_partition_pairs(self):
+        yield from []
+
+    def get_participating_variables(self, encoding):
+        assert not encoding
+        yield from []
+
+    def get_branch_encodings(self, current_string=""):
+        yield from []
diff --git a/entmoot/models/mean_models/tree_ensemble.py b/entmoot/models/mean_models/tree_ensemble.py
new file mode 100644
index 0000000..fa9e151
--- /dev/null
+++ b/entmoot/models/mean_models/tree_ensemble.py
@@ -0,0 +1,665 @@
+from entmoot.models.base_model import BaseModel
+from entmoot.models.mean_models.lgbm_utils import read_lgbm_tree_model_dict
+from entmoot.models.mean_models.meta_tree_ensemble import MetaTreeModel
+import warnings
+import numpy as np
+
+
+class TreeEnsemble(BaseModel):
+    def __init__(self, problem_config, params=None):
+
+        if params is None:
+            params = {}
+
+        self._problem_config = problem_config
+        self._train_lib = params.get("train_lib", "lgbm")
+        self._rnd_seed = problem_config.rnd_seed
+
+        assert self._train_lib in ("lgbm", "catboost", "xgboost"), (
+            "Parameter 'train_lib' for tree ensembles needs to be "
+            "in '('lgbm', 'catboost', 'xgboost')'."
+        )
+
+        if "train_params" not in params:
+            # default training params
+            warnings.warn(
+                "No 'train_params' for tree ensemble training specified. "
+                "Switch training to default params!"
+            )
+
+            self._train_params = {
+                "objective": "regression",
+                "metric": "rmse",
+                "boosting": "gbdt",
+                "num_boost_round": 100,
+                "max_depth": 3,
+                "min_data_in_leaf": 1,
+                "min_data_per_group": 1,
+                "verbose": -1,
+            }
+
+            if self._rnd_seed is not None:
+                self._train_params["random_state"] = self._rnd_seed
+        else:
+            self._train_params = params["train_params"]
+
+        self._tree_dict = None
+        self._meta_tree_dict = {}
+        self._min_y, self._max_y = None, None
+
+    @property
+    def tree_dict(self):
+        assert (
+            self._tree_dict is not None
+        ), "No tree model is trained yet. Call '.fit(X, y)' first."
+        return self._tree_dict
+
+    @property
+    def meta_tree_dict(self):
+        assert (
+            len(self._meta_tree_dict) > 0
+        ), "No tree model is trained yet. Call '.fit(X, y)' first."
+        return self._meta_tree_dict
+
+    @property
+    def min_y(self):
+        return self._min_y
+
+    @property
+    def max_y(self):
+        return self._max_y
+
+    def fit(self, X, y):
+        # train tree models for every objective
+        if self._tree_dict is None:
+            self._tree_dict = {}
+
+        # cache min/max vals for normalization
+        self._min_y = np.min(y, axis=0)
+        self._max_y = np.max(y, axis=0)
+
+        # train tree models for every objective
+        for i, obj in enumerate(self._problem_config.obj_list):
+            if self._train_lib == "lgbm":
+                tree_model = self._train_lgbm(X, y[:, i])
+            elif self._train_lib == "catboost":
+                raise NotImplementedError()
+            elif self._train_lib == "xgboost":
+                raise NotImplementedError()
+            else:
+                raise IOError(
+                    "Parameter 'train_lib' for tree ensembles needs to be "
+                    "in '('lgbm', 'catboost', 'xgboost')'."
+                )
+            self._tree_dict[obj.name] = tree_model
+
+    def _train_lgbm(self, X, y):
+        import lightgbm as lgb
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+
+            if self._problem_config.cat_idx:
+                # train for categorial vars
+                train_data = lgb.Dataset(
+                    X,
+                    label=y,
+                    categorical_feature=self._problem_config.cat_idx,
+                    free_raw_data=False,
+                    params={"verbose": -1},
+                )
+
+                tree_model = lgb.train(
+                    self._train_params,
+                    train_data,
+                    categorical_feature=self._problem_config.cat_idx,
+                    verbose_eval=False,
+                )
+            else:
+                # train for non-categorical vars
+                train_data = lgb.Dataset(X, label=y, params={"verbose": -1})
+
+                tree_model = lgb.train(
+                    self._train_params, train_data, verbose_eval=False
+                )
+        return tree_model
+
+    def predict(self, X):
+        # predict vals
+        tree_pred = []
+        for obj in self._problem_config.obj_list:
+            tree_pred.append(self.tree_dict[obj.name].predict(X))
+        return np.column_stack(tree_pred)
+
+    def _update_meta_tree_dict(self):
+        self._meta_tree_dict = {}
+
+        # get model information
+        for obj in self._problem_config.obj_list:
+
+            if self._train_lib == "lgbm":
+                lib_out = self.tree_dict[obj.name].dump_model()
+            elif self._train_lib == "catboost":
+                raise NotImplementedError()
+            elif self._train_lib == "xgboost":
+                raise NotImplementedError()
+            else:
+                raise IOError(
+                    "Parameter 'train_lib' for tree ensembles needs to be "
+                    "in '('lgbm', 'catboost', 'xgboost')'."
+                )
+
+            # order tree_model_dict
+            ordered_tree_model_dict = read_lgbm_tree_model_dict(
+                lib_out, cat_idx=self._problem_config.cat_idx
+            )
+
+            # populate meta_tree_model
+            self._meta_tree_dict[obj.name] = MetaTreeModel(ordered_tree_model_dict)
+
+    def add_to_gurobipy_model(self, model, add_mu_var=True, normalize_mean=False):
+        from gurobipy import GRB, quicksum
+
+        self._update_meta_tree_dict()
+
+        # attach tree info
+        model._num_trees = lambda obj_name: self._meta_tree_dict[obj_name].num_trees
+
+        # attach leaf info
+        model._leaves = lambda obj_name, tree: tuple(
+            self._meta_tree_dict[obj_name].get_leaf_encodings(tree)
+        )
+        model._leaf_weight = lambda obj_name, tree, leaf: self._meta_tree_dict[
+            obj_name
+        ].get_leaf_weight(tree, leaf)
+        model._leaf_vars = lambda obj_name, tree, leaf: tuple(
+            var
+            for var in self._meta_tree_dict[obj_name].get_participating_variables(
+                tree, leaf
+            )
+        )
+
+        # attach breakpoint info
+        var_break_points = [
+            tree_model.get_var_break_points()
+            for tree_model in self._meta_tree_dict.values()
+        ]
+        model._breakpoints = {}
+        model._breakpoint_index = []
+
+        for idx, feat in enumerate(self._problem_config.feat_list):
+            if feat.is_cat():
+                continue
+            else:
+                splits = set()
+                for vb in var_break_points:
+                    try:
+                        splits = splits.union(set(vb[idx]))
+                    except KeyError:
+                        pass
+                if splits:
+                    model._breakpoints[idx] = sorted(splits)
+                    model._breakpoint_index.append(idx)
+
+        # define indexing helper functions
+        def tree_index(model_obj):
+            for obj in self._problem_config.obj_list:
+                for tree in range(model_obj._num_trees(obj.name)):
+                    yield obj.name, tree
+
+        def leaf_index(model_obj):
+            for label, tree in tree_index(model_obj):
+                for leaf in model_obj._leaves(label, tree):
+                    yield label, tree, leaf
+
+        def interval_index(model_obj):
+            for var in model_obj._breakpoint_index:
+                for j in range(len(model_obj._breakpoints[var])):
+                    yield var, j
+
+        def split_index(model_obj, meta_tree_dict):
+            for label, tree in tree_index(model_obj):
+                for encoding in meta_tree_dict[label].get_branch_encodings(tree):
+                    yield label, tree, encoding
+
+        # add leaf variables
+        model._z = model.addVars(
+            leaf_index(model), lb=0, ub=GRB.INFINITY, name="z", vtype="C"
+        )
+
+        # add split variables
+        model._nu = model.addVars(interval_index(model), name="nu", vtype=GRB.BINARY)
+
+        # add single leaf constraints
+        def single_leaf_rule(model_obj, label, tree):
+            z, leaves = model_obj._z, model_obj._leaves
+            return quicksum(z[label, tree, leaf] for leaf in leaves(label, tree)) == 1
+
+        model.addConstrs(
+            (
+                single_leaf_rule(model, label, tree)
+                for (label, tree) in tree_index(model)
+            ),
+            name="single_leaf",
+        )
+
+        # add left split constraints
+        def left_split_r(model_obj, meta_tree_dict, label, tree, split_enc):
+            split_var, split_val = meta_tree_dict[label].get_branch_partition_pair(
+                tree, split_enc
+            )
+
+            if not isinstance(split_val, list):
+                # non-cat vars
+                nu_val = model_obj._breakpoints[split_var].index(split_val)
+                return (
+                    quicksum(
+                        model_obj._z[label, tree, leaf]
+                        for leaf in meta_tree_dict[label].get_left_leaves(
+                            tree, split_enc
+                        )
+                    )
+                    <= model_obj._nu[split_var, nu_val]
+                )
+            else:
+                # cat vars
+                return quicksum(
+                    model_obj._z[label, tree, leaf]
+                    for leaf in meta_tree_dict[label].get_left_leaves(tree, split_enc)
+                ) <= quicksum(model_obj._all_feat[split_var][cat] for cat in split_val)
+
+        model.addConstrs(
+            (
+                left_split_r(model, self.meta_tree_dict, label, tree, encoding)
+                for (label, tree, encoding) in split_index(model, self.meta_tree_dict)
+            ),
+            name="left_split",
+        )
+
+        # add right split constraints
+        def right_split_r(model_obj, meta_tree_dict, label, tree, split_enc):
+            split_var, split_val = meta_tree_dict[label].get_branch_partition_pair(
+                tree, split_enc
+            )
+
+            if not isinstance(split_val, list):
+                # non-cat vars
+                nu_val = model_obj._breakpoints[split_var].index(split_val)
+                return (
+                    quicksum(
+                        model_obj._z[label, tree, leaf]
+                        for leaf in meta_tree_dict[label].get_right_leaves(
+                            tree, split_enc
+                        )
+                    )
+                    <= 1 - model_obj._nu[split_var, nu_val]
+                )
+            else:
+                # cat vars
+                return quicksum(
+                    model_obj._z[label, tree, leaf]
+                    for leaf in meta_tree_dict[label].get_right_leaves(tree, split_enc)
+                ) <= 1 - quicksum(
+                    model_obj._all_feat[split_var][cat] for cat in split_val
+                )
+
+        model.addConstrs(
+            (
+                right_split_r(model, self.meta_tree_dict, label, tree, encoding)
+                for (label, tree, encoding) in split_index(model, self.meta_tree_dict)
+            ),
+            name="right_split",
+        )
+
+        # add split order constraints
+        def y_order_r(model_obj, i, j):
+            return model_obj._nu[i, j] <= model_obj._nu[i, j + 1]
+
+        model.addConstrs(
+            (
+                y_order_r(model, var, j)
+                for (var, j) in interval_index(model)
+                if j != len(model._breakpoints[var]) - 1
+            ),
+            name="y_order",
+        )
+
+        # add cat var constraints
+        def cat_sums(model_obj, i):
+            return (
+                quicksum(model_obj._all_feat[i][cat] for cat in model_obj._all_feat[i])
+                == 1
+            )
+
+        model.addConstrs(
+            (cat_sums(model, var) for var in self._problem_config.cat_idx),
+            name="cat_sums",
+        )
+
+        # add split / space linking constraints
+        def var_lower_r(model_obj, lb, i, j):
+            j_bound = model_obj._breakpoints[i][j]
+            return model_obj._all_feat[i] >= lb + (j_bound - lb) * (
+                1 - model_obj._nu[i, j]
+            )
+
+        def var_upper_r(model_obj, ub, i, j):
+            j_bound = model_obj._breakpoints[i][j]
+            return model_obj._all_feat[i] <= ub + (j_bound - ub) * (model_obj._nu[i, j])
+
+        model.addConstrs(
+            (
+                var_lower_r(model, self._problem_config.feat_list[var].lb, var, j)
+                for (var, j) in interval_index(model)
+            ),
+            name="var_lower",
+        )
+
+        model.addConstrs(
+            (
+                var_upper_r(model, self._problem_config.feat_list[var].ub, var, j)
+                for (var, j) in interval_index(model)
+            ),
+            name="var_upper",
+        )
+
+        if add_mu_var:
+            # add mu variables for all objectives
+            def obj_leaf_index(model_obj, obj_name):
+                for tree in range(model_obj._num_trees(obj_name)):
+                    for leaf in model_obj._leaves(obj_name, tree):
+                        yield tree, leaf
+
+            model._aux_mu = []
+            model._unscaled_mu = []
+
+            for idx, obj in enumerate(self._problem_config.obj_list):
+                weighted_sum = quicksum(
+                    model._leaf_weight(obj.name, tree, leaf)
+                    * model._z[obj.name, tree, leaf]
+                    for tree, leaf in obj_leaf_index(model, obj.name)
+                )
+
+                model._aux_mu.append(
+                    model.addVar(
+                        lb=-GRB.INFINITY,
+                        ub=GRB.INFINITY,
+                        name=f"aux_mean_obj_{obj.name}",
+                        vtype="C",
+                    )
+                )
+
+                model._unscaled_mu.append(
+                    model.addVar(
+                        lb=-GRB.INFINITY,
+                        ub=GRB.INFINITY,
+                        name=f"unscaled_mean_obj_{obj.name}",
+                        vtype="C",
+                    )
+                )
+
+                if normalize_mean:
+                    shift, scale = self.min_y[idx], self.max_y[idx] - self.min_y[idx]
+                else:
+                    shift, scale = 0.0, 1.0
+
+                model.addConstr(
+                    model._aux_mu[-1] == (weighted_sum - shift) / scale,
+                    name=f"aux_mean_obj_{obj.name}_tree_link",
+                )
+
+                model.addConstr(
+                    model._unscaled_mu[-1] == (model._aux_mu[-1] * scale) + shift,
+                    name=f"unscaled_mean_obj_{obj.name}_tree_link",
+                )
+
+        model.update()
+
+    def add_to_pyomo_model(
+        self, model, add_mu_var: bool = True, normalize_mean: bool = False
+    ):
+        import pyomo.environ as pyo
+
+        self._update_meta_tree_dict()
+
+        # attach tree info
+        model._num_trees = lambda obj_name: self._meta_tree_dict[obj_name].num_trees
+
+        # attach leaf info
+        model._leaves = lambda obj_name, tree: tuple(
+            self._meta_tree_dict[obj_name].get_leaf_encodings(tree)
+        )
+        model._leaf_weight = lambda obj_name, tree, leaf: self._meta_tree_dict[
+            obj_name
+        ].get_leaf_weight(tree, leaf)
+        model._leaf_vars = lambda obj_name, tree, leaf: tuple(
+            var
+            for var in self._meta_tree_dict[obj_name].get_participating_variables(
+                tree, leaf
+            )
+        )
+
+        # attach breakpoint info
+        var_break_points = [
+            tree_model.get_var_break_points()
+            for tree_model in self._meta_tree_dict.values()
+        ]
+        model._breakpoints = {}
+        model._breakpoint_index = []
+
+        for idx, feat in enumerate(self._problem_config.feat_list):
+            if feat.is_cat():
+                continue
+            else:
+                splits = set()
+                for vb in var_break_points:
+                    try:
+                        splits = splits.union(set(vb[idx]))
+                    except KeyError:
+                        pass
+                if splits:
+                    model._breakpoints[idx] = sorted(splits)
+                    model._breakpoint_index.append(idx)
+
+        # define indexing helper functions
+        def tree_index(model_obj):
+            for obj in self._problem_config.obj_list:
+                for tree in range(model_obj._num_trees(obj.name)):
+                    yield obj.name, tree
+
+        def leaf_index(model_obj):
+            for label, tree in tree_index(model_obj):
+                for leaf in model_obj._leaves(label, tree):
+                    yield label, tree, leaf
+
+        def interval_index(model_obj):
+            for var in model_obj._breakpoint_index:
+                for j in range(len(model_obj._breakpoints[var])):
+                    yield var, j
+
+        def split_index(model_obj, meta_tree_dict):
+            for label, tree in tree_index(model_obj):
+                for encoding in meta_tree_dict[label].get_branch_encodings(tree):
+                    yield label, tree, encoding
+
+        # add leaf variables
+        model._z = pyo.Var(leaf_index(model), domain=pyo.NonNegativeReals)
+
+        # add split variables
+        model._nu = pyo.Var(interval_index(model), domain=pyo.Binary)
+
+        def single_leaf_rule(model_obj, label, tree):
+            z, leaves = model_obj._z, model_obj._leaves
+            return sum(z[label, tree, leaf] for leaf in leaves(label, tree)) == 1
+
+        model.single_leaf_constraints = pyo.Constraint(
+            [(label, tree) for (label, tree) in tree_index(model)],
+            rule=single_leaf_rule,
+        )
+
+        # We hide the dictionary in the model data since we will need it some functions (e.g.left_split_r()), but
+        # the function is not allowed to have the dict as an argument since otherwise Pyomo would not accept it as rule.
+        model.meta_tree_dict = self.meta_tree_dict
+
+        # add left split constraints
+        def left_split_r(model_obj, label, tree, split_enc):
+            meta_tree_dict = model.meta_tree_dict
+            split_var, split_val = meta_tree_dict[label].get_branch_partition_pair(
+                tree, split_enc
+            )
+
+            if not isinstance(split_val, list):
+                # non-cat vars
+                nu_val = model_obj._breakpoints[split_var].index(split_val)
+                return (
+                    sum(
+                        model_obj._z[label, tree, leaf]
+                        for leaf in meta_tree_dict[label].get_left_leaves(
+                            tree, split_enc
+                        )
+                    )
+                    <= model_obj._nu[split_var, nu_val]
+                )
+            else:
+                # cat vars
+                return sum(
+                    model_obj._z[label, tree, leaf]
+                    for leaf in meta_tree_dict[label].get_left_leaves(tree, split_enc)
+                ) <= sum(model_obj._all_feat[split_var][cat] for cat in split_val)
+
+        model.left_split_constraints = pyo.Constraint(
+            [
+                (label, tree, encoding)
+                for (label, tree, encoding) in split_index(model, self.meta_tree_dict)
+            ],
+            rule=left_split_r,
+        )
+
+        # add right split constraints
+        def right_split_r(model_obj, label, tree, split_enc):
+            meta_tree_dict = model_obj.meta_tree_dict
+            split_var, split_val = meta_tree_dict[label].get_branch_partition_pair(
+                tree, split_enc
+            )
+
+            if not isinstance(split_val, list):
+                # non-cat vars
+                nu_val = model_obj._breakpoints[split_var].index(split_val)
+                return (
+                    sum(
+                        model_obj._z[label, tree, leaf]
+                        for leaf in meta_tree_dict[label].get_right_leaves(
+                            tree, split_enc
+                        )
+                    )
+                    <= 1 - model_obj._nu[split_var, nu_val]
+                )
+            else:
+                # cat vars
+                return sum(
+                    model_obj._z[label, tree, leaf]
+                    for leaf in meta_tree_dict[label].get_right_leaves(tree, split_enc)
+                ) <= 1 - sum(model_obj._all_feat[split_var][cat] for cat in split_val)
+
+        model.right_split_constraints = pyo.Constraint(
+            [
+                (label, tree, encoding)
+                for (label, tree, encoding) in split_index(model, self.meta_tree_dict)
+            ],
+            rule=right_split_r,
+        )
+
+        # add split order constraints
+        def y_order_r(model_obj, i, j):
+            return model_obj._nu[i, j] <= model_obj._nu[i, j + 1]
+
+        model.y_order_constraints = pyo.Constraint(
+            [
+                (var, j)
+                for (var, j) in interval_index(model)
+                if j != len(model._breakpoints[var]) - 1
+            ],
+            rule=y_order_r,
+        )
+
+        # add cat var constraints
+        def cat_sums(model_obj, i):
+            return (
+                sum(model_obj._all_feat[i][cat] for cat in model_obj._all_feat[i]) == 1
+            )
+
+        model.categorical_sum_constraints = pyo.Constraint(
+            [var for var in self._problem_config.cat_idx], rule=cat_sums
+        )
+
+        # add split / space linking constraints
+        def var_lower_r(model_obj, i, j):
+            lb = model_obj._all_feat[i].lb
+            j_bound = model_obj._breakpoints[i][j]
+            return model_obj._all_feat[i] >= lb + (j_bound - lb) * (
+                1 - model_obj._nu[i, j]
+            )
+
+        def var_upper_r(model_obj, i, j):
+            ub = model_obj._all_feat[i].ub
+            j_bound = model_obj._breakpoints[i][j]
+            return model_obj._all_feat[i] <= ub + (j_bound - ub) * (model_obj._nu[i, j])
+
+        model.linking_constraints_lower = pyo.Constraint(
+            [(var, j) for (var, j) in interval_index(model)], rule=var_lower_r
+        )
+
+        model.linking_constraints_upper = pyo.Constraint(
+            [(var, j) for (var, j) in interval_index(model)], rule=var_upper_r
+        )
+
+        if add_mu_var:
+            # add mu variables for all objectives
+            def obj_leaf_index(model_obj, obj_name):
+                for tree in range(model_obj._num_trees(obj_name)):
+                    for leaf in model_obj._leaves(obj_name, tree):
+                        yield tree, leaf
+
+            objective_names = [obj.name for obj in self._problem_config.obj_list]
+
+            model._aux_mu = pyo.Var(objective_names, domain=pyo.Reals)
+            model._unscaled_mu = pyo.Var(objective_names, domain=pyo.Reals)
+
+            weighted_sum = {}
+            shift = {}
+            scale = {}
+
+            for idx, obj_name in enumerate(objective_names):
+                weighted_sum[obj_name] = sum(
+                    model._leaf_weight(obj_name, tree, leaf)
+                    * model._z[obj_name, tree, leaf]
+                    for tree, leaf in obj_leaf_index(model, obj_name)
+                )
+
+                if normalize_mean:
+                    shift[obj_name], scale[obj_name] = (
+                        self.min_y[idx],
+                        self.max_y[idx] - self.min_y[idx],
+                    )
+                else:
+                    shift[obj_name], scale[obj_name] = 0.0, 1.0
+
+            def mu_objectives(model_obj, name):
+                return (
+                    model_obj._aux_mu[name]
+                    == (weighted_sum[name] - shift[name]) / scale[name]
+                )
+
+            model.constraints_mu_objectives = pyo.Constraint(
+                objective_names, rule=mu_objectives
+            )
+
+            def link_unscaled_mu_auxmu(model_obj: pyo.ConcreteModel, name: str):
+                return (
+                    model_obj._unscaled_mu[name]
+                    == (model._aux_mu[name] * scale[name]) + shift[name]
+                )
+
+            model.constraints_link_unscaled_mu_auxmu = pyo.Constraint(
+                objective_names, rule=link_unscaled_mu_auxmu
+            )
diff --git a/entmoot/models/uncertainty_models/__init__.py b/entmoot/models/uncertainty_models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/entmoot/models/uncertainty_models/base_distance.py b/entmoot/models/uncertainty_models/base_distance.py
new file mode 100644
index 0000000..421a9fb
--- /dev/null
+++ b/entmoot/models/uncertainty_models/base_distance.py
@@ -0,0 +1,162 @@
+from entmoot.models.base_model import BaseModel
+import numpy as np
+
+
+class NonCatDistance(BaseModel):
+    def __init__(self, problem_config, acq_sense, dist_trafo):
+        self._problem_config = problem_config
+        self._acq_sense = acq_sense
+        self._dist_trafo = dist_trafo
+
+        self._shift, self._scale = None, None
+        self._x_trafo = None
+
+    @property
+    def shift(self):
+        return self._shift
+
+    @property
+    def scale(self):
+        return self._scale
+
+    @property
+    def x_trafo(self):
+        assert (
+            self._x_trafo is not None
+        ), f"Uncertainty model needs fit function call before it can predict."
+        return self._x_trafo
+
+    def get_big_m(self):
+        lb = self._trafo(self._problem_config.non_cat_lb)
+        ub = self._trafo(self._problem_config.non_cat_ub)
+        return self._get_distance(lb, ub)
+
+    def _get_distance(self, x_left, x_right):
+        raise NotImplementedError()
+
+    def _trafo(self, non_cat_x):
+        return (non_cat_x - self._shift) / self._scale
+
+    def predict(self, xi):
+        non_cat_x = self._trafo(np.atleast_2d(xi)[:, self._problem_config.non_cat_idx])
+        return self._get_distance(self.x_trafo, non_cat_x)
+
+    def _array_predict(self, X):
+        raise NotImplementedError()
+
+    def fit(self, X):
+        non_cat_x = X[:, self._problem_config.non_cat_idx]
+
+        # define shift and scalar values for non-cat feats
+        if self._dist_trafo == "normal":
+            self._shift = np.asarray(self._problem_config.non_cat_lb)
+            self._scale = np.asarray(self._problem_config.non_cat_bnd_diff)
+        elif self._dist_trafo == "standard":
+            self._shift = np.mean(np.asarray(non_cat_x), axis=0)
+            self._scale = np.std(np.asarray(non_cat_x), axis=0)
+        else:
+            raise IOError(
+                "Parameter 'dist_trafo' for uncertainty model needs to be "
+                "in '('normal', 'standard')'."
+            )
+
+        self._x_trafo = self._trafo(non_cat_x)
+
+    def get_gurobipy_model_constr(self, model_core):
+        raise NotImplementedError()
+
+    def get_pyomo_model_constr(self, model_core):
+        raise NotImplementedError()
+
+
+class CatDistance(BaseModel):
+    def __init__(self, problem_config, acq_sense):
+        self._problem_config = problem_config
+        self._acq_sense = acq_sense
+
+        self._cat_x = None
+        self._sim_map = None
+        self._cache_x = None
+
+    @property
+    def cache_x(self):
+        assert (
+            self._cache_x is not None
+        ), f"Uncertainty model needs fit function call before it can predict."
+        return self._cache_x
+
+    @property
+    def sim_map(self):
+        assert (
+            self._sim_map is not None
+        ), f"Uncertainty model needs fit function call before it can predict."
+        return self._sim_map
+
+    def get_big_m(self):
+        # categorical contribution has maximum of one
+        return len(self._problem_config.cat_idx)
+
+    def predict(self, xi):
+        # iterate through every evaluation row
+        sim_vec = np.zeros(len(self.cache_x))
+
+        for idx in self._problem_config.cat_idx:
+            pred_cat = int(xi[idx])
+            cached_cats = (np.rint(self.cache_x[:, idx])).astype(int)
+            sim_vec += self.sim_map[idx][pred_cat, cached_cats]
+
+        # compute distance based on similarity
+        dist_vec = (-1) * sim_vec + len(self._problem_config.cat_idx)
+        return dist_vec
+
+    def _sim_mat_rule(self, x_left, x_right, cat_idx):
+        raise NotImplementedError()
+
+    def fit(self, X):
+        self._cache_x = X
+
+        # generate similarity matrix for all data points
+        self._sim_map = {}
+
+        for idx in self._problem_config.cat_idx:
+            feat = self._problem_config.feat_list[idx]
+            all_cats = feat.enc_cat_list
+
+            # creates similarity entries for all categories of all categorical features
+            mat = np.fromfunction(
+                np.vectorize(self._sim_mat_rule, otypes=[float]),
+                (len(all_cats), len(all_cats)),
+                dtype=int,
+                cat_idx=idx,
+            )
+            self._sim_map[idx] = mat
+
+    def get_gurobipy_model_constr_terms(self, model):
+        feat = model._all_feat
+        constr_list = []
+        for xi in self.cache_x:
+            constr = 0
+
+            # iterate through all categories to check which distances are active
+            for idx in self._problem_config.cat_idx:
+                for cat in self._problem_config.feat_list[idx].enc_cat_list:
+                    sim = self.sim_map[idx][cat, int(xi[idx])]
+                    constr += (1 - sim) * feat[idx][cat]
+
+            constr_list.append(constr)
+        return constr_list
+
+    def get_pyomo_model_constr_terms(self, model):
+        feat = model._all_feat
+        constr_list = []
+        for xi in self.cache_x:
+            constr = 0
+
+            # iterate through all categories to check which distances are active
+            for idx in self._problem_config.cat_idx:
+                for cat in self._problem_config.feat_list[idx].enc_cat_list:
+                    sim = self.sim_map[idx][cat, int(xi[idx])]
+                    constr += (1 - sim) * feat[idx][cat]
+
+            constr_list.append(constr)
+        return constr_list
diff --git a/entmoot/models/uncertainty_models/distance_based_uncertainty.py b/entmoot/models/uncertainty_models/distance_based_uncertainty.py
new file mode 100644
index 0000000..f4c2c36
--- /dev/null
+++ b/entmoot/models/uncertainty_models/distance_based_uncertainty.py
@@ -0,0 +1,358 @@
+from entmoot.models.base_model import BaseModel
+from entmoot.models.uncertainty_models.euclidean_squared_distance import (
+    EuclideanSquaredDistance,
+)
+from entmoot.models.uncertainty_models.l1_distance import L1Distance
+from entmoot.models.uncertainty_models.l2_distance import L2Distance
+
+from entmoot.models.uncertainty_models.overlap_distance import OverlapDistance
+from entmoot.models.uncertainty_models.goodall4_distance import Goodall4Distance
+from entmoot.models.uncertainty_models.of_distance import OfDistance
+
+import numpy as np
+
+
+class DistanceBasedUncertainty(BaseModel):
+    def __init__(self, problem_config, params):
+        self._problem_config = problem_config
+
+        dist_metric = params.get("dist_metric", "euclidean_squared")
+        dist_trafo = params.get("dist_trafo", "normal")
+        cat_metric = params.get("cat_metric", "overlap")
+        acq_sense = params.get("acq_sense", "exploration")
+
+        self._non_cat_x, self._cat_x = None, None
+        self._dist_bound = None
+        self._dist_metric = dist_metric
+        self._num_cache_x = None
+        self._acq_sense = acq_sense
+
+        if dist_trafo == "standard":
+            assert (
+                len(self._problem_config.obj_list) == 1
+            ), "Distance transformation 'standard' can only be used for single objective problems."
+
+            assert (
+                len(self._problem_config.cat_idx) == 0
+            ), "Distance transformation 'standard' can only be used for non-categorical problems."
+
+            self._dist_has_var_bound = False if self._acq_sense == "penalty" else True
+            self._bound_coeff = params.get("bound_coeff", 0.5)
+            self._dist_coeff = 1.0
+        elif dist_trafo == "normal":
+            self._dist_has_var_bound = False
+            self._bound_coeff = None
+            self._dist_coeff = 1 / len(self._problem_config.feat_list)
+        else:
+            raise IOError(
+                f"Pick 'dist_trafo' '{dist_trafo}' in '('normal', 'standard')'."
+            )
+
+        # pick distance metric for non-cat features
+        if dist_metric == "euclidean_squared":
+            self.non_cat_unc_model = EuclideanSquaredDistance(
+                problem_config=self._problem_config,
+                acq_sense=acq_sense,
+                dist_trafo=dist_trafo,
+            )
+        elif dist_metric == "l1":
+            self.non_cat_unc_model = L1Distance(
+                problem_config=self._problem_config,
+                acq_sense=acq_sense,
+                dist_trafo=dist_trafo,
+            )
+        elif dist_metric == "l2":
+            self.non_cat_unc_model = L2Distance(
+                problem_config=self._problem_config,
+                acq_sense=acq_sense,
+                dist_trafo=dist_trafo,
+            )
+        else:
+            raise IOError(
+                f"Non-categorical uncertainty metric '{dist_metric}' for "
+                f"{self.__class__.__name__} model is not supported. "
+                f"Check 'params['uncertainty_type']'."
+            )
+
+        # pick distance metric for cat features
+        if cat_metric == "overlap":
+            self.cat_unc_model = OverlapDistance(
+                problem_config=self._problem_config, acq_sense=acq_sense
+            )
+        elif cat_metric == "of":
+            self.cat_unc_model = OfDistance(
+                problem_config=self._problem_config, acq_sense=acq_sense
+            )
+        elif cat_metric == "goodall4":
+            self.cat_unc_model = Goodall4Distance(
+                problem_config=self._problem_config, acq_sense=acq_sense
+            )
+        else:
+            raise IOError(
+                f"Categorical uncertainty metric '{cat_metric}' for {self.__class__.__name__} "
+                f"model is not supported. Check 'params['uncertainty_type']'."
+            )
+
+    @property
+    def num_cache_x(self):
+        assert (
+            self._num_cache_x is not None
+        ), f"Uncertainty model needs fit function call before it can predict."
+        return self._num_cache_x
+
+    def fit(self, X, y):
+        if self._dist_has_var_bound:
+            self._dist_bound = abs(np.var(y) * self._bound_coeff)
+
+        self._num_cache_x = len(X)
+
+        self.non_cat_unc_model.fit(X)
+        self.cat_unc_model.fit(X)
+
+    def predict(self, X):
+        comb_pred = []
+        for xi in X:
+            non_cat_pred = self.non_cat_unc_model.predict(xi)
+            cat_pred = self.cat_unc_model.predict(xi)
+            dist_pred = np.min(non_cat_pred + cat_pred) * self._dist_coeff
+
+            # the standard trafo case has a bound on the prediction
+            if self._dist_has_var_bound:
+                if dist_pred > self._dist_bound:
+                    dist_pred = self._dist_bound
+
+            comb_pred.append(dist_pred)
+        return np.asarray(comb_pred)
+
+    def add_to_gurobipy_model(self, model):
+        from gurobipy import GRB
+
+        # define main uncertainty variables
+        if self._dist_has_var_bound:
+            dist_bound = self._dist_bound
+        else:
+            dist_bound = GRB.INFINITY
+
+        model._unc = model.addVar(lb=0.0, ub=dist_bound, name="uncertainty", vtype="C")
+
+        # get constr terms for non-categorical and categorical contributions
+        non_cat_term_list = self.non_cat_unc_model.get_gurobipy_model_constr_terms(
+            model
+        )
+        cat_term_list = self.cat_unc_model.get_gurobipy_model_constr_terms(model)
+
+        # add binaries for big_m penalty constraints
+        if self._acq_sense == "penalty":
+            model._bin_penalty = []
+            big_m = self.non_cat_unc_model.get_big_m() + self.cat_unc_model.get_big_m()
+
+        for i, (non_cat_term, cat_term) in enumerate(
+            zip(non_cat_term_list, cat_term_list)
+        ):
+            # check if penalty term is needed
+            if self._acq_sense == "penalty":
+
+                model._bin_penalty.append(
+                    model.addVar(name=f"bin_penalty_{i}", vtype="B")
+                )
+
+                big_m_term = big_m * (1 - model._bin_penalty[-1])
+
+                if self._dist_metric == "l2":
+                    # take sqrt for l2 distance
+                    aux_non_cat_unc = model.addVar(
+                        lb=0.0, ub=dist_bound, name=f"aux_non_cat_unc_x_{i}", vtype="C"
+                    )
+
+                    model.addQConstr(
+                        aux_non_cat_unc * aux_non_cat_unc >= non_cat_term,
+                        name=f"unc_x_{i}",
+                    )
+
+                    model.addQConstr(
+                        model._unc + big_m_term
+                        >= (aux_non_cat_unc + cat_term) * self._dist_coeff,
+                        name=f"unc_x_{i}",
+                    )
+                    model.params.NonConvex = 2
+                else:
+                    # self._dist_metric == "l1" or "euclidean_squared"
+                    model.addQConstr(
+                        model._unc + big_m_term
+                        >= (non_cat_term + cat_term) * self._dist_coeff,
+                        name=f"unc_x_{i}",
+                    )
+            elif self._acq_sense =="exploration":
+                if self._dist_metric == "l2":
+
+                    # take sqrt for l2 distance
+                    aux_non_cat_unc = model.addVar(
+                        lb=0.0, ub=dist_bound, name=f"aux_non_cat_unc_x_{i}", vtype="C"
+                    )
+
+                    model.addQConstr(
+                        aux_non_cat_unc * aux_non_cat_unc <= non_cat_term,
+                        name=f"unc_x_{i}",
+                    )
+
+                    model.addQConstr(
+                        model._unc * model._unc
+                        <= (aux_non_cat_unc + cat_term) * self._dist_coeff,
+                        name=f"unc_x_{i}",
+                    )
+                else:
+                    model.addQConstr(
+                        model._unc <= (non_cat_term + cat_term) * self._dist_coeff,
+                        name=f"unc_x_{i}",
+                    )
+
+        if self._acq_sense == "exploration":
+            model.params.NonConvex = 2
+
+        model.update()
+
+    def add_to_pyomo_model(self, model):
+        import pyomo.environ as pyo
+
+        # define main uncertainty variables
+        if self._dist_has_var_bound:
+            dist_bound = self._dist_bound
+        else:
+            dist_bound = float("inf")
+
+        model._unc = pyo.Var(bounds=(0, dist_bound), domain=pyo.Reals)
+
+        # get constr terms for non-categorical and categorical contributions
+        non_cat_term_list = self.non_cat_unc_model.get_pyomo_model_constr_terms(model)
+        cat_term_list = self.cat_unc_model.get_pyomo_model_constr_terms(model)
+
+        model.terms_constrs_cat_noncat_contr = list(
+            zip(non_cat_term_list, cat_term_list)
+        )
+        model.indices_constrs_cat_noncat_contr = list(
+            range(len(model.terms_constrs_cat_noncat_contr))
+        )
+
+        # add binaries for big_m penalty constraints
+        if self._acq_sense == "penalty":
+            big_m = self.non_cat_unc_model.get_big_m() + self.cat_unc_model.get_big_m()
+
+            model._bin_penalty = pyo.Var(
+                model.indices_constrs_cat_noncat_contr, domain=pyo.Binary
+            )
+
+            big_m_term = {
+                i: big_m * (1 - model._bin_penalty[i])
+                for i in model.indices_constrs_cat_noncat_contr
+            }
+
+            if self._dist_metric == "l2":
+                # take sqrt for l2 distance
+                model.aux_non_cat_unc = pyo.Var(
+                    model.indices_constrs_cat_noncat_contr, bounds=(0, dist_bound)
+                )
+
+                def constrs_unc_x_aux(model_obj, k):
+                    return (
+                        model_obj.aux_non_cat_unc[k] * model_obj.aux_non_cat_unc[k]
+                        >= non_cat_term_list[k]
+                    )
+
+                model.constrs_unc_x_aux = pyo.Constraint(
+                    model.indices_constrs_cat_noncat_contr, rule=constrs_unc_x_aux
+                )
+
+                def constrs_unc_x(model_obj, k):
+                    return (
+                        model_obj._unc + big_m_term[k]
+                        >= (model_obj.aux_non_cat_unc[k] + cat_term_list[k])
+                        * self._dist_coeff
+                    )
+
+                model.constrs_unc_x_l2 = pyo.Constraint(
+                    model.indices_constrs_cat_noncat_contr, rule=constrs_unc_x
+                )
+
+            elif self._dist_metric == "l1" or "euclidean_squared":
+
+                def constrs_unc_x_l1l2squared(model_obj, k):
+                    return (
+                        model_obj._unc + big_m_term[k]
+                        >= (non_cat_term_list[k] + cat_term_list[k]) * self._dist_coeff
+                    )
+
+                model.constrs_unc_x_l1l2squared = pyo.Constraint(
+                    model.indices_constrs_cat_noncat_contr,
+                    rule=constrs_unc_x_l1l2squared,
+                )
+
+            else:
+                # Illegal value for self._dist_metric
+                raise ValueError
+        elif self._acq_sense == "exploration":
+            if self._dist_metric == "l2":
+
+                # take sqrt for l2 distance
+                model.aux_non_cat_unc = pyo.Var(
+                    model.indices_constrs_cat_noncat_contr, bounds=(0, dist_bound)
+                )
+
+                def constrs_unc_x_aux_expl(model_obj, k):
+                    return (
+                        model_obj.aux_non_cat_unc[k] * model_obj.aux_non_cat_unc[k]
+                        <= non_cat_term_list[k]
+                    )
+
+                model.constrs_unc_x_aux_expl = pyo.Constraint(
+                    model.indices_constrs_cat_noncat_contr, rule=constrs_unc_x_aux_expl
+                )
+
+                def constrs_unc_x_expl(model_obj, k):
+                    return (
+                        model_obj._unc * model_obj._unc
+                        <= (model_obj.aux_non_cat_unc[k] + cat_term_list[k])
+                        * self._dist_coeff
+                    )
+
+                model.constrs_unc_x_expl = pyo.Constraint(
+                    model.indices_constrs_cat_noncat_contr, rule=constrs_unc_x_expl
+                )
+
+            elif self._dist_metric == "l1" or "euclidean_squared":
+
+                def constrs_unc_x_l1l2squared_expl(model_obj, k):
+                    return (
+                        model_obj._unc
+                        <= (non_cat_term_list[k] + cat_term_list[k]) * self._dist_coeff
+                    )
+
+                model.constrs_unc_x_l1l2squared = pyo.Constraint(
+                    model.indices_constrs_cat_noncat_contr,
+                    rule=constrs_unc_x_l1l2squared_expl,
+                )
+
+            else:
+                # Illegal value for self._dist_metric
+                raise ValueError
+
+        def rule_constr_cat_noncat_quadr(model_obj, i):
+            # take sqrt for l2 distance
+            non_cat_term, cat_term = model.terms_constrs_cat_noncat_contr[i]
+            return (
+                model._unc * model._unc
+                <= (non_cat_term + cat_term) * self._dist_coeff**2
+            )
+
+        def rule_constr_cat_noncat(model_obj, i):
+            non_cat_term, cat_term = model.terms_constrs_cat_noncat_contr[i]
+            return model._unc <= (non_cat_term + cat_term) * self._dist_coeff
+
+        if self._dist_metric == "l2":
+            model.constrs_cat_noncat_contr = pyo.Constraint(
+                model.indices_constrs_cat_noncat_contr,
+                rule=rule_constr_cat_noncat_quadr,
+            )
+        else:
+            model.constrs_cat_noncat_contr = pyo.Constraint(
+                model.indices_constrs_cat_noncat_contr, rule=rule_constr_cat_noncat
+            )
diff --git a/entmoot/models/uncertainty_models/euclidean_squared_distance.py b/entmoot/models/uncertainty_models/euclidean_squared_distance.py
new file mode 100644
index 0000000..b4d8015
--- /dev/null
+++ b/entmoot/models/uncertainty_models/euclidean_squared_distance.py
@@ -0,0 +1,45 @@
+from entmoot.models.uncertainty_models.base_distance import NonCatDistance
+import numpy as np
+
+
+class EuclideanSquaredDistance(NonCatDistance):
+    def _get_distance(self, x_left, x_right):
+        if x_left.ndim == 1:
+            dist = np.sum((x_left - x_right) ** 2)
+        else:
+            dist = np.sum((x_left - x_right) ** 2, axis=1)
+        return dist
+
+    def _array_predict(self, X):
+        raise NotImplementedError()
+
+    def get_gurobipy_model_constr_terms(self, model):
+
+        from gurobipy import quicksum
+
+        feat = model._all_feat
+
+        constr_list = []
+        for xi in self.x_trafo:
+            constr = quicksum(
+                (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                * (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                for i, idx in enumerate(self._problem_config.non_cat_idx)
+            )
+            constr_list.append(constr)
+        return constr_list
+
+    def get_pyomo_model_constr_terms(self, model):
+
+        feat = model._all_feat
+
+        constr_list = []
+
+        for xi in self.x_trafo:
+            constr = sum(
+                (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                * (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                for i, idx in enumerate(self._problem_config.non_cat_idx)
+            )
+            constr_list.append(constr)
+        return constr_list
diff --git a/entmoot/models/uncertainty_models/goodall4_distance.py b/entmoot/models/uncertainty_models/goodall4_distance.py
new file mode 100644
index 0000000..42a5ac5
--- /dev/null
+++ b/entmoot/models/uncertainty_models/goodall4_distance.py
@@ -0,0 +1,16 @@
+from entmoot.models.uncertainty_models.base_distance import CatDistance
+import numpy as np
+
+
+class Goodall4Distance(CatDistance):
+    def _get_pk2(self, cat_rows, cat):
+        count_cat = np.sum(cat_rows == cat)
+        n_rows = len(cat_rows)
+        return (count_cat * (count_cat - 1)) / (n_rows * (n_rows - 1))
+
+    def _sim_mat_rule(self, x_left, x_right, cat_idx):
+        return (
+            self._get_pk2(self._cache_x[:, cat_idx], x_left)
+            if x_left == x_right
+            else 0.0
+        )
diff --git a/entmoot/models/uncertainty_models/l1_distance.py b/entmoot/models/uncertainty_models/l1_distance.py
new file mode 100644
index 0000000..0f392b9
--- /dev/null
+++ b/entmoot/models/uncertainty_models/l1_distance.py
@@ -0,0 +1,144 @@
+from entmoot.models.uncertainty_models.base_distance import NonCatDistance
+import numpy as np
+
+
+class L1Distance(NonCatDistance):
+    def _get_distance(self, x_left, x_right):
+        if x_left.ndim == 1:
+            dist = np.sum(np.abs(x_left - x_right))
+        else:
+            dist = np.sum(np.abs(x_left - x_right), axis=1)
+        return dist
+
+    def _array_predict(self, X):
+        raise NotImplementedError()
+
+    def get_gurobipy_model_constr_terms(self, model):
+        from gurobipy import GRB, quicksum
+
+        feat = model._all_feat
+
+        # define auxiliary variables
+        feat_dict = {
+            idx: self._problem_config.feat_list[idx]
+            for idx in self._problem_config.non_cat_idx
+        }
+
+        aux_pos = model.addVars(len(self.x_trafo), feat_dict, name="aux_pos", vtype="C")
+
+        aux_neg = model.addVars(len(self.x_trafo), feat_dict, name="aux_neg", vtype="C")
+
+        # define distance constraints
+        constr_list = []
+        for data_idx, xi in enumerate(self.x_trafo):
+            for i, idx in enumerate(self._problem_config.non_cat_idx):
+                # capture positive and negative contributions
+                model.addConstr(
+                    (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                    == aux_pos[data_idx, idx] - aux_neg[data_idx, idx],
+                    name=f"unc_aux_({data_idx},{idx})",
+                )
+
+                # make sure only one variable is non-zero
+                model.addSOS(
+                    GRB.SOS_TYPE1, [aux_pos[data_idx, idx], aux_neg[data_idx, idx]]
+                )
+
+            # define individual sum terms
+            constr = quicksum(
+                aux_pos[data_idx, idx] + aux_neg[data_idx, idx]
+                for idx in self._problem_config.non_cat_idx
+            )
+
+            constr_list.append(constr)
+
+        model.update()
+        return constr_list
+
+    def get_pyomo_model_constr_terms(self, model):
+        import pyomo.environ as pyo
+
+        feat = model._all_feat
+
+        # define auxiliary variables
+        feat_dict = {
+            idx: self._problem_config.feat_list[idx]
+            for idx in self._problem_config.non_cat_idx
+        }
+
+        indices_aux_vars = [(i, j) for i in range(len(self.x_trafo)) for j in feat_dict]
+        model.aux_pos = pyo.Var(indices_aux_vars, domain=pyo.NonNegativeReals)
+        model.aux_neg = pyo.Var(indices_aux_vars, domain=pyo.NonNegativeReals)
+
+        indices_l1_constraints = [
+            (data_idx, i, idx)
+            for data_idx, _ in enumerate(self.x_trafo)
+            for i, idx in enumerate(self._problem_config.non_cat_idx)
+        ]
+
+        def rule_contrs_l1_pos_neg_contr(modelobj, data_idx, i, idx):
+            xi = self.x_trafo[data_idx]
+            return (
+                xi[i] - (feat[idx] - self.shift[i]) / self.scale[i]
+            ) == modelobj.aux_pos[data_idx, idx] - modelobj.aux_neg[data_idx, idx]
+
+        model.contrs_l1_pos_neg_contr = pyo.Constraint(
+            indices_l1_constraints, rule=rule_contrs_l1_pos_neg_contr
+        )
+
+        model.l1_bigm = {
+            (data_idx, idx): max(
+                abs(
+                    self.x_trafo[data_idx][i]
+                    - (feat[idx].ub - self.shift[i]) / self.scale[i]
+                ),
+                abs(
+                    self.x_trafo[data_idx][i]
+                    - (feat[idx].lb - self.shift[i]) / self.scale[i]
+                ),
+            )
+            for (data_idx, i, idx) in indices_l1_constraints
+        }
+
+        model.indices_bigm = list(model.l1_bigm.keys())
+
+        model.y_pos_bigm = pyo.Var(model.indices_bigm, domain=pyo.Binary)
+        model.y_neg_bigm = pyo.Var(model.indices_bigm, domain=pyo.Binary)
+
+        def bigm_ub_pos(modelobj, data_idx, idx):
+            return (
+                modelobj.aux_pos[data_idx, idx]
+                <= model.l1_bigm[data_idx, idx] * model.y_pos_bigm[data_idx, idx]
+            )
+
+        def bigm_ub_neg(modelobj, data_idx, idx):
+            return (
+                modelobj.aux_neg[data_idx, idx]
+                <= model.l1_bigm[data_idx, idx] * model.y_neg_bigm[data_idx, idx]
+            )
+
+        model.constrs_l1_bigm_ub_pos = pyo.Constraint(
+            model.indices_bigm, rule=bigm_ub_pos
+        )
+        model.constrs_l1_bigm_ub_neg = pyo.Constraint(
+            model.indices_bigm, rule=bigm_ub_neg
+        )
+
+        def binaries_mutually_exclusive(modelobj, data_idx, idx):
+            return (
+                model.y_pos_bigm[data_idx, idx] + model.y_neg_bigm[data_idx, idx] == 1
+            )
+
+        model.constrs_l1_bigm_mut_excl = pyo.Constraint(
+            model.indices_bigm, rule=binaries_mutually_exclusive
+        )
+
+        constr_list = [
+            sum(
+                model.aux_pos[data_idx, idx] + model.aux_neg[data_idx, idx]
+                for idx in self._problem_config.non_cat_idx
+            )
+            for (data_idx, _) in enumerate(self.x_trafo)
+        ]
+
+        return constr_list
diff --git a/entmoot/models/uncertainty_models/l2_distance.py b/entmoot/models/uncertainty_models/l2_distance.py
new file mode 100644
index 0000000..f8f6a32
--- /dev/null
+++ b/entmoot/models/uncertainty_models/l2_distance.py
@@ -0,0 +1,43 @@
+from entmoot.models.uncertainty_models.base_distance import NonCatDistance
+import numpy as np
+
+
+class L2Distance(NonCatDistance):
+    def _get_distance(self, x_left, x_right):
+        if x_left.ndim == 1:
+            dist = np.sqrt(np.sum((x_left - x_right) ** 2))
+        else:
+            dist = np.sqrt(np.sum((x_left - x_right) ** 2, axis=1))
+        return dist
+
+    def _array_predict(self, X):
+        raise NotImplementedError()
+
+    def get_gurobipy_model_constr_terms(self, model):
+        from gurobipy import quicksum
+
+        feat = model._all_feat
+
+        constr_list = []
+        for xi in self.x_trafo:
+            constr = quicksum(
+                (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                * (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                for i, idx in enumerate(self._problem_config.non_cat_idx)
+            )
+            constr_list.append(constr)
+        return constr_list
+
+    def get_pyomo_model_constr_terms(self, model):
+
+        feat = model._all_feat
+
+        constr_list = []
+        for xi in self.x_trafo:
+            constr = sum(
+                (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                * (xi[i] - (feat[idx] - self.shift[i]) / self.scale[i])
+                for i, idx in enumerate(self._problem_config.non_cat_idx)
+            )
+            constr_list.append(constr)
+        return constr_list
diff --git a/entmoot/models/uncertainty_models/of_distance.py b/entmoot/models/uncertainty_models/of_distance.py
new file mode 100644
index 0000000..639518f
--- /dev/null
+++ b/entmoot/models/uncertainty_models/of_distance.py
@@ -0,0 +1,18 @@
+from entmoot.models.uncertainty_models.base_distance import CatDistance
+import numpy as np
+from math import log
+
+
+class OfDistance(CatDistance):
+    def _get_of_frac(self, cat_rows, cat_left, cat_right):
+        count_cat_left = np.sum(cat_rows == cat_left)
+        count_cat_right = np.sum(cat_rows == cat_right)
+        n_rows = len(cat_rows)
+        return 1 / (1 + log(n_rows / count_cat_left) * log(n_rows / count_cat_right))
+
+    def _sim_mat_rule(self, x_left, x_right, cat_idx):
+        return (
+            self._get_of_frac(self._cache_x[:, cat_idx], x_left, x_right)
+            if x_left != x_right
+            else 1.0
+        )
diff --git a/entmoot/models/uncertainty_models/overlap_distance.py b/entmoot/models/uncertainty_models/overlap_distance.py
new file mode 100644
index 0000000..c9a9abc
--- /dev/null
+++ b/entmoot/models/uncertainty_models/overlap_distance.py
@@ -0,0 +1,6 @@
+from entmoot.models.uncertainty_models.base_distance import CatDistance
+
+
+class OverlapDistance(CatDistance):
+    def _sim_mat_rule(self, x_left, x_right, cat_idx):
+        return 1.0 if x_left == x_right else 0.0
diff --git a/entmoot/optimizer/__init__.py b/entmoot/optimizer/__init__.py
deleted file mode 100644
index 88ac37f..0000000
--- a/entmoot/optimizer/__init__.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from .optimizer import Optimizer
-from .entmoot_minimize import entmoot_minimize
-
-__all__ = [
-    "Optimizer","entmoot_minimize"
-]
diff --git a/entmoot/optimizer/entmoot_minimize.py b/entmoot/optimizer/entmoot_minimize.py
deleted file mode 100644
index 94e27a0..0000000
--- a/entmoot/optimizer/entmoot_minimize.py
+++ /dev/null
@@ -1,399 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-NOTE: Changes were made to the scikit-optimize source code included here. 
-For the most recent version of scikit-optimize we refer to:
-https://github.com/scikit-optimize/scikit-optimize/
-
-Copyright (c) 2019-2020 Alexander Thebelt.
-"""
-
-from entmoot.optimizer.optimizer import Optimizer
-
-import copy
-import inspect
-import numbers
-
-try:
-    from collections.abc import Iterable
-except ImportError:
-    from collections import Iterable
-
-def entmoot_minimize(
-    func,
-    dimensions,
-    n_calls=60,
-    batch_size=None,
-    batch_strategy="cl_mean",
-    n_points=10000,
-    base_estimator="GBRT",
-    n_initial_points=50, 
-    initial_point_generator="random",
-    acq_func="LCB", 
-    acq_optimizer="global",
-    x0=None,
-    y0=None,  
-    random_state=None, 
-    acq_func_kwargs=None, 
-    acq_optimizer_kwargs=None,
-    base_estimator_kwargs=None,
-    model_queue_size=None,
-    verbose=False,
-):
-    """Run bayesian optimisation loop.
-
-    An `Optimizer` is initialized which represents the steps of a bayesian 
-    optimisation loop. To use it you need to provide your own loop mechanism. 
-
-    Parameters specify how `Optimizer` is initialized, which function is 
-    optimized and how each iteration is conducted.
-
-    Parameters
-    ----------
-    func : BenchmarkFunction instance,
-        Use class entmoot.benchmarks.BenchmarkFunction as a base to incorporate
-        own black-box functions
-
-    dimensions : list, shape (n_dims,)
-        List of search space dimensions.
-        Each search dimension can be defined either as
-
-        - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-          dimensions),
-        - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-          dimensions),
-        - as a list of categories (for `Categorical` dimensions), or
-        - an instance of a `Dimension` object (`Real`, `Integer` or
-          `Categorical`).
-
-        *NOTE: `Integer` and `Categorical` variables are not yet supported and
-        will be added in the next release.
-
-    n_calls : int, default: 60
-        Budget of func calls that shall be used. `n_calls` includes initial
-        points as well as suggested model-based points.
-
-    batch_size : int, default: None
-        Number of suggestions given per sequential iteration. For `None` only
-        one suggestion per iteration is given.
-
-    batch_strategy : string, default: "cl_mean"
-        In case `batch_size` is not None, `batch_strategy` specifies the 
-        strategy used to suggest multiple points per iteration. Available
-        strategies are ["cl_min", "cl_mean", "cl_max"]
-
-    n_points : int, default: 10,000
-        For `acq_optimizer` = "sampling", `n_points` specifies the number of
-        random points used to sample the acquisition function to derive its
-        minimum.
-
-    base_estimator : string, default: "GBRT",
-        A default LightGBM surrogate model of the corresponding type is used  
-        minimize `func`.
-        The following model types are available:
-        
-        - "GBRT" for gradient-boosted trees
-
-        - "RF" for random forests
-
-        *NOTE: More model types will be added in the future
-
-    std_estimator : string, default: "BDD",
-        A model is used to estimate uncertainty of `base_estimator`.
-        Different types can be classified as exploration measures, i.e. move
-        as far as possible from reference points, and penalty measures, i.e.
-        stay as close as possible to reference points. Within these types, the 
-        following uncertainty estimators are available:
-        
-        - exploration:
-            - "BDD" for bounded-data distance, which uses squared euclidean
-              distance to standardized data points
-            - "L1BDD" for bounded-data distance, which uses manhattan
-              distance to standardized data points
-        
-        - penalty:
-            - "DDP" for data distance, which uses squared euclidean
-              distance to standardized data points
-            - "L1DDP" for data distance, which uses manhattan
-              distance to standardized data points
-
-        *NOTE: More model uncertainty estimators will be added in the future
-
-    n_initial_points : int, default: 50
-        Number of evaluations of `func` with initialization points
-        before approximating it with `base_estimator`. Initial point
-        generator can be changed by setting `initial_point_generator`. For 
-        `n_initial_points` <= 20 we need to set `min_child_samples` <= 20 in 
-        `base_estimator_kwargs` so LightGBM can train tree models based on small
-        data sets.
-
-    initial_point_generator : str, InitialPointGenerator instance, \
-            default: "random"
-        Sets a initial points generator. Can be either
-
-        - "random" for uniform random numbers,
-        - "sobol" for a Sobol sequence,
-        - "halton" for a Halton sequence,
-        - "hammersly" for a Hammersly sequence,
-        - "lhs" for a latin hypercube sequence,
-        - "grid" for a uniform grid sequence
-
-    acq_func : string, default: "LCB"
-        Function to minimize over the posterior distribution. Can be either
-
-        - "LCB" for lower confidence bound.
-
-    acq_optimizer : string, default: "sampling"
-        Method to minimize the acquisition function. The fit model
-        is updated with the optimal value obtained by optimizing `acq_func`
-        with `acq_optimizer`.
-
-        - If set to "sampling", then `acq_func` is optimized by computing
-          `acq_func` at `n_points` randomly sampled points.
-        - If set to "global", then `acq_optimizer` is optimized by using
-          global solver to find minimum of `acq_func`.
-
-    x0 : list, list of lists or `None`
-        Initial input points.
-        - If it is a list of lists, use it as a list of input points.
-        - If it is a list, use it as a single initial input point.
-        - If it is `None`, no initial input points are used.
-
-    y0 : list, scalar or `None`
-        Evaluation of initial input points.
-        - If it is a list, then it corresponds to evaluations of the function
-          at each element of `x0` : the i-th element of `y0` corresponds
-          to the function evaluated at the i-th element of `x0`.
-        - If it is a scalar, then it corresponds to the evaluation of the
-          function at `x0`.
-        - If it is None and `x0` is provided, then the function is evaluated
-          at each element of `x0`.
-
-    random_state : int, RandomState instance, or None (default)
-        Set random state to something other than None for reproducible
-        results.
-
-    acq_func_kwargs : dict
-        Additional arguments to be passed to the acquisition function.
-
-    acq_optimizer_kwargs : dict
-        Additional arguments to be passed to the acquisition optimizer.
-
-    base_estimator_kwargs : dict
-        Additional arguments to be passed to the base_estimator.
-
-    std_estimator_kwargs : dict
-        Additional arguments to be passed to the std_estimator.
-
-    model_queue_size : int or None, default: None
-        Keeps list of models only as long as the argument given. In the
-        case of None, the list has no capped length.
-
-    verbose : bool or int:
-        - If it is `True`, general solver information is printed at every 
-          iteration
-        - If it is `False`, no output is printed
-
-        - If it is 0, same as if `False`
-        - If it is 1, same as if `True`
-        - If it is 2, general solver information is printed at every iteration
-          as well as detailed solver information, i.e. gurobi log
-
-    Returns
-    -------
-    result : `OptimizeResult`, scipy object
-        OptimizeResult instance with the required information.
-    """
-
-    specs = {"args": copy.copy(inspect.currentframe().f_locals),
-             "function": inspect.currentframe().f_code.co_name}
-
-    if acq_optimizer_kwargs is None:
-        acq_optimizer_kwargs = {}
-
-    acq_optimizer_kwargs["n_points"] = n_points
-
-    # Initialize optimization
-    # Suppose there are points provided (x0 and y0), record them
-
-    # check x0: list-like, requirement of minimal points
-    if x0 is None:
-        x0 = []
-    elif not isinstance(x0[0], (list, tuple)):
-        x0 = [x0]
-    if not isinstance(x0, list):
-        raise ValueError("`x0` should be a list, but got %s" % type(x0))
-
-    if n_initial_points <= 0 and not x0:
-        raise ValueError("Either set `n_initial_points` > 0,"
-                         " or provide `x0`")
-    # check y0: list-like, requirement of maximal calls
-    if isinstance(y0, Iterable):
-        y0 = list(y0)
-    elif isinstance(y0, numbers.Number):
-        y0 = [y0]
-    required_calls = n_initial_points + (len(x0) if not y0 else 0)
-    if n_calls < required_calls:
-        raise ValueError(
-            "Expected `n_calls` >= %d, got %d" % (required_calls, n_calls))
-    # calculate the total number of initial points
-    n_initial_points = n_initial_points + len(x0)
-
-    # Build optimizer
-
-    # create optimizer class
-    optimizer = Optimizer(
-        dimensions, 
-        base_estimator=base_estimator,
-        n_initial_points=n_initial_points, 
-        initial_point_generator=initial_point_generator,
-        acq_func=acq_func, 
-        acq_optimizer=acq_optimizer,
-        random_state=random_state, 
-        acq_func_kwargs=acq_func_kwargs, 
-        acq_optimizer_kwargs=acq_optimizer_kwargs,
-        base_estimator_kwargs=base_estimator_kwargs,
-        model_queue_size=model_queue_size,
-        verbose=verbose
-    )
-
-    # Record provided points
-
-    # create return object
-    result = None
-    # evaluate y0 if only x0 is provided
-    if x0 and y0 is None:
-        y0 = list(map(func, x0))
-        n_calls -= len(y0)
-    # record through tell function
-    if x0:
-        if not (isinstance(y0, Iterable) or isinstance(y0, numbers.Number)):
-            raise ValueError(
-                "`y0` should be an iterable or a scalar, got %s" % type(y0))
-        if len(x0) != len(y0):
-            raise ValueError("`x0` and `y0` should have the same length")
-        result = optimizer.tell(x0, y0)
-        result.specs = specs
-
-    # Handle solver output
-    if not isinstance(verbose, (int, type(None))):
-        raise TypeError("verbose should be an int of [0,1,2] or bool, "
-                        "got {}".format(type(verbose)))
-
-    if isinstance(verbose, bool):
-        if verbose:
-            verbose = 1
-        else:
-            verbose = 0
-    elif isinstance(verbose, int):
-        if verbose not in [0,1,2]:
-            raise TypeError("if verbose is int, it should in [0,1,2], "
-                        "got {}".format(verbose))
-
-    # Optimize
-    _n_calls = n_calls
-
-    itr = 1
-
-    if verbose > 0:
-        print("")
-        print("SOLVER: start solution process...")
-        print("")
-        print(f"SOLVER: generate \033[1m {n_initial_points}\033[0m initial points...")
-
-
-    while _n_calls>0:
-
-        # check if optimization is performed in batches
-        if batch_size is not None:
-            _batch_size = min([_n_calls,batch_size])
-            _n_calls -= _batch_size
-            next_x = optimizer.ask(_batch_size,strategy=batch_strategy)
-        else:
-            _n_calls -= 1
-            next_x = optimizer.ask(strategy=batch_strategy)
-            
-        next_y = func(next_x)
-
-        # first iteration uses next_y as best point instead of min of next_y
-        if itr == 1:
-            if batch_size is None:
-                best_fun = next_y
-            else:
-                best_fun = min(next_y)
-
-        # handle output print at every iteration
-        if verbose > 0:
-            print("")
-            print(f"\033[1m itr_{itr}\033[0m")
-
-            if isinstance(next_y,Iterable):
-                # in case of batch optimization, print all new proposals and 
-                # mark improvements of objectives with (*)
-                print_str = []
-                for y in next_y:
-                    if y <= best_fun:
-                        print_str.append(f"\033[1m{round(y,5)}\033[0m (*)")
-                    else:
-                        print_str.append(str(round(y,5)))
-                print(f"   new points obj.: {print_str[0]}")
-                for y_str in print_str[1:]:
-                    print(f"                    {y_str}")    
-            else:
-                # in case of single point sequential optimization, print new 
-                # point proposal
-                if next_y <= best_fun:
-                    print_str = f"\033[1m itr_{round(next_y,5)}\033[0m"
-                else:
-                    print_str = str(round(next_y,5))
-                print(f"   new point obj.: {round(next_y,5)}")
-
-            # print best obj until (not including) current iteration
-            print(f"   best obj.:       {round(best_fun,5)}")
-
-        itr += 1
-
-        optimizer.tell(
-            next_x, next_y
-        )
-
-        result = optimizer.get_result()
-
-        best_fun = result.fun  
-
-        result.specs = specs
-    
-    # print end of solve once convergence criteria is met
-    if verbose>0:
-        print("")
-        print("SOLVER: finished run!")
-        print(f"SOLVER: best obj.: {round(result.fun,5)}")
-        print("")
-
-    return result
\ No newline at end of file
diff --git a/entmoot/optimizer/gurobi_utils.py b/entmoot/optimizer/gurobi_utils.py
deleted file mode 100644
index 8b6456e..0000000
--- a/entmoot/optimizer/gurobi_utils.py
+++ /dev/null
@@ -1,571 +0,0 @@
-from gurobipy import GRB, quicksum
-import gurobipy as gp
-import numpy as np
-
-def get_opt_core_copy(opt_core):
-    new_opt_core = opt_core.copy()
-    new_opt_core._n_feat = opt_core._n_feat
-
-    # transfer var dicts
-    new_opt_core._cont_var_dict = {}
-    new_opt_core._cat_var_dict = {}
-
-    ## transfer cont_var_dict
-    for var in opt_core._cont_var_dict.keys():
-        var_name = opt_core._cont_var_dict[var].VarName
-
-        new_opt_core._cont_var_dict[var] = \
-            new_opt_core.getVarByName(var_name)
-
-    ## transfer cat_var_dict
-    for var in opt_core._cat_var_dict.keys():
-        for cat in opt_core._cat_var_dict[var].keys():
-            var_name = opt_core._cat_var_dict[var][cat].VarName
-
-            if var not in new_opt_core._cat_var_dict.keys():
-                new_opt_core._cat_var_dict[var] = {}
-
-            new_opt_core._cat_var_dict[var][cat] = \
-                new_opt_core.getVarByName(var_name)
-
-    return new_opt_core
-
-def get_core_gurobi_model(space, add_model_core=None, env=None):
-    """Add core to gurobi model, i.e. bounds, variables and parameters.
-
-    Parameters
-    ----------
-    space : scikit-optimize object
-        Captures the feature space
-    model : gurobipy.Model
-        Model to which the core is added.
-    env : gurobipy.Env
-        Env to use. If not given Gurobi will use the default env.
-
-    Returns
-    -------
-    -
-    """
-    from entmoot.utils import get_cat_idx
-    from entmoot.space.space import Integer
-
-    if add_model_core is None:
-        if env is None:
-            model = gp.Model()
-        else:
-            model = gp.Model(env=env)
-            
-        cat_idx = get_cat_idx(space)
-
-        x_lb = {}
-        x_ub = {}
-
-        model._cont_var_dict = {}
-        model._cat_var_dict = {}
-
-        all_names = [dim.name for dim in space.dimensions]
-
-        for idx,dim in enumerate(space.dimensions):
-            if dim.name:
-                temp_var_name = str(dim.name)
-            else:
-                temp_var_name = '_'.join(['x', str(idx)])
-
-            assert all_names.count(temp_var_name) <= 1,\
-                f"please pick unique names for dimensions, i.e. '{temp_var_name}' appears multiple times"
-
-            if idx not in cat_idx:
-                # store bounds of numerical vars
-                x_lb[idx] = dim.transformed_bounds[0]
-                x_ub[idx] = dim.transformed_bounds[1]
-
-                if isinstance(dim, Integer) and dim.prior in ["uniform"]:
-                    model._cont_var_dict[idx] = \
-                        model.addVar(lb=x_lb[idx], ub=x_ub[idx], 
-                            name=temp_var_name, vtype='I')
-                else:
-                    model._cont_var_dict[idx] = \
-                        model.addVar(lb=x_lb[idx], ub=x_ub[idx], 
-                            name=temp_var_name, vtype='C')
-            else:
-                # store categories of categorical vars
-                model._cat_var_dict[idx] = {}
-                for cat in dim.transform(dim.bounds):
-                    model._cat_var_dict[idx][cat] = \
-                        model.addVar(name=f"{temp_var_name}_{cat}", vtype=GRB.BINARY)
-            
-        model._n_feat = len(model._cont_var_dict) + len(model._cat_var_dict)
-        model.update()
-
-        return model
-    else:
-        # validate model core input
-        assert type(add_model_core) is gp.Model, \
-            "wrong model core type given, please create model core using: " + \
-                "entmoot.optimizer.gurobi_utils.add_core_to_gurobi_model(..."
-
-        check_attributes = \
-            hasattr(add_model_core,"_cont_var_dict") and \
-            hasattr(add_model_core,"_cat_var_dict") and \
-            hasattr(add_model_core,"_n_feat")
-
-        assert check_attributes, \
-            "model core was not configured correctly, please create model core using: " + \
-                "entmoot.optimizer.gurobi_utils.add_core_to_gurobi_model(..."
-
-        # create copy of the original model core, so we can keep reusing it
-        model_core_copy = get_opt_core_copy(add_model_core)
-
-        return model_core_copy
-
-def add_std_to_gurobi_model(est, model):
-    """Adds standard estimator formulation to gurobi model.
-
-    Parameters
-    ----------
-    est : EntingRegressor object
-        Has both tree model and uncertainty estimate embedded
-    model : gurobipy.Model,
-        Model to which the core is added.
-
-    Returns
-    -------
-    -
-    """
-    est.std_estimator.add_to_gurobi_model(model)
-    model.update()
-
-def set_gurobi_init_to_ref(est, model):
-    """Sets intial values of gurobi model variables to the reference points.
-
-    Parameters
-    ----------
-    est : EntingRegressor object
-        Has both tree model and uncertainty estimate embedded
-    model : gurobipy.Model,
-        Model to which the core is added.
-
-    Returns
-    -------
-    -
-    """
-    ref_points_unscaled = \
-        est.std_estimator.Xi
-
-    best_val = est.predict(ref_points_unscaled[0].reshape(1, -1))
-    best_ref = ref_points_unscaled[0]
-
-    for ref_point in ref_points_unscaled:
-        temp_val = est.predict(ref_point.reshape(1, -1))
-
-        if best_val > temp_val:
-            best_val = temp_val
-            best_ref = ref_point
-
-    for i in model._cont_var_dict.keys():
-        model._cont_var_dict[i].start = best_ref[i]
-
-    for i in model._cat_var_dict.keys():
-        model._cat_var_dict[i][int(best_ref[i])].start = 1
-
-    model._alpha.start = 0.0
-    model.update()
-
-def get_gbm_model_multi_obj_mu(opt_model, yi):
-    # normalize mu based on targets
-    y_min = np.min(yi, axis=0)
-    y_max = np.max(yi, axis=0)
-    y_range = y_max - y_min
-
-    # collect weighted sum for every label
-    ob_expr = {}
-    for label in opt_model._gbm_set:
-        weighted_sum = quicksum(
-            opt_model._leaf_weight(label, tree, leaf) *
-            opt_model._z_l[label, tree, leaf]
-            for tree, leaf in label_leaf_index(opt_model, label)
-        )
-        # normalized mu contribution
-        ob_expr[label] = (weighted_sum - y_min[label]) / y_range[label]
-    return ob_expr
-
-def get_gbm_model_mu(opt_model, yi, norm=False):
-    # normalize mu based on targets
-    y_min = np.min(yi)
-    y_max = np.max(yi)
-    y_range = y_max - y_min
-
-    weighted_sum = quicksum(
-        opt_model._leaf_weight(label, tree, leaf) * \
-        opt_model._z_l[label, tree, leaf]
-        for label, tree, leaf in leaf_index(opt_model)
-    )
-    if norm:
-        return (weighted_sum - y_min) / y_range
-    else:
-        return weighted_sum
-
-def add_acq_to_gurobi_model(model, model_mu, model_unc,
-                            est,
-                            acq_func="LCB",
-                            acq_func_kwargs=None,
-                            num_obj=1,
-                            weights=None):
-    """Sets gurobi model objective function to acquisition function.
-
-    Parameters
-    ----------
-    model : gurobipy.Model
-        Model to which the core is added.
-    est : EntingRegressor object
-        Has both tree model and uncertainty estimate embedded
-    acq_func : str
-        Type of acquisition function used for gurobi model objective
-    acq_func_kwargs : dict
-        Allows additional parameter settings for acquisition function
-
-    Returns
-    -------
-    -
-    """
-    # check inputs
-    if acq_func_kwargs is None:
-        acq_func_kwargs = {}
-
-    kappa = acq_func_kwargs.get("kappa", 1.96)
-
-
-    if num_obj > 1:
-        model._mu = model.addVar(
-            name='mu', vtype='C', lb=-GRB.INFINITY, ub=GRB.INFINITY)
-
-        for obj_idx in model_mu.keys():
-            temp_constr = \
-                model.addConstr(
-                    model._mu >= weights[obj_idx] * model_mu[obj_idx],
-                    name=f"multi-obj_{obj_idx}"
-                )
-            model.update()
-        proc_mu = model._mu
-    else:
-        proc_mu = model_mu
-
-    if acq_func == "LCB":
-        if model_unc is not None:
-            ob_expr = quicksum((proc_mu, kappa * model_unc))
-        else:
-            ob_expr = proc_mu
-        model.setObjective(ob_expr, GRB.MINIMIZE)
-
-    elif acq_func == "HLCB":
-        obj_cost = acq_func_kwargs.get("obj_cost", 0.5)
-        model.setObjectiveN(proc_mu, 0, 1, reltol=obj_cost)
-        model.setObjectiveN(model_unc, 1, 0)
-
-    model.update()
-
-def get_gurobi_obj(model, est, return_std=False, acq_func_kwargs=None):
-    """Returns gurobi model objective contribution of tree model and std
-    estimator.
-
-    Returns tree model objective contribution if `return_std` is set to False.
-    If `return_std` is set to True, then return tree model and std estimator 
-    contribution to gurobi model objective.
-
-    Parameters
-    ----------
-    model : gurobipy.Model
-        Model to which the core is added.
-    est : EntingRegressor object
-        Has both tree model and uncertainty estimate embedded
-    return_std : bool
-        Set `True` to return both tree model and std estimator contribution
-
-    Returns
-    -------
-    mean or (mean, std): gurobipy.expr, or tuple(gurobipy.expr, gurobipy.expr), 
-    depending on value of `return_std`.
-    
-    """
-
-    mean = get_gbm_obj(model)
-
-    if return_std:
-        std = est.std_estimator.get_gurobi_obj(model)
-        return mean, std
-
-    return mean
-
-def get_gbm_obj(model):
-    """Returns objective of `gbm_model` specified in gurobi model.
-
-    Parameters
-    ----------
-    model : gurobipy.Model
-        Model to which the core is added.
-
-    Returns
-    -------
-    ob_expr: gurobipy.expr
-        Defines the gurobipy expression corresponding to the tree model objective
-        contribution.
-    
-    """
-    weighted_sum = quicksum(
-        model._leaf_weight(label, tree, leaf) * \
-            model._z_l[label, tree, leaf]
-        for label, tree, leaf in leaf_index(model)
-    )
-    ob_expr = weighted_sum
-    return ob_expr
-
-def get_gbm_obj_from_model(model, label):
-    temp_sum = sum(model._leaf_weight(label, tree, leaf) * \
-        round(model._z_l[label, tree, leaf].x,1)
-        for temp_label, tree, leaf in leaf_index(model) if temp_label == label)
-    return temp_sum
-
-def get_gbm_multi_obj_from_model(model):
-    temp_sum = []
-    for label in model._gbm_set:
-        temp_sum.append(sum(model._leaf_weight(label, tree, leaf) * \
-            round(model._z_l[label, tree, leaf].x,1)
-            for temp_label, tree, leaf in leaf_index(model) if temp_label == label))
-    return temp_sum
-
-
-### GBT HANDLER
-## gbt model helper functions
-def label_leaf_index(model, label):
-    for tree in range(model._num_trees(label)):
-        for leaf in model._leaves(label, tree):
-            yield (tree, leaf)
-
-def tree_index(model):
-    for label in model._gbm_set:
-        for tree in range(model._num_trees(label)):
-            yield (label, tree)
-
-tree_index.dimen = 2
-
-def leaf_index(model):
-    for label, tree in tree_index(model):
-        for leaf in model._leaves(label, tree):
-            yield (label, tree, leaf)
-
-leaf_index.dimen = 3
-
-def misic_interval_index(model):
-    for var in model._breakpoint_index:
-        for j in range(len(model._breakpoints(var))):
-            yield (var, j)
-
-misic_interval_index.dimen = 2
-
-def misic_split_index(model):
-    gbm_models = model._gbm_models
-    for label, tree in tree_index(model):
-        for encoding in gbm_models[label].get_branch_encodings(tree):
-            yield (label, tree, encoding)
-
-misic_split_index.dimen = 3
-
-def alt_interval_index(model):
-    for var in model.breakpoint_index:
-        for j in range(1, len(model.breakpoints[var])+1):
-            yield (var, j)
-
-alt_interval_index.dimen = 2
-
-def add_gbm_to_gurobi_model(space, gbm_model_dict, model):
-    from entmoot.utils import get_cat_idx
-    cat_idx = get_cat_idx(space)
-
-    add_gbm_parameters(cat_idx, gbm_model_dict, model)
-    add_gbm_variables(model)
-    add_gbm_constraints(cat_idx, model)
-
-def add_gbm_parameters(cat_idx, gbm_model_dict, model):
-    model._gbm_models = gbm_model_dict
-
-    model._gbm_set = set(gbm_model_dict.keys())
-    model._num_trees = lambda label: \
-        gbm_model_dict[label].n_trees
-    
-    model._leaves = lambda label, tree: \
-            tuple(gbm_model_dict[label].get_leaf_encodings(tree))
-
-    model._leaf_weight = lambda label, tree, leaf: \
-            gbm_model_dict[label].get_leaf_weight(tree, leaf)
-
-    vbs = [v.get_var_break_points() for v in gbm_model_dict.values()]
-
-    all_breakpoints = {}
-    for i in range(model._n_feat):
-        if i in cat_idx:
-            continue
-        else:
-            s = set()
-            for vb in vbs:
-                try:
-                    s = s.union(set(vb[i]))
-                except KeyError:
-                    pass
-            if s:
-                all_breakpoints[i] = sorted(s)
-
-    model._breakpoint_index = list(all_breakpoints.keys())
-
-    model._breakpoints = lambda i: all_breakpoints[i]
-
-    model._leaf_vars = lambda label, tree, leaf: \
-            tuple(i
-                for i in gbm_model_dict[label].get_participating_variables(
-                    tree, leaf))
-
-def add_gbm_variables(model):
-    model._z_l = model.addVars(
-                    leaf_index(model), 
-                    lb=0,
-                    ub=GRB.INFINITY, 
-                    name="z_l", vtype='C'
-                )
-                            
-    model._y = model.addVars(
-                    misic_interval_index(model), 
-                    name="y", 
-                    vtype=GRB.BINARY
-                )
-    model.update()
-
-def add_gbm_constraints(cat_idx, model):
-
-    def single_leaf_rule(model_, label, tree):
-        z_l, leaves = model_._z_l, model_._leaves
-        return (quicksum(z_l[label, tree, leaf] 
-                for leaf in leaves(label, tree))
-                == 1)
-
-    model.addConstrs(
-        (single_leaf_rule(model, label, tree)
-        for (label,tree) in tree_index(model)),
-        name="single_leaf"
-    )
-
-    def left_split_r(model_, label, tree, split_enc):
-        gbt = model_._gbm_models[label]
-        split_var, split_val = gbt.get_branch_partition_pair(
-            tree,
-            split_enc
-        )
-        y_var = split_var
-
-        if not isinstance(split_val, list):
-            # for conti vars
-            y_val = model_._breakpoints(y_var).index(split_val)
-            return \
-                quicksum(
-                    model_._z_l[label, tree, leaf]
-                    for leaf in gbt.get_left_leaves(tree, split_enc)
-                ) <= \
-                    model_._y[y_var, y_val]
-        else:
-            # for cat vars
-            return \
-                quicksum(
-                    model_._z_l[label, tree, leaf]
-                    for leaf in gbt.get_left_leaves(tree, split_enc)
-                ) <= \
-                    quicksum(
-                        model_._cat_var_dict[split_var][cat]
-                        for cat in split_val
-                    )
-
-    def right_split_r(model_, label, tree, split_enc):
-        gbt = model_._gbm_models[label]
-        split_var, split_val = gbt.get_branch_partition_pair(
-            tree,
-            split_enc
-        )
-        y_var = split_var
-        if not isinstance(split_val, list):
-            # for conti vars
-            y_val = model_._breakpoints(y_var).index(split_val)
-            return \
-                quicksum(
-                    model_._z_l[label, tree, leaf]
-                    for leaf in gbt.get_right_leaves(tree, split_enc)
-            ) <= \
-                1 - model_._y[y_var, y_val]
-        else:
-            # for cat vars
-            return \
-                quicksum(
-                    model_._z_l[label, tree, leaf]
-                    for leaf in gbt.get_right_leaves(tree, split_enc)
-            ) <= 1 - \
-                quicksum(
-                    model_._cat_var_dict[split_var][cat]
-                    for cat in split_val
-                )
-
-    def y_order_r(model_, i, j):
-        if j == len(model_._breakpoints(i)):
-            return Constraint.Skip
-        return model_._y[i, j] <= model_._y[i, j+1]
-
-    def cat_sums(model_, i):
-        return quicksum(
-            model_._cat_var_dict[i][cat]
-            for cat in model_._cat_var_dict[i].keys()
-        ) == 1
-
-    def var_lower_r(model_, i, j):
-        lb = model_._cont_var_dict[i].lb
-        j_bound = model_._breakpoints(i)[j]
-        return model_._cont_var_dict[i] >= lb + (j_bound - lb)*(1-model_._y[i, j])
-
-    def var_upper_r(model_, i, j):
-        ub = model_._cont_var_dict[i].ub
-        j_bound = model_._breakpoints(i)[j]
-        return model_._cont_var_dict[i] <= ub + (j_bound - ub)*(model_._y[i, j])
-    
-    model.addConstrs(
-        (left_split_r(model, label, tree, encoding)
-        for (label, tree, encoding) in misic_split_index(model)),
-        name="left_split"
-    )
-
-    model.addConstrs(
-        (right_split_r(model, label, tree, encoding)
-        for (label, tree, encoding) in misic_split_index(model)),
-        name="right_split"
-    )
-
-    # for conti vars
-    model.addConstrs(
-        (y_order_r(model, var, j)
-        for (var, j) in misic_interval_index(model)
-        if j != len(model._breakpoints(var))-1),
-        name="y_order"
-    )
-
-    # for cat vars
-    model.addConstrs(
-        (cat_sums(model, var)
-        for var in cat_idx),
-        name="cat_sums"
-    )    
-
-
-    model.addConstrs(
-        (var_lower_r(model, var, j)
-        for (var, j) in misic_interval_index(model)),
-        name="var_lower"
-    )
-
-    model.addConstrs(
-        (var_upper_r(model, var, j)
-        for (var, j) in misic_interval_index(model)),
-        name="var_upper"
-    )
\ No newline at end of file
diff --git a/entmoot/optimizer/optimizer.py b/entmoot/optimizer/optimizer.py
deleted file mode 100644
index 984d91e..0000000
--- a/entmoot/optimizer/optimizer.py
+++ /dev/null
@@ -1,813 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-NOTE: Changes were made to the scikit-optimize source code included here. 
-For the most recent version of scikit-optimize we refer to:
-https://github.com/scikit-optimize/scikit-optimize/
-
-Copyright (c) 2019-2020 Alexander Thebelt.
-"""
-from typing import Optional
-import warnings
-from numbers import Number
-import numpy as np
-
-from entmoot.acquisition import _gaussian_acquisition
-
-
-class Optimizer(object):
-    """
-    This class is used to run the main BO loop. 
-    Optimizer objects define all BO settings and store the current data set.
-    
-    The self.ask function provides new input proposals, and resulting black-box
-    values can be added through self.tell. This procedure is automated via
-    self.run, where a callable function is provided as an input.
-    
-    Functions self.predict_with_est and self.predict_with_acq use the current
-    surrogate model to predict inputs.
-    
-    :param dimensions: list
-        list of search-space variables, i.e.
-            If (lower: float, upper: float), gives continuous variable
-            If (lower: int, upper: int), gives discrete variable
-            If list of categories, gives categorical variable
-    :param base_estimator: str
-        currently available estimator types are in ["ENTING"]
-    :param n_initial_points: int
-        number of initial points sampled before surrogate model is trained
-    :param initial_point_generator: str
-        currently supported sampling generators are in
-        ["sobol", "halton", "hammersly", "lhs", "random", "grid"]
-    :param num_obj: int
-        gives the number of objectives that the black-box is being optimized for
-    :param acq_func: str
-        acquisition function type that is used for exploitation vs. exploration
-        trade-off, i.e. currently supported ["LCB"] and ["LCB", "HLCB"] for
-        num_obj == 1
-    :param acq_optimizer: str
-        optimization method used to minimize the acquisition function, i.e.
-        currently supports ["sampling", "global"]
-    :param random_state: Optional[int]
-        fixed random seed to generate reproducible results
-    :param acq_func_kwargs: Optional[dict]
-        define additional params for acquisition function, i.e.
-            "kappa": int, influences acquisition function "min mu - kappa * alpha"
-    :param acq_optimizer_kwargs: Optional[dict]
-        define additional params for acqu. optimizer "sampling":
-            "n_points": int, number of points to minimize the acquisition function
-        define additional params for acqu. optimizer "global":
-            "env": GRBEnv, defines Gurobi environment for cluster computations
-            "gurobi_timelimit": int, optimization time limit in sec
-            "add_model_core": GRBModel, Gurobi optimization model that includes 
-                additional constraints
-    :param base_estimator_kwargs: Optional[dict]
-        defines additional params that influence the base_estimator behavior
-            "lgbm_params": dict, additional parameters that are passed to lightgbm
-            "ensemble_type": str, options in ['GBRT', 'RF'],
-                "GBRT": uses gradient-boosted tree regressor
-                "RF": uses random forest
-            "unc_metric": str, options in ['exploration', 'penalty'], i.e.
-                negative or positive alpha contribution in "min mu \pm kappa * alpha"
-            "unc_scaling": str, options in ["standard", "normalize"], i.e.
-                scaling used for the distance-based uncertainty metric
-            "dist_metric": str, options in ["manhattan",'squared_euclidean'], i.e.
-                metric used to define non-categorical distance for uncertainty
-            "cat_dist_metric": str, options in ["overlap", "goodall4", "of"], i.e.
-                metric used to define categorical distance for uncertainty
-    :param model_queue_size: Optional[int]
-        defines number of previous models that are stored in self.models
-    :param verbose: bool
-        defines the verbosity level of the output
-    """
-
-    def __init__(
-        self,
-        dimensions: list,
-        base_estimator: str = "ENTING",
-        n_initial_points: int = 50,
-        initial_point_generator: str = "random",
-        num_obj: int = 1,
-        acq_func: str = "LCB",
-        acq_optimizer: str = "global",
-        random_state: Optional[int] = None,
-        acq_func_kwargs: Optional[dict] = None,
-        acq_optimizer_kwargs: Optional[dict] = None,
-        base_estimator_kwargs: Optional[dict] = None,
-        model_queue_size: Optional[int] = None,
-        verbose: bool = False
-    ):
-
-        from entmoot.utils import is_supported
-        from entmoot.utils import cook_estimator
-        from entmoot.utils import cook_initial_point_generator
-        from sklearn.utils import check_random_state
-
-        # define random state
-        self.rng = check_random_state(random_state)
-
-        # store and create acquisition function set
-        self.acq_func = acq_func
-        if acq_func_kwargs is None:
-            self.acq_func_kwargs = {}
-        else:
-            self.acq_func_kwargs = acq_func_kwargs
-
-        allowed_acq_funcs = ["LCB","HLCB"]
-        if self.acq_func not in allowed_acq_funcs:
-            raise ValueError("expected acq_func to be in %s, got %s" %
-                             (",".join(allowed_acq_funcs), self.acq_func))
-
-        # configure counter of points
-        if n_initial_points < 0:
-            raise ValueError(
-                "Expected `n_initial_points` >= 0, got %d" % n_initial_points)
-        self._n_initial_points = n_initial_points
-        self.n_initial_points_ = n_initial_points
-        
-        # initialize search space and output
-        from entmoot.space.space import Space
-
-        self.space = Space(dimensions)
-
-        self._initial_samples = None
-        self._initial_point_generator = \
-            cook_initial_point_generator(initial_point_generator)
-
-        if self._initial_point_generator is not None:
-            transformer = self.space.get_transformer()
-            self._initial_samples = self._initial_point_generator.generate(
-                self.space.dimensions, n_initial_points,
-                random_state=self.rng.randint(0, np.iinfo(np.int32).max))
-            self.space.set_transformer(transformer)
-
-        self.num_obj = num_obj
-
-        # create base_estimator
-        self.base_estimator_kwargs = {} if base_estimator_kwargs is None else base_estimator_kwargs
-
-        from entmoot.learning.tree_model import EntingRegressor,MisicRegressor
-
-        if type(base_estimator) not in [EntingRegressor,MisicRegressor]:
-            if type(base_estimator) in [str]:
-                # define random_state of estimator
-                est_random_state = self.rng.randint(0, np.iinfo(np.int32).max)
-
-                # check support of base_estimator if exists
-                if not is_supported(base_estimator):
-                    raise ValueError("Estimator type: %s is not supported." % base_estimator)
-
-                # build base_estimator
-                base_estimator = cook_estimator(
-                    self.space,
-                    base_estimator,
-                    self.base_estimator_kwargs,
-                    num_obj=self.num_obj,
-                    random_state=est_random_state)
-            else:
-                raise ValueError("Estimator type: %s is not supported." % base_estimator)
-
-        self.base_estimator_ = base_estimator
-
-        # Configure Optimizer
-        self.acq_optimizer = acq_optimizer
-
-        # record other arguments
-        if acq_optimizer_kwargs is None:
-            acq_optimizer_kwargs = {}
-
-        self.acq_optimizer_kwargs = acq_optimizer_kwargs
-        self.n_points = acq_optimizer_kwargs.get("n_points", 10000)
-        self.gurobi_env = acq_optimizer_kwargs.get("env", None)
-        self.gurobi_timelimit = acq_optimizer_kwargs.get("gurobi_timelimit", None)
-
-        ## Initialize storage for optimization
-        if not isinstance(model_queue_size, (int, type(None))):
-            raise TypeError("model_queue_size should be an int or None, "
-                            "got {}".format(type(model_queue_size)))
-
-        # model cache
-        self.max_model_queue_size = model_queue_size
-        self.models = []
-
-        # data set cache
-        self.Xi = []
-        self.yi = []
-
-        # model_mu and model_std cache
-        self.model_mu = []
-        self.model_std = []
-
-        # global opti metrics
-        self.gurobi_mipgap = []
-
-        # Initialize cache for `ask` method responses
-        # This ensures that multiple calls to `ask` with n_points set
-        # return same sets of points. Reset to {} at every call to `tell`.
-        self.cache_ = {}
-
-        # Handle solver print output
-        if not isinstance(verbose, (int, type(None))):
-            raise TypeError("verbose should be an int of [0,1,2] or bool, "
-                            "got {}".format(type(verbose)))
-
-        if isinstance(verbose, bool):
-            if verbose:
-                self.verbose = 1
-            else:
-                self.verbose = 0
-        elif isinstance(verbose, int):
-            if verbose not in [0,1,2]:
-                raise TypeError("if verbose is int, it should in [0,1,2], "
-                            "got {}".format(verbose))
-            else:
-                self.verbose = verbose
-
-        # printed_switch_to_model defines if notification of switch from 
-        # intitial point generation to model-based point generation has been
-        # printed yet
-        if self.verbose > 0:
-            self.printed_switch_to_model = False
-        else:
-            self.printed_switch_to_model = True
-
-    def copy(self, random_state: Optional[int] = None):
-        """Create a shallow copy of an instance of the optimizer.
-
-        :param random_state: Optional[int]
-        :return optimizer: Optimizer
-        """
-
-        optimizer = Optimizer(
-            dimensions=self.space.dimensions,
-            base_estimator=self.base_estimator_,
-            n_initial_points=self.n_initial_points_,
-            initial_point_generator=self._initial_point_generator,
-            acq_func=self.acq_func,
-            acq_optimizer=self.acq_optimizer,
-            acq_func_kwargs=self.acq_func_kwargs,
-            acq_optimizer_kwargs=self.acq_optimizer_kwargs,
-            random_state=random_state,
-            verbose=self.verbose
-        )
-
-        optimizer._initial_samples = self._initial_samples
-        optimizer.printed_switch_to_model = self.printed_switch_to_model
-        optimizer.gurobi_env = self.gurobi_env
-
-        if self.Xi:
-            optimizer._tell(self.Xi, self.yi)
-
-        return optimizer
-
-    def ask(
-        self,
-        n_points: int = None,
-        strategy: str = "cl_min",
-        weights: Optional[list] = None,
-        add_model_core=None,
-        gurobi_env=None,
-    ) -> list:
-        """
-        Computes the next point (or multiple points) at which the objective
-        should be evaluated.
-
-        :param n_points: int = None,
-            gives the number of points that should be returned
-        :param strategy: str = "cl_min",
-            determines the liar strategy, i.e.
-            (https://hal.archives-ouvertes.fr/hal-00732512/document)
-            used for single-objective batch proposals, i.e.
-                "cl_min": str, uses minimum of observations as lie
-                "cl_mean": str, uses mean of observations as lie
-                "cl_max": str, uses maximum of observations as lie
-        :param weights: Optional[list] = None,
-            1D list of weights of size num_obj for single multi-objective proposal;
-            2D list of weights of size shape(n_points,num_obj) for
-                batch multi-objectiveproposals
-        :param add_model_core: GRBModel = None,
-            Gurobi optimization model that includes additional constraints
-        :param gurobi_env: Gurobi Env = None,
-            Gurobi environment used for computation of the next proposal
-
-        :return next_x: list, next proposal of shape(n_points, n_dims)
-        """
-
-        # update gurobi_env attribute
-        if gurobi_env:
-            self.gurobi_env = gurobi_env
-
-        # check if single point or batch of point is returned
-        if n_points is None and weights is None:
-            return self._ask(add_model_core=add_model_core)
-        elif self.num_obj > 1:
-            if weights is not None:
-                n_points = len(weights)
-
-            X = []
-            for i in range(n_points):
-                self._next_x = None
-
-                # use the weights if provided
-                if weights:
-                    w = weights[i]
-                    assert len(w) == self.num_obj, \
-                        f"The {i}'th provided weight has dim '{len(w)}' but " \
-                        f"number of objectives is '{self.num_obj}'."
-                    next_x = self._ask(weight=w, add_model_core=add_model_core)
-                else:
-                    next_x = self._ask(add_model_core=add_model_core)
-                X.append(next_x)
-            return X if len(X) > 1 else X[0]
-
-        supported_strategies = ["cl_min", "cl_mean", "cl_max"]
-
-        if not (isinstance(n_points, int) and n_points > 0):
-            raise ValueError(
-                "n_points should be int > 0, got " + str(n_points)
-            )
-
-        if strategy not in supported_strategies:
-            raise ValueError(
-                "Expected parallel_strategy to be one of " +
-                str(supported_strategies) + ", " + "got %s" % strategy
-            )
-
-        # caching the result with n_points not None. If some new parameters
-        # are provided to the ask, the cache_ is not used.
-        if (n_points, strategy) in self.cache_:
-            return self.cache_[(n_points, strategy)]
-
-        # Copy of the optimizer is made in order to manage the
-        # deletion of points with "lie" objective (the copy of
-        # optimizer is simply discarded)
-        opt = self.copy(
-            random_state=self.rng.randint(0,np.iinfo(np.int32).max))
-
-        X = []
-        for i in range(n_points):
-            
-            x = opt.ask()
-            X.append(x)
-
-            if strategy == "cl_min":
-                y_lie = np.min(opt.yi) if opt.yi else 0.0  # CL-min lie
-            elif strategy == "cl_mean":
-                y_lie = np.mean(opt.yi) if opt.yi else 0.0  # CL-mean lie
-            else:
-                y_lie = np.max(opt.yi) if opt.yi else 0.0  # CL-max lie
-
-            opt._tell(x, y_lie)
-
-
-        self.printed_switch_to_model = opt.printed_switch_to_model
-
-        self.cache_ = {(n_points, strategy): X}  # cache_ the result
-
-        return X
-
-    def _ask(
-        self,
-        weight: Optional[list] = None,
-        add_model_core = None
-    ):
-        """
-        Computes the next point at which the objective should be evaluated.
-
-        :param weight: Optional[list] = None,
-            list of weights of size num_obj to trade-off between different
-            objective contributions
-        :param add_model_core: GRBModel = None,
-            Gurobi optimization model that includes additional constraints
-
-        :return next_x: list, next proposal of shape(n_points, n_dims)
-        """
-
-        if self._n_initial_points > 0 or self.base_estimator_ is None:
-            # this will not make a copy of `self.rng` and hence keep advancing
-            # our random state.
-            if self._initial_samples is None:
-                return self.space.rvs(random_state=self.rng)[0]
-            else:
-                # The samples are evaluated starting form initial_samples[0]
-                return self._initial_samples[len(self._initial_samples) - self._n_initial_points]
-
-        elif self._next_x is not None:
-            # return self._next_x if optimizer hasn't learned anything new
-            return self._next_x
-
-        else:
-            # after being "told" n_initial_points we switch from sampling
-            # random points to using a surrogate model
-
-            # create fresh copy of base_estimator
-            est = self.base_estimator_.copy()
-            self.base_estimator_ = est
-
-            # esimator is fitted using a generic fit function
-            est.fit(self.space.transform(self.Xi), self.yi)
-
-            # we cache the estimator in model_queue
-            if self.max_model_queue_size is None:
-                self.models.append(est)
-            elif len(self.models) < self.max_model_queue_size:
-                self.models.append(est)
-            else:
-                # Maximum list size obtained, remove oldest model.
-                self.models.pop(0)
-                self.models.append(est)
-
-            if not self.printed_switch_to_model:
-                print("")
-                print("SOLVER: initial points exhausted")
-                print("   -> switch to model-based optimization")
-                self.printed_switch_to_model = True
-
-            # this code provides a heuristic solution that uses sampling as the optimization strategy
-            if self.acq_optimizer == "sampling":
-                # sample a large number of points and then pick the best ones as
-                # starting points
-                X = self.space.transform(self.space.rvs(
-                    n_samples=self.n_points, random_state=self.rng))
-
-                values = _gaussian_acquisition(
-                    X=X, model=est,
-                    acq_func=self.acq_func,
-                    acq_func_kwargs=self.acq_func_kwargs)
-                # Find the minimum of the acquisition function by randomly
-                # sampling points from the space
-                next_x = X[np.argmin(values)]
-
-                # derive model mu and std
-                # next_xt = self.space.transform([next_x])[0]
-                next_model_mu, next_model_std = \
-                    self.models[-1].predict(
-                        X=np.asarray(next_x).reshape(1, -1),
-                        return_std=True)
-
-                model_mu = next_model_mu[0]
-                model_std = next_model_std[0]
-
-            # acquisition function is optimized globally
-            elif self.acq_optimizer == "global":
-                try:
-                    import gurobipy as gp
-                except ModuleNotFoundError:
-                    ImportError("GurobiNotFoundError: "
-                          "To run `aqu_optimizer='global'` "
-                          "please install the Gurobi solver "
-                          "(https://www.gurobi.com/) and its interface "
-                          "gurobipy. "
-                          "Alternatively, change `aqu_optimizer='sampling'`.")
-
-                if add_model_core is None:
-                    add_model_core = \
-                        self.acq_optimizer_kwargs.get("add_model_core", None)
-
-                next_x, model_mu, model_std, gurobi_mipgap = \
-                    self.models[-1].get_global_next_x(acq_func=self.acq_func,
-                                                      acq_func_kwargs=self.acq_func_kwargs,
-                                                      add_model_core=add_model_core,
-                                                      weight=weight,
-                                                      verbose=self.verbose,
-                                                      gurobi_env=self.gurobi_env,
-                                                      gurobi_timelimit=self.gurobi_timelimit)
-
-                self.gurobi_mipgap.append(gurobi_mipgap)
-
-            # note the need for [0] at the end
-            self._next_x = self.space.inverse_transform(
-                next_x.reshape((1, -1)))[0]
-
-            from entmoot.utils import get_cat_idx
-
-            for idx, xi in enumerate(self._next_x):
-                if idx not in get_cat_idx(self.space):
-                    self._next_x[idx] = xi
-
-                    # enforce variable bounds
-                    if self._next_x[idx] > self.space.bounds[idx][1]:
-                        self._next_x[idx] = self.space.bounds[idx][1]
-                    elif self._next_x[idx] < self.space.bounds[idx][0]:
-                        self._next_x[idx] = self.space.bounds[idx][0]
-
-            self._model_mu = model_mu
-            self._model_std = model_std
-
-            if self.models:
-                self.model_mu.append(self._model_mu)
-                self.model_std.append(self._model_std)
-
-
-            # check how far the new next_x is away from existing data
-            next_x = self._next_x
-            min_delta_x = min([self.space.distance(next_x, xi)
-                               for xi in self.Xi])
-            if abs(min_delta_x) <= 1e-8:
-                warnings.warn("The objective has been evaluated "
-                              "at this point before.")
-
-            # return point computed from last call to tell()
-            return next_x
-
-    def tell(self, x: list, y: list):
-        """
-        Checks that both x and y are valid points for the given search space.
-
-        :param x: list, locations of new data points
-        :param y: list, target value of new data points
-        """
-
-        from entmoot.utils import check_x_in_space
-        check_x_in_space(x, self.space)
-        self._check_y_is_valid(x, y)
-        self._tell(x, y)
-
-    def _tell(self, x, y):
-        """
-        Adds the new data points to the data set.
-
-        :param x: list, locations of new data points
-        :param y: list, target value of new data points
-        """
-
-        from entmoot.utils import is_listlike
-        from entmoot.utils import is_2Dlistlike
-
-        # if y isn't a scalar it means we have been handed a batch of points
-        if is_listlike(y) and is_2Dlistlike(x):
-            self.Xi.extend(x)
-            self.yi.extend(y)
-            self._n_initial_points -= len(y)
-        elif is_listlike(x):
-            self.Xi.append(x)
-            self.yi.append(y)
-            self._n_initial_points -= 1
-        else:
-            raise ValueError("Type of arguments `x` (%s) and `y` (%s) "
-                             "not compatible." % (type(x), type(y)))
-            
-        # optimizer learned something new - discard cache
-        self.cache_ = {}
-
-        # set self._next_x to None to indicate that the solver has learned something new
-        self._next_x = None
-
-    def _check_y_is_valid(self, x, y):
-        """check if the shape and types of x and y are consistent."""
-
-        from entmoot.utils import is_listlike
-        from entmoot.utils import is_2Dlistlike
-
-        # single objective checks for scalar values
-        if self.num_obj == 1:
-            # if y isn't a scalar it means we have been handed a batch of points
-            if is_listlike(y) and is_2Dlistlike(x):
-                for y_value in y:
-                    if not isinstance(y_value, Number):
-                        raise ValueError("expected y to be a list of scalars")
-
-            elif is_listlike(x):
-                if not isinstance(y, Number):
-                    raise ValueError("`func` should return a scalar")
-
-            else:
-                raise ValueError("Type of arguments `x` (%s) and `y` (%s) "
-                                 "not compatible." % (type(x), type(y)))
-        else:
-            # if y isn't a scalar it means we have been handed a batch of points
-            if is_listlike(y[0]) and is_2Dlistlike(x):
-                for y_value in y:
-                    if len(y_value) != self.num_obj:
-                        raise ValueError(f"expected y to be of size {self.num_obj}")
-                    for yi in y_value:
-                        if not isinstance(yi, Number):
-                            raise ValueError(f"expected y to be a list of list-like items of length {self.num_obj}")
-            elif is_listlike(x):
-                if len(y) != self.num_obj:
-                    raise ValueError(
-                        f"`func` should return a list-like item of length {self.num_obj}")
-                for yi in y:
-                    if not isinstance(yi, Number):
-                        raise ValueError(f"`func` should return a list-like item of length {self.num_obj}")
-            else:
-                raise ValueError("Type of arguments `x` (%s) and `y` (%s) "
-                                 "not compatible." % (type(x), type(y)))
-
-
-    def run(self, func, n_iter=1, no_progress_bar=False, update_min=False):
-        """Execute ask() + tell() `n_iter` times"""
-        from tqdm import tqdm
-
-        for itr in tqdm(range(n_iter),disable=no_progress_bar):
-            x = self.ask()
-            self.tell(x, func(x))
-
-            if no_progress_bar and update_min:
-                print(f"Min. obj.: {round(min(self.yi),2)} at itr.: {itr+1} / {n_iter}", end="\r")
-
-        from entmoot.utils import create_result
-
-        result = create_result(self.Xi, self.yi, self.space, self.rng,
-                               models=self.models, 
-                               model_mu=self.model_mu,
-                               model_std=self.model_std,
-                               gurobi_mipgap=self.gurobi_mipgap)
-
-        if no_progress_bar and update_min:
-            print(f"Min. obj.: {round(min(self.yi),2)} at itr.: {n_iter} / {n_iter}")
-        return result
-
-    def update_next(self):
-        """Updates the value returned by opt.ask(). Useful if a parameter
-        was updated after ask was called."""
-        self.cache_ = {}
-        # Ask for a new next_x.
-        # We only need to overwrite _next_x if it exists.
-        if hasattr(self, '_next_x'):
-            opt = self.copy(random_state=self.rng)
-            self._next_x = opt._next_x
-
-    def get_result(self):
-        """Returns the same result that would be returned by opt.tell()
-        but without calling tell
-
-        Parameters
-        ----------
-        -
-
-        Returns
-        -------
-        res : `OptimizeResult`, scipy object
-            OptimizeResult instance with the required information.
-
-        """
-        from entmoot.utils import create_result
-
-        result = create_result(self.Xi, self.yi, self.space, self.rng,
-                               models=self.models, 
-                               model_mu=self.model_mu,
-                               model_std=self.model_std,
-                               gurobi_mipgap=self.gurobi_mipgap)
-        return result
-
-    def predict_with_est(self, x, return_std=True):
-        from entmoot.utils import is_2Dlistlike
-
-        if is_2Dlistlike(x):
-            next_x = np.asarray(
-                self.space.transform(x)
-            )
-        else:
-            next_x = np.asarray(
-                self.space.transform([x])[0]
-            ).reshape(1, -1)
-
-        est = self.base_estimator_
-        est.fit(self.space.transform(self.Xi), self.yi)
-        temp_mu, temp_std = \
-            est.predict(
-                X=next_x,
-                return_std=True)
-        
-        if is_2Dlistlike(x):
-            if return_std:
-                return temp_mu, temp_std
-            else:
-                return temp_mu
-        else:
-            if return_std:
-                return temp_mu[0], temp_std[0]
-            else:
-                return temp_mu[0]
-
-    def predict_with_acq(self, x):
-        from entmoot.utils import is_2Dlistlike
-
-        if is_2Dlistlike(x):
-            next_x = np.asarray(
-                self.space.transform(x)
-            )
-        else:
-            next_x = np.asarray(
-                self.space.transform([x])[0]
-            ).reshape(1, -1)
-
-        if self.models:
-            temp_val = _gaussian_acquisition(
-                X=next_x, model=self.models[-1], 
-                y_opt=np.min(self.yi),
-                acq_func=self.acq_func,
-                acq_func_kwargs=self.acq_func_kwargs
-            )
-        else:
-            est = self.base_estimator_
-            est.fit(self.space.transform(self.Xi), self.yi)
-
-            temp_val = _gaussian_acquisition(
-                X=next_x, model=est, 
-                y_opt=np.min(self.yi),
-                acq_func=self.acq_func,
-                acq_func_kwargs=self.acq_func_kwargs
-            )
-            
-        if is_2Dlistlike(x):
-            return temp_val
-        else:
-            return temp_val[0]
-
-    def predict_pareto(
-        self,
-        sampling_strategy: str = 'random',
-        num_samples: int = 10,
-        num_levels: int = 10,
-        add_model_core = None,
-        gurobi_env=None,
-    ):
-        """
-        Computes the next point at which the objective should be evaluated.
-
-        :param sampling_strategy: str = 'grid'
-            picks the strategy to sample weights for the muli-objective function
-                'random': gives 'num_samples' randomly drawn weights that sum to 1
-                'grid': defines ordered grid of samples depending on 'num_levels'
-        :param num_samples: int = 10,
-            defines number of samples for 'random' sampling strategy
-        :param num_levels: int = 10,
-            defines levels per dimension for 'grid' sampling strategy
-        :param add_model_core: GRBModel = None,
-            Gurobi optimization model that includes additional constraints
-        :param gurobi_env: Gurobi Env = None,
-            Gurobi environment used for computation of the next proposal
-
-        :return pareto_x: list, next pareto point predictions of
-            shape(n_points, n_dims)
-        """
-
-        # update gurobi_env attribute
-        if gurobi_env:
-            self.gurobi_env = gurobi_env
-
-        assert self.num_obj > 1, \
-            f"Number of objectives needs to be > 1 to" \
-            f" compute Pareto frontiers."
-
-        from opti.sampling.simplex import sample, grid
-
-        # pick the sampling strategy
-        if sampling_strategy == 'random':
-            weights = sample(self.num_obj, num_samples)
-        elif sampling_strategy == 'grid':
-            weights = grid(self.num_obj, num_levels)
-        else:
-            raise ValueError("'sampling_type' must be in ['random', 'grid'")
-
-        # fit current model
-        est = self.base_estimator_
-        est.fit(self.space.transform(self.Xi), self.yi)
-
-        # add model constraints if necessary
-        if add_model_core is None:
-            add_model_core = \
-                self.acq_optimizer_kwargs.get("add_model_core", None)
-
-        # compute pareto points based on weight vector
-        pareto = []
-        for w in weights:
-            temp_x, temp_mu, model_std, gurobi_mipgap = \
-                est.get_global_next_x(acq_func=self.acq_func,
-                                      acq_func_kwargs=self.acq_func_kwargs,
-                                      acq_optimizer_kwargs=self.acq_optimizer_kwargs,
-                                      add_model_core=add_model_core,
-                                      weight=w,
-                                      verbose=self.verbose,
-                                      gurobi_env=self.gurobi_env,
-                                      gurobi_timelimit=self.gurobi_timelimit)
-            pareto.append((temp_x, temp_mu))
-        return pareto
diff --git a/entmoot/optimizers/__init__.py b/entmoot/optimizers/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/entmoot/optimizers/base_opt.py b/entmoot/optimizers/base_opt.py
new file mode 100644
index 0000000..670127d
--- /dev/null
+++ b/entmoot/optimizers/base_opt.py
@@ -0,0 +1,9 @@
+class BaseOptimizer:
+    def __init__(self, space, params):
+        raise NotImplementedError()
+
+    def solve(model):
+        raise NotImplementedError()
+
+    def sample_feas(num_points):
+        raise NotImplementedError()
diff --git a/entmoot/optimizers/gurobi_opt.py b/entmoot/optimizers/gurobi_opt.py
new file mode 100644
index 0000000..e4fb7d1
--- /dev/null
+++ b/entmoot/optimizers/gurobi_opt.py
@@ -0,0 +1,113 @@
+from collections import namedtuple
+from entmoot import Enting, ProblemConfig
+from entmoot.utils import OptResult
+import gurobipy as gur
+import os
+
+
+class GurobiOptimizer:
+    def __init__(self, problem_config: ProblemConfig, params: dict = None) -> float:
+        self._params = {} if params is None else params
+        self._problem_config = problem_config
+        self._curr_sol = None
+        self._active_leaves = None
+
+    def get_curr_sol(self) -> list:
+        """
+        returns current solution (i.e. optimal points) from optimization run
+        """
+        assert self._curr_sol is not None, "No solution was generated yet."
+        return self._curr_sol
+
+    def get_active_leaf_sol(self) -> list:
+        """
+        returns current solution (i.e. optimal points) from optimization based
+        """
+        assert self._curr_sol is not None, "No solution was generated yet."
+        return self._active_leaves
+
+    def solve(
+            self,
+            tree_model: Enting,
+            model_core: gur.Model = None,
+            weights: tuple = None,
+            use_env: bool = False,
+    ) -> namedtuple:
+        """
+        Solves the Gurobi optimization model
+        """
+
+        if model_core is None:
+            if use_env:
+                if "CLOUDACCESSID" in os.environ:
+                    # Use Gurobi Cloud
+                    connection_params_cld = {
+                        "CLOUDACCESSID": os.getenv("CLOUDACCESSID"),
+                        "CLOUDSECRETKEY": os.getenv("CLOUDSECRETKEY"),
+                        "CLOUDPOOL": os.getenv("CLOUDPOOL"),
+                    }
+                    env_cld = gur.Env(params=connection_params_cld)
+                    env_cld.start()
+                    opt_model = self._problem_config.get_gurobi_model_core(env=env_cld)
+
+            else:
+                opt_model = self._problem_config.get_gurobi_model_core()
+        else:
+            # create model copy to not overwrite original one
+            opt_model = self._problem_config.copy_gurobi_model_core(model_core)
+
+        # check weights
+        if weights is not None:
+            assert len(weights) == len(self._problem_config.obj_list), (
+                f"Number of 'weights' is '{len(weights)}', number of objectives "
+                f"is '{len(self._problem_config.obj_list)}'."
+            )
+            assert sum(weights) == 1.0, "weights don't add up to 1.0"
+
+        # set solver parameters
+        for param, param_val in self._params.items():
+            opt_model.setParam(param, param_val)
+
+        # build gurobi model
+        tree_model.add_to_gurobipy_model(opt_model, weights=weights)
+
+        # Solve optimization model
+        opt_model.optimize()
+
+        # update current solution
+        self._curr_sol, self._active_leaves = self._get_sol(opt_model)
+
+        return OptResult(
+            self.get_curr_sol(),
+            opt_model.ObjVal,
+            [x.X for x in opt_model._unscaled_mu],
+        )
+
+    def _get_sol(self, solved_model: gur.Model) -> list:
+        # extract solutions from conti and discrete variables
+        res = []
+        for idx, feat in enumerate(self._problem_config.feat_list):
+            curr_var = solved_model._all_feat[idx]
+            if feat.is_cat():
+                # find active category
+                sol_cat = [
+                    int(round(curr_var[enc_cat].x)) for enc_cat in feat.enc_cat_list
+                ].index(1)
+                res.append(sol_cat)
+            else:
+                res.append(curr_var.x)
+
+        # extract active leaves of solution
+        def obj_leaf_index(model_obj, obj_name):
+            # this function is the same as in 'tree_ensemble.py', TODO: put this in a tree_utils?
+            for tree in range(model_obj._num_trees(obj_name)):
+                for leaf in model_obj._leaves(obj_name, tree):
+                    yield tree, leaf
+
+        act_leaves = []
+        for idx, obj in enumerate(self._problem_config.obj_list):
+            act_leaves.append(
+                [(tree_id, leaf_enc) for tree_id, leaf_enc in obj_leaf_index(solved_model, obj.name)
+                 if round(solved_model._z[obj.name, tree_id, leaf_enc].x) == 1.0])
+
+        return self._problem_config.decode([res]), act_leaves
diff --git a/entmoot/optimizers/pyomo_opt.py b/entmoot/optimizers/pyomo_opt.py
new file mode 100644
index 0000000..61b3c50
--- /dev/null
+++ b/entmoot/optimizers/pyomo_opt.py
@@ -0,0 +1,100 @@
+from collections import namedtuple
+from entmoot import Enting, ProblemConfig
+from entmoot.utils import OptResult
+import pyomo.environ as pyo
+
+
+class PyomoOptimizer:
+    def __init__(self, problem_config: ProblemConfig, params: dict = None):
+        self._params = {} if params is None else params
+        self._problem_config = problem_config
+        self._curr_sol = None
+        self._active_leaves = None
+
+    @property
+    def get_curr_sol(self) -> list:
+        """
+        returns current solution (i.e. optimal points) from optimization run
+        """
+        assert self._curr_sol is not None, "No solution was generated yet."
+        return self._curr_sol
+
+    def get_active_leaf_sol(self) -> list:
+        """
+        returns current solution (i.e. optimal points) from optimization based
+        """
+        assert self._curr_sol is not None, "No solution was generated yet."
+        return self._active_leaves
+
+    def solve(
+        self, tree_model: Enting, model_core: pyo.ConcreteModel = None, weights: tuple = None
+    ) -> namedtuple:
+        """
+        Solves the Pyomo optimization model
+        """
+        if model_core is None:
+            opt_model = self._problem_config.get_pyomo_model_core()
+        else:
+            # create model copy to not overwrite original one
+            opt_model = self._problem_config.copy_pyomo_model_core(model_core)
+
+        # check weights
+        if weights is not None:
+            assert len(weights) == len(self._problem_config.obj_list), (
+                f"Number of 'weights' is '{len(weights)}', number of objectives "
+                f"is '{len(self._problem_config.obj_list)}'."
+            )
+            assert sum(weights) == 1.0, "weights don't add up to 1.0"
+
+        # choose solver
+        opt = pyo.SolverFactory(self._params["solver_name"])
+
+        # set solver parameters
+        if "solver_options" in self._params:
+            for k, v in self._params["solver_options"].items():
+                opt.options[k] = v
+
+        # build pyomo model using information from tree model
+        tree_model.add_to_pyomo_model(opt_model)
+
+        # Solve optimization model
+        opt.solve(opt_model, tee=True)
+
+        # update current solution
+        self._curr_sol, self._active_leaves = self._get_sol(opt_model)
+
+        return OptResult(
+            self.get_curr_sol,
+            pyo.value(opt_model.obj),
+            [opt_model._unscaled_mu[k].value for k in opt_model._unscaled_mu],
+        )
+
+    def _get_sol(self, solved_model: pyo.ConcreteModel) -> list:
+        # extract solutions from conti and discrete variables
+        res = []
+        for idx, feat in enumerate(self._problem_config.feat_list):
+            curr_var = solved_model._all_feat[idx]
+            if feat.is_cat():
+                # find active category
+                sol_cat = [
+                    int(round(pyo.value(curr_var[enc_cat])))
+                    for enc_cat in feat.enc_cat_list
+                ].index(1)
+                res.append(sol_cat)
+            else:
+                res.append(pyo.value(curr_var))
+
+        # extract active leaves of solution
+        def obj_leaf_index(model_obj, obj_name):
+            # this function is the same as in 'tree_ensemble.py', TODO: put this in a tree_utils?
+            for tree in range(model_obj._num_trees(obj_name)):
+                for leaf in model_obj._leaves(obj_name, tree):
+                    yield tree, leaf
+
+        act_leaves = []
+        for idx, obj in enumerate(self._problem_config.obj_list):
+            act_leaves.append(
+                [(tree_id, leaf_enc) for tree_id, leaf_enc in obj_leaf_index(solved_model, obj.name)
+                 if round(pyo.value(solved_model._z[obj.name, tree_id, leaf_enc])) == 1.0])
+
+        return self._problem_config.decode([res]), act_leaves
diff --git a/entmoot/optimizers/sampling_opt.py b/entmoot/optimizers/sampling_opt.py
new file mode 100644
index 0000000..83704f3
--- /dev/null
+++ b/entmoot/optimizers/sampling_opt.py
@@ -0,0 +1,9 @@
+class SamplingOptimizer:
+    def __init__(self, space, params):
+        raise NotImplementedError()
+
+    def solve(model):
+        raise NotImplementedError()
+
+    def sample_feas(num_points):
+        raise NotImplementedError()
diff --git a/entmoot/problem_config.py b/entmoot/problem_config.py
new file mode 100644
index 0000000..2b7c468
--- /dev/null
+++ b/entmoot/problem_config.py
@@ -0,0 +1,452 @@
+from typing import Tuple, List
+import numpy as np
+import random
+
+
+class ProblemConfig:
+    def __init__(self, rnd_seed: int = None):
+        self._feat_list = []
+        self._obj_list = []
+        self._rnd_seed = rnd_seed
+
+        if rnd_seed is not None:
+            assert isinstance(
+                rnd_seed, int
+            ), f"Argument 'rnd_seed' needs to be an integer. Got type '{type(rnd_seed)}'."
+            np.random.seed(rnd_seed)
+            random.seed(rnd_seed)
+
+    @property
+    def cat_idx(self):
+        return tuple([i for i, feat in enumerate(self.feat_list) if feat.is_cat()])
+
+    @property
+    def non_cat_idx(self):
+        return tuple([i for i, feat in enumerate(self.feat_list) if not feat.is_cat()])
+
+    @property
+    def non_cat_lb(self):
+        return tuple(
+            [feat.lb for i, feat in enumerate(self.feat_list) if not feat.is_cat()]
+        )
+
+    @property
+    def non_cat_ub(self):
+        return tuple(
+            [feat.ub for i, feat in enumerate(self.feat_list) if not feat.is_cat()]
+        )
+
+    @property
+    def non_cat_bnd_diff(self):
+        return tuple(
+            [
+                feat.ub - feat.lb
+                for i, feat in enumerate(self.feat_list)
+                if not feat.is_cat()
+            ]
+        )
+
+    @property
+    def feat_list(self):
+        return self._feat_list
+
+    @property
+    def obj_list(self):
+        return self._obj_list
+
+    @property
+    def rnd_seed(self):
+        return self._rnd_seed
+
+    def _encode_xi(self, xi: List):
+        return np.asarray(
+            [feat.encode(xi[idx]) for idx, feat in enumerate(self.feat_list)]
+        )
+
+    def _decode_xi(self, xi: List):
+        return [feat.decode(xi[idx]) for idx, feat in enumerate(self.feat_list)]
+
+    def encode(self, X: List):
+        if len(X) == 0:
+            return []
+        elif len(X) == 1:
+            return self._encode_xi(X[0])
+        else:
+            enc = [self._encode_xi(xi) for xi in X]
+            return np.asarray(enc)
+
+    def decode(self, X: List):
+        if len(X) == 0:
+            return []
+        elif len(X) == 1:
+            return self._decode_xi(X[0])
+        else:
+            dec = [self._decode_xi(xi) for xi in X]
+            return np.asarray(dec)
+
+    def add_feature(self, feat_type: str, bounds: Tuple = None, name: str = None):
+        if name is None:
+            name = f"feat_{len(self.feat_list)}"
+
+        if bounds is None and feat_type in ("real", "integer", "categorical"):
+            raise IOError(
+                "Please provide bounds for feature types in '(real, integer, categorical)'"
+            )
+
+        # perform basic input checks and define features of space
+        if feat_type in ("real", "integer"):
+            assert len(bounds) == 2, (
+                f"Too many values for bounds for feat_type '{feat_type}'. Expected '2', "
+                f"got '{len(bounds)}'. Check feature '{name}'"
+            )
+
+            lb, ub = bounds
+
+            if feat_type == "real":
+                if isinstance(lb, int):
+                    lb = float(lb)
+                if isinstance(ub, int):
+                    ub = float(ub)
+
+                assert isinstance(lb, float) and isinstance(ub, float), (
+                    f"Wrong type for bounds of feat_type '{feat_type}'. Expected '(float, "
+                    f"float)', got '({type(lb)}, {type(ub)})'. Check feature '{name}'"
+                )
+
+                assert lb < ub, (
+                    f"Lower bound '{lb}' of feat_type '{feat_type}' needs to be "
+                    f"smaller than upper bound. Check feature '{name}'."
+                )
+
+                self._feat_list.append(Real(lb=lb, ub=ub, name=name))
+            else:
+                assert isinstance(lb, int) and isinstance(ub, int), (
+                    f"Wrong type for bounds of feat_type '{feat_type}'. Expected '(int, int)', "
+                    f"got '({type(lb)}, {type(ub)})'. Check feature '{name}'"
+                )
+
+                assert lb < ub, (
+                    f"Lower bound '{lb}' of feat_type '{feat_type}' needs to be "
+                    f"smaller than upper bound. Check feature '{name}'."
+                )
+
+                self._feat_list.append(Integer(lb=lb, ub=ub, name=name))
+
+        elif feat_type == "binary":
+            self._feat_list.append(Binary(name=name))
+
+        elif feat_type == "categorical":
+            assert len(bounds) > 1, (
+                f"Not enough categories specified for bounds of feat_type '{feat_type}'. "
+                f"Expected 'len(bounds) > 1', got 'len(bounds) = {len(bounds)}'. "
+                f"Check feature '{name}'."
+            )
+
+            assert all(isinstance(item, str) for item in bounds), (
+                f"Wrong type for bounds of feat_type '{feat_type}'. Expected 'str' for all "
+                f"categories."
+            )
+
+            assert len(bounds) == len(
+                set(bounds)
+            ), f"Categories of feat_type '{feat_type}' are not all unique."
+
+            self._feat_list.append(Categorical(cat_list=bounds, name=name))
+
+        else:
+            raise IOError(
+                f"No support for feat_type '{feat_type}'. Check feature '{name}'."
+            )
+
+    def add_min_objective(self, name: str = None):
+        if name is None:
+            name = f"obj_{len(self.obj_list)}"
+
+        self._obj_list.append(MinObjective(name=name))
+
+    def get_rnd_sample_numpy(self, num_samples):
+        # returns np.array for faster processing
+        array_list = []
+        for feat in self.feat_list:
+            if feat.is_real():
+                array_list.append(
+                    np.random.uniform(low=feat.lb, high=feat.ub, size=num_samples)
+                )
+            elif feat.is_cat():
+                array_list.append(
+                    np.random.randint(len(feat.cat_list), size=num_samples)
+                )
+            elif feat.is_int() or feat.is_bin():
+                array_list.append(
+                    np.random.random_integers(
+                        low=feat.lb, high=feat.ub, size=num_samples
+                    )
+                )
+        return np.squeeze(np.column_stack(array_list))
+
+    def get_rnd_sample_list(self, num_samples=1, cat_enc=False):
+        # returns list of tuples
+        sample_list = []
+        for _ in range(num_samples):
+            sample = []
+            for feat in self.feat_list:
+                if feat.is_real():
+                    sample.append(random.uniform(feat.lb, feat.ub))
+                elif feat.is_cat():
+                    if cat_enc:
+                        sample.append(random.randrange(0, len(feat.cat_list)))
+                    else:
+                        sample.append(random.choice(feat.cat_list))
+                elif feat.is_int() or feat.is_bin():
+                    sample.append(random.randint(feat.lb, feat.ub))
+            sample_list.append(tuple(sample))
+        return sample_list if len(sample_list) > 1 else sample_list[0]
+
+    def get_gurobi_model_core(self, env=None):
+        import gurobipy as gur
+
+        # initialize gurobi model
+        if env is None:
+            model = gur.Model()
+        else:
+            model = gur.Model(env=env)
+
+        # initalize var space
+        model._all_feat = []
+
+        for i, feat in enumerate(self.feat_list):
+            if feat.is_real():
+                model._all_feat.append(
+                    model.addVar(lb=feat.lb, ub=feat.ub, name=feat.name, vtype="C")
+                )
+            elif feat.is_cat():
+                model._all_feat.append(dict())
+                for enc, cat in zip(feat.enc_cat_list, feat.cat_list):
+                    comb_name = f"{feat.name}_{cat}"
+                    model._all_feat[i][enc] = model.addVar(name=comb_name, vtype="B")
+            elif feat.is_int():
+                model._all_feat.append(
+                    model.addVar(lb=feat.lb, ub=feat.ub, name=feat.name, vtype="I")
+                )
+            elif feat.is_bin():
+                model._all_feat.append(model.addVar(name=feat.name, vtype="B"))
+
+        model.update()
+        return model
+
+    def copy_gurobi_model_core(self, model_core):
+        """
+        Computes a copy of a Gurobi model which is decoupled from the original, i.e. you can modify the copy without
+        changing the original model.
+        :param model_core: The gurobi model we want to copy
+        :type model_core: gurobi.Model
+        :return:
+        :rtype: gurobi.Model
+        """
+        copy_model_core = model_core.copy()
+        # Since Gurobi's .copy() method does not copy attributes like ._all_feat we have to copy them manually
+        copy_model_core._all_feat = []
+
+        # transfer feature var list to model copy
+        for i, feat in enumerate(self.feat_list):
+            if feat.is_cat():
+                copy_model_core._all_feat.append(dict())
+                for enc, cat in zip(feat.enc_cat_list, feat.cat_list):
+                    var_name = model_core._all_feat[i][enc].VarName
+                    copy_model_core._all_feat[i][enc] = copy_model_core.getVarByName(
+                        var_name
+                    )
+            else:
+                var_name = model_core._all_feat[i].VarName
+                copy_model_core._all_feat.append(copy_model_core.getVarByName(var_name))
+        return copy_model_core
+
+    def copy_pyomo_model_core(self, model_core):
+        """
+        Computes a copy of a Pyomo model which is decoupled from the original, i.e. you can modify the copy without
+        changing the original model.
+        :param model_core: A Pyomo model
+        :type model_core: pyomo.environ.ConcreteModel
+        :return: A copy of the Pyomo model
+        :rtype: pyomo.environ.ConcreteModel
+        """
+        return model_core.clone()
+
+    def get_pyomo_model_core(self):
+        import pyomo.environ as pyo
+
+        # initialize Pyomo model
+        model = pyo.ConcreteModel()
+
+        # Initialize var space
+        model._all_feat = []
+
+        # Define some auxiliary dictionaries
+        index_to_var_domain = {}
+        index_to_var_bounds = {}
+        for i, feat in enumerate(self.feat_list):
+            if feat.is_real():
+                index_to_var_domain[i] = pyo.Reals
+                index_to_var_bounds[i] = (feat.lb, feat.ub)
+            elif feat.is_cat():
+                for enc, cat in zip(feat.enc_cat_list, feat.cat_list):
+                    # We encode the index of this variable by (i, enc, cat), where 'i' is the position in the list of
+                    # features, 'enc' is the corresponding encoded numerical value and 'cat' is the category
+                    index_to_var_domain[i, enc, cat] = pyo.Binary
+            elif feat.is_int():
+                index_to_var_domain[i] = pyo.Integers
+                index_to_var_bounds[i] = (feat.lb, feat.ub)
+            elif feat.is_bin():
+                index_to_var_domain[i] = pyo.Binary
+
+        # Build Pyomo index set from dictionary keys
+        # Why strings?? Well, pyomo doesn't support lists where some indices are multidimensional, but this is handy
+        # for encoded categorical features
+        model.indices_features = pyo.Set(
+            initialize=[str(x) for x in index_to_var_domain]
+        )
+
+        # These auxiliary functions are needed, because Pyomo does not accept the dictionaries instead. The 'model'
+        # argument is needed as well, but not explicitly used.
+        # We hve the eval-statement eval(i) instead of just i, because the indices are strings for reasons explained
+        # above :(
+        def i_to_dom(model, i):
+            return index_to_var_domain[eval(i)]
+
+        def i_to_bounds(model, i):
+            if eval(i) in index_to_var_bounds:
+                return index_to_var_bounds[eval(i)]
+            else:
+                return (None, None)
+
+        # Define decision variables for Pyomo model
+        model.x = pyo.Var(model.indices_features, domain=i_to_dom, bounds=i_to_bounds)
+        # We store the decision variables corresponding to the features in the list _all_feat in order to follow a
+        # similar strategy to the gurobi model. Can be replaced in the future by a more direct approach.
+        for i_str in model.x:
+            i = eval(i_str)
+            if type(i) is int:
+                model._all_feat.append(model.x[i_str])
+            elif type(i) is tuple:
+                if i[1] == 0:
+                    model._all_feat.append({})
+                    model._all_feat[-1][i[1]] = model.x[i_str]
+                elif i[1] > 0:
+                    model._all_feat[-1][i[1]] = model.x[i_str]
+                else:
+                    raise TypeError
+            else:
+                raise TypeError
+
+        return model
+
+    def __str__(self):
+        out_str = list(["\nPROBLEM SUMMARY"])
+        out_str.append(len(out_str[-1][:-1]) * "-")
+        out_str.append("features:")
+        for feat in self.feat_list:
+            if feat.is_cat():
+                out_str.append(
+                    f"{feat.name} :: {feat.__class__.__name__} :: {feat.cat_list} "
+                )
+            else:
+                out_str.append(
+                    f"{feat.name} :: {feat.__class__.__name__} :: ({feat.lb}, {feat.ub}) "
+                )
+
+        out_str.append("\nobjectives:")
+        for obj in self.obj_list:
+            out_str.append(f"{obj.name} :: {obj.__class__.__name__}")
+        return "\n".join(out_str)
+
+
+class FeatureType:
+    def is_real(self):
+        return False
+
+    def is_cat(self):
+        return False
+
+    def is_int(self):
+        return False
+
+    def is_bin(self):
+        return False
+
+    def encode(self, xi):
+        return xi
+
+    def decode(self, xi):
+        return xi
+
+
+class Real(FeatureType):
+    def __init__(self, lb, ub, name):
+        self.lb = lb
+        self.ub = ub
+        self.name = name
+
+    def is_real(self):
+        return True
+
+
+class Categorical(FeatureType):
+    def __init__(self, cat_list, name):
+        self._cat_list = cat_list
+        self.name = name
+
+        # encode categories
+        self._enc2str, self._str2enc = {}, {}
+        self._enc_cat_list = []
+        for enc, cat in enumerate(cat_list):
+            self._enc_cat_list.append(enc)
+            self._enc2str[enc] = cat
+            self._str2enc[cat] = enc
+
+    @property
+    def cat_list(self):
+        return self._cat_list
+
+    @property
+    def enc_cat_list(self):
+        return self._enc_cat_list
+
+    def encode(self, xi):
+        return self._str2enc[xi]
+
+    def decode(self, xi):
+        return self._enc2str[xi]
+
+    def is_cat(self):
+        return True
+
+
+class Integer(FeatureType):
+    def __init__(self, lb, ub, name):
+        self.lb = lb
+        self.ub = ub
+        self.name = name
+
+    def is_int(self):
+        return True
+
+    def decode(self, xi):
+        return int(xi)
+
+
+class Binary(FeatureType):
+    def __init__(self, name):
+        self.lb = 0
+        self.ub = 1
+        self.name = name
+
+    def decode(self, xi):
+        return abs(int(xi))
+
+    def is_bin(self):
+        return True
+
+
+class MinObjective:
+    def __init__(self, name):
+        self.name = name
diff --git a/entmoot/sampler/__init__.py b/entmoot/sampler/__init__.py
deleted file mode 100644
index c48efef..0000000
--- a/entmoot/sampler/__init__.py
+++ /dev/null
@@ -1,46 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-"""
-Utilities for generating initial sequences
-"""
-from entmoot.sampler.lhs import Lhs
-from entmoot.sampler.sobol import Sobol
-from entmoot.sampler.halton import Halton
-from entmoot.sampler.hammersly import Hammersly
-from entmoot.sampler.grid import Grid
-from entmoot.sampler.base import InitialPointGenerator
-
-
-__all__ = [
-    "Lhs", "Sobol",
-    "Halton", "Hammersly",
-    "Grid", "InitialPointGenerator"
-]
diff --git a/entmoot/sampler/base.py b/entmoot/sampler/base.py
deleted file mode 100644
index b5ec0ef..0000000
--- a/entmoot/sampler/base.py
+++ /dev/null
@@ -1,57 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-from collections import defaultdict
-
-
-class InitialPointGenerator(object):
-    def generate(self, dimensions, n_samples, random_state=None):
-        raise NotImplemented
-
-    def set_params(self, **params):
-        """
-        Set the parameters of this initial point generator.
-
-        Parameters
-        ----------
-        **params : dict
-            Generator parameters.
-        Returns
-        -------
-        self : object
-            Generator instance.
-        """
-        if not params:
-            # Simple optimization to gain speed (inspect is slow)
-            return self
-        for key, value in params.items():
-            setattr(self, key, value)
-
-        return self
diff --git a/entmoot/sampler/grid.py b/entmoot/sampler/grid.py
deleted file mode 100644
index 40034ec..0000000
--- a/entmoot/sampler/grid.py
+++ /dev/null
@@ -1,198 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-Inspired by https://github.com/jonathf/chaospy/blob/master/chaospy/
-distributions/sampler/sequences/grid.py
-"""
-import numpy as np
-from entmoot.sampler.base import InitialPointGenerator
-from entmoot.space.space import Space
-from sklearn.utils import check_random_state
-
-
-def _quadrature_combine(args):
-    args = [np.asarray(arg).reshape(len(arg), -1) for arg in args]
-    shapes = [arg.shape for arg in args]
-
-    size = np.prod(shapes, 0)[0] * np.sum(shapes, 0)[1]
-    if size > 10 ** 9:
-        raise MemoryError("Too large sets")
-
-    out = args[0]
-    for arg in args[1:]:
-        out = np.hstack([
-            np.tile(out, len(arg)).reshape(-1, out.shape[1]),
-            np.tile(arg.T, len(out)).reshape(arg.shape[1], -1).T,
-        ])
-    return out
-
-
-def _create_uniform_grid_exclude_border(n_dim, order):
-    assert order > 0
-    assert n_dim > 0
-    x_data = np.arange(1, order + 1) / (order + 1.)
-    x_data = _quadrature_combine([x_data] * n_dim)
-    return x_data
-
-
-def _create_uniform_grid_include_border(n_dim, order):
-    assert order > 1
-    assert n_dim > 0
-    x_data = np.arange(0, order) / (order - 1.)
-    x_data = _quadrature_combine([x_data] * n_dim)
-    return x_data
-
-
-def _create_uniform_grid_only_border(n_dim, order):
-    assert n_dim > 0
-    assert order > 1
-    x = [[0., 1.]] * (n_dim - 1)
-    x.append(list(np.arange(0, order) / (order - 1.)))
-    x_data = _quadrature_combine(x)
-    return x_data
-
-
-class Grid(InitialPointGenerator):
-    """Generate samples from a regular grid.
-
-    Parameters
-    ----------
-    border : str, default='exclude'
-        defines how the samples are generated:
-        - 'include' : Includes the border into the grid layout
-        - 'exclude' : Excludes the border from the grid layout
-        - 'only' : Selects only points at the border of the dimension
-    use_full_layout : boolean, default=True
-        When True, a  full factorial design is generated and
-        missing points are taken from the next larger full factorial
-        design, depending on `append_border`
-        When False, the next larger  full factorial design is
-        generated and points are randomly selected from it.
-    append_border : str, default="only"
-        When use_full_layout is True, this parameter defines how the missing
-        points will be generated from the next larger grid layout:
-        - 'include' : Includes the border into the grid layout
-        - 'exclude' : Excludes the border from the grid layout
-        - 'only' : Selects only points at the border of the dimension
-    """
-
-    def __init__(self, border="exclude", use_full_layout=True,
-                 append_border="only"):
-        self.border = border
-        self.use_full_layout = use_full_layout
-        self.append_border = append_border
-
-    def generate(self, dimensions, n_samples, random_state=None):
-        """Creates samples from a regular grid.
-
-        Parameters
-        ----------
-        dimensions : list, shape (n_dims,)
-            List of search space dimensions.
-            Each search dimension can be defined either as
-
-            - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-              dimensions),
-            - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-              dimensions),
-            - as a list of categories (for `Categorical` dimensions), or
-            - an instance of a `Dimension` object (`Real`, `Integer` or
-              `Categorical`).
-
-        n_samples : int
-            The order of the Halton sequence. Defines the number of samples.
-        random_state : int, RandomState instance, or None (default)
-            Set random state to something other than None for reproducible
-            results.
-
-        Returns
-        -------
-        np.array, shape=(n_dim, n_samples)
-            grid set
-        """
-        rng = check_random_state(random_state)
-        space = Space(dimensions)
-        n_dim = space.n_dims
-        transformer = space.get_transformer()
-        space.set_transformer("normalize")
-
-        if self.border == "include":
-            if self.use_full_layout:
-                order = int(np.floor(np.sqrt(n_samples)))
-            else:
-                order = int(np.ceil(np.sqrt(n_samples)))
-            if order < 2:
-                order = 2
-            h = _create_uniform_grid_include_border(n_dim, order)
-        elif self.border == "exclude":
-            if self.use_full_layout:
-                order = int(np.floor(np.sqrt(n_samples)))
-            else:
-                order = int(np.ceil(np.sqrt(n_samples)))
-            if order < 1:
-                order = 1
-            h = _create_uniform_grid_exclude_border(n_dim, order)
-        elif self.border == "only":
-            if self.use_full_layout:
-                order = int(np.floor(n_samples / 2.))
-            else:
-                order = int(np.ceil(n_samples / 2.))
-            if order < 2:
-                order = 2
-            h = _create_uniform_grid_exclude_border(n_dim, order)
-        else:
-            raise ValueError("Wrong value for border")
-        if np.size(h, 0) > n_samples:
-            rng.shuffle(h)
-            h = h[:n_samples, :]
-        elif np.size(h, 0) < n_samples:
-            if self.append_border == "only":
-                order = int(np.ceil((n_samples - np.size(h, 0)) / 2.))
-                if order < 2:
-                    order = 2
-                h2 = _create_uniform_grid_only_border(n_dim, order)
-            elif self.append_border == "include":
-                order = int(np.ceil(np.sqrt(n_samples - np.size(h, 0))))
-                if order < 2:
-                    order = 2
-                h2 = _create_uniform_grid_include_border(n_dim, order)
-            elif self.append_border == "exclude":
-                order = int(np.ceil(np.sqrt(n_samples - np.size(h, 0))))
-                if order < 1:
-                    order = 1
-                h2 = _create_uniform_grid_exclude_border(n_dim, order)
-            else:
-                raise ValueError("Wrong value for append_border")
-            h = np.vstack((h, h2[:(n_samples - np.size(h, 0))]))
-            rng.shuffle(h)
-        else:
-            rng.shuffle(h)
-        h = space.inverse_transform(h)
-        space.set_transformer(transformer)
-        return h
diff --git a/entmoot/sampler/halton.py b/entmoot/sampler/halton.py
deleted file mode 100644
index f0dcb99..0000000
--- a/entmoot/sampler/halton.py
+++ /dev/null
@@ -1,197 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-Inspired by https://github.com/jonathf/chaospy/blob/master/chaospy/
-distributions/sampler/sequences/halton.py
-"""
-
-import numpy as np
-from .base import InitialPointGenerator
-from ..space import Space
-from sklearn.utils import check_random_state
-from typing import Union
-
-
-class Halton(InitialPointGenerator):
-    """Creates `Halton` sequence samples.
-    In statistics, Halton sequences are sequences used to generate
-    points in space for numerical methods such as Monte Carlo simulations.
-    Although these sequences are deterministic, they are of low discrepancy,
-    that is, appear to be random
-    for many purposes. They were first introduced in 1960 and are an example
-    of a quasi-random number sequence. They generalise the one-dimensional
-    van der Corput sequences.
-    For ``dim == 1`` the sequence falls back to Van Der Corput sequence.
-    Parameters
-    ----------
-    min_skip : int
-        Minimum skipped seed number. When `min_skip != max_skip`
-        a random number is picked.
-    max_skip : int
-        Maximum skipped seed number. When `min_skip != max_skip`
-        a random number is picked.
-    primes : tuple, default=None
-        The (non-)prime base to calculate values along each axis. If
-        empty or None, growing prime values starting from 2 will be used.
-    """
-    def __init__(self, min_skip: int = 0, max_skip: int = 0, primes=None):
-        self.primes = primes
-        self.min_skip = min_skip
-        self.max_skip = max_skip
-
-    def generate(self, dimensions, n_samples: int, random_state=None):
-        """Creates samples from Halton set.
-        Parameters
-        ----------
-        dimensions : list, shape (n_dims,)
-            List of search space dimensions.
-            Each search dimension can be defined either as
-            - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-              dimensions),
-            - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-              dimensions),
-            - as a list of categories (for `Categorical` dimensions), or
-            - an instance of a `Dimension` object (`Real`, `Integer` or
-              `Categorical`).
-        n_samples : int
-            The order of the Halton sequence. Defines the number of samples.
-        random_state : int, RandomState instance, or None (default)
-            Set random state to something other than None for reproducible
-            results.
-        Returns
-        -------
-        np.array, shape=(n_dim, n_samples)
-            Halton set.
-        """
-        rng = check_random_state(random_state)
-        if self.primes is None:
-            primes = []
-        else:
-            primes = list(self.primes)
-        space = Space(dimensions)
-        n_dim = space.n_dims
-        transformer = space.get_transformer()
-        space.set_transformer("normalize")
-        if len(primes) < n_dim:
-            prime_order = 10 * n_dim
-            while len(primes) < n_dim:
-                primes = _create_primes(prime_order)
-                prime_order *= 2
-
-            primes = primes[:n_dim]
-        assert len(primes) == n_dim, "not enough primes"
-
-        if self.min_skip == self.max_skip:
-            skip = self.min_skip
-        elif self.min_skip < 0 and self.max_skip < 0:
-            skip = max(primes)
-        elif self.min_skip < 0 or self.max_skip < 0:
-            skip = np.max(self.min_skip, self.max_skip)
-        else:
-            skip = rng.randint(self.min_skip, self.max_skip)
-
-        out = np.empty((n_dim, n_samples))
-        indices = [idx + skip for idx in range(n_samples)]
-        for dim_ in range(n_dim):
-            out[dim_] = _van_der_corput_samples(
-                indices, number_base=primes[dim_])
-        out = space.inverse_transform(np.transpose(out))
-        space.set_transformer(transformer)
-        return out
-
-
-def _van_der_corput_samples(idx: Union[int, np.ndarray], number_base: int = 2):
-    """Create `Van Der Corput` low discrepancy sequence samples.
-    A van der Corput sequence is an example of the simplest one-dimensional
-    low-discrepancy sequence over the unit interval; it was first described in
-    1935 by the Dutch mathematician J. G. van der Corput. It is constructed by
-    reversing the base-n representation of the sequence of natural numbers
-    (1, 2, 3, ...).
-    In practice, use Halton sequence instead of Van Der Corput, as it is the
-    same, but generalized to work in multiple dimensions.
-    Parameters
-    ----------
-    idx (int, numpy.ndarray):
-        The index of the sequence. If array is provided, all values in
-        array is returned.
-    number_base : int
-        The numerical base from where to create the samples from.
-    Returns
-    -------
-    float, numpy.ndarray
-        Van der Corput samples.
-    """
-    assert number_base > 1
-
-    idx = np.asarray(idx).flatten()
-    out = np.zeros(len(idx), dtype=float)
-
-    base = float(number_base)
-    active = np.ones(len(idx), dtype=bool)
-    while np.any(active):
-        out[active] += (idx[active] % number_base)/base
-        idx //= number_base
-        base *= number_base
-        active = idx > 0
-    return out
-
-
-def _create_primes(threshold: int):
-    """
-    Generate prime values using sieve of Eratosthenes method.
-    Parameters
-    ----------
-    threshold : int
-        The upper bound for the size of the prime values.
-    Returns
-    ------
-    List
-        All primes from 2 and up to ``threshold``.
-    """
-    if threshold == 2:
-        return [2]
-
-    elif threshold < 2:
-        return []
-
-    numbers = list(range(3, threshold+1, 2))
-    root_of_threshold = threshold ** 0.5
-    half = int((threshold+1)/2-1)
-    idx = 0
-    counter = 3
-    while counter <= root_of_threshold:
-        if numbers[idx]:
-            idy = int((counter*counter-3)/2)
-            numbers[idy] = 0
-            while idy < half:
-                numbers[idy] = 0
-                idy += counter
-        idx += 1
-        counter = 2*idx+3
-    return [2] + [number for number in numbers if number]
diff --git a/entmoot/sampler/hammersly.py b/entmoot/sampler/hammersly.py
deleted file mode 100644
index 605187e..0000000
--- a/entmoot/sampler/hammersly.py
+++ /dev/null
@@ -1,114 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-# -*- coding: utf-8 -*-
-""" Inspired by https://github.com/jonathf/chaospy/blob/master/chaospy/
-distributions/sampler/sequences/hammersley.py
-"""
-
-import numpy as np
-from .halton import Halton
-from ..space import Space
-from .base import InitialPointGenerator
-from sklearn.utils import check_random_state
-
-
-class Hammersly(InitialPointGenerator):
-    """Creates `Hammersley` sequence samples.
-    The Hammersley set is equivalent to the Halton sequence, except for one
-    dimension is replaced with a regular grid. It is not recommended to
-    generate a Hammersley sequence with more than 10 dimension.
-    For ``dim == 1`` the sequence falls back to Van Der Corput sequence.
-    References
-    ----------
-    T-T. Wong, W-S. Luk, and P-A. Heng, "Sampling with Hammersley and Halton
-    Points," Journal of Graphics Tools, vol. 2, no. 2, 1997, pp. 9 - 24.
-    Parameters
-    ----------
-    min_skip : int, default=-1
-        Minimum skipped seed number. When `min_skip != max_skip` and
-        both are > -1, a random number is picked.
-    max_skip : int, default=-1
-        Maximum skipped seed number. When `min_skip != max_skip` and
-        both are > -1, a random number is picked.
-    primes : tuple, default=None
-        The (non-)prime base to calculate values along each axis. If
-        empty, growing prime values starting from 2 will be used.
-    """
-    def __init__(self, min_skip: int = 0, max_skip: int = 0, primes=None):
-        self.primes = primes
-        self.min_skip = min_skip
-        self.max_skip = max_skip
-
-    def generate(self, dimensions, n_samples: int, random_state=None):
-        """Creates samples from Hammersly set.
-        Parameters
-        ----------
-        dimensions : list, shape (n_dims,)
-            List of search space dimensions.
-            Each search dimension can be defined either as
-            - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-              dimensions),
-            - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-              dimensions),
-            - as a list of categories (for `Categorical` dimensions), or
-            - an instance of a `Dimension` object (`Real`, `Integer` or
-              `Categorical`).
-        n_samples : int
-            The order of the Hammersley sequence.
-            Defines the number of samples.
-        random_state : int, RandomState instance, or None (default)
-            Set random state to something other than None for reproducible
-            results.
-        Returns
-        -------
-        np.array, shape=(n_dim, n_samples)
-            Hammersley set.
-        """
-        rng = check_random_state(random_state)
-        halton = Halton(min_skip=self.min_skip, max_skip=self.max_skip,
-                        primes=self.primes)
-        space = Space(dimensions)
-        n_dim = space.n_dims
-        transformer = space.get_transformer()
-        space.set_transformer("normalize")
-        if n_dim == 1:
-            out = halton.generate(dimensions, n_samples,
-                                  random_state=rng)
-        else:
-            out = np.empty((n_dim, n_samples), dtype=float)
-            out[:n_dim - 1] = np.array(halton.generate(
-                [(0., 1.), ] * (n_dim - 1), n_samples,
-                random_state=rng)).T
-
-            out[n_dim - 1] = np.linspace(0, 1, n_samples + 1)[:-1]
-            out = space.inverse_transform(out.T)
-        space.set_transformer(transformer)
-        return out
diff --git a/entmoot/sampler/lhs.py b/entmoot/sampler/lhs.py
deleted file mode 100644
index a4bf056..0000000
--- a/entmoot/sampler/lhs.py
+++ /dev/null
@@ -1,173 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-Lhs functions are inspired by
-https://github.com/clicumu/pyDOE2/blob/
-master/pyDOE2/doe_lhs.py
-"""
-import numpy as np
-from sklearn.utils import check_random_state
-from scipy import spatial
-from entmoot.sampler.base import InitialPointGenerator
-from entmoot.space.space import Space
-from typing import Optional
-
-
-def _random_permute_matrix(h, random_state=None):
-    rng = check_random_state(random_state)
-    h_rand_perm = np.zeros_like(h)
-    samples, n = h.shape
-    for j in range(n):
-        order = rng.permutation(range(samples))
-        h_rand_perm[:, j] = h[order, j]
-    return h_rand_perm
-
-
-class Lhs(InitialPointGenerator):
-    """Latin hypercube sampling
-
-    Parameters
-    ----------
-    lhs_type : str, default='classic'
-        - 'classic' - a small random number is added
-        - 'centered' - points are set uniformly in each interval
-
-    criterion : str or None, default='maximin'
-        When set to None, the LHS is not optimized
-
-        - 'correlation' : optimized LHS by minimizing the correlation
-        - 'maximin' : optimized LHS by maximizing the minimal pdist
-        - 'ratio' : optimized LHS by minimizing the ratio
-          `max(pdist) / min(pdist)`
-
-    iterations : int
-        Defines the number of iterations for optimizing LHS
-    """
-    def __init__(self, lhs_type: str = "classic", criterion: Optional[str] = "maximin", iterations=1000):
-        self.lhs_type = lhs_type
-        self.criterion = criterion
-        self.iterations = iterations
-
-    def generate(self, dimensions, n_samples: int, random_state: Optional[int] = None):
-        """Creates latin hypercube samples.
-
-        Parameters
-        ----------
-        dimensions : list, shape (n_dims,)
-            List of search space dimensions.
-            Each search dimension can be defined either as
-
-            - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-              dimensions),
-            - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-              dimensions),
-            - as a list of categories (for `Categorical` dimensions), or
-            - an instance of a `Dimension` object (`Real`, `Integer` or
-              `Categorical`).
-
-        n_samples : int
-            The order of the LHS sequence. Defines the number of samples.
-        random_state : int, RandomState instance, or None (default)
-            Set random state to something other than None for reproducible
-            results.
-
-        Returns
-        -------
-        np.array, shape=(n_dim, n_samples)
-            LHS set
-        """
-        rng = check_random_state(random_state)
-        space = Space(dimensions)
-        transformer = space.get_transformer()
-        n_dim = space.n_dims
-        space.set_transformer("normalize")
-        if self.criterion is None or n_samples == 1:
-            h = self._lhs_normalized(n_dim, n_samples, rng)
-            h = space.inverse_transform(h)
-            space.set_transformer(transformer)
-            return h
-        else:
-            h_opt = self._lhs_normalized(n_dim, n_samples, rng)
-            h_opt = space.inverse_transform(h_opt)
-            if self.criterion == "correlation":
-                mincorr = np.inf
-                for i in range(self.iterations):
-                    # Generate a random LHS
-                    h = self._lhs_normalized(n_dim, n_samples, rng)
-                    r = np.corrcoef(np.array(h).T)
-                    if len(np.abs(r[r != 1])) > 0 and \
-                            np.max(np.abs(r[r != 1])) < mincorr:
-                        mincorr = np.max(np.abs(r - np.eye(r.shape[0])))
-                        h_opt = h.copy()
-                        h_opt = space.inverse_transform(h_opt)
-            elif self.criterion == "maximin":
-                maxdist = 0
-                # Maximize the minimum distance between points
-                for i in range(self.iterations):
-                    h = self._lhs_normalized(n_dim, n_samples, rng)
-                    d = spatial.distance.pdist(np.array(h), 'euclidean')
-                    if maxdist < np.min(d):
-                        maxdist = np.min(d)
-                        h_opt = h.copy()
-                        h_opt = space.inverse_transform(h_opt)
-            elif self.criterion == "ratio":
-                minratio = np.inf
-
-                # Maximize the minimum distance between points
-                for i in range(self.iterations):
-                    h = self._lhs_normalized(n_dim, n_samples, rng)
-                    p = spatial.distance.pdist(np.array(h), 'euclidean')
-                    if np.min(p) == 0:
-                        ratio = np.max(p) / 1e-8
-                    else:
-                        ratio = np.max(p) / np.min(p)
-                    if minratio > ratio:
-                        minratio = ratio
-                        h_opt = h.copy()
-                        h_opt = space.inverse_transform(h_opt)
-            else:
-                raise ValueError("Wrong criterion."
-                                 "Got {}".format(self.criterion))
-            space.set_transformer(transformer)
-            return h_opt
-
-    def _lhs_normalized(self, n_dim, n_samples: int, random_state: Optional[int]):
-        rng = check_random_state(random_state)
-        x = np.linspace(0, 1, n_samples + 1)
-        u = rng.rand(n_samples, n_dim)
-        h = np.zeros_like(u)
-        if self.lhs_type == "centered":
-            for j in range(n_dim):
-                h[:, j] = np.diff(x) / 2.0 + x[:n_samples]
-        elif self.lhs_type == "classic":
-            for j in range(n_dim):
-                h[:, j] = u[:, j] * np.diff(x) + x[:n_samples]
-        else:
-            raise ValueError("Wrong lhs_type. Got ".format(self.lhs_type))
-        return _random_permute_matrix(h, random_state=rng)
diff --git a/entmoot/sampler/sobol.py b/entmoot/sampler/sobol.py
deleted file mode 100644
index b75463d..0000000
--- a/entmoot/sampler/sobol.py
+++ /dev/null
@@ -1,425 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-  Authors:
-    Original FORTRAN77 version of i4_sobol by Bennett Fox.
-    MATLAB version by John Burkardt.
-    PYTHON version by Corrado Chisari
-
-    Original Python version of is_prime by Corrado Chisari
-
-    Original MATLAB versions of other functions by John Burkardt.
-    PYTHON versions by Corrado Chisari
-
-    Modified Python version by Holger Nahrstaedt
-
-    Original code is available from
-    http://people.sc.fsu.edu/~jburkardt/py_src/sobol/sobol.html
-"""
-
-from __future__ import division
-
-import warnings
-
-import numpy as np
-from .base import InitialPointGenerator
-from ..space import Space
-from sklearn.utils import check_random_state
-from typing import Optional
-
-
-class Sobol(InitialPointGenerator):
-    """Generates a new quasirandom Sobol' vector with each call.
-    Parameters
-    ----------
-    skip : int
-        Skipped seed number.
-    randomize : bool, default=False
-        When set to True, random shift is applied.
-    Notes
-    -----
-    Sobol' sequences [1]_ provide :math:`n=2^m` low discrepancy points in
-    :math:`[0,1)^{dim}`. Scrambling them makes them suitable for singular
-    integrands, provides a means of error estimation, and can improve their
-    rate of convergence.
-    There are many versions of Sobol' sequences depending on their
-    'direction numbers'. Here, the maximum number of dimension is 40.
-    The routine adapts the ideas of Antonov and Saleev [2]_.
-    .. warning::
-       Sobol' sequences are a quadrature rule and they lose their balance
-       properties if one uses a sample size that is not a power of 2, or skips
-       the first point, or thins the sequence [5]_.
-       If :math:`n=2^m` points are not enough then one should take :math:`2^M`
-       points for :math:`M>m`. When scrambling, the number R of independent
-       replicates does not have to be a power of 2.
-       Sobol' sequences are generated to some number :math:`B` of bits. Then
-       after :math:`2^B` points have been generated, the sequence will repeat.
-       Currently :math:`B=30`.
-    References
-    ----------
-    .. [1] I. M. Sobol. The distribution of points in a cube and the accurate
-       evaluation of integrals. Zh. Vychisl. Mat. i Mat. Phys., 7:784-802,
-       1967.
-    .. [2] Antonov, Saleev,
-       USSR Computational Mathematics and Mathematical Physics,
-       Volume 19, 1980, pages 252 - 256.
-    .. [3] Paul Bratley, Bennett Fox,
-       Algorithm 659:
-       Implementing Sobol's Quasirandom Sequence Generator,
-       ACM Transactions on Mathematical Software,
-       Volume 14, Number 1, pages 88-100, 1988.
-    .. [4] Bennett Fox,
-       Algorithm 647:
-       Implementation and Relative Efficiency of Quasirandom
-       Sequence Generators,
-    .. [5] Art B. Owen. On dropping the first Sobol' point. arXiv 2008.08051,
-       2020.
-    """
-    def __init__(self, skip: int = 0, randomize: bool = True):
-
-        if not (skip & (skip - 1) == 0):
-            raise ValueError("The balance properties of Sobol' points require"
-                             " skipping a power of 2.")
-        if skip != 0:
-            warnings.warn(f"{skip} points have been skipped: "
-                          f"{skip} points can be generated before the "
-                          f"sequence repeats.")
-        self.skip = skip
-
-        self.num_generated = 0
-        self.randomize = randomize
-
-        self.dim_max = 40
-        self.log_max = 30
-        self.atmost = 2 ** self.log_max - 1
-        self.lastq = None
-        self.maxcol = None
-        self.poly = None
-        self.recipd = None
-        self.seed_save = -1
-        self.v = np.zeros((self.dim_max, self.log_max))
-        self.dim_num_save = -1
-
-    def init(self, dim_num: int):
-        self.dim_num_save = dim_num
-        self.v = np.zeros((self.dim_max, self.log_max))
-        self.v[0:40, 0] = np.transpose([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
-                                        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
-                                        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
-                                        1, 1, 1, 1])
-
-        self.v[2:40, 1] = np.transpose([1, 3, 1, 3, 1, 3, 3, 1, 3, 1, 3, 1,
-                                        3, 1, 1, 3, 1, 3, 1, 3, 1, 3, 3, 1,
-                                        3, 1, 3, 1, 3, 1, 1, 3, 1, 3, 1, 3,
-                                        1, 3])
-
-        self.v[3:40, 2] = np.transpose([7, 5, 1, 3, 3, 7, 5, 5, 7, 7, 1,
-                                        3, 3, 7, 5, 1, 1, 5, 3, 3, 1, 7, 5,
-                                        1, 3, 3, 7, 5, 1, 1, 5, 7, 7, 5, 1,
-                                        3, 3])
-
-        self.v[5:40, 3] = np.transpose([1, 7,  9,  13, 11, 1, 3,  7,  9,
-                                        5,  13, 13, 11, 3,  15, 5, 3, 15,
-                                        7,  9,  13, 9,  1,  11, 7, 5, 15,
-                                        1,  15, 11, 5,  3,  1,  7,  9])
-
-        self.v[7:40, 4] = np.transpose([9,  3,  27, 15, 29, 21, 23, 19,
-                                        11, 25, 7,  13, 17, 1,  25, 29,
-                                        3,  31, 11, 5,  23, 27, 19, 21,
-                                        5,  1,  17, 13, 7,  15, 9,  31, 9])
-
-        self.v[13:40, 5] = np.transpose([37, 33, 7,  5,  11, 39, 63, 27,
-                                         17, 15, 23, 29, 3,  21, 13, 31,
-                                         25, 9,  49, 33, 19, 29, 11, 19,
-                                         27, 15, 25])
-
-        self.v[19:40, 6] = np.transpose([13, 33, 115, 41, 79, 17, 29, 119,
-                                         75, 73, 105, 7,  59,  65, 21, 3,
-                                         113, 61,  89, 45, 107])
-
-        self.v[37:40, 7] = np.transpose([7, 23, 39])
-
-        #  Set POLY.
-        self.poly = [1, 3, 7, 11, 13, 19, 25, 37, 59, 47, 61, 55, 41, 67, 97,
-                     91, 109, 103, 115, 131, 193, 137, 145, 143, 241, 157,
-                     185, 167, 229, 171, 213, 191, 253, 203, 211, 239, 247,
-                     285, 369, 299]
-
-        #  Find the number of bits in ATMOST.
-        self.maxcol = _bit_hi1(self.atmost)
-
-        #  Initialize row 1 of V.
-        self.v[0, 0:self.maxcol] = 1
-
-        #  Check parameters.
-        if dim_num < 1 or self.dim_max < dim_num:
-            raise ValueError(f'I4_SOBOL - Fatal error!\n'
-                             f'  The spatial dimension DIM_NUM should '
-                             f'satisfy:\n'
-                             f'  1 <= DIM_NUM <= {self.dim_max}\n'
-                             f'  But this input value is DIM_NUM = {dim_num}')
-
-        #  Initialize the remaining rows of V.
-        for i in range(2, dim_num + 1):
-
-            #  The bits of the integer POLY(I) gives the form of polynomial I.
-            #  Find the degree of polynomial I from binary encoding.
-            j = self.poly[i - 1]
-            m = 0
-            j //= 2
-            while j > 0:
-                j //= 2
-                m += 1
-
-            #  Expand this bit pattern to separate components
-            #  of the logical array INCLUD.
-            j = self.poly[i - 1]
-            includ = np.zeros(m)
-            for k in range(m, 0, -1):
-                j2 = j // 2
-                includ[k - 1] = (j != 2 * j2)
-                j = j2
-
-            #  Calculate the remaining elements of row I as explained
-            #  in Bratley and Fox, section 2.
-            for j in range(m + 1, self.maxcol + 1):
-                newv = self.v[i - 1, j - m - 1]
-                p2 = 1
-                for k in range(1, m + 1):
-                    p2 *= 2
-                    if includ[k - 1]:
-                        newv = np.bitwise_xor(
-                            int(newv), int(p2 * self.v[i - 1, j - k - 1]))
-                self.v[i - 1, j - 1] = newv
-        #  Multiply columns of V by appropriate power of 2.
-        p2 = 1
-        for j in range(self.maxcol - 1, 0, -1):
-            p2 *= 2
-            self.v[0:dim_num, j - 1] = self.v[0:dim_num, j - 1] * p2
-
-        #  RECIPD is 1/(common denominator of the elements in V).
-        self.recipd = 1.0 / (2 * p2)
-        self.lastq = np.zeros(dim_num)
-
-    def generate(self, dimensions, n_samples: int, random_state: Optional[int] = None):
-        """Creates samples from Sobol' set.
-        Parameters
-        ----------
-        dimensions : list, shape (n_dims,)
-            List of search space dimensions.
-            Each search dimension can be defined either as
-            - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-              dimensions),
-            - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-              dimensions),
-            - as a list of categories (for `Categorical` dimensions), or
-            - an instance of a `Dimension` object (`Real`, `Integer` or
-              `Categorical`).
-        n_samples : int
-            The order of the Sobol' sequence. Defines the number of samples.
-        random_state : int, RandomState instance, or None (default)
-            Set random state to something other than None for reproducible
-            results.
-        Returns
-        -------
-        sample : array_like (n_samples, dim)
-            Sobol' set.
-        """
-        total_n_samples = self.num_generated + n_samples
-        if not (total_n_samples & (total_n_samples - 1) == 0):
-            warnings.warn("The balance properties of Sobol' points require "
-                          "n to be a power of 2. {0} points have been "
-                          "previously generated, then: n={0}+{1}={2}. "
-                          .format(self.num_generated, n_samples,
-                                  total_n_samples))
-        if self.skip != 0 and total_n_samples > self.skip:
-            raise ValueError(f"{self.skip} points have been skipped: "
-                             f"generating "
-                             f"{n_samples} more points would cause the "
-                             f"sequence to repeat.")
-
-        rng = check_random_state(random_state)
-        space = Space(dimensions)
-        n_dim = space.n_dims
-        transformer = space.get_transformer()
-        space.set_transformer("normalize")
-        r = np.full((n_samples, n_dim), np.nan)
-
-        seed = self.skip
-        for j in range(n_samples):
-            r[j, 0:n_dim], seed = self._sobol(n_dim, seed)
-
-        if self.randomize:
-            r = _random_shift(r, rng)
-
-        r = space.inverse_transform(r)
-        space.set_transformer(transformer)
-
-        self.num_generated += n_samples
-
-        return r
-
-    def _sobol(self, dim_num: int, seed: int):
-        """Generates a new quasirandom Sobol' vector with each call.
-        Parameters
-        ----------
-        dim_num : int
-          Number of spatial dimensions.
-          `dim_num` must satisfy 1 <= DIM_NUM <= 40.
-        seed : int
-          the `seed` for the sequence.
-          This is essentially the index in the sequence of the quasirandom
-          value to be generated. On output, `seed` has been set to the
-          appropriate next value, usually simply `seed`+1.
-          If `seed` is less than 0 on input, it is treated as though it were 0.
-          An input value of 0 requests the first (0-th) element of
-          the sequence.
-        Returns
-        -------
-        vector, seed : np.array (n_dim,), int
-            The next quasirandom vector and the seed of its next vector.
-        """
-        #  Things to do only if the dimension changed.
-        if dim_num != self.dim_num_save:
-            self.init(dim_num)
-
-        seed = int(np.floor(seed))
-
-        if seed < 0:
-            seed = 0
-
-        pos_lo0 = 1
-        if seed == 0:
-            self.lastq = np.zeros(dim_num)
-
-        elif seed == self.seed_save + 1:
-
-            #  Find the position of the right-hand zero in SEED.
-            pos_lo0 = _bit_lo0(seed)
-
-        elif seed <= self.seed_save:
-
-            self.seed_save = 0
-            self.lastq = np.zeros(dim_num)
-
-            for seed_temp in range(int(self.seed_save), int(seed)):
-                pos_lo0 = _bit_lo0(seed_temp)
-                for i in range(1, dim_num + 1):
-                    self.lastq[i - 1] = np.bitwise_xor(
-                        int(self.lastq[i - 1]),
-                        int(self.v[i - 1, pos_lo0 - 1]))
-
-            pos_lo0 = _bit_lo0(seed)
-
-        elif self.seed_save + 1 < seed:
-
-            for seed_temp in range(int(self.seed_save + 1), int(seed)):
-                pos_lo0 = _bit_lo0(seed_temp)
-                for i in range(1, dim_num + 1):
-                    self.lastq[i - 1] = np.bitwise_xor(
-                        int(self.lastq[i - 1]),
-                        int(self.v[i - 1, pos_lo0 - 1]))
-
-            pos_lo0 = _bit_lo0(seed)
-
-        #  Check that the user is not calling too many times!
-        if self.maxcol < pos_lo0:
-            raise ValueError(f'I4_SOBOL - Fatal error!\n'
-                             f' Too many calls!\n'
-                             f' MAXCOL = {self.maxcol}\n'
-                             f' L =      {pos_lo0}\n')
-
-        #  Calculate the new components of QUASI.
-        quasi = np.zeros(dim_num)
-        for i in range(1, dim_num + 1):
-            quasi[i - 1] = self.lastq[i - 1] * self.recipd
-            self.lastq[i - 1] = np.bitwise_xor(
-                int(self.lastq[i - 1]), int(self.v[i - 1, pos_lo0 - 1]))
-
-        self.seed_save = seed
-        seed += 1
-
-        return [quasi, seed]
-
-
-def _bit_hi1(n: int):
-    """Returns the position of the high 1 bit base 2 in an integer.
-    Parameters
-    ----------
-    n : int
-        Input, should be positive.
-    """
-    bin_repr = np.binary_repr(n)
-    most_left_one = bin_repr.find('1')
-    if most_left_one == -1:
-        return 0
-    else:
-        return len(bin_repr) - most_left_one
-
-
-def _bit_lo0(n: int):
-    """Returns the position of the low 0 bit base 2 in an integer.
-    Parameters
-    ----------
-    n : int
-        Input, should be positive.
-    """
-    bin_repr = np.binary_repr(n)
-    most_right_zero = bin_repr[::-1].find('0')
-    if most_right_zero == -1:
-        most_right_zero = len(bin_repr)
-    return most_right_zero + 1
-
-
-def _random_shift(dm: np.array, random_state: Optional[int] = None):
-    """Random shifting of a vector.
-    Randomization of the quasi-MC samples can be achieved in the easiest manner
-    by random shift (or the Cranley-Patterson rotation).
-    References
-    -----------
-    .. [1] C. Lemieux, "Monte Carlo and Quasi-Monte Carlo Sampling," Springer
-       Series in Statistics 692, Springer Science+Business Media, New York,
-       2009
-    Parameters
-    ----------
-    dm : array, shape(n, d)
-        Input matrix.
-    random_state : int, RandomState instance, or None (default)
-        Set random state to something other than None for reproducible
-        results.
-    Returns
-    -------
-    dm :  array, shape(n, d)
-        Randomized Sobol' design matrix.
-    """
-    rng = check_random_state(random_state)
-    # Generate random shift matrix from uniform distribution
-    shift = np.repeat(rng.rand(1, dm.shape[1]), dm.shape[0], axis=0)
-    # Return the shifted Sobol' design
-    return (dm + shift) % 1
diff --git a/entmoot/space/__init__.py b/entmoot/space/__init__.py
deleted file mode 100644
index a969a45..0000000
--- a/entmoot/space/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from .space import *
-
diff --git a/entmoot/space/space.py b/entmoot/space/space.py
deleted file mode 100644
index a1fc589..0000000
--- a/entmoot/space/space.py
+++ /dev/null
@@ -1,1219 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-import numbers
-import numpy as np
-import yaml
-
-from scipy.stats.distributions import randint
-from scipy.stats.distributions import rv_discrete
-from scipy.stats.distributions import uniform
-
-from sklearn.utils import check_random_state
-from sklearn.utils.fixes import sp_version
-
-from entmoot.space.transformers import CategoricalEncoder
-from entmoot.space.transformers import StringEncoder
-from entmoot.space.transformers import LabelEncoder
-from entmoot.space.transformers import Normalize
-from entmoot.space.transformers import Identity
-from entmoot.space.transformers import LogN
-from entmoot.space.transformers import Pipeline
-from entmoot.space.transformers import Logit
-from entmoot.space.transformers import Bilog
-
-# helper class to be able to print [1, ..., 4] instead of [1, '...', 4]
-class _Ellipsis:
-    def __repr__(self):
-        return '...'
-
-
-def _transpose_list_array(x):
-    """Transposes a list matrix
-    """
-
-    n_dims = len(x)
-    assert n_dims > 0
-    n_samples = len(x[0])
-    rows = [None] * n_samples
-    for i in range(n_samples):
-        r = [None] * n_dims
-        for j in range(n_dims):
-            r[j] = x[j][i]
-        rows[i] = r
-    return rows
-
-
-def check_dimension(dimension, transform=None):
-    """Turn a provided dimension description into a dimension object.
-
-    Checks that the provided dimension falls into one of the
-    supported types. For a list of supported types, look at
-    the documentation of ``dimension`` below.
-
-    If ``dimension`` is already a ``Dimension`` instance, return it.
-
-    Parameters
-    ----------
-    dimension : Dimension
-        Search space Dimension.
-        Each search dimension can be defined either as
-
-        - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-          dimensions),
-        - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-          dimensions),
-        - as a list of categories (for `Categorical` dimensions), or
-        - an instance of a `Dimension` object (`Real`, `Integer` or
-          `Categorical`).
-
-    transform : "identity", "normalize", "string", "label", "onehot" optional
-        - For `Categorical` dimensions, the following transformations are
-          supported.
-
-          - "onehot" (default) one-hot transformation of the original space.
-          - "label" integer transformation of the original space
-          - "string" string transformation of the original space.
-          - "identity" same as the original space.
-
-        - For `Real` and `Integer` dimensions, the following transformations
-          are supported.
-
-          - "identity", (default) the transformed space is the same as the
-            original space.
-          - "normalize", the transformed space is scaled to be between 0 and 1.
-
-    Returns
-    -------
-    dimension : Dimension
-        Dimension instance.
-    """
-    if isinstance(dimension, Dimension):
-        return dimension
-
-    if not isinstance(dimension, (list, tuple, np.ndarray)):
-        raise ValueError("Dimension has to be a list or tuple.")
-
-    # A `Dimension` described by a single value is assumed to be
-    # a `Categorical` dimension. This can be used in `BayesSearchCV`
-    # to define subspaces that fix one value, e.g. to choose the
-    # model type, see "sklearn-gridsearchcv-replacement.py"
-    # for examples.
-    if len(dimension) == 1:
-        return Categorical(dimension, transform=transform)
-
-    if len(dimension) == 2:
-        if any([isinstance(d, (str, bool)) or isinstance(d, np.bool_)
-                for d in dimension]):
-            return Categorical(dimension, transform=transform)
-        elif all([isinstance(dim, numbers.Integral) for dim in dimension]):
-            return Integer(*dimension, transform=transform)
-        elif any([isinstance(dim, numbers.Real) for dim in dimension]):
-            return Real(*dimension, transform=transform)
-        else:
-            raise ValueError("Invalid dimension {}. Read the documentation for"
-                             " supported types.".format(dimension))
-
-    if len(dimension) == 3:
-        if (any([isinstance(dim, int) for dim in dimension[:2]]) and
-                dimension[2] in ["uniform", "log-uniform", "bilog-uniform"]):
-            return Integer(*dimension, transform=transform)
-        elif (any([isinstance(dim, (float, int)) for dim in dimension[:2]]) and
-              dimension[2] in ["uniform", "log-uniform", "logit-uniform", "bilog-uniform"]):
-            return Real(*dimension, transform=transform)
-        else:
-            return Categorical(dimension, transform=transform)
-
-    if len(dimension) == 4:
-        # covers the case if log-uniform is picked and base provided
-        if (any([isinstance(dim, int) for dim in dimension[:2]]) and
-                dimension[2] == "log-uniform" and isinstance(dimension[3],
-                                                             int)):
-            return Integer(*dimension, transform=transform)
-        elif (any([isinstance(dim, (float, int)) for dim in dimension[:2]]) and
-              dimension[2] == "log-uniform" and isinstance(dimension[3], int)):
-            return Real(*dimension, transform=transform)
-
-    if len(dimension) > 3:
-        return Categorical(dimension, transform=transform)
-
-    raise ValueError("Invalid dimension {}. Read the documentation for "
-                     "supported types.".format(dimension))
-
-
-class Dimension(object):
-    """Base class for search space dimensions."""
-
-    prior = None
-
-    def rvs(self, n_samples=1, random_state=None):
-        """Draw random samples.
-
-        Parameters
-        ----------
-        n_samples : int or None
-            The number of samples to be drawn.
-
-        random_state : int, RandomState instance, or None (default)
-            Set random state to something other than None for reproducible
-            results.
-        """
-        rng = check_random_state(random_state)
-        samples = self._rvs.rvs(size=n_samples, random_state=rng)
-        return self.inverse_transform(samples)
-
-    def transform(self, X):
-        """Transform samples form the original space to a warped space."""
-        return self.transformer.transform(X)
-
-    def inverse_transform(self, Xt):
-        """Inverse transform samples from the warped space back into the
-           original space.
-        """
-        return self.transformer.inverse_transform(Xt)
-
-    def set_transformer(self):
-        raise NotImplementedError
-
-    @property
-    def size(self):
-        return 1
-
-    @property
-    def transformed_size(self):
-        return 1
-
-    @property
-    def bounds(self):
-        raise NotImplementedError
-
-    @property
-    def is_constant(self):
-        raise NotImplementedError
-
-    @property
-    def transformed_bounds(self):
-        raise NotImplementedError
-
-    @property
-    def name(self):
-        return self._name
-
-    @name.setter
-    def name(self, value):
-        if isinstance(value, str) or value is None:
-            self._name = value
-        else:
-            raise ValueError("Dimension's name must be either string or None.")
-
-
-def _uniform_inclusive(loc=0.0, scale=1.0):
-    # like scipy.stats.distributions but inclusive of `high`
-    # XXX scale + 1. might not actually be a float after scale if
-    # XXX scale is very large.
-    return uniform(loc=loc, scale=np.nextafter(scale, scale + 1.))
-
-
-class Real(Dimension):
-    """Search space dimension that can take on any real value.
-
-    Parameters
-    ----------
-    low : float
-        Lower bound (inclusive).
-
-    high : float
-        Upper bound (inclusive).
-
-    prior : "uniform" or "log-uniform", default="uniform"
-        Distribution to use when sampling random points for this dimension.
-
-        - If `"uniform"`, points are sampled uniformly between the lower
-          and upper bounds.
-        - If `"log-uniform"`, points are sampled uniformly between
-          `log(lower, base)` and `log(upper, base)` where log
-          has base `base`.
-
-    base : int
-        The logarithmic base to use for a log-uniform prior.
-        - Default 10, otherwise commonly 2.
-
-    transform : "identity", "normalize", optional
-        The following transformations are supported.
-
-        - "identity", (default) the transformed space is the same as the
-          original space.
-        - "normalize", the transformed space is scaled to be between
-          0 and 1.
-
-    name : str or None
-        Name associated with the dimension, e.g., "learning rate".
-
-    dtype : str or dtype, default=np.float
-        float type which will be used in inverse_transform,
-        can be float.
-
-    """
-    def __init__(self, low, high, prior="uniform", base=10, transform=None,
-                 name=None, dtype=np.float):
-        if high <= low:
-            raise ValueError("the lower bound {} has to be less than the"
-                             " upper bound {}".format(low, high))
-        self.low = low
-        self.high = high
-        self.prior = prior
-        self.base = base
-        self.log_base = np.log10(base)
-        self.name = name
-        self.dtype = dtype
-        self._rvs = None
-        self.transformer = None
-        self.transform_ = transform
-        if isinstance(self.dtype, str) and self.dtype\
-                not in ['float', 'float16', 'float32', 'float64']:
-            raise ValueError("dtype must be 'float', 'float16', 'float32'"
-                             "or 'float64'"
-                             " got {}".format(self.dtype))
-        elif isinstance(self.dtype, type) and self.dtype\
-                not in [float, np.float, np.float16, np.float32, np.float64]:
-            raise ValueError("dtype must be float, np.float"
-                             " got {}".format(self.dtype))
-
-        if transform is None:
-            transform = "identity"
-        self.set_transformer(transform)
-
-    def set_transformer(self, transform="identitiy"):
-        """Define rvs and transformer spaces.
-
-        Parameters
-        ----------
-        transform : str
-           Can be 'normalize' or 'identity'
-
-        """
-        self.transform_ = transform
-
-        if self.transform_ not in ["normalize", "identity"]:
-            raise ValueError("transform should be 'normalize' or 'identity'"
-                             " got {}".format(self.transform_))
-
-        # XXX: The _rvs is for sampling in the transformed space.
-        # The rvs on Dimension calls inverse_transform on the points sampled
-        # using _rvs
-        if self.transform_ == "normalize":
-            # set upper bound to next float after 1. to make the numbers
-            # inclusive of upper edge
-            self._rvs = _uniform_inclusive(0., 1.)
-            if self.prior == "uniform":
-                self.transformer = Pipeline(
-                    [Identity(), Normalize(self.low, self.high)])
-            elif self.prior == "log-uniform":
-                self.transformer = Pipeline(
-                    [LogN(self.base),
-                     Normalize(np.log10(self.low) / self.log_base,
-                               np.log10(self.high) / self.log_base)]
-                )
-            elif self.prior == "logit-uniform":
-                self.transformer = \
-                    [
-                        Logit(), 
-                        Normalize(
-                            Logit.transform(self.low), 
-                            Logit.transform(self.high)
-                        )
-                    ]
-            elif self.prior == "bilog-uniform":
-                self.transformer = \
-                    [
-                        Bilog(), 
-                        Normalize(
-                            Bilog.transform(self.low), 
-                            Bilog.transform(self.high)
-                        )
-                    ]
-        else:
-            if self.prior == "uniform":
-                self._rvs = _uniform_inclusive(self.low, self.high - self.low)
-                self.transformer = Identity()
-            elif self.prior == "log-uniform":
-                self._rvs = _uniform_inclusive(
-                    np.log10(self.low) / self.log_base,
-                    np.log10(self.high) / self.log_base -
-                    np.log10(self.low) / self.log_base)
-                self.transformer = LogN(self.base)
-            elif self.prior == "logit-uniform":
-                self._rvs = _uniform_inclusive(
-                    Logit.transform(self.low),
-                    Logit.transform(self.high) -
-                    Logit.transform(self.low))
-                self.transformer = Logit()
-            elif self.prior == "bilog-uniform":
-                self._rvs = _uniform_inclusive(
-                    Bilog.transform(self.low),
-                    Bilog.transform(self.high) -
-                    Bilog.transform(self.low))
-                self.transformer = Bilog()  
-
-    def __eq__(self, other):
-        return (type(self) is type(other) and
-                np.allclose([self.low], [other.low]) and
-                np.allclose([self.high], [other.high]) and
-                self.prior == other.prior and
-                self.transform_ == other.transform_)
-
-    def __repr__(self):
-        return "Real(low={}, high={}, prior='{}', transform='{}')".format(
-            self.low, self.high, self.prior, self.transform_)
-
-    def inverse_transform(self, Xt):
-        """Inverse transform samples from the warped space back into the
-           original space.
-        """
-
-        inv_transform = super(Real, self).inverse_transform(Xt)
-        if isinstance(inv_transform, list):
-            inv_transform = np.array(inv_transform)
-        inv_transform = np.clip(inv_transform,
-                                self.low, self.high).astype(self.dtype)
-        if self.dtype == float or self.dtype == 'float':
-            # necessary, otherwise the type is converted to a numpy type
-            return getattr(inv_transform, "tolist", lambda: value)()
-        else:
-            return inv_transform
-
-    @property
-    def bounds(self):
-        return (self.low, self.high)
-
-    @property
-    def is_constant(self):
-        return self.low == self.high
-
-    def __contains__(self, point):
-        if isinstance(point, list):
-            point = np.array(point)
-        return self.low <= point <= self.high
-
-    @property
-    def transformed_bounds(self):
-        if self.transform_ == "normalize":
-            return 0.0, 1.0
-        else:
-            if self.prior == "uniform":
-                return self.low, self.high
-            elif self.prior == "log-uniform":
-                return np.log10(self.low), np.log10(self.high)
-            elif self.prior == "logit-uniform":
-                return Logit.transform(self.low), Logit.transform(self.high)
-            elif self.prior == "bilog-uniform":
-                return Bilog.transform(self.low), Bilog.transform(self.high) 
-
-    def distance(self, a, b):
-        """Compute distance between point `a` and `b`.
-
-        Parameters
-        ----------
-        a : float
-            First point.
-
-        b : float
-            Second point.
-        """
-        if not (a in self and b in self):
-            raise RuntimeError("Can only compute distance for values within "
-                               "the space, not %s and %s." % (a, b))
-        return abs(a - b)
-
-
-class Integer(Dimension):
-    """Search space dimension that can take on integer values.
-
-    Parameters
-    ----------
-    low : int
-        Lower bound (inclusive).
-
-    high : int
-        Upper bound (inclusive).
-
-    prior : "uniform" or "log-uniform", default="uniform"
-        Distribution to use when sampling random integers for
-        this dimension.
-
-        - If `"uniform"`, integers are sampled uniformly between the lower
-          and upper bounds.
-        - If `"log-uniform"`, integers are sampled uniformly between
-          `log(lower, base)` and `log(upper, base)` where log
-          has base `base`.
-
-    base : int
-        The logarithmic base to use for a log-uniform prior.
-
-        - Default 10, otherwise commonly 2.
-
-    transform : "identity", "normalize", optional
-        The following transformations are supported.
-
-        - "identity", (default) the transformed space is the same as the
-          original space.
-        - "normalize", the transformed space is scaled to be between
-          0 and 1.
-
-    name : str or None
-        Name associated with dimension, e.g., "number of trees".
-
-    dtype : str or dtype, default=np.int64
-        integer type which will be used in inverse_transform,
-        can be int, np.int16, np.uint32, np.int32, np.int64 (default).
-        When set to int, `inverse_transform` returns a list instead of
-        a numpy array
-
-    """
-    def __init__(self, low, high, prior="uniform", base=10, transform=None,
-                 name=None, dtype=np.int64):
-        if high <= low:
-            raise ValueError("the lower bound {} has to be less than the"
-                             " upper bound {}".format(low, high))
-        self.low = low
-        self.high = high
-        self.prior = prior
-        self.base = base
-        self.log_base = np.log10(base)
-        self.name = name
-        self.dtype = dtype
-        self.transform_ = transform
-        self._rvs = None
-        self.transformer = None
-
-        if isinstance(self.dtype, str) and self.dtype\
-            not in ['int', 'int8', 'int16', 'int32', 'int64',
-                    'uint8', 'uint16', 'uint32', 'uint64']:
-            raise ValueError("dtype must be 'int', 'int8', 'int16',"
-                             "'int32', 'int64', 'uint8',"
-                             "'uint16', 'uint32', or"
-                             "'uint64', but got {}".format(self.dtype))
-        elif isinstance(self.dtype, type) and self.dtype\
-                not in [int, np.int8, np.int16, np.int32, np.int64,
-                        np.uint8, np.uint16, np.uint32, np.uint64]:
-            raise ValueError("dtype must be 'int', 'np.int8', 'np.int16',"
-                             "'np.int32', 'np.int64', 'np.uint8',"
-                             "'np.uint16', 'np.uint32', or"
-                             "'np.uint64', but got {}".format(self.dtype))
-
-        if transform is None:
-            transform = "identity"
-        self.set_transformer(transform)
-
-    def set_transformer(self, transform="identitiy"):
-        """Define _rvs and transformer spaces.
-
-        Parameters
-        ----------
-        transform : str
-           Can be 'normalize' or 'identity'
-
-        """
-        self.transform_ = transform
-
-        if transform not in ["normalize", "identity"]:
-            raise ValueError("transform should be 'normalize' or 'identity'"
-                             " got {}".format(self.transform_))
-
-        if self.transform_ == "normalize":
-            self._rvs = _uniform_inclusive(0.0, 1.0)
-            if self.prior == "uniform":
-                self.transformer = Pipeline(
-                    [Identity(), Normalize(self.low, self.high, is_int=True)])
-            elif self.prior == "log-uniform":
-                self.transformer = Pipeline(
-                    [LogN(self.base),
-                     Normalize(np.log10(self.low) / self.log_base,
-                               np.log10(self.high) / self.log_base)]
-                )
-            elif self.prior == "bilog-uniform":
-                self.transformer = \
-                    [
-                        Bilog(), 
-                        Normalize(
-                            Bilog.transform(self.low), 
-                            Bilog.transform(self.high)
-                        )
-                    ]
-        else:
-            if self.prior == "uniform":
-                self._rvs = randint(self.low, self.high + 1)
-                self.transformer = Identity()
-            elif self.prior == "log-uniform":
-                self._rvs = _uniform_inclusive(
-                    np.log10(self.low) / self.log_base,
-                    np.log10(self.high) / self.log_base -
-                    np.log10(self.low) / self.log_base)
-                self.transformer = LogN(self.base)
-            elif self.prior == "bilog-uniform":
-                self._rvs = _uniform_inclusive(
-                    Bilog.transform(self.low),
-                    Bilog.transform(self.high) -
-                    Bilog.transform(self.low))
-                self.transformer = Bilog()  
-
-    def __eq__(self, other):
-        return (type(self) is type(other) and
-                np.allclose([self.low], [other.low]) and
-                np.allclose([self.high], [other.high]))
-
-    def __repr__(self):
-        return "Integer(low={}, high={}, prior='{}', transform='{}')".format(
-            self.low, self.high, self.prior, self.transform_)
-
-    def inverse_transform(self, Xt):
-        """Inverse transform samples from the warped space back into the
-           original space.
-        """
-        # The concatenation of all transformed dimensions makes Xt to be
-        # of type float, hence the required cast back to int.
-        inv_transform = super(Integer, self).inverse_transform(Xt)
-        if isinstance(inv_transform, list):
-            inv_transform = np.array(inv_transform)
-        inv_transform = np.clip(inv_transform,
-                                self.low, self.high)
-        if self.dtype == int or self.dtype == 'int':
-            # necessary, otherwise the type is converted to a numpy type
-            return getattr(np.round(inv_transform).astype(self.dtype),
-                           "tolist", lambda: value)()
-        else:
-            return np.round(inv_transform).astype(self.dtype)
-
-    @property
-    def bounds(self):
-        return (self.low, self.high)
-
-    @property
-    def is_constant(self):
-        return self.low == self.high
-
-    def __contains__(self, point):
-        if isinstance(point, list):
-            point = np.array(point)
-        return self.low <= point <= self.high
-
-    @property
-    def transformed_bounds(self):
-        if self.transform_ == "normalize":
-            return 0., 1.
-        else:
-            if self.prior == "uniform":
-                return self.low, self.high
-            elif self.prior == "log-uniform":
-                return np.log10(self.low), np.log10(self.high)
-            elif self.prior == "bilog-uniform":
-                return Bilog.transform(self.low), Bilog.transform(self.high) 
-
-    def distance(self, a, b):
-        """Compute distance between point `a` and `b`.
-
-        Parameters
-        ----------
-        a : int
-            First point.
-
-        b : int
-            Second point.
-        """
-        if not (a in self and b in self):
-            raise RuntimeError("Can only compute distance for values within "
-                               "the space, not %s and %s." % (a, b))
-        return abs(a - b)
-
-
-class Categorical(Dimension):
-    """Search space dimension that can take on categorical values.
-
-    Parameters
-    ----------
-    categories : list, shape=(n_categories,)
-        Sequence of possible categories.
-
-    prior : list, shape=(categories,), default=None
-        Prior probabilities for each category. By default all categories
-        are equally likely.
-
-    transform : "onehot", "string", "identity", "label", default="onehot"
-
-        - "identity", the transformed space is the same as the original
-          space.
-        - "string",  the transformed space is a string encoded
-          representation of the original space.
-        - "label", the transformed space is a label encoded
-          representation (integer) of the original space.
-        - "onehot", the transformed space is a one-hot encoded
-          representation of the original space.
-
-    name : str or None
-        Name associated with dimension, e.g., "colors".
-
-    """
-    def __init__(self, categories, prior=None, transform=None, name=None):
-        self.categories = tuple(categories)
-
-        self.name = name
-
-        if transform is None:
-            transform = "label"
-        self.transform_ = transform
-        self.transformer = None
-        self._rvs = None
-        self.prior = prior
-
-        if prior is None:
-            self.prior_ = np.tile(1. / len(self.categories),
-                                  len(self.categories))
-        else:
-            self.prior_ = prior
-        self.set_transformer(transform)
-
-    def set_transformer(self, transform="label"):
-        """Define _rvs and transformer spaces.
-
-        Parameters
-        ----------
-        transform : str
-           Can be 'normalize', 'onehot', 'string', 'label', or 'identity'
-
-        """
-        self.transform_ = transform
-        if transform not in ["identity", "onehot", "string", "normalize",
-                             "label"]:
-            raise ValueError("Expected transform to be 'identity', 'string',"
-                             "'label' or 'onehot' got {}".format(transform))
-        if transform == "onehot":
-            self.transformer = CategoricalEncoder()
-            self.transformer.fit(self.categories)
-        elif transform == "string":
-            self.transformer = StringEncoder()
-            self.transformer.fit(self.categories)
-        elif transform == "label":
-            self.transformer = LabelEncoder()
-            self.transformer.fit(self.categories)
-        elif transform == "normalize":
-            self.transformer = Pipeline(
-                [LabelEncoder(list(self.categories)),
-                 Normalize(0, len(self.categories) - 1, is_int=True)])
-        else:
-            self.transformer = Identity()
-            self.transformer.fit(self.categories)
-        if transform == "normalize":
-            self._rvs = _uniform_inclusive(0.0, 1.0)
-        else:
-            # XXX check that sum(prior) == 1
-            self._rvs = rv_discrete(
-                values=(range(len(self.categories)), self.prior_)
-            )
-
-    def __eq__(self, other):
-        return (type(self) is type(other) and
-                self.categories == other.categories and
-                np.allclose(self.prior_, other.prior_))
-
-    def __repr__(self):
-        if len(self.categories) > 7:
-            cats = self.categories[:3] + (_Ellipsis(),) + self.categories[-3:]
-        else:
-            cats = self.categories
-
-        if self.prior is not None and len(self.prior) > 7:
-            prior = self.prior[:3] + [_Ellipsis()] + self.prior[-3:]
-        else:
-            prior = self.prior
-
-        return "Categorical(categories={}, prior={})".format(cats, prior)
-
-    def inverse_transform(self, Xt):
-        """Inverse transform samples from the warped space back into the
-           original space.
-        """
-        # The concatenation of all transformed dimensions makes Xt to be
-        # of type float, hence the required cast back to int.
-        inv_transform = super(Categorical, self).inverse_transform(Xt)
-        if isinstance(inv_transform, list):
-            inv_transform = np.array(inv_transform)
-        return inv_transform
-
-    def rvs(self, n_samples=None, random_state=None):
-        choices = self._rvs.rvs(size=n_samples, random_state=random_state)
-
-        if isinstance(choices, numbers.Integral):
-            return self.categories[choices]
-        elif self.transform_ == "normalize" and isinstance(choices, float):
-            return self.inverse_transform([(choices)])
-        elif self.transform_ == "normalize":
-            return self.inverse_transform(list(choices))
-        else:
-            return [self.categories[c] for c in choices]
-
-    @property
-    def transformed_size(self):
-        if self.transform_ == "onehot":
-            size = len(self.categories)
-            # when len(categories) == 2, CategoricalEncoder outputs a
-            # single value
-            return size if size != 2 else 1
-        return 1
-
-    @property
-    def bounds(self):
-        return self.categories
-
-    @property
-    def is_constant(self):
-        return len(self.categories) <= 1
-
-    def __contains__(self, point):
-        return point in self.categories
-
-    @property
-    def transformed_bounds(self):
-        if self.transformed_size == 1:
-            return 0.0, 1.0
-        else:
-            return [(0.0, 1.0) for i in range(self.transformed_size)]
-
-    def distance(self, a, b):
-        """Compute distance between category `a` and `b`.
-
-        As categories have no order the distance between two points is one
-        if a != b and zero otherwise.
-
-        Parameters
-        ----------
-        a : category
-            First category.
-
-        b : category
-            Second category.
-        """
-        if not (a in self and b in self):
-            raise RuntimeError("Can only compute distance for values within"
-                               " the space, not {} and {}.".format(a, b))
-        return 1 if a != b else 0
-
-
-class Space(object):
-    """Initialize a search space from given specifications.
-
-    Parameters
-    ----------
-    dimensions : list, shape=(n_dims,)
-        List of search space dimensions.
-        Each search dimension can be defined either as
-
-        - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-          dimensions),
-        - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-          dimensions),
-        - as a list of categories (for `Categorical` dimensions), or
-        - an instance of a `Dimension` object (`Real`, `Integer` or
-          `Categorical`).
-
-        .. note::
-            The upper and lower bounds are inclusive for `Integer`
-            dimensions.
-    """
-    def __init__(self, dimensions):
-        self.dimensions = [check_dimension(dim) for dim in dimensions]
-
-    def __eq__(self, other):
-        return all([a == b for a, b in zip(self.dimensions, other.dimensions)])
-
-    def __repr__(self):
-        if len(self.dimensions) > 31:
-            dims = self.dimensions[:15] + [_Ellipsis()] + self.dimensions[-15:]
-        else:
-            dims = self.dimensions
-        return "Space([{}])".format(',\n       '.join(map(str, dims)))
-
-    def __iter__(self):
-        return iter(self.dimensions)
-
-    @property
-    def dimension_names(self):
-        """
-        Names of all the dimensions in the search-space.
-        """
-        index = 0
-        names = []
-        for dim in self.dimensions:
-            if dim.name is None:
-                names.append("X_%d" % index)
-            else:
-                names.append(dim.name)
-            index += 1
-        return names
-
-    @property
-    def is_real(self):
-        """
-        Returns true if all dimensions are Real
-        """
-        return all([isinstance(dim, Real) for dim in self.dimensions])
-
-    @classmethod
-    def from_yaml(cls, yml_path, namespace=None):
-        """Create Space from yaml configuration file
-
-        Parameters
-        ----------
-        yml_path : str
-            Full path to yaml configuration file, example YaML below:
-            Space:
-
-            - Integer:
-              low: -5
-              high: 5
-            - Categorical:
-              categories:
-              - a
-              - b
-            - Real:
-              low: 1.0
-              high: 5.0
-              prior: log-uniform
-
-        namespace : str, default=None
-           Namespace within configuration file to use, will use first
-           namespace if not provided
-
-        Returns
-        -------
-        space : Space
-           Instantiated Space object
-
-        """
-        with open(yml_path, 'rb') as f:
-            config = yaml.safe_load(f)
-
-        dimension_classes = {'real': Real,
-                             'integer': Integer,
-                             'categorical': Categorical}
-
-        # Extract space options for configuration file
-        if isinstance(config, dict):
-            if namespace is None:
-                options = next(iter(config.values()))
-            else:
-                options = config[namespace]
-        elif isinstance(config, list):
-            options = config
-        else:
-            raise TypeError('YaML does not specify a list or dictionary')
-
-        # Populate list with Dimension objects
-        dimensions = []
-        for option in options:
-            key = next(iter(option.keys()))
-            # Make configuration case insensitive
-            dimension_class = key.lower()
-            values = {k.lower(): v for k, v in option[key].items()}
-            if dimension_class in dimension_classes:
-                # Instantiate Dimension subclass and add it to the list
-                dimension = dimension_classes[dimension_class](**values)
-                dimensions.append(dimension)
-
-        space = cls(dimensions=dimensions)
-
-        return space
-
-    def rvs(self, n_samples=1, random_state=None):
-        """Draw random samples.
-
-        The samples are in the original space. They need to be transformed
-        before being passed to a model or minimizer by `space.transform()`.
-
-        Parameters
-        ----------
-        n_samples : int, default=1
-            Number of samples to be drawn from the space.
-
-        random_state : int, RandomState instance, or None (default)
-            Set random state to something other than None for reproducible
-            results.
-
-        Returns
-        -------
-        points : list of lists, shape=(n_points, n_dims)
-           Points sampled from the space.
-        """
-        rng = check_random_state(random_state)
-
-        # Draw
-        columns = []
-
-        for dim in self.dimensions:
-            columns.append(dim.rvs(n_samples=n_samples, random_state=rng))
-        
-        # Transpose
-        return _transpose_list_array(columns)
-
-    def set_transformer(self, transform):
-        """Sets the transformer of all dimension objects to `transform`
-
-        Parameters
-        ----------
-        transform : str or list of str
-           Sets all transformer,, when `transform`  is a string.
-           Otherwise, transform must be a list with strings with
-           the same length as `dimensions`
-        """
-        # Transform
-        for j in range(self.n_dims):
-            if isinstance(transform, list):
-                self.dimensions[j].set_transformer(transform[j])
-            else:
-                self.dimensions[j].set_transformer(transform)
-
-    def set_transformer_by_type(self, transform, dim_type):
-        """Sets the transformer of `dim_type` objects to `transform`
-
-        Parameters
-        ----------
-        transform : str
-           Sets all transformer of type `dim_type` to `transform`
-        dim_type : type
-            Can be `skopt.space.Real`, `skopt.space.Integer` or
-             `skopt.space.Categorical`
-        """
-        # Transform
-        for j in range(self.n_dims):
-            if isinstance(self.dimensions[j], dim_type):
-                self.dimensions[j].set_transformer(transform)
-
-    def get_transformer(self):
-        """Returns all transformers as list"""
-        return [self.dimensions[j].transform_ for j in range(self.n_dims)]
-
-    def transform(self, X):
-        """Transform samples from the original space into a warped space.
-
-        Note: this transformation is expected to be used to project samples
-              into a suitable space for numerical optimization.
-
-        Parameters
-        ----------
-        X : list of lists, shape=(n_samples, n_dims)
-            The samples to transform.
-
-        Returns
-        -------
-        Xt : array of floats, shape=(n_samples, transformed_n_dims)
-            The transformed samples.
-        """
-        # Pack by dimension
-        columns = []
-        for dim in self.dimensions:
-            columns.append([])
-
-        for i in range(len(X)):
-            for j in range(self.n_dims):
-                columns[j].append(X[i][j])
-
-        # Transform
-        for j in range(self.n_dims):
-            columns[j] = self.dimensions[j].transform(columns[j])
-
-        # Repack as an array
-        Xt = np.hstack([np.asarray(c).reshape((len(X), -1)) for c in columns])
-
-        return Xt
-
-    def inverse_transform(self, Xt):
-        """Inverse transform samples from the warped space back to the
-           original space.
-
-        Parameters
-        ----------
-        Xt : array of floats, shape=(n_samples, transformed_n_dims)
-            The samples to inverse transform.
-
-        Returns
-        -------
-        X : list of lists, shape=(n_samples, n_dims)
-            The original samples.
-        """
-        # Inverse transform
-        columns = []
-        start = 0
-        Xt = np.asarray(Xt)
-        for j in range(self.n_dims):
-            dim = self.dimensions[j]
-            offset = dim.transformed_size
-
-            if offset == 1:
-                columns.append(dim.inverse_transform(Xt[:, start]))
-            else:
-                columns.append(
-                    dim.inverse_transform(Xt[:, start:start + offset]))
-
-            start += offset
-
-        # Transpose
-        return _transpose_list_array(columns)
-
-    @property
-    def n_dims(self):
-        """The dimensionality of the original space."""
-        return len(self.dimensions)
-
-    @property
-    def transformed_n_dims(self):
-        """The dimensionality of the warped space."""
-        return sum([dim.transformed_size for dim in self.dimensions])
-
-    @property
-    def bounds(self):
-        """The dimension bounds, in the original space."""
-        b = []
-
-        for dim in self.dimensions:
-            if dim.size == 1:
-                b.append(dim.bounds)
-            else:
-                b.extend(dim.bounds)
-
-        return b
-
-    def __contains__(self, point):
-        """Check that `point` is within the bounds of the space."""
-        for component, dim in zip(point, self.dimensions):
-            if component not in dim:
-                return False
-        return True
-
-    def __getitem__(self, dimension_names):
-        """
-        Lookup and return the search-space dimension with the given name.
-
-        This allows for dict-like lookup of dimensions, for example:
-        `space['foo']` returns the dimension named 'foo' if it exists,
-        otherwise `None` is returned.
-
-        It also allows for lookup of a list of dimension-names, for example:
-        `space[['foo', 'bar']]` returns the two dimensions named
-        'foo' and 'bar' if they exist.
-
-        Parameters
-        ----------
-        dimension_names : str or list(str)
-            Name of a single search-space dimension (str).
-            List of names for search-space dimensions (list(str)).
-
-        Returns
-        -------
-        dims tuple (index, Dimension), list(tuple(index, Dimension)), \
-                (None, None)
-            A single search-space dimension with the given name,
-            or a list of search-space dimensions with the given names.
-        """
-
-        def _get(dimension_name):
-            """Helper-function for getting a single dimension."""
-            index = 0
-            # Get the index of the search-space dimension using its name.
-            for dim in self.dimensions:
-                if dimension_name == dim.name:
-                    return (index, dim)
-                elif dimension_name == index:
-                    return (index, dim)
-                index += 1
-            return (None, None)
-
-        if isinstance(dimension_names, (str, int)):
-            # Get a single search-space dimension.
-            dims = _get(dimension_name=dimension_names)
-        elif isinstance(dimension_names, (list, tuple)):
-            # Get a list of search-space dimensions.
-            # Note that we do not check whether the names are really strings.
-            dims = [_get(dimension_name=name) for name in dimension_names]
-        else:
-            msg = "Dimension name should be either string or" \
-                  "list of strings, but got {}."
-            raise ValueError(msg.format(type(dimension_names)))
-
-        return dims
-
-    @property
-    def transformed_bounds(self):
-        """The dimension bounds, in the warped space."""
-        b = []
-
-        for dim in self.dimensions:
-            if dim.transformed_size == 1:
-                b.append(dim.transformed_bounds)
-            else:
-                b.extend(dim.transformed_bounds)
-
-        return b
-
-    @property
-    def is_categorical(self):
-        """Space contains exclusively categorical dimensions"""
-        return all([isinstance(dim, Categorical) for dim in self.dimensions])
-
-    @property
-    def is_partly_categorical(self):
-        """Space contains any categorical dimensions"""
-        return any([isinstance(dim, Categorical) for dim in self.dimensions])
-
-    @property
-    def n_constant_dimensions(self):
-        """Returns the number of constant dimensions which have zero degree of
-        freedom, e.g. an Integer dimensions with (0., 0.) as bounds.
-        """
-        n = 0
-        for dim in self.dimensions:
-            if dim.is_constant:
-                n += 1
-        return n
-
-    def distance(self, point_a, point_b):
-        """Compute distance between two points in this space.
-
-        Parameters
-        ----------
-        point_a : array
-            First point.
-
-        point_b : array
-            Second point.
-        """
-        distance = 0.
-        for a, b, dim in zip(point_a, point_b, self.dimensions):
-            distance += dim.distance(a, b)
-
-        return distance
diff --git a/entmoot/space/transformers.py b/entmoot/space/transformers.py
deleted file mode 100644
index 64d499c..0000000
--- a/entmoot/space/transformers.py
+++ /dev/null
@@ -1,373 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-from __future__ import division
-import numpy as np
-from sklearn.preprocessing import LabelBinarizer
-from sklearn.utils import column_or_1d
-
-
-class Transformer(object):
-    """Base class for all 1-D transformers.
-    """
-    def fit(self, X):
-        return self
-
-    def transform(self, X):
-        raise NotImplementedError
-
-    def inverse_transform(self, X):
-        raise NotImplementedError
-
-
-class Identity(Transformer):
-    """Identity transform.
-    """
-
-    def transform(self, X):
-        return X
-
-    def inverse_transform(self, Xt):
-        return Xt
-
-
-class StringEncoder(Transformer):
-    """StringEncoder transform.
-       The transform will cast everything to a
-       string and the inverse transform will cast to the type defined in dtype.
-    """
-
-    def __init__(self, dtype=str):
-        super(StringEncoder, self).__init__()
-        self.dtype = dtype
-
-    def fit(self, X):
-        """Fit a list or array of categories. All elements must be from the
-        same type.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_categories,)
-            List of categories.
-        """
-        if len(X) > 0:
-            self.dtype = type(X[0])
-
-    def transform(self, X):
-        """Transform an array of categories to a string encoded representation.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_samples,)
-            List of categories.
-
-        Returns
-        -------
-        Xt : array-like, shape=(n_samples,)
-            The string encoded categories.
-        """
-        return [str(x) for x in X]
-
-    def inverse_transform(self, Xt):
-        """Inverse transform string encoded categories back to their original
-           representation.
-
-        Parameters
-        ----------
-        Xt : array-like, shape=(n_samples,)
-            String encoded categories.
-
-        Returns
-        -------
-        X : array-like, shape=(n_samples,)
-            The original categories.
-        """
-        return [self.dtype(x) for x in Xt]
-
-
-class LogN(Transformer):
-    """Base N logarithm transform."""
-
-    def __init__(self, base):
-        self._base = base
-
-    def transform(self, X):
-        return np.log10(np.asarray(X, dtype=np.float)) / np.log10(self._base)
-
-    def inverse_transform(self, Xt):
-        return self._base ** np.asarray(Xt, dtype=np.float)
-
-
-class Logit(Transformer):
-    """uses logit transformation, see: 
-    https://github.com/uber/bayesmark/blob/master/bayesmark/space.py"""
-
-    def __init__(self):
-        pass
-
-    @staticmethod
-    def transform(X):
-        from scipy.special import logit
-        return logit(X)
-
-    @staticmethod
-    def inverse_transform(Xt):
-        from scipy.special import expit as logistic
-        return logistic(Xt)
-
-class Bilog(Transformer):
-    """uses bilog transformation, see: 
-    https://github.com/uber/bayesmark/blob/master/bayesmark/space.py"""
-
-    def __init__(self):
-        pass
-
-    @staticmethod
-    def transform(X):
-        return np.sign(X) * np.log(1.0 + np.abs(X))
-
-    @staticmethod
-    def inverse_transform(Xt):
-        return np.sign(Xt) * (np.exp(np.abs(Xt)) - 1.0)
-
-
-class CategoricalEncoder(Transformer):
-    """OneHotEncoder that can handle categorical variables."""
-
-    def __init__(self):
-        """Convert labeled categories into one-hot encoded features."""
-        self._lb = LabelBinarizer()
-
-    def fit(self, X):
-        """Fit a list or array of categories.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_categories,)
-            List of categories.
-        """
-        self.mapping_ = {v: i for i, v in enumerate(X)}
-        self.inverse_mapping_ = {i: v for v, i in self.mapping_.items()}
-        self._lb.fit([self.mapping_[v] for v in X])
-        self.n_classes = len(self._lb.classes_)
-
-        return self
-
-    def transform(self, X):
-        """Transform an array of categories to a one-hot encoded representation.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_samples,)
-            List of categories.
-
-        Returns
-        -------
-        Xt : array-like, shape=(n_samples, n_categories)
-            The one-hot encoded categories.
-        """
-        return self._lb.transform([self.mapping_[v] for v in X])
-
-    def inverse_transform(self, Xt):
-        """Inverse transform one-hot encoded categories back to their original
-           representation.
-
-        Parameters
-        ----------
-        Xt : array-like, shape=(n_samples, n_categories)
-            One-hot encoded categories.
-
-        Returns
-        -------
-        X : array-like, shape=(n_samples,)
-            The original categories.
-        """
-        Xt = np.asarray(Xt)
-        return [
-            self.inverse_mapping_[i] for i in self._lb.inverse_transform(Xt)
-        ]
-
-
-class LabelEncoder(Transformer):
-    """LabelEncoder that can handle categorical variables."""
-    def __init__(self, X=None):
-        if X is not None:
-            self.fit(X)
-
-    def fit(self, X):
-        """Fit a list or array of categories.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_categories,)
-            List of categories.
-        """
-        X = np.asarray(X)
-        if X.dtype == object:
-            self.mapping_ = {v: i for i, v in enumerate(X)}
-        else:
-            i = 0
-            self.mapping_ = {}
-            for v in np.unique(X):
-                self.mapping_[v] = i
-                i += 1
-        self.inverse_mapping_ = {i: v for v, i in self.mapping_.items()}
-        return self
-
-    def transform(self, X):
-        """Transform an array of categories to a one-hot encoded
-        representation.
-
-        Parameters
-        ----------
-        X : array-like, shape=(n_samples,)
-            List of categories.
-
-        Returns
-        -------
-        Xt : array-like, shape=(n_samples, n_categories)
-            The integer categories.
-        """
-        X = np.asarray(X)
-        return [self.mapping_[v] for v in X]
-
-    def inverse_transform(self, Xt):
-        """Inverse transform integer categories back to their original
-           representation.
-
-        Parameters
-        ----------
-        Xt : array-like, shape=(n_samples, n_categories)
-            Integer categories.
-
-        Returns
-        -------
-        X : array-like, shape=(n_samples,)
-            The original categories.
-        """
-        if isinstance(Xt, (float, np.float64)):
-            Xt = [Xt]
-        else:
-            Xt = np.asarray(Xt)
-        return [
-            self.inverse_mapping_[int(np.round(i))] for i in Xt
-        ]
-
-
-class Normalize(Transformer):
-    """
-    Scales each dimension into the interval [0, 1].
-
-    Parameters
-    ----------
-    low : float
-        Lower bound.
-
-    high : float
-        Higher bound.
-
-    is_int : bool, default=True
-        Round and cast the return value of `inverse_transform` to integer. Set
-        to `True` when applying this transform to integers.
-    """
-    def __init__(self, low, high, is_int=False):
-        self.low = float(low)
-        self.high = float(high)
-        self.is_int = is_int
-        self._eps = 1e-8
-
-    def transform(self, X):
-        X = np.asarray(X)
-        if self.is_int:
-            if np.any(np.round(X) > self.high):
-                raise ValueError("All integer values should"
-                                 "be less than %f" % self.high)
-            if np.any(np.round(X) < self.low):
-                raise ValueError("All integer values should"
-                                 "be greater than %f" % self.low)
-        else:
-            if np.any(X > self.high + self._eps):
-                raise ValueError("All values should"
-                                 "be less than %f" % self.high)
-            if np.any(X < self.low - self._eps):
-                raise ValueError("All values should"
-                                 "be greater than %f" % self.low)
-        if (self.high - self.low) == 0.:
-            return X * 0.
-        if self.is_int:
-            return (np.round(X).astype(np.int) - self.low) /\
-                   (self.high - self.low)
-        else:
-            return (X - self.low) / (self.high - self.low)
-
-    def inverse_transform(self, X):
-        X = np.asarray(X)
-        if np.any(X > 1.0 + self._eps):
-            raise ValueError("All values should be less than 1.0")
-        if np.any(X < 0.0 - self._eps):
-            raise ValueError("All values should be greater than 0.0")
-        X_orig = X * (self.high - self.low) + self.low
-        if self.is_int:
-            return np.round(X_orig).astype(np.int)
-        return X_orig
-
-
-class Pipeline(Transformer):
-    """
-    A lightweight pipeline to chain transformers.
-
-    Parameters
-    ----------
-    transformers : list
-        A list of Transformer instances.
-    """
-    def __init__(self, transformers):
-        self.transformers = list(transformers)
-        for transformer in self.transformers:
-            if not isinstance(transformer, Transformer):
-                raise ValueError(
-                    "Provided transformers should be a Transformer "
-                    "instance. Got %s" % transformer
-                )
-
-    def fit(self, X):
-        for transformer in self.transformers:
-            transformer.fit(X)
-        return self
-
-    def transform(self, X):
-        for transformer in self.transformers:
-            X = transformer.transform(X)
-        return X
-
-    def inverse_transform(self, X):
-        for transformer in self.transformers[::-1]:
-            X = transformer.inverse_transform(X)
-        return X
diff --git a/entmoot/utils.py b/entmoot/utils.py
index 94eca0f..b936c74 100644
--- a/entmoot/utils.py
+++ b/entmoot/utils.py
@@ -1,463 +1,74 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-NOTE: Changes were made to the scikit-optimize source code included here. 
-For the most recent version of scikit-optimize we refer to:
-https://github.com/scikit-optimize/scikit-optimize/
-
-Copyright (c) 2019-2020 Alexander Thebelt.
-"""
-
-import os
-from gurobipy import Env
+from collections import namedtuple
 import numpy as np
-from numpy.typing import ArrayLike
-from scipy.optimize import OptimizeResult
-from joblib import dump as dump_
-from joblib import load as load_
-from .space import Space, Dimension
-from typing import Optional
-
-
-def get_gurobi_env():
-    """Return a Gurobi CloudEnv if environment variables are set, else None."""
-    if "GRB_CLOUDPOOL" in os.environ:
-        return Env.CloudEnv(
-            logfilename="gurobi.log",
-            accessID=os.environ["GRB_CLOUDACCESSID"],
-            secretKey=os.environ["GRB_CLOUDKEY"],
-            pool=os.environ["GRB_CLOUDPOOL"],
-        )
-    return None
-
-
-def is_supported(base_estimator):
-    return base_estimator in ['ENTING']
-
-
-def get_cat_idx(space):
-    from entmoot.space.space import Space, Categorical, Integer, Real, Dimension
-
-    if space is None:
-        return []
-
-    # collect indices of categorical features
-    cat_idx = []
-    for idx,dim in enumerate(space):
-        if isinstance(dim, Categorical):
-            cat_idx.append(idx)
-    return cat_idx
-
+from scipy.special import comb
 
-def cook_estimator(
-        space: Space,
-        base_estimator: str,
-        base_estimator_kwargs: Optional[dict] = None,
-        num_obj: int = 1,
-        random_state: Optional[int] = None):
-    """Cook an estimator that used for mean and uncertainty prediction.
 
-    :param space: Space, defines the search space of variables
-    :param base_estimator: str
-        currently available estimator types are in ["ENTING"]
-    :param base_estimator_kwargs: Optional[dict]
-        defines additional params that influence the base_estimator behavior
-            "lgbm_params": dict, additional parameters that are passed to lightgbm
-            "ensemble_type": str, options in ['GBRT', 'RF'],
-                "GBRT": uses gradient-boosted tree regressor
-                "RF": uses random forest
-            "unc_metric": str, options in ['exploration', 'penalty'], i.e.
-                negative or positive alpha contribution in "min mu \pm kappa * alpha"
-            "unc_scaling": str, options in ["standard", "normalize"], i.e.
-                scaling used for the distance-based uncertainty metric
-            "dist_metric": str, options in ["manhattan",'squared_euclidean'], i.e.
-                metric used to define non-categorical distance for uncertainty
-            "cat_dist_metric": str, options in ["overlap", "goodall4", "of"], i.e.
-                metric used to define categorical distance for uncertainty
-    :param num_obj: int
-        gives the number of objectives that the black-box is being optimized for
-    :param random_state: Optional[int]
-        fixed random seed to generate reproducible results
-    """
-
-    from lightgbm import LGBMRegressor
-    from entmoot.learning.tree_model import EntingRegressor, MisicRegressor
-
-    # collect indices of categorical features
-    cat_idx = get_cat_idx(space)
-
-    if not is_supported(base_estimator):
-        raise ValueError("base_estimator is not supported.")
-
-    # build enting estimator with distance-based uncertainty
-    if base_estimator == "ENTING":
+OptResult = namedtuple("OptResult", ["opt_point", "opt_val", "mu_unscaled"])
 
-        # check uncertainty metric parameters
-        unc_metric = base_estimator_kwargs.get("unc_metric", "exploration")
-        unc_scaling = base_estimator_kwargs.get("unc_scaling", "standard")
-        dist_metric = base_estimator_kwargs.get("dist_metric", "squared_euclidean")
-        cat_dist_metric = base_estimator_kwargs.get("cat_dist_metric", "goodall4")
 
-        unc_estimator = cook_unc_estimator(space,
-                                           unc_metric=unc_metric,
-                                           unc_scaling=unc_scaling,
-                                           dist_metric=dist_metric,
-                                           cat_dist_metric=cat_dist_metric,
-                                           num_obj=num_obj)
+def grid(dimension: int, levels: int) -> np.ndarray:
+    """Construct a regular grid on the unit simplex.
 
-        ensemble_type = base_estimator_kwargs.get("ensemble_type", "GBRT")
-        lgbm_params = base_estimator_kwargs.get("lgbm_params", {})
+    The number of grid points is L = (n+m-1)!/(n!*(m-1)!)
+    where m is the dimension and n+1 is the number of levels.
 
-        if unc_estimator.std_type == 'distance':
-            tree_reg = EntingRegressor
-        elif unc_estimator.std_type == 'proximity':
-            tree_reg = MisicRegressor
+    Args:
+        dimension (int): Number of variables.
+        levels (int): Number of levels for each variable.
 
-        if ensemble_type == "GBRT":
-            gbrt = []
-            for _ in range(num_obj):
-                gbrt.append(
-                    LGBMRegressor(boosting_type='gbdt',
-                                  objective='regression',
-                                  verbose=-1))
-            base_estimator = tree_reg(space,
-                                    gbrt,
-                                    unc_estimator,
-                                    random_state=random_state,
-                                    cat_idx=cat_idx)
-        elif ensemble_type == "RF":
-            rf = []
-            for _ in range(num_obj):
-                rf.append(
-                    LGBMRegressor(boosting_type='random_forest',
-                                  objective='regression',
-                                  verbose=0,
-                                  subsample_freq=1,
-                                  subsample=0.9,
-                                  bagging_seed=random_state))
-            base_estimator = tree_reg(space,
-                                    rf,
-                                    unc_estimator,
-                                    random_state=random_state,
-                                    cat_idx=cat_idx)
+    Returns:
+       array: Regularily spaced points on the unit simplex.
 
-    if base_estimator_kwargs is not None:
-        base_estimator.set_params(**lgbm_params)
+    Examples:
+        --> simplex_grid(3, 3)
+        array([
+            [0. , 0. , 1. ],
+            [0. , 0.5, 0.5],
+            [0. , 1. , 0. ],
+            [0.5, 0. , 0.5],
+            [0.5, 0.5, 0. ],
+            [1. , 0. , 0. ]
+        ])
 
-    return base_estimator
-
-
-def cook_unc_estimator(space: Space,
-                    unc_metric: str,
-                    unc_scaling: str,
-                    dist_metric: str,
-                    cat_dist_metric: str,
-                    num_obj: int = 1):
-    """Cook a distance-based uncertainty estimator.
-
-    :param space: Space, defines the search space of variables
-    :param unc_metric: str, options in ['exploration', 'penalty'], i.e.
-        negative or positive alpha contribution in "min mu \pm kappa * alpha"
-    :param unc_scaling: str, options in ["standard", "normalize"], i.e.
-        scaling used for the distance-based uncertainty metric
-    :param dist_metric: str, options in ["manhattan",'squared_euclidean'], i.e.
-        metric used to define non-categorical distance for uncertainty
-    :param cat_dist_metric: str, options in ["overlap", "goodall4", "of"], i.e.
-        metric used to define categorical distance for uncertainty
-    :param num_obj: int
-        gives the number of objectives that the black-box is being optimized for
-    :param random_state: Optional[int]
-        fixed random seed to generate reproducible results
+    References:
+        Nijenhuis and Wilf, Combinatorial Algorithms, Chapter 5, Academic Press, 1978.
     """
+    m = dimension
+    n = levels - 1
+    L = int(comb(dimension - 1 + levels - 1, dimension - 1, exact=True))
 
-    from entmoot.learning.distance_based_std import \
-        DistanceBasedExploration, DistanceBasedPenalty
-
-    from entmoot.learning.proximit_based_std import \
-        ProximityMetric
+    x = np.zeros(m, dtype=int)
+    x[-1] = n
 
-    # normalize distances if multi-obj
-    if num_obj>1: unc_scaling="normalize"
+    out = np.empty((L, m), dtype=int)
+    out[0] = x
 
-    if unc_metric == "exploration":
-        unc_estimator = DistanceBasedExploration(
-                space,
-                unc_scaling=unc_scaling,
-                dist_metric=dist_metric,
-                cat_dist_metric=cat_dist_metric)
+    h = m
+    for i in range(1, L):
+        h -= 1
 
-    elif unc_metric == "penalty":
-        unc_estimator = DistanceBasedPenalty(
-                space,
-                unc_scaling=unc_scaling,
-                dist_metric=dist_metric,
-                cat_dist_metric=cat_dist_metric)
+        val = x[h]
+        x[h] = 0
+        x[h - 1] += 1
+        x[-1] = val - 1
 
-    elif unc_metric == "prox":
-        if num_obj > 1:
-            raise NotImplementedError(
-                "unc_metric 'prox' is currently unavailable and will be adapted for "
-                "multi-objective usage.")
-        unc_estimator = \
-            ProximityMetric(space)
+        if val != 1:
+            h = m
 
-    return unc_estimator
+        out[i] = x
 
+    return out / n
 
-def cook_initial_point_generator(generator, **kwargs):
-    """Cook a default initial point generator.
 
-    For the special generator called "random" the return value is None.
+def sample(dimension: int, n_samples: int = 1) -> np.ndarray:
+    """Sample uniformly from the unit simplex.
 
-    Parameters
-    ----------
-    generator : "lhs", "sobol", "halton", "hammersly", "grid", "random" \
-            or InitialPointGenerator instance"
-        Should inherit from `skopt.sampler.InitialPointGenerator`.
+    Args:
+        dimension (int): Number of dimensions.
+        n_samples (int): Number of samples to draw.
 
-    kwargs : dict
-        Extra parameters provided to the generator at init time.
+    Returns:
+        array, shape=(n_samples, dimesnion): Random samples from the unit simplex.
     """
-    from entmoot.sampler import Sobol, Lhs, Hammersly, Halton, Grid
-    from entmoot.sampler import InitialPointGenerator
-
-    if generator is None:
-        generator = "random"
-    elif isinstance(generator, str):
-        generator = generator.lower()
-        if generator not in ["sobol", "halton", "hammersly", "lhs", "random",
-                             "grid"]:
-            raise ValueError("Valid strings for the generator parameter "
-                             " are: 'sobol', 'lhs', 'halton', 'hammersly',"
-                             "'random', or 'grid' not "
-                             "%s." % generator)
-    elif not isinstance(generator, InitialPointGenerator):
-        raise ValueError("generator has to be an InitialPointGenerator."
-                         "Got %s" % (str(type(generator))))
-
-    if isinstance(generator, str):
-        if generator == "sobol":
-            generator = Sobol()
-        elif generator == "halton":
-            generator = Halton()
-        elif generator == "hammersly":
-            generator = Hammersly()
-        elif generator == "lhs":
-            generator = Lhs()
-        elif generator == "grid":
-            generator = Grid()
-        elif generator == "random":
-            return None
-    generator.set_params(**kwargs)
-    return generator
-
-
-def is_listlike(x):
-    return isinstance(x, (list, tuple))
-
-
-def is_2Dlistlike(x):
-    return np.all([is_listlike(xi) for xi in x])
-
-
-def check_x_in_space(x, space):
-    if is_2Dlistlike(x):
-        if not np.all([p in space for p in x]):
-            raise ValueError("Not all points are within the bounds of"
-                             " the space.")
-        if any([len(p) != len(space.dimensions) for p in x]):
-            raise ValueError("Not all points have the same dimensions as"
-                             " the space.")
-    elif is_listlike(x):
-        if x not in space:
-            raise ValueError("Point (%s) is not within the bounds of"
-                             " the space (%s)."
-                             % (x, space.bounds))
-        if len(x) != len(space.dimensions):
-            raise ValueError("Dimensions of point (%s) and space (%s) do not match"
-                             % (x, space.bounds))
-
-
-def create_result(
-        Xi: list, yi: ArrayLike, space=None, rng=None, specs: Optional[dict] = None, models: Optional[list] = None,
-        model_mu=None, model_std=None, gurobi_mipgap=None):
-    """
-    Initialize an `OptimizeResult` object.
-
-    Parameters
-    ----------
-    Xi : list of lists, shape (n_iters, n_features)
-        Location of the minimum at every iteration.
-
-    yi : array-like, shape (n_iters,)
-        Minimum value obtained at every iteration.
-
-    space : Space instance, optional
-        Search space.
-
-    rng : RandomState instance, optional
-        State of the random state.
-
-    specs : dict, optional
-        Call specifications.
-
-    models : list, optional
-        List of fit surrogate models.
-
-    Returns
-    -------
-    res : `OptimizeResult`, scipy object
-        OptimizeResult instance with the required information.
-    """
-    res = OptimizeResult()
-    yi = np.asarray(yi)
-    if np.ndim(yi) == 2:
-        res.log_time = np.ravel(yi[:, 1])
-        yi = np.ravel(yi[:, 0])
-    best = np.argmin(yi)
-    res.x = Xi[best]
-    res.fun = yi[best]
-    res.func_vals = yi
-    res.x_iters = Xi
-    res.models = models
-    res.model_mu = model_mu
-    res.model_std = model_std
-    res.gurobi_mipgap = gurobi_mipgap
-    res.space = space
-    res.random_state = rng
-    res.specs = specs
-    return res
-
-
-def dump(res, filename, store_objective=True, **kwargs):
-    """
-    Store an skopt optimization result into a file.
-
-    Parameters
-    ----------
-    res : `OptimizeResult`, scipy object
-        Optimization result object to be stored.
-
-    filename : string or `pathlib.Path`
-        The path of the file in which it is to be stored. The compression
-        method corresponding to one of the supported filename extensions ('.z',
-        '.gz', '.bz2', '.xz' or '.lzma') will be used automatically.
-
-    store_objective : boolean, default=True
-        Whether the objective function should be stored. Set `store_objective`
-        to `False` if your objective function (`.specs['args']['func']`) is
-        unserializable (i.e. if an exception is raised when trying to serialize
-        the optimization result).
-
-        Notice that if `store_objective` is set to `False`, a deep copy of the
-        optimization result is created, potentially leading to performance
-        problems if `res` is very large. If the objective function is not
-        critical, one can delete it before calling `skopt.dump()` and thus
-        avoid deep copying of `res`.
-
-    **kwargs : other keyword arguments
-        All other keyword arguments will be passed to `joblib.dump`.
-    """
-    if store_objective:
-        dump_(res, filename, **kwargs)
-
-    elif 'func' in res.specs['args']:
-        # If the user does not want to store the objective and it is indeed
-        # present in the provided object, then create a deep copy of it and
-        # remove the objective function before dumping it with joblib.dump.
-        res_without_func = deepcopy(res)
-        del res_without_func.specs['args']['func']
-        dump_(res_without_func, filename, **kwargs)
-
-    else:
-        # If the user does not want to store the objective and it is already
-        # missing in the provided object, dump it without copying.
-        dump_(res, filename, **kwargs)
-
-
-def load(filename, **kwargs):
-    """
-    Reconstruct a skopt optimization result from a file
-    persisted with skopt.dump.
-
-    .. note::
-        Notice that the loaded optimization result can be missing
-        the objective function (`.specs['args']['func']`) if `skopt.dump`
-        was called with `store_objective=False`.
-
-    Parameters
-    ----------
-    filename : string or `pathlib.Path`
-        The path of the file from which to load the optimization result.
-
-    **kwargs : other keyword arguments
-        All other keyword arguments will be passed to `joblib.load`.
-
-    Returns
-    -------
-    res : `OptimizeResult`, scipy object
-        Reconstructed OptimizeResult instance.
-    """
-    return load_(filename, **kwargs)
-
-def normalize_dimensions(dimensions: list):
-    """Create a ``Space`` where all dimensions are normalized to unit range.
-    This is particularly useful for Gaussian process based regressors and is
-    used internally by ``gp_minimize``.
-    Parameters
-    ----------
-    dimensions : list, shape (n_dims,)
-        List of search space dimensions.
-        Each search dimension can be defined either as
-        - a `(lower_bound, upper_bound)` tuple (for `Real` or `Integer`
-          dimensions),
-        - a `(lower_bound, upper_bound, "prior")` tuple (for `Real`
-          dimensions),
-        - as a list of categories (for `Categorical` dimensions), or
-        - an instance of a `Dimension` object (`Real`, `Integer` or
-          `Categorical`).
-         NOTE: The upper and lower bounds are inclusive for `Integer`
-         dimensions.
-    """
-    space = Space(dimensions)
-    transformed_dimensions = []
-    for dimension in space.dimensions:
-        # check if dimension is of a Dimension instance
-        if isinstance(dimension, Dimension):
-            # Change the transformer to normalize
-            # and add it to the new transformed dimensions
-            dimension.set_transformer("normalize")
-            transformed_dimensions.append(
-                dimension
-            )
-        else:
-            raise RuntimeError("Unknown dimension type "
-                               "(%s)" % type(dimension))
-
-    return Space(transformed_dimensions)
+    s = np.random.standard_exponential((n_samples, dimension))
+    return (s.T / s.sum(axis=1)).T
diff --git a/requirements-dev.txt b/requirements-dev.txt
new file mode 100644
index 0000000..dc76545
--- /dev/null
+++ b/requirements-dev.txt
@@ -0,0 +1,8 @@
+numpy>=1.18.4
+lightgbm>=2.3.1
+tqdm
+pyomo
+gurobipy
+pytest-cov
+IPython
+nbsphinx
diff --git a/requirements.txt b/requirements.txt
index e7fb89f..c0c65d8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,3 @@
 numpy>=1.18.4
-scikit-learn>=0.21.3
-PyYAML>=5.3.1
 lightgbm>=2.3.1
 tqdm
-mopti
diff --git a/setup.cfg b/setup.cfg
index d6c01b6..1782bdb 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -4,4 +4,7 @@ include_trailing_comma=True
 force_grid_wrap=0
 use_parentheses=True
 line_length=88
-known_first_party=entmoot
\ No newline at end of file
+known_first_party=entmoot
+
+[flake8]
+max-line-length = 121
diff --git a/setup.py b/setup.py
index e4e91fb..6bfbf08 100644
--- a/setup.py
+++ b/setup.py
@@ -4,28 +4,24 @@
 project_root = Path(__file__).resolve().parent
 
 about = {}
-version_path = project_root / 'entmoot' / '__version__.py'
+version_path = project_root / "entmoot" / "__version__.py"
 with version_path.open() as f:
     exec(f.read(), about)
 
 setup(
-    name='entmoot',
-    author=about['__author__'],
-    author_email=about['__author_email__'],
-    license=about['__license__'],
-    version=about['__version__'],
-    url='https://github.com/cog-imperial/entmoot',
-    packages=find_packages(exclude=['tests','docs']),
+    name="entmoot",
+    author=about["__author__"],
+    author_email=about["__author_email__"],
+    license=about["__license__"],
+    version=about["__version__"],
+    url="https://github.com/cog-imperial/entmoot",
+    packages=find_packages(exclude=["tests", "docs"]),
     install_requires=[
-        'numpy>=1.18.4',
-        'scikit-learn>=0.21.3',
-        'pyYAML>=5.3.1',
-        'lightgbm>=2.3.1'
+        "numpy>=1.18.4",
+        "lightgbm>=2.3.1",
     ],
     setup_requires=[
-        'numpy>=1.18.4',
-        'scikit-learn>=0.21.3',
-        'pyYAML>=5.3.1',
-        'lightgbm>=2.3.1'
-    ]
+        "numpy>=1.18.4",
+        "lightgbm>=2.3.1",
+    ],
 )


Test patch:
diff --git a/tests/test_curr.py b/tests/test_curr.py
new file mode 100644
index 0000000..9177a48
--- /dev/null
+++ b/tests/test_curr.py
@@ -0,0 +1,209 @@
+from entmoot import Enting, ProblemConfig, GurobiOptimizer, PyomoOptimizer
+from entmoot.benchmarks import (
+    build_multi_obj_categorical_problem,
+    eval_multi_obj_cat_testfunc,
+)
+
+import numpy as np
+import pytest
+import random
+
+
+def test_core_model_copy():
+    # define problem
+    problem_config = ProblemConfig(rnd_seed=73)
+    # number of objectives
+    number_objectives = 2
+    build_multi_obj_categorical_problem(problem_config, n_obj=number_objectives)
+
+    core_model_gurobi = problem_config.get_gurobi_model_core()
+    core_model_gurobi_copy = problem_config.copy_gurobi_model_core(core_model_gurobi)
+
+    assert len(core_model_gurobi.getVars()) == len(core_model_gurobi_copy.getVars())
+    assert len(core_model_gurobi._all_feat) == len(core_model_gurobi_copy._all_feat)
+
+    core_model_pyomo = problem_config.get_pyomo_model_core()
+    core_model_pyomo_copy = problem_config.copy_pyomo_model_core(core_model_pyomo)
+
+    assert len(core_model_pyomo.x) == len(core_model_pyomo_copy.x)
+    assert len(core_model_pyomo._all_feat) == len(core_model_pyomo_copy._all_feat)
+
+
+@pytest.mark.fast_test
+def test_multiobj_constraints():
+    # define problem
+    problem_config = ProblemConfig(rnd_seed=73)
+    # number of objectives
+    number_objectives = 2
+    build_multi_obj_categorical_problem(problem_config, n_obj=number_objectives)
+
+    # sample data
+    rnd_sample = problem_config.get_rnd_sample_list(num_samples=20)
+    testfunc_evals = eval_multi_obj_cat_testfunc(rnd_sample, n_obj=number_objectives)
+
+    params = {"unc_params": {"dist_metric": "l1", "acq_sense": "exploration"}}
+    enting = Enting(problem_config, params=params)
+    # fit tree ensemble
+    enting.fit(rnd_sample, testfunc_evals)
+
+    # Add constraints
+
+    # Gurobi version
+    # get optimization model
+    model_gur = problem_config.get_gurobi_model_core()
+    # extract decision variables
+    x = model_gur._all_feat[3]
+    y = model_gur._all_feat[4]
+    z = model_gur._all_feat[5]
+    # add constraint that all variables should coincide
+    model_gur.addConstr(x == y)
+    model_gur.addConstr(y == z)
+    # Update model
+    model_gur.update()
+
+    # Build GurobiOptimizer object and solve optimization problem
+    params_gurobi = {"MIPGap": 0}
+    opt_gur = GurobiOptimizer(problem_config, params=params_gurobi)
+
+    res_gur = opt_gur.solve(enting, model_core=model_gur)
+    x_opt, y_opt, z_opt = res_gur.opt_point[3:]
+
+    assert round(x_opt, 5) == round(y_opt, 5) and round(y_opt, 5) == round(z_opt, 5)
+
+    # Pyomo version
+    import pyomo.environ as pyo
+
+    model_pyo = problem_config.get_pyomo_model_core()
+    # extract decision variables
+    x = model_pyo._all_feat[3]
+    y = model_pyo._all_feat[4]
+    z = model_pyo._all_feat[5]
+    # add constraint that all variables should coincide
+    model_pyo.xy_equal_constr = pyo.Constraint(expr=x == y)
+    model_pyo.yz_equal_constr = pyo.Constraint(expr=y == z)
+
+    # Build GurobiOptimizer object and solve optimization problem
+    params_pyomo = {"solver_name": "gurobi"}
+    opt_pyo = PyomoOptimizer(problem_config, params=params_pyomo)
+
+    res_pyo = opt_pyo.solve(enting, model_core=model_pyo)
+    x_opt, y_opt, z_opt = res_pyo.opt_point[3:]
+
+    assert round(x_opt, 5) == round(y_opt, 5) and round(y_opt, 5) == round(z_opt, 5)
+
+
+def test_simple_test():
+    def my_func(x: float) -> float:
+        return x**2 + 1 + random.uniform(-0.2, 0.2)
+
+    # Define a one-dimensional minimization problem with one real variable bounded by -2 and 3
+    problem_config = ProblemConfig()
+    problem_config.add_feature("real", (-2, 3))
+    problem_config.add_min_objective()
+
+    # Create training data using the randomly disturbed function f(x) = x^2 + 1 + eps
+    X_train = np.reshape(np.linspace(-2, 3, 10), (-1, 1))
+    y_train = np.reshape([my_func(x) for x in X_train], (-1, 1))
+
+    # Define enting object and corresponding parameters
+    params = {"unc_params": {"dist_metric": "l1", "acq_sense": "penalty"}}
+    enting = Enting(problem_config, params=params)
+    # Fit tree model
+    enting.fit(X_train, y_train)
+
+    # Build PyomoOptimizer object with Gurobi as solver and solve optimization problem
+    params_pyo = {"solver_name": "gurobi"}
+    opt_pyo = PyomoOptimizer(problem_config, params=params_pyo)
+    res = opt_pyo.solve(enting)
+
+    assert round(res.opt_point[0]) == 0
+
+
+@pytest.mark.fast_test
+def test_compare_pyomo_gurobipy_multiobj():
+    """
+    Ensures for a multi objective example with l1  and l2 uncertainty metric and mixed feature types that optimization
+    results for Gurobipy model and Pyomo model with Gurobi as optimizer coincide.
+    """
+
+    # define problem
+    problem_config = ProblemConfig(rnd_seed=73)
+    # number of objectives
+    number_objectives = 2
+    build_multi_obj_categorical_problem(problem_config, n_obj=number_objectives)
+
+    # sample data
+    rnd_sample = problem_config.get_rnd_sample_list(num_samples=20)
+    testfunc_evals = eval_multi_obj_cat_testfunc(rnd_sample, n_obj=number_objectives)
+
+
+    # The (l2, exploration) combination does not work which is probably due to a rounding error
+    for metric in ["l1", "euclidean_squared"]:
+        for acq_sense in ["penalty"]:
+            params = {"unc_params": {"dist_metric": metric, "acq_sense": acq_sense}}
+            enting = Enting(problem_config, params=params)
+            # fit tree ensemble
+            enting.fit(rnd_sample, testfunc_evals)
+
+            # Build GurobiOptimizer object and solve optimization problem
+            params_gurobi = {"NonConvex": 2, "MIPGap": 0}
+            opt_gur = GurobiOptimizer(problem_config, params=params_gurobi)
+            res_gur = opt_gur.solve(enting, weights=(0.4, 0.6))
+
+            # Build PyomoOptimizer object with Gurobi as solver and solve optimization problem
+            params_pyo = {
+                "solver_name": "gurobi",
+                "solver_options": {"NonConvex": 2, "MIPGap": 0},
+            }
+            opt_pyo = PyomoOptimizer(problem_config, params=params_pyo)
+            res_pyo = opt_pyo.solve(enting, weights=(0.4, 0.6))
+
+            # Assert that active leaves coincide for both models
+            assert round(res_gur.opt_val, 2) == round(res_pyo.opt_val, 2)
+
+            assert opt_gur._active_leaves == opt_pyo._active_leaves
+
+
+@pytest.mark.fast_test
+def test_compare_pyomo_gurobipy_singleobj():
+    """
+    Ensures for a single objective example with l1  and l2 uncertainty metric and mixed feature types that optimization
+    results for Gurobipy model and Pyomo model with Gurobi as optimizer coincide.
+    """
+
+    # define problem
+    problem_config = ProblemConfig(rnd_seed=73)
+    # number of objectives
+    number_objectives = 1
+    build_multi_obj_categorical_problem(problem_config, n_obj=number_objectives)
+
+    # sample data
+    rnd_sample = problem_config.get_rnd_sample_list(num_samples=20)
+    testfunc_evals = eval_multi_obj_cat_testfunc(rnd_sample, n_obj=number_objectives)
+
+    for metric in ["l1", "l2", "euclidean_squared"]:
+        for acq_sense in ["exploration", "penalty"]:
+            params = {"unc_params": {"dist_metric": metric, "acq_sense": acq_sense}}
+            enting = Enting(problem_config, params=params)
+            # fit tree ensemble
+            enting.fit(rnd_sample, testfunc_evals)
+
+            # Build GurobiOptimizer object and solve optimization problem
+            params_gurobi = {"NonConvex": 2, "MIPGap": 0}
+            opt_gur = GurobiOptimizer(problem_config, params=params_gurobi)
+            res_gur = opt_gur.solve(enting)
+
+            # Build PyomoOptimizer object with Gurobi as solver and solve optimization problem
+            params_pyo = {
+                "solver_name": "gurobi",
+                "solver_options": {"NonConvex": 2, "MIPGap": 0},
+            }
+            opt_pyo = PyomoOptimizer(problem_config, params=params_pyo)
+            res_pyo = opt_pyo.solve(enting)
+
+            # Compare optimal values (e.g. objective values) ...
+            assert round(res_gur.opt_val / res_pyo.opt_val, 5) <= 1.0001
+            # ... and optimal points (e.g. feature variables)
+            assert [round(x) for x in res_gur.opt_point[2:]] == [
+                round(x) for x in res_pyo.opt_point[2:]
+            ]
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
deleted file mode 100644
index 5f578e8..0000000
--- a/tests/test_optimizer.py
+++ /dev/null
@@ -1,317 +0,0 @@
-
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-NOTE: Changes were made to the scikit-optimize source code included here. 
-For the most recent version of scikit-optimize we refer to:
-https://github.com/scikit-optimize/scikit-optimize/
-
-Copyright (c) 2019-2020 Alexander Thebelt.
-"""
-
-import numpy as np
-import pytest
-
-from numpy.testing import assert_array_equal
-from numpy.testing import assert_equal
-from numpy.testing import assert_raises
-
-from entmoot.benchmarks import Rosenbrock, SimpleCat
-
-from entmoot.optimizer import Optimizer
-
-ESTIMATOR_STRINGS = ("GBRT", "RF")
-
-STD_ESTIMATORS = ("BDD","L1BDD","DDP","L1DDP")
-ACQ_FUNCS = ["LCB"]
-ACQ_OPTIMIZER = ["sampling", "global"]
-
-@pytest.mark.fast_test
-def test_multiple_asks():
-    # calling ask() multiple times without a tell() inbetween should
-    # be a "no op"
-    bench1 = Rosenbrock()
-
-    opt = Optimizer(bench1.get_bounds(2), "ENTING", n_initial_points=10,
-                    acq_optimizer="sampling",
-                    base_estimator_kwargs={"min_child_samples":2})
-
-    opt.run(bench1, n_iter=13)
-    # ask() computes the next point ready for the next call to ask()
-    # hence there are three after three iterations
-    assert_equal(len(opt.models), 3)
-    assert_equal(len(opt.Xi), 13)
-    opt.ask()
-    # training happens in ask(), that's why there's one more model now
-    assert_equal(len(opt.models), 4)
-    assert_equal(len(opt.Xi), 13)
-    assert_equal(opt.ask(), opt.ask())
-    opt.update_next()
-    assert_equal(opt.ask(), opt.ask())
-
-@pytest.mark.fast_test
-def test_model_queue_size():
-    # Check if model_queue_size limits the model queue size
-    bench1 = Rosenbrock()
-
-    opt = Optimizer(bench1.get_bounds(2), "ENTING", n_initial_points=10,
-                    acq_optimizer="sampling", model_queue_size=2,
-                    base_estimator_kwargs={"min_child_samples":2})
-
-    opt.run(bench1, n_iter=13)
-    # tell() computes the next point ready for the next call to ask()
-    # hence there are three after three iterations
-    assert_equal(len(opt.models), 2)
-    assert_equal(len(opt.Xi), 13)
-    opt.ask()
-    assert_equal(len(opt.models), 2)
-    assert_equal(len(opt.Xi), 13)
-    assert_equal(opt.ask(), opt.ask())
-    opt.update_next()
-    assert_equal(opt.ask(), opt.ask())
-
-@pytest.mark.fast_test
-def test_invalid_tell_arguments():
-    bench1 = Rosenbrock()
-
-    opt = Optimizer(bench1.get_bounds(2), "ENTING", n_initial_points=10,
-                    acq_optimizer="sampling", model_queue_size=2,
-                    base_estimator_kwargs={"min_child_samples":2})
-
-    # can't have single point and multiple values for y
-    assert_raises(ValueError, opt.tell, [1.], [1., 1.])
-
-@pytest.mark.fast_test
-def test_invalid_tell_arguments_list():
-    bench1 = Rosenbrock()
-
-    opt = Optimizer(bench1.get_bounds(2), "ENTING", n_initial_points=10,
-                    acq_optimizer="sampling", model_queue_size=2,
-                    base_estimator_kwargs={"min_child_samples":2})
-
-    assert_raises(ValueError, opt.tell, [[1.], [2.]], [1., None])
-
-@pytest.mark.fast_test
-def test_bounds_checking_1D():
-    low = -2.
-    high = 2.
-
-    opt = Optimizer([(low, high)], "ENTING", n_initial_points=1,
-                    acq_optimizer="sampling")
-
-    assert_raises(ValueError, opt.tell, [high + 0.5], 2.)
-    assert_raises(ValueError, opt.tell, [low - 0.5], 2.)
-    # feed two points to tell() at once
-    assert_raises(ValueError, opt.tell, [high + 0.5, high], (2., 3.))
-    assert_raises(ValueError, opt.tell, [low - 0.5, high], (2., 3.))
-
-@pytest.mark.fast_test
-def test_bounds_checking_2D():
-    low = -2.
-    high = 2.
-
-    opt = Optimizer([(low, high)], "ENTING", n_initial_points=1,
-                    acq_optimizer="sampling")
-
-    assert_raises(ValueError, opt.tell, [high + 0.5, high + 4.5], 2.)
-    assert_raises(ValueError, opt.tell, [low - 0.5, low - 4.5], 2.)
-
-    # first out, second in
-    assert_raises(ValueError, opt.tell, [high + 0.5, high + 0.5], 2.)
-    assert_raises(ValueError, opt.tell, [low - 0.5, high + 0.5], 2.)
-
-
-@pytest.mark.fast_test
-def test_bounds_checking_2D_multiple_points():
-    low = -2.
-    high = 2.
-
-    opt = Optimizer([(low, high)], "ENTING", n_initial_points=1,
-                    acq_optimizer="sampling")
-
-    # first component out, second in
-    assert_raises(ValueError, opt.tell,
-                  [(high + 0.5, high + 0.5), (high + 0.5, high + 0.5)],
-                  [2., 3.])
-    assert_raises(ValueError, opt.tell,
-                  [(low - 0.5, high + 0.5), (low - 0.5, high + 0.5)],
-                  [2., 3.])
-
-
-@pytest.mark.fast_test
-def test_dimension_checking_1D():
-    low = -2
-    high = 2
-    opt = Optimizer([(low, high)], "ENTING", n_initial_points=1,
-                    acq_optimizer="sampling")
-
-    with pytest.raises(ValueError) as e:
-        # within bounds but one dimension too high
-        opt.tell([low+1, low+1], 2.)
-    assert "Dimensions of point " in str(e.value)
-
-
-@pytest.mark.fast_test
-def test_dimension_checking_2D():
-    low = -2
-    high = 2
-    opt = Optimizer([(low, high), (low, high)], "ENTING", n_initial_points=10,
-                    acq_optimizer="sampling")
-
-    # within bounds but one dimension too little
-    with pytest.raises(ValueError) as e:
-        opt.tell([low+1, ], 2.)
-    assert "Dimensions of point " in str(e.value)
-    # within bounds but one dimension too much
-    with pytest.raises(ValueError) as e:
-        opt.tell([low+1, low+1, low+1], 2.)
-    assert "Dimensions of point " in str(e.value)
-
-
-@pytest.mark.fast_test
-def test_dimension_checking_2D_multiple_points():
-    low = -2
-    high = 2
-    opt = Optimizer([(low, high), (low, high)])
-    # within bounds but one dimension too little
-    with pytest.raises(ValueError) as e:
-        opt.tell([[low+1, ], [low+1, low+2], [low+1, low+3]], 2.)
-    assert "dimensions as the space" in str(e.value)
-    # within bounds but one dimension too much
-    with pytest.raises(ValueError) as e:
-        opt.tell([[low + 1, low + 1, low + 1], [low + 1, low + 2],
-                  [low + 1, low + 3]], 2.)
-    assert "dimensions as the space" in str(e.value)
-
-@pytest.mark.consistent
-@pytest.mark.parametrize("acq_optimizer", ACQ_OPTIMIZER)
-def test_result_rosenbrock(acq_optimizer):
-    bench1 = Rosenbrock()
-
-    opt = Optimizer(bench1.get_bounds(2), "ENTING", n_initial_points=10,
-                    acq_optimizer=acq_optimizer,
-                    base_estimator_kwargs={"min_child_samples":2},
-                    random_state=100)
-
-    opt.run(bench1, n_iter=20, no_progress_bar=True)
-
-    if acq_optimizer == "sampling":
-        res_known = np.array(
-            [[0.701053475292035, -0.36025795213264167],
-            [0.10806258681567815, -1.4393145750513043],
-            [-1.4061114159677395, -1.2842304177627843],
-            [-1.1873986869551771, -0.19357735788859287],
-            [1.5161042350797533, -1.7871624483849504],
-            [0.5091814944214517, 0.09563236039209677],
-            [0.25516550013357264, -2.0241727838227086],
-            [-0.7887945257988347, 1.843954950307714],
-            [-1.5292242379168393, -1.7244656733329853],
-            [-0.7726975624367143, 0.5422431772370664],
-            [-0.26851521090965313, -0.27168208839378893],
-            [0.18559023118643525, -0.30091277104839054],
-            [2.037449372031385, -0.003908949645985427],
-            [0.3618262062035891, -0.886575853445807],
-            [-0.70099499312817, -0.013753379624784401],
-            [0.3074119029928717, 1.3227888228213858],
-            [0.43370695792140035, -0.5401762255031577],
-            [0.060723595879015324, 2.0360116729103517],
-            [-0.8891589013732526, -0.2756841022715071],
-            [-0.5803197996462619, -0.6720516369715639]]
-        )
-    elif acq_optimizer == "global":
-        res_known = np.array(
-            [[0.701053475292035, -0.36025795213264167],
-            [0.10806258681567815, -1.4393145750513043],
-            [-1.4061114159677395, -1.2842304177627843],
-            [-1.1873986869551771, -0.19357735788859287],
-            [1.5161042350797533, -1.7871624483849504],
-            [0.5091814944214517, 0.09563236039209677],
-            [0.25516550013357264, -2.0241727838227086],
-            [-0.7887945257988347, 1.843954950307714],
-            [-1.5292242379168393, -1.7244656733329853],
-            [-0.7726975624367143, 0.5422431772370664],
-            [-0.27264284694027374, -0.27692000000000005],
-            [0.191375558305514, -0.31858999999999993],
-            [2.048, 0.0],
-            [0.3821699999999999, -0.896229816148446],
-            [-0.7085110933262221, 0.0],
-            [-0.8912055180959243, -0.23524999999999974],
-            [0.0, 1.2022767998385142],
-            [0.31867, 2.048],
-            [0.4353459945606931, -0.5351577612613215],
-            [-0.05614795431522378, 0.4412061957063193]]
-        )
-
-    res = opt.get_result()
-    assert_array_equal(res.x_iters, res_known)
-
-
-@pytest.mark.consistent
-@pytest.mark.parametrize("acq_optimizer", ACQ_OPTIMIZER)
-def test_result_rosenbrock(acq_optimizer):
-
-    bench1 = SimpleCat()
-
-    opt = Optimizer(bench1.get_bounds(), "ENTING", n_initial_points=5,
-                        acq_optimizer=acq_optimizer,
-                        base_estimator_kwargs={"min_child_samples":2},
-                        random_state=100)
-
-    opt.run(bench1, n_iter=10, no_progress_bar=True)
-
-    if acq_optimizer == "sampling":
-        res_known = np.array(
-            [[0.6846225344648778, -0.35181440637953276, 'pow2'],
-            [-1.4055806396985395, -1.3731556796559956, 'mult6'],
-            [-1.159569030229665, -0.18904038856307892, 'pow2'],
-            [-1.7452758285009282, 0.4972475531459488, 'pow2'],
-            [0.24918505872419194, -1.9767312342018637, 'mult6'],
-            [1.9533151312376096, 1.9911019960902978, 'mult6'],
-            [-0.30869536630063044, -0.8116308097286711, 'mult6'],
-            [1.9832830176095446, -1.6137030366593792, 'pow2'],
-            [-0.7534228429864178, -1.9474210829020842, 'mult6'],
-            [-1.9397005641914857, -1.9937770993390025, 'mult6']]
-        )
-
-    elif acq_optimizer == "global":
-        res_known = np.array(
-            [[0.6846225344648778, -0.35181440637953276, 'pow2'],
-            [-1.4055806396985395, -1.3731556796559956, 'mult6'],
-            [-1.159569030229665, -0.18904038856307892, 'pow2'],
-            [-1.7452758285009282, 0.4972475531459488, 'pow2'],
-            [0.24918505872419194, -1.9767312342018637, 'mult6'],
-            [2.0, 2.0, 'mult6'],
-            [-0.26728431683174625, -0.8793390813532928, 'mult6'],
-            [2.0, -1.6258926741796051, 'pow2'],
-            [-0.7295094146491492, -2.0, 'mult6'],
-            [-2.0, -2.0, 'mult6']]
-        )
-
-    res = opt.get_result()
-    assert_array_equal(res.x_iters, res_known)
\ No newline at end of file
diff --git a/tests/test_sampler.py b/tests/test_sampler.py
deleted file mode 100644
index 9a20c67..0000000
--- a/tests/test_sampler.py
+++ /dev/null
@@ -1,316 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-import pytest
-import numbers
-import numpy as np
-import os
-import yaml
-from tempfile import NamedTemporaryFile
-
-from numpy.testing import assert_array_almost_equal
-from numpy.testing import assert_almost_equal
-from numpy.testing import assert_array_equal
-from numpy.testing import assert_equal
-from numpy.testing import assert_raises
-from scipy import spatial
-from entmoot import Optimizer
-from entmoot.space import Space
-from entmoot.space import Real
-from entmoot.space import Integer
-from entmoot.space import Categorical
-from entmoot.space import check_dimension as space_check_dimension
-from entmoot.sampler.sobol import _bit_lo0, _bit_hi1
-from entmoot.sampler.halton import _van_der_corput_samples, _create_primes
-from entmoot.sampler import Hammersly, Halton, Lhs, Sobol, Grid
-from entmoot.sampler import InitialPointGenerator
-from entmoot.sampler.grid import _create_uniform_grid_include_border
-from entmoot.sampler.grid import _create_uniform_grid_exclude_border
-from entmoot.sampler.grid import _quadrature_combine
-from entmoot.sampler.grid import _create_uniform_grid_only_border
-from entmoot.utils import cook_initial_point_generator
-
-
-LHS_TYPE = ["classic", "centered"]
-CRITERION = ["maximin", "ratio", "correlation", None]
-SAMPLER = ["lhs", "halton", "sobol", "hammersly", "grid"]
-
-
-@pytest.mark.fast_test
-def test_lhs_centered():
-    lhs = Lhs(lhs_type="centered")
-    samples = lhs.generate([(0., 1.), ] * 3, 3)
-    assert_almost_equal(np.sum(samples), 4.5)
-
-
-@pytest.mark.parametrize("samlper", SAMPLER)
-def test_sampler(samlper):
-    s = cook_initial_point_generator(samlper)
-    samples = s.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
-    assert isinstance(s, InitialPointGenerator)
-
-    samples = s.generate([("a", "b", "c")], 3)
-    assert samples[0][0] in ["a", "b", "c"]
-
-    samples = s.generate([("a", "b", "c"), (0, 1)], 1)
-    assert samples[0][0] in ["a", "b", "c"]
-    assert samples[0][1] in [0, 1]
-
-    samples = s.generate([("a", "b", "c"), (0, 1)], 3)
-    assert samples[0][0] in ["a", "b", "c"]
-    assert samples[0][1] in [0, 1]
-
-
-@pytest.mark.parametrize("lhs_type", LHS_TYPE)
-@pytest.mark.parametrize("criterion", CRITERION)
-def test_lhs_criterion(lhs_type, criterion):
-    lhs = Lhs(lhs_type=lhs_type, criterion=criterion, iterations=100)
-    samples = lhs.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
-    samples = lhs.generate([("a", "b", "c")], 3)
-    assert samples[0][0] in ["a", "b", "c"]
-
-    samples = lhs.generate([("a", "b", "c"), (0, 1)], 1)
-    assert samples[0][0] in ["a", "b", "c"]
-    assert samples[0][1] in [0, 1]
-
-    samples = lhs.generate([("a", "b", "c"), (0, 1)], 3)
-    assert samples[0][0] in ["a", "b", "c"]
-    assert samples[0][1] in [0, 1]
-
-
-def test_lhs_pdist():
-    n_dim = 2
-    n_samples = 20
-    lhs = Lhs()
-
-    h = lhs._lhs_normalized(n_dim, n_samples, 0)
-    d_classic = spatial.distance.pdist(np.array(h), 'euclidean')
-    lhs = Lhs(criterion="maximin", iterations=100)
-    h = lhs.generate([(0., 1.), ] * n_dim, n_samples, random_state=0)
-    d = spatial.distance.pdist(np.array(h), 'euclidean')
-    assert np.min(d) > np.min(d_classic)
-
-
-@pytest.mark.parametrize("criterion", CRITERION)
-def test_lhs_random_state(criterion):
-    n_dim = 2
-    n_samples = 20
-    lhs = Lhs()
-
-    h = lhs._lhs_normalized(n_dim, n_samples, 0)
-    h2 = lhs._lhs_normalized(n_dim, n_samples, 0)
-    assert_array_equal(h, h2)
-    lhs = Lhs(criterion=criterion, iterations=100)
-    h = lhs.generate([(0., 1.), ] * n_dim, n_samples, random_state=0)
-    h2 = lhs.generate([(0., 1.), ] * n_dim, n_samples, random_state=0)
-    assert_array_equal(h, h2)
-
-
-@pytest.mark.fast_test
-def test_bit():
-    X = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
-    res = [2, 1, 3, 1, 2, 1, 4, 1, 2, 1]
-    for i in range(len(X)):
-        assert _bit_lo0(X[i]) == res[i]
-
-    X = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
-    res = [1, 2, 2, 3, 3, 3, 3, 4, 4, 4]
-    for i in range(len(X)):
-        assert _bit_hi1(X[i]) == res[i]
-
-
-@pytest.mark.fast_test
-def test_sobol():
-    sobol = Sobol()
-    x, seed = sobol._sobol(3, 0)
-    assert_array_equal(x, [0., 0., 0.])
-    x, seed = sobol._sobol(3, 1)
-    assert_array_equal(x, [0.5, 0.5, 0.5])
-    x, seed = sobol._sobol(3, 2)
-    assert_array_equal(x, [0.75, 0.25, 0.75])
-    x, seed = sobol._sobol(3, 3)
-    assert_array_equal(x, [0.25, 0.75, 0.25])
-    x, seed = sobol._sobol(3, 4)
-    assert_array_equal(x, [0.375, 0.375, 0.625])
-    x, seed = sobol._sobol(3, 5)
-    assert_array_equal(x, [0.875, 0.875, 0.125])
-    x, seed = sobol._sobol(3, 6)
-    assert_array_equal(x, [0.625, 0.125, 0.375])
-
-
-@pytest.mark.fast_test
-def test_generate():
-    sobol = Sobol(randomize=False)
-    x = sobol.generate([(0., 1.), ] * 3, 4)
-    x = np.array(x)
-    assert_array_equal(x[0, :], [0., 0., 0.])
-    assert_array_equal(x[1, :], [0.5, 0.5, 0.5])
-    assert_array_equal(x[2, :], [0.75, 0.25, 0.75])
-    assert_array_equal(x[3, :], [0.25, 0.75, 0.25])
-
-    sobol.set_params(skip=2)
-    assert sobol.skip == 2
-    assert isinstance(sobol, InitialPointGenerator)
-
-
-@pytest.mark.fast_test
-def test_van_der_corput():
-    x = _van_der_corput_samples(range(12), number_base=10)
-    y = [0., 0.1,  0.2,  0.3,  0.4,  0.5,  0.6,  0.7,  0.8,  0.9, 0.01, 0.11]
-    assert_array_equal(x, y)
-
-    x = _van_der_corput_samples(range(9), number_base=2)
-    y = [0., 0.5, 0.25, 0.75, 0.125, 0.625, 0.375, 0.875, 0.0625]
-    assert_array_equal(x, y)
-
-
-@pytest.mark.fast_test
-def test_halton():
-    h = Halton()
-    x = h.generate([(0., 1.), ], 9)
-    y = _van_der_corput_samples(range(9), number_base=2)
-    assert_array_almost_equal(np.array(x).flatten(), y)
-
-    h = Halton()
-    x = h.generate([(0., 1.), ] * 2, 6)
-    y = np.array([[0, 0], [1 / 2, 1 / 3], [1 / 4, 2 / 3], [3 / 4, 1 / 9],
-                  [1 / 8, 4 / 9], [5 / 8, 7 / 9]])
-    assert_array_almost_equal(x, y)
-
-    h = Halton(min_skip=0, max_skip=3)
-    x = h.generate([(0., 1.), ] * 2, 4, random_state=12345)
-    assert_array_almost_equal(x, y[2:])
-
-    samples = h.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
-
-
-@pytest.mark.fast_test
-def test_hammersly():
-    h = Hammersly()
-    x = h.generate([(0., 1.), ] * 2, 4)
-    y = np.array([[0, 0], [1 / 2, 0.25], [1 / 4, 0.5], [3 / 4, 0.75]])
-    assert_almost_equal(x, y)
-
-    samples = h.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
-
-
-@pytest.mark.fast_test
-def test_primes():
-
-    x = _create_primes(1)
-    assert_array_equal(x, [])
-    x = _create_primes(2)
-    assert_array_equal(x, [2])
-    x = _create_primes(3)
-    assert_array_equal(x, [2, 3])
-    x = _create_primes(20)
-    assert_array_equal(x, [2, 3, 5, 7, 11, 13, 17, 19])
-
-
-@pytest.mark.fast_test
-def test_quadrature_combine():
-    a = [1, 2]
-    b = [[4, 4], [5, 6]]
-    x = [[1, 4, 4], [1, 5, 6], [2, 4, 4], [2, 5, 6]]
-    x_test = _quadrature_combine([a, b])
-    assert_array_equal(x_test, x)
-
-
-@pytest.mark.fast_test
-def test_uniform_grid():
-    x = _create_uniform_grid_exclude_border(1, 2)
-    assert_array_equal(x, [[1./3.], [2./3.]])
-    x = _create_uniform_grid_include_border(1, 2)
-    assert_array_equal(x, [[0.], [1.]])
-    x = _create_uniform_grid_only_border(1, 2)
-    assert_array_equal(x, [[0.], [1.]])
-
-    x = _create_uniform_grid_exclude_border(1, 3)
-    assert_array_equal(x, [[1./4.], [2./4.], [3./4.]])
-    x = _create_uniform_grid_include_border(1, 3)
-    assert_array_equal(x, [[0./2.], [1./2.], [2./2.]])
-    x = _create_uniform_grid_only_border(1, 3)
-    assert_array_equal(x, [[0./2.], [1./2.], [2./2.]])
-
-    x = _create_uniform_grid_exclude_border(1, 5)
-    assert_array_equal(x, [[1./6.], [2./6.], [3./6.], [4./6.], [5./6.]])
-    x = _create_uniform_grid_include_border(1, 5)
-    assert_array_equal(x, [[0./4.], [1./4.], [2./4.], [3./4.], [4./4.]])
-    x = _create_uniform_grid_only_border(1, 5)
-    assert_array_equal(x, [[0./4.], [1./4.], [2./4.], [3./4.], [4./4.]])
-
-    x = _create_uniform_grid_exclude_border(2, 2)
-    assert_array_equal(x, [[1. / 3., 1./3.], [1. / 3., 2. / 3.],
-                           [2. / 3., 1. / 3.], [2. / 3., 2. / 3.]])
-    x = _create_uniform_grid_include_border(2, 2)
-    assert_array_equal(x, [[0., 0.], [0., 1.],
-                           [1., 0.], [1., 1.]])
-    x = _create_uniform_grid_only_border(2, 3)
-    assert_array_equal(x, [[0., 0.], [0., 0.5],
-                           [0., 1.], [1., 0.],
-                           [1., 0.5], [1., 1.]])
-
-    assert_raises(AssertionError, _create_uniform_grid_exclude_border, 1, 0)
-    assert_raises(AssertionError, _create_uniform_grid_exclude_border, 0, 1)
-    assert_raises(AssertionError, _create_uniform_grid_include_border, 1, 0)
-    assert_raises(AssertionError, _create_uniform_grid_include_border, 0, 1)
-    assert_raises(AssertionError, _create_uniform_grid_only_border, 1, 1)
-    assert_raises(AssertionError, _create_uniform_grid_only_border, 0, 2)
-
-
-@pytest.mark.fast_test
-def test_grid():
-    grid = Grid()
-    samples = grid.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
-
-    grid = Grid(border="include")
-    samples = grid.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
-
-    grid = Grid(use_full_layout=False)
-    samples = grid.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
-
-    grid = Grid(use_full_layout=True, append_border="include")
-    samples = grid.generate([(0., 1.), ] * 2, 200)
-    assert len(samples) == 200
-    assert len(samples[0]) == 2
\ No newline at end of file
diff --git a/tests/test_space.py b/tests/test_space.py
deleted file mode 100644
index c750c7a..0000000
--- a/tests/test_space.py
+++ /dev/null
@@ -1,828 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-import pytest
-import numbers
-import numpy as np
-import os
-import yaml
-from tempfile import NamedTemporaryFile
-
-from numpy.testing import assert_array_almost_equal
-from numpy.testing import assert_array_equal
-from numpy.testing import assert_equal
-from numpy.testing import assert_raises_regex
-
-from entmoot import Optimizer
-from entmoot.space import Space
-from entmoot.space import Real
-from entmoot.space import Integer
-from entmoot.space import Categorical
-from entmoot.space import check_dimension as space_check_dimension
-from entmoot.utils import normalize_dimensions
-
-
-def check_dimension(Dimension, vals, random_val):
-    x = Dimension(*vals)
-    assert_equal(x, Dimension(*vals))
-    assert x != Dimension(vals[0], vals[1] + 1)
-    assert x != Dimension(vals[0] + 1, vals[1])
-    y = x.rvs(random_state=1)
-    if isinstance(y, list):
-        y = np.array(y)
-    assert_equal(y, random_val)
-
-
-def check_categorical(vals, random_val):
-    x = Categorical(vals)
-    assert_equal(x, Categorical(vals))
-    assert x != Categorical(vals[:-1] + ("zzz",))
-    assert_equal(x.rvs(random_state=1), random_val)
-
-
-def check_limits(value, low, high):
-    # check if low <= value <= high
-    if isinstance(value, list):
-        value = np.array(value)
-    assert np.all(low <= value)
-    assert np.all(high >= value)
-
-
-@pytest.mark.fast_test
-def test_dimensions():
-    check_dimension(Real, (1., 4.), 2.251066014107722)
-    check_dimension(Real, (1, 4), 2.251066014107722)
-    check_dimension(Integer, (1, 4), 2)
-    check_dimension(Integer, (1., 4.), 2)
-    check_dimension(Integer, (1, 4), 2)
-    check_categorical(("a", "b", "c", "d"), "b")
-    check_categorical((1., 2., 3., 4.), 2.)
-    check_categorical((1, 2, 3, 4), 2)
-
-
-@pytest.mark.fast_test
-def test_real_log_sampling_in_bounds():
-    dim = Real(low=1, high=32, prior='log-uniform', transform='normalize')
-
-    # round trip a value that is within the bounds of the space
-    #
-    # x = dim.inverse_transform(dim.transform(31.999999999999999))
-    for n in (32., 31.999999999999999):
-        round_tripped = dim.inverse_transform(dim.transform([n]))
-        assert np.allclose([n], round_tripped)
-        assert n in dim
-        assert round_tripped in dim
-
-
-@pytest.mark.fast_test
-def test_real():
-    a = Real(1, 25)
-    for i in range(50):
-        r = a.rvs(random_state=i)
-        check_limits(r, 1, 25)
-        assert r in a
-
-    random_values = a.rvs(random_state=0, n_samples=10)
-    assert len(random_values) == 10
-    assert_array_equal(a.transform(random_values), random_values)
-    assert_array_equal(a.inverse_transform(random_values), random_values)
-
-    log_uniform = Real(10**-5, 10**5, prior="log-uniform")
-    assert log_uniform != Real(10**-5, 10**5)
-    for i in range(50):
-        random_val = log_uniform.rvs(random_state=i)
-        check_limits(random_val, 10**-5, 10**5)
-    random_values = log_uniform.rvs(random_state=0, n_samples=10)
-    assert len(random_values) == 10
-    transformed_vals = log_uniform.transform(random_values)
-    assert_array_equal(transformed_vals, np.log10(random_values))
-    assert_array_equal(
-        log_uniform.inverse_transform(transformed_vals), random_values)
-
-
-@pytest.mark.fast_test
-def test_real_bounds():
-    # should give same answer as using check_limits() but this is easier
-    # to read
-    a = Real(1., 2.1)
-    assert 0.99 not in a
-    assert 1. in a
-    assert 2.09 in a
-    assert 2.1 in a
-    assert np.nextafter(2.1, 3.) not in a
-
-
-@pytest.mark.fast_test
-def test_integer():
-    a = Integer(1, 10)
-    for i in range(50):
-        r = a.rvs(random_state=i)
-        assert 1 <= r
-        assert 11 >= r
-        assert r in a
-
-    random_values = a.rvs(random_state=0, n_samples=10)
-    assert_array_equal(random_values.shape, (10))
-    assert_array_equal(a.transform(random_values), random_values)
-    assert_array_equal(a.inverse_transform(random_values), random_values)
-
-
-@pytest.mark.fast_test
-def test_categorical_transform():
-    categories = ["apple", "orange", "banana", None, True, False, 3]
-    cat = Categorical(categories,transform="onehot")
-
-    apple = [1., 0., 0., 0., 0., 0., 0.]
-    orange = [0., 1.0, 0.0, 0.0, 0., 0., 0.]
-    banana = [0., 0., 1., 0., 0., 0., 0.]
-    none = [0., 0., 0., 1., 0., 0., 0.]
-    true = [0., 0., 0., 0., 1., 0., 0.]
-    false = [0., 0., 0., 0., 0., 1., 0.]
-    three = [0., 0., 0., 0., 0., 0., 1.]
-
-    assert_equal(cat.transformed_size, 7)
-    assert_equal(cat.transformed_size, cat.transform(["apple"]).size)
-    assert_array_equal(
-        cat.transform(categories),
-        [apple, orange, banana, none, true, false, three]
-        )
-    assert_array_equal(cat.transform(["apple", "orange"]), [apple, orange])
-    assert_array_equal(cat.transform(["apple", "banana"]), [apple, banana])
-    assert_array_equal(cat.inverse_transform([apple, orange]),
-                       ["apple", "orange"])
-    assert_array_equal(cat.inverse_transform([apple, banana]),
-                       ["apple", "banana"])
-    ent_inverse = cat.inverse_transform(
-        [apple, orange, banana, none, true, false, three])
-    assert_array_equal(ent_inverse, categories)
-
-
-@pytest.mark.fast_test
-def test_categorical_transform_binary():
-    categories = ["apple", "orange"]
-    cat = Categorical(categories,transform="onehot")
-
-    apple = [0.]
-    orange = [1.]
-
-    assert_equal(cat.transformed_size, 1)
-    assert_equal(cat.transformed_size, cat.transform(["apple"]).size)
-    assert_array_equal(cat.transform(categories), [apple, orange])
-    assert_array_equal(cat.transform(["apple", "orange"]), [apple, orange])
-    assert_array_equal(cat.inverse_transform([apple, orange]),
-                       ["apple", "orange"])
-    ent_inverse = cat.inverse_transform([apple, orange])
-    assert_array_equal(ent_inverse, categories)
-
-
-@pytest.mark.fast_test
-def test_categorical_repr():
-    small_cat = Categorical([1, 2, 3, 4, 5])
-    assert (small_cat.__repr__() ==
-            "Categorical(categories=(1, 2, 3, 4, 5), prior=None)")
-
-    big_cat = Categorical([1, 2, 3, 4, 5, 6, 7, 8])
-    assert (big_cat.__repr__() ==
-            'Categorical(categories=(1, 2, 3, ..., 6, 7, 8), prior=None)')
-
-
-@pytest.mark.fast_test
-def test_space_consistency():
-    # Reals (uniform)
-
-    s1 = Space([Real(0.0, 1.0)])
-    s2 = Space([Real(0.0, 1.0)])
-    s3 = Space([Real(0, 1)])
-    s4 = Space([(0.0, 1.0)])
-    s5 = Space([(0.0, 1.0, "uniform")])
-    s6 = Space([(0, 1.0)])
-    s7 = Space([(np.float64(0.0), 1.0)])
-    s8 = Space([(0, np.float64(1.0))])
-    a1 = s1.rvs(n_samples=10, random_state=0)
-    a2 = s2.rvs(n_samples=10, random_state=0)
-    a3 = s3.rvs(n_samples=10, random_state=0)
-    a4 = s4.rvs(n_samples=10, random_state=0)
-    a5 = s5.rvs(n_samples=10, random_state=0)
-    assert_equal(s1, s2)
-    assert_equal(s1, s3)
-    assert_equal(s1, s4)
-    assert_equal(s1, s5)
-    assert_equal(s1, s6)
-    assert_equal(s1, s7)
-    assert_equal(s1, s8)
-    assert_array_equal(a1, a2)
-    assert_array_equal(a1, a3)
-    assert_array_equal(a1, a4)
-    assert_array_equal(a1, a5)
-
-    # Reals (log-uniform)
-    s1 = Space([Real(10**-3.0, 10**3.0, prior="log-uniform", base=10)])
-    s2 = Space([Real(10**-3.0, 10**3.0, prior="log-uniform", base=10)])
-    s3 = Space([Real(10**-3, 10**3, prior="log-uniform", base=10)])
-    s4 = Space([(10**-3.0, 10**3.0, "log-uniform", 10)])
-    s5 = Space([(np.float64(10**-3.0), 10**3.0, "log-uniform", 10)])
-    a1 = s1.rvs(n_samples=10, random_state=0)
-    a2 = s2.rvs(n_samples=10, random_state=0)
-    a3 = s3.rvs(n_samples=10, random_state=0)
-    a4 = s4.rvs(n_samples=10, random_state=0)
-    assert_equal(s1, s2)
-    assert_equal(s1, s3)
-    assert_equal(s1, s4)
-    assert_equal(s1, s5)
-    assert_array_equal(a1, a2)
-    assert_array_equal(a1, a3)
-    assert_array_equal(a1, a4)
-
-    # Integers
-    s1 = Space([Integer(1, 5)])
-    s2 = Space([Integer(1.0, 5.0)])
-    s3 = Space([(1, 5)])
-    s4 = Space([(np.int64(1.0), 5)])
-    s5 = Space([(1, np.int64(5.0))])
-    a1 = s1.rvs(n_samples=10, random_state=0)
-    a2 = s2.rvs(n_samples=10, random_state=0)
-    a3 = s3.rvs(n_samples=10, random_state=0)
-    assert_equal(s1, s2)
-    assert_equal(s1, s3)
-    assert_equal(s1, s4)
-    assert_equal(s1, s5)
-    assert_array_equal(a1, a2)
-    assert_array_equal(a1, a3)
-
-    # Integers (log-uniform)
-    s1 = Space([Integer(16, 512, prior="log-uniform", base=2)])
-    s2 = Space([Integer(16.0, 512.0, prior="log-uniform", base=2)])
-    s3 = Space([(16, 512, "log-uniform", 2)])
-    s4 = Space([(np.int64(16.0), 512, "log-uniform", 2)])
-    s5 = Space([(16, np.int64(512.0), "log-uniform", 2)])
-    a1 = s1.rvs(n_samples=10, random_state=0)
-    a2 = s2.rvs(n_samples=10, random_state=0)
-    a3 = s3.rvs(n_samples=10, random_state=0)
-    assert_equal(s1, s2)
-    assert_equal(s1, s3)
-    assert_equal(s1, s4)
-    assert_equal(s1, s5)
-    assert_array_equal(a1, a2)
-    assert_array_equal(a1, a3)
-
-    # Categoricals
-    s1 = Space([Categorical(["a", "b", "c"])])
-    s2 = Space([Categorical(["a", "b", "c"])])
-    s3 = Space([["a", "b", "c"]])
-    a1 = s1.rvs(n_samples=10, random_state=0)
-    a2 = s2.rvs(n_samples=10, random_state=0)
-    a3 = s3.rvs(n_samples=10, random_state=0)
-    assert_equal(s1, s2)
-    assert_array_equal(a1, a2)
-    assert_equal(s1, s3)
-    assert_array_equal(a1, a3)
-
-    s1 = Space([(True, False)])
-    s2 = Space([Categorical([True, False])])
-    s3 = Space([np.array([True, False])])
-    assert s1 == s2 == s3
-
-    # Categoricals Integer
-    s1 = Space([Categorical([1, 2, 3])])
-    s2 = Space([Categorical([1, 2, 3])])
-    s3 = Space([[1, 2, 3]])
-    a1 = s1.rvs(n_samples=10, random_state=0)
-    a2 = s2.rvs(n_samples=10, random_state=0)
-    a3 = s3.rvs(n_samples=10, random_state=0)
-    assert_equal(s1, s2)
-    assert_array_equal(a1, a2)
-    assert_equal(s1, s3)
-    assert_array_equal(a1, a3)
-
-    s1 = Space([(True, False)])
-    s2 = Space([Categorical([True, False])])
-    s3 = Space([np.array([True, False])])
-    assert s1 == s2 == s3
-
-@pytest.mark.fast_test
-def test_space_api():
-    space = Space([(0.0, 1.0), (-5, 5),
-                   Categorical(["a", "b", "c"], transform="onehot"), 
-                   (1.0, 5.0, "log-uniform"), ("e", "f")])
-
-    cat_space = Space([(1, "r"), (1.0, "r")])
-    assert isinstance(cat_space.dimensions[0], Categorical)
-    assert isinstance(cat_space.dimensions[1], Categorical)
-
-    assert_equal(len(space.dimensions), 5)
-    assert isinstance(space.dimensions[0], Real)
-    assert isinstance(space.dimensions[1], Integer)
-    assert isinstance(space.dimensions[2], Categorical)
-    assert isinstance(space.dimensions[3], Real)
-    assert isinstance(space.dimensions[4], Categorical)
-
-    samples = space.rvs(n_samples=10, random_state=0)
-    assert_equal(len(samples), 10)
-    assert_equal(len(samples[0]), 5)
-
-    assert isinstance(samples, list)
-    for n in range(4):
-        assert isinstance(samples[n], list)
-
-    assert isinstance(samples[0][0], numbers.Real)
-    assert isinstance(samples[0][1], numbers.Integral)
-    assert isinstance(samples[0][2], str)
-    assert isinstance(samples[0][3], numbers.Real)
-    assert isinstance(samples[0][4], str)
-
-    samples_transformed = space.transform(samples)
-    assert_equal(samples_transformed.shape[0], len(samples))
-    assert_equal(samples_transformed.shape[1], 1 + 1 + 3 + 1 + 1)
-
-    # our space contains mixed types, this means we can't use
-    # `array_allclose` or similar to check points are close after a round-trip
-    # of transformations
-    for orig, round_trip in zip(samples,
-                                space.inverse_transform(samples_transformed)):
-        assert space.distance(orig, round_trip) < 1.e-8
-
-    samples = space.inverse_transform(samples_transformed)
-    assert isinstance(samples[0][0], numbers.Real)
-    assert isinstance(samples[0][1], numbers.Integral)
-    assert isinstance(samples[0][2], str)
-    assert isinstance(samples[0][3], numbers.Real)
-    assert isinstance(samples[0][4], str)
-
-    for b1, b2 in zip(space.bounds,
-                      [(0.0, 1.0), (-5, 5),
-                       np.asarray(["a", "b", "c"]), (1.0, 5.0),
-                       np.asarray(["e", "f"])]):
-        assert_array_equal(b1, b2)
-
-    for b1, b2 in zip(space.transformed_bounds,
-                      [(0.0, 1.0), (-5, 5), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0),
-                       (np.log10(1.0), np.log10(5.0)), (0.0, 1.0)]):
-        assert_array_equal(b1, b2)
-
-
-@pytest.mark.fast_test
-def test_space_from_space():
-    # can you pass a Space instance to the Space constructor?
-    space = Space([(0.0, 1.0), (-5, 5),
-                   ("a", "b", "c"), (1.0, 5.0, "log-uniform"), ("e", "f")])
-
-    space2 = Space(space)
-
-    assert_equal(space, space2)
-
-
-@pytest.mark.fast_test
-def test_constant_property():
-    space = Space([(0.0, 1.0), (1,),
-                   ("a", "b", "c"), (1.0, 5.0, "log-uniform"), ("e",)])
-    assert space.n_constant_dimensions == 2
-    for i in [1, 4]:
-        assert space.dimensions[i].is_constant
-    for i in [0, 2, 3]:
-        assert not space.dimensions[i].is_constant
-
-
-@pytest.mark.fast_test
-def test_set_get_transformer():
-    # can you pass a Space instance to the Space constructor?
-    space = Space([(0.0, 1.0), (-5, 5),
-                   ("a", "b", "c"), (1.0, 5.0, "log-uniform"), ("e", "f")])
-
-    transformer = space.get_transformer()
-    assert_array_equal(["identity", "identity", "label",
-                        "identity", "label"], transformer)
-    space.set_transformer("normalize")
-    transformer = space.get_transformer()
-    assert_array_equal(["normalize"] * 5, transformer)
-    space.set_transformer(transformer)
-    assert_array_equal(transformer, space.get_transformer())
-
-    space.set_transformer_by_type("label", Categorical)
-    assert space.dimensions[2].transform(["a"]) == [0]
-
-
-@pytest.mark.fast_test
-def test_normalize():
-    # can you pass a Space instance to the Space constructor?
-    space = Space([(0.0, 1.0), (-5, 5),
-                   ("a", "b", "c"), (1.0, 5.0, "log-uniform"), ("e", "f")])
-    space.set_transformer("normalize")
-    X = [[0., -5, 'a', 1., 'e']]
-    Xt = np.zeros((1, 5))
-    assert_array_equal(space.transform(X), Xt)
-    assert_array_equal(space.inverse_transform(Xt), X)
-    assert_array_equal(space.inverse_transform(space.transform(X)), X)
-
-
-@pytest.mark.fast_test
-def test_normalize_types():
-    # can you pass a Space instance to the Space constructor?
-    space = Space([(0.0, 1.0), Integer(-5, 5, dtype=int), (True, False)])
-    space.set_transformer("normalize")
-    X = [[0., -5, False]]
-    Xt = np.zeros((1, 3))
-    assert_array_equal(space.transform(X), Xt)
-    assert_array_equal(space.inverse_transform(Xt), X)
-    assert_array_equal(space.inverse_transform(space.transform(X)), X)
-    assert isinstance(space.inverse_transform(Xt)[0][0], float)
-    assert isinstance(space.inverse_transform(Xt)[0][1], int)
-    assert isinstance(space.inverse_transform(Xt)[0][2], (np.bool_, bool))
-
-@pytest.mark.fast_test
-def test_normalize_real():
-
-    a = Real(2.0, 30.0, transform="normalize")
-    for i in range(50):
-        check_limits(a.rvs(random_state=i), 2, 30)
-
-    rng = np.random.RandomState(0)
-    X = rng.randn(100)
-    X = 28 * (X - X.min()) / (X.max() - X.min()) + 2
-
-    # Check transformed values are in [0, 1]
-    assert np.all(a.transform(X) <= np.ones_like(X))
-    assert np.all(np.zeros_like(X) <= a.transform(X))
-
-    # Check inverse transform
-    assert_array_almost_equal(a.inverse_transform(a.transform(X)), X)
-
-    # log-uniform prior
-    a = Real(10**2.0, 10**4.0, prior="log-uniform", transform="normalize")
-    for i in range(50):
-        check_limits(a.rvs(random_state=i), 10**2, 10**4)
-
-    rng = np.random.RandomState(0)
-    X = np.clip(10**3 * rng.randn(100), 10**2.0, 10**4.0)
-
-    # Check transform
-    assert np.all(a.transform(X) <= np.ones_like(X))
-    assert np.all(np.zeros_like(X) <= a.transform(X))
-
-    # Check inverse transform
-    assert_array_almost_equal(a.inverse_transform(a.transform(X)), X)
-
-    a = Real(0, 1, transform="normalize", dtype=float)
-    for i in range(50):
-        check_limits(a.rvs(random_state=i), 0, 1)
-    assert_array_equal(a.transformed_bounds, (0, 1))
-
-    X = rng.rand()
-    # Check transformed values are in [0, 1]
-    assert np.all(a.transform(X) <= np.ones_like(X))
-    assert np.all(np.zeros_like(X) <= a.transform(X))
-
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, float)
-    assert_array_equal(X_orig, X)
-
-    a = Real(0, 1, transform="normalize", dtype='float64')
-    X = np.float64(rng.rand())
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, np.float64)
-
-    a = Real(0, 1, transform="normalize", dtype=np.float64)
-    X = np.float64(rng.rand())
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, np.float64)
-
-    a = Real(0, 1, transform="normalize", dtype='float64')
-    X = np.float64(rng.rand())
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, np.float64)
-
-
-@pytest.mark.fast_test
-def test_normalize_integer():
-    a = Integer(2, 30, transform="normalize")
-    for i in range(50):
-        check_limits(a.rvs(random_state=i), 2, 30)
-    assert_array_equal(a.transformed_bounds, (0, 1))
-    rng = np.random.RandomState(0)
-    X = rng.randint(2, 31, dtype=np.int64)
-    # Check transformed values are in [0, 1]
-    assert np.all(a.transform(X) <= np.ones_like(X))
-    assert np.all(np.zeros_like(X) <= a.transform(X))
-
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, np.int64)
-    assert_array_equal(X_orig, X)
-
-    a = Integer(2, 30, transform="normalize", dtype=int)
-    X = rng.randint(2, 31, dtype=int)
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, int)
-
-    a = Integer(2, 30, transform="normalize", dtype='int')
-    X = rng.randint(2, 31, dtype=int)
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, int)
-
-    a = Integer(2, 30, prior="log-uniform", base=2, transform="normalize",
-                dtype=int)
-    for i in range(50):
-        check_limits(a.rvs(random_state=i), 2, 30)
-    assert_array_equal(a.transformed_bounds, (0, 1))
-
-    X = rng.randint(2, 31, dtype=int)
-    # Check transformed values are in [0, 1]
-    assert np.all(a.transform(X) <= np.ones_like(X))
-    assert np.all(np.zeros_like(X) <= a.transform(X))
-
-    # Check inverse transform
-    X_orig = a.inverse_transform(a.transform(X))
-    assert isinstance(X_orig, int)
-    assert_array_equal(X_orig, X)
-
-
-@pytest.mark.fast_test
-def test_normalize_categorical():
-    categories = ["cat", "dog", "rat"]
-    a = Categorical(categories, transform="normalize")
-    for i in range(len(categories)):
-        assert a.rvs(random_state=i)[0] in categories
-    assert a.inverse_transform([0.]) == [categories[0]]
-    assert a.inverse_transform([0.5]) == [categories[1]]
-    assert a.inverse_transform([1.0]) == [categories[2]]
-    assert_array_equal(categories, a.inverse_transform([0., 0.5, 1]))
-
-    categories = [1, 2, 3]
-    a = Categorical(categories, transform="normalize")
-    assert_array_equal(categories, np.sort(np.unique(a.rvs(100,
-                                                           random_state=1))))
-    assert_array_equal(categories, a.inverse_transform([0., 0.5, 1.]))
-
-    categories = [1., 2., 3.]
-    a = Categorical(categories, transform="normalize")
-    assert_array_equal(categories, np.sort(np.unique(a.rvs(100,
-                                                           random_state=1))))
-    assert_array_equal(categories, a.inverse_transform([0., 0.5, 1.]))
-
-    categories = [1, 2, 3]
-    a = Categorical(categories, transform="string")
-    a.set_transformer("normalize")
-    assert_array_equal(categories, np.sort(np.unique(a.rvs(100,
-                                                           random_state=1))))
-    assert_array_equal(categories, a.inverse_transform([0., 0.5, 1.]))
-
-@pytest.mark.fast_test
-def test_normalize_integer():
-    for dtype in ['int', 'int8', 'int16', 'int32', 'int64',
-                  'uint8', 'uint16', 'uint32', 'uint64']:
-        a = Integer(2, 30, transform="normalize", dtype=dtype)
-        for X in range(2, 31):
-            X_orig = a.inverse_transform(a.transform(X))
-            assert_array_equal(X_orig, X)
-    for dtype in [int, np.int8, np.int16, np.int32, np.int64,
-                  np.uint8, np.uint16, np.uint32, np.uint64]:
-        a = Integer(2, 30, transform="normalize", dtype=dtype)
-        for X in range(2, 31):
-            X_orig = a.inverse_transform(a.transform(X))
-            assert_array_equal(X_orig, X)
-            assert isinstance(X_orig, dtype)
-
-
-def check_valid_transformation(klass):
-    assert klass(2, 30, transform="normalize")
-    assert klass(2, 30, transform="identity")
-    assert_raises_regex(ValueError, "should be 'normalize' or 'identity'",
-                        klass, 2, 30, transform='not a valid transform name')
-
-
-@pytest.mark.fast_test
-def test_valid_transformation():
-    check_valid_transformation(Integer)
-    check_valid_transformation(Real)
-
-
-@pytest.mark.fast_test
-def test_invalid_dimension():
-    assert_raises_regex(ValueError, "has to be a list or tuple",
-                        space_check_dimension, "23")
-    # single value fixes dimension of space
-    space_check_dimension((23,))
-
-
-@pytest.mark.fast_test
-def test_categorical_identity():
-    categories = ["cat", "dog", "rat"]
-    cat = Categorical(categories, transform="identity")
-    samples = cat.rvs(100)
-    assert all([t in categories for t in cat.rvs(100)])
-    transformed = cat.transform(samples)
-    assert_array_equal(transformed, samples)
-    assert_array_equal(samples, cat.inverse_transform(transformed))
-
-
-@pytest.mark.fast_test
-def test_categorical_string():
-    categories = [1, 2, 3]
-    categories_transformed = ["1", "2", "3"]
-    cat = Categorical(categories, transform="string")
-    samples = cat.rvs(100)
-    assert all([t in categories for t in cat.rvs(100)])
-    transformed = cat.transform(samples)
-    assert all([t in categories_transformed for t in transformed])
-    assert_array_equal(samples, cat.inverse_transform(transformed))
-
-
-@pytest.mark.fast_test
-def test_categorical_distance():
-    categories = ['car', 'dog', 'orange']
-    cat = Categorical(categories)
-    for cat1 in categories:
-        for cat2 in categories:
-            delta = cat.distance(cat1, cat2)
-            if cat1 == cat2:
-                assert delta == 0
-            else:
-                assert delta == 1
-
-
-@pytest.mark.fast_test
-def test_integer_distance():
-    ints = Integer(1, 10)
-    for i in range(1, 10+1):
-        assert_equal(ints.distance(4, i), abs(4 - i))
-
-
-@pytest.mark.fast_test
-def test_integer_distance_out_of_range():
-    ints = Integer(1, 10)
-    assert_raises_regex(RuntimeError, "compute distance for values within",
-                        ints.distance, 11, 10)
-
-
-@pytest.mark.fast_test
-def test_real_distance_out_of_range():
-    ints = Real(1, 10)
-    assert_raises_regex(RuntimeError, "compute distance for values within",
-                        ints.distance, 11, 10)
-
-
-@pytest.mark.fast_test
-def test_real_distance():
-    reals = Real(1, 10)
-    for i in range(1, 10+1):
-        assert_equal(reals.distance(4.1234, i), abs(4.1234 - i))
-
-
-@pytest.mark.parametrize("dimension, bounds",
-                         [(Real, (2, 1)), (Integer, (2, 1)),
-                          (Real, (2, 2)), (Integer, (2, 2))])
-def test_dimension_bounds(dimension, bounds):
-    with pytest.raises(ValueError) as exc:
-        dim = dimension(*bounds)
-        assert "has to be less than the upper bound " in exc.value.args[0]
-
-
-@pytest.mark.parametrize("dimension, name",
-                         [(Real(1, 2, name="learning_rate"), "learning_rate"),
-                          (Integer(1, 100, name="n_trees"), "n_trees"),
-                          (Categorical(["red, blue"], name="colors"), "colors")])
-def test_dimension_name(dimension, name):
-    assert dimension.name == name
-
-
-def test_dimension_name():
-    notnames = [1, 1., True]
-    for n in notnames:
-        with pytest.raises(ValueError) as exc:
-            real = Real(1, 2, name=n)
-            assert("Dimension's name must be either string or"
-                   "None." == exc.value.args[0])
-    s = Space([Real(1, 2, name="a"),
-               Integer(1, 100, name="b"),
-               Categorical(["red, blue"], name="c")])
-    assert s["a"] == (0, s.dimensions[0])
-    assert s["a", "c"] == [(0, s.dimensions[0]), (2, s.dimensions[2])]
-    assert s[["a", "c"]] == [(0, s.dimensions[0]), (2, s.dimensions[2])]
-    assert s[("a", "c")] == [(0, s.dimensions[0]), (2, s.dimensions[2])]
-    assert s[0] == (0, s.dimensions[0])
-    assert s[0, "c"] == [(0, s.dimensions[0]), (2, s.dimensions[2])]
-    assert s[0, 2] == [(0, s.dimensions[0]), (2, s.dimensions[2])]
-
-
-@pytest.mark.parametrize("dimension",
-                         [Real(1, 2), Integer(1, 100), Categorical(["red, blue"])])
-def test_dimension_name_none(dimension):
-    assert dimension.name is None
-
-
-@pytest.mark.fast_test
-def test_space_from_yaml():
-    with NamedTemporaryFile(delete=False) as tmp:
-        tmp.write(b"""
-        Space:
-            - Real:
-                low: 0.0
-                high: 1.0
-            - Integer:
-                low: -5
-                high: 5
-            - Categorical:
-                categories:
-                - a
-                - b
-                - c
-            - Real:
-                low: 1.0
-                high: 5.0
-                prior: log-uniform
-            - Categorical:
-                categories:
-                - e
-                - f
-        """)
-        tmp.flush()
-
-        space = Space([(0.0, 1.0),
-                       (-5, 5),
-                       ("a", "b", "c"),
-                       (1.0, 5.0, "log-uniform"),
-                       ("e", "f")])
-
-        space2 = Space.from_yaml(tmp.name)
-        assert_equal(space, space2)
-        tmp.close()
-        os.unlink(tmp.name)
-
-@pytest.mark.parametrize("name", [1, 1., True])
-def test_dimension_with_invalid_names(name):
-    with pytest.raises(ValueError) as exc:
-        Real(1, 2, name=name)
-    assert("Dimension's name must be either string or None." ==
-           exc.value.args[0])
-
-
-
-@pytest.mark.fast_test
-def test_purely_categorical_space():
-    # Test reproduces the bug in #908, make sure it doesn't come back
-    dims = [Categorical(['a', 'b', 'c']), Categorical(['A', 'B', 'C'])]
-    optimizer = Optimizer(dims, n_initial_points=2, random_state=3, acq_optimizer="sampling")
-
-    for _ in range(2):
-        x = optimizer.ask()
-        # before the fix this call raised an exception
-        optimizer.tell(x, np.random.uniform())
-
-
-@pytest.mark.fast_test
-def test_partly_categorical_space():
-    dims = Space([Categorical(['a', 'b', 'c']), Categorical(['A', 'B', 'C'])])
-    assert dims.is_partly_categorical
-    dims = Space([Categorical(['a', 'b', 'c']), Integer(1, 2)])
-    assert dims.is_partly_categorical
-    assert not dims.is_categorical
-    dims = Space([Integer(1, 2), Integer(1, 2)])
-    assert not dims.is_partly_categorical
-
-
-@pytest.mark.fast_test
-def test_normalize_bounds():
-    bounds = [(-999, 189000), Categorical((True, False))]
-    space = Space(normalize_dimensions(bounds))
-    for a in np.linspace(1e-9, 0.4999, 1000):
-        x = space.inverse_transform([[a, a]])
-        check_limits(x[0][0], -999, 189000)
-        y = space.transform(x)
-        check_limits(y, 0., 1.)
-    for a in np.linspace(0.50001, 1e-9 + 1., 1000):
-        x = space.inverse_transform([[a, a]])
-        check_limits(x[0][0], -999, 189000)
-        y = space.transform(x)
-        check_limits(y, 0., 1.)
diff --git a/tests/test_transformer.py b/tests/test_transformer.py
deleted file mode 100644
index 6f322a8..0000000
--- a/tests/test_transformer.py
+++ /dev/null
@@ -1,150 +0,0 @@
-"""
-Copyright (c) 2016-2020 The scikit-optimize developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-"""
-
-import pytest
-import numbers
-import numpy as np
-from numpy.testing import assert_raises
-from numpy.testing import assert_array_equal
-from numpy.testing import assert_equal
-from numpy.testing import assert_raises_regex
-from entmoot.space import LogN, Normalize
-from entmoot.space.transformers import StringEncoder, LabelEncoder, Identity
-
-
-@pytest.mark.fast_test
-def test_logn2_integer():
-
-    transformer = LogN(2)
-    for X in range(2, 31):
-        X_orig = transformer.inverse_transform(transformer.transform(X))
-        assert_array_equal(int(np.round(X_orig)), X)
-
-
-@pytest.mark.fast_test
-def test_logn10_integer():
-
-    transformer = LogN(2)
-    for X in range(2, 31):
-        X_orig = transformer.inverse_transform(transformer.transform(X))
-        assert_array_equal(int(np.round(X_orig)), X)
-
-
-@pytest.mark.fast_test
-def test_integer_encoder():
-
-    transformer = LabelEncoder()
-    X = [1, 5, 9]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), [0, 1, 2])
-    assert_array_equal(transformer.inverse_transform([0, 1, 2]), X)
-
-    transformer = LabelEncoder(X)
-    assert_array_equal(transformer.transform(X), [0, 1, 2])
-    assert_array_equal(transformer.inverse_transform([0, 1, 2]), X)
-
-    X = ["a", "b", "c"]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), [0, 1, 2])
-    assert_array_equal(transformer.inverse_transform([0, 1, 2]), X)
-
-    transformer = LabelEncoder(X)
-    assert_array_equal(transformer.transform(X), [0, 1, 2])
-    assert_array_equal(transformer.inverse_transform([0, 1, 2]), X)
-
-
-@pytest.mark.fast_test
-def test_string_encoder():
-
-    transformer = StringEncoder()
-    X = [1, 5, 9]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), ['1', '5', '9'])
-    assert_array_equal(transformer.inverse_transform(['1', '5', '9']), X)
-
-    X = ['a', True, 1]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), ['a', 'True', '1'])
-    assert_array_equal(transformer.inverse_transform(['a', 'True', '1']), X)
-
-    X = ["a", "b", "c"]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), X)
-    assert_array_equal(transformer.inverse_transform(X), X)
-
-
-@pytest.mark.fast_test
-def test_identity_encoder():
-
-    transformer = Identity()
-    X = [1, 5, 9, 9, 5, 1]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), X)
-    assert_array_equal(transformer.inverse_transform(X), X)
-
-    X = ['a', True, 1, 'a', True, 1]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), X)
-    assert_array_equal(transformer.inverse_transform(X), X)
-
-    X = ["a", "b", "c", "a", "b", "c"]
-    transformer.fit(X)
-    assert_array_equal(transformer.transform(X), X)
-    assert_array_equal(transformer.inverse_transform(X), X)
-
-
-@pytest.mark.fast_test
-def test_normalize_integer():
-    transformer = Normalize(1, 20, is_int=True)
-    assert transformer.transform(19.8) == 1.0
-    assert transformer.transform(20.2) == 1.0
-    assert transformer.transform(1.2) == 0.0
-    assert transformer.transform(0.9) == 0.0
-    assert_raises(ValueError, transformer.transform, 20.6)
-    assert_raises(ValueError, transformer.transform, 0.4)
-
-    assert transformer.inverse_transform(0.99) == 20
-    assert transformer.inverse_transform(0.01) == 1
-    assert_raises(ValueError, transformer.inverse_transform, 1. + 1e-6)
-    assert_raises(ValueError, transformer.transform, 0. - 1e-6)
-    transformer = Normalize(0, 20, is_int=True)
-    assert transformer.transform(-0.2) == 0.0
-    assert_raises(ValueError, transformer.transform, -0.6)
-
-
-@pytest.mark.fast_test
-def test_normalize():
-    transformer = Normalize(1, 20, is_int=False)
-    assert transformer.transform(20.) == 1.0
-    assert transformer.transform(1.) == 0.0
-    assert_raises(ValueError, transformer.transform, 20. + 1e-6)
-    assert_raises(ValueError, transformer.transform, 1.0 - 1e-6)
-    assert_raises(ValueError, transformer.inverse_transform, 1. + 1e-6)
-    assert_raises(ValueError, transformer.transform, 0. - 1e-6)
\ No newline at end of file