Repository: Systems-Theory-in-Systems-Biology/EPI (language: python)

Issue:
Scale all data to have mean 0 and variance 1 to improve numerical stability for high data dimensionality
**Is your feature request related to a problem? Please describe.**
Right now, EPI doesn't work properly for the stock example if the distributions are too wide. This seems to stem from an numerical issue, where the variances of the data dimensions differ by multiple orders of magnitude, leading to rounding errors stacking up to wrong results. The large support of the pdfs in the data space also leads to extremely small densities, adding to the numerical problems.

**Describe the solution you'd like**
A possible fix, as suggested by @VivianeKlingel for this problem would be to scale all data dimensions to have mean zero and variance one. This transformation wouldn't affect much of the computational base, as it only adds a constant factor to the parameter distribution (see below). However, some additional discussion between the main contributors is required where exactly the transformation should be applied (@TheVincentWagner , @kaiserls ). 

![image](https://user-images.githubusercontent.com/63878559/231768967-b5090105-0890-4d72-857c-95fc6da53aab.png)


**Describe alternatives you've considered**
The alternative of just scaling the data distribution by a constant factor does not seem promising as it doesn't address the issue of large differences of variance in the data space. 

**Additional context**
Estimation results often look like this:

![image](https://user-images.githubusercontent.com/63878559/231774856-bf13a429-7a92-4e61-954b-fae991f20de4.png)


Use principal component analysis (PCA) to reduce number of data dimesions
**Is your feature request related to a problem? Please describe.**
High data dimensions cause issues with the distance based kernel-density estimation (see, for instance, https://homes.cs.washington.edu/%7Epedrod/papers/cacm12.pdf). This is likely one of the causes for the problem described in #38 .

**Describe the solution you'd like**
 If the number of parameter dimensions is lower, EPI could use PCA to reduce the number data dimensions (for instance s.t. the number of data dimensions is the number of parameter dimensions). This should retain most of the information contained in the data (especially, the assumption that performing a PCA usually should not break injectivity of the model should still hold).

**Describe alternatives you've considered**
The approach of just normalizing the data to reduce numerical errors as described in #38 did not fully solve the issue.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3eb52fba..43755202 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,19 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+## [0.3.0] - 2023-04-27
+
+### Added
+
+- Added data normalization to improve numerics and performance in the inference method.
+- Added possiblity to use no PCA transformation or completely custom transformation for inference.
+- More tests to validate EPI
+
+### Fixed
+
+- Fixed grid results beeing unusable due to numerical issues
+- Fixed sbml models ignoring parameter_names argument
+
 ## [0.2.1] - 2023-04-21
 
 ### Fixed
diff --git a/docs/source/eulerpi.core.data_transformation.rst b/docs/source/eulerpi.core.data_transformation.rst
new file mode 100644
index 00000000..6d7927ab
--- /dev/null
+++ b/docs/source/eulerpi.core.data_transformation.rst
@@ -0,0 +1,7 @@
+eulerpi.core.data\_transformation module
+========================================
+
+.. automodule:: eulerpi.core.data_transformation
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.data_transformation_types.rst b/docs/source/eulerpi.core.data_transformation_types.rst
new file mode 100644
index 00000000..3f117e0e
--- /dev/null
+++ b/docs/source/eulerpi.core.data_transformation_types.rst
@@ -0,0 +1,7 @@
+eulerpi.core.data\_transformation\_types module
+===============================================
+
+.. automodule:: eulerpi.core.data_transformation_types
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.rst b/docs/source/eulerpi.core.rst
index d763cfad..31406d15 100644
--- a/docs/source/eulerpi.core.rst
+++ b/docs/source/eulerpi.core.rst
@@ -7,6 +7,8 @@ Submodules
 .. toctree::
    :maxdepth: 4
 
+   Data Transformation<eulerpi.core.data_transformation>
+   Data Transformation Types<eulerpi.core.data_transformation_types>
    Dense Grid<eulerpi.core.dense_grid>
    Dense Grid Types<eulerpi.core.dense_grid_types>
    Inference<eulerpi.core.inference>
diff --git a/eulerpi/core/data_transformation.py b/eulerpi/core/data_transformation.py
new file mode 100644
index 00000000..765a3026
--- /dev/null
+++ b/eulerpi/core/data_transformation.py
@@ -0,0 +1,225 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Union
+
+import jax.numpy as jnp
+from jax import jit, tree_util
+from sklearn.decomposition import PCA
+
+
+class DataTransformation(ABC):
+    """Class for normalizing data. The data is normalized by subtracting the mean and multiplying by the inverse of the Cholesky decomposition of the covariance matrix."""
+
+    @abstractmethod
+    def transform(self, data: jnp.ndarray) -> jnp.ndarray:
+        """Transform the given data."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def jacobian(self, data: jnp.ndarray) -> jnp.ndarray:
+        """Calculate the jacobian of the transformation at the given data point(s).
+
+        Args:
+            data (jnp.ndarray): The data at which the jacobian should be evaluated.
+
+        Raises:
+            NotImplementedError: Raised if the jacobian is not implemented in the subclass.
+
+        Returns:
+            jnp.ndarray: The jacobian of the transformation at the given data point(s).
+        """
+        raise NotImplementedError()
+
+
+class DataIdentity(DataTransformation):
+    """The identity transformation. Does not change the data."""
+
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    def transform(self, data: jnp.ndarray) -> jnp.ndarray:
+        """Returns the data unchanged.
+
+        Args:
+            data (jnp.ndarray): The data which should be transformed.
+
+        Returns:
+            jnp.ndarray: The data unchanged.
+        """
+        return data
+
+    def jacobian(self, data: jnp.ndarray) -> jnp.ndarray:
+        """Returns the identity matrix.
+
+        Args:
+            data (jnp.ndarray): The data at which the jacobian should be evaluated.
+
+        Returns:
+            jnp.ndarray: The identity matrix.
+        """
+        return jnp.eye(data.shape[-1])
+
+
+class DataNormalizer(DataTransformation):
+    """Class for normalizing data. The data is normalized by subtracting the mean and multiplying by the inverse of the Cholesky decomposition of the covariance matrix."""
+
+    def __init__(
+        self,
+        normalizing_matrix: Optional[jnp.ndarray] = None,
+        mean_vector: Optional[jnp.ndarray] = None,
+    ):
+        """Initialize a DataNormalizer object.
+
+        Args:
+            normalizing_matrix (Optional[jnp.ndarray], optional): The normalizing matrix. Defaults to None.
+            mean_vector (Optional[jnp.ndarray], optional): The mean / shift vector. Defaults to None.
+        """
+        super().__init__()
+
+        self.normalizing_matrix = normalizing_matrix
+        self.mean_vector = mean_vector
+
+    @classmethod
+    def from_data(cls, data: jnp.ndarray) -> "DataTransformation":
+        """Initialize a DataTransformation object by calculating the mean vector and normalizing matrix from the given data.
+
+        Args:
+            data (jnp.ndarray): The data from which to calculate the mean vector and normalizing matrix.
+
+        Returns:
+            DataTransformation: The initialized DataTransformation object.
+        """
+        instance = cls()
+        instance.mean_vector = jnp.mean(data, axis=0)
+        cov = jnp.cov(data, rowvar=False)
+        L = jnp.linalg.cholesky(jnp.atleast_2d(cov))
+        instance.normalizing_matrix = jnp.linalg.inv(L)
+
+        if instance.normalizing_matrix.shape == (1, 1):
+            instance.normalizing_matrix = instance.normalizing_matrix[0, 0]
+
+        return instance
+
+    @classmethod
+    def from_transformation(
+        cls,
+        mean_vector: jnp.ndarray,
+        normalizing_matrix: jnp.ndarray,
+    ) -> "DataTransformation":
+        """Initialize a DataTransformation object from the given mean vector, normalizing matrix and determinant.
+
+        Args:
+            mean_vector (jnp.ndarray): The vector to shift the data by.
+            normalizing_matrix (jnp.ndarray): The matrix to multiply the data by.
+
+        Returns:
+            DataTransformation: The initialized DataTransformation object.
+        """
+        instance = cls()
+        instance.mean_vector = mean_vector
+        instance.normalizing_matrix = normalizing_matrix
+        return instance
+
+    @jit
+    def transform(
+        self, data: Union[jnp.double, jnp.ndarray]
+    ) -> Union[jnp.double, jnp.ndarray]:
+        """Normalize the given data.
+
+        Args:
+            data (Union[jnp.double, jnp.ndarray]): The data to be normalized. Columns correspond to different dimensions. Rows correspond to different observations.
+
+        Returns:
+            Union[jnp.double, jnp.ndarray]: The normalized data.
+        """
+        # possible shapes
+        # normalizing matrix: (d, d)
+        # data: (n, d)
+        # data: (d)
+        # The correct output shape is (n, d) or (d) depending on the input shape.
+        return jnp.inner(data - self.mean_vector, self.normalizing_matrix)
+
+    def jacobian(self, data: jnp.ndarray) -> jnp.ndarray:
+        return self.normalizing_matrix
+
+    def _tree_flatten(self):
+        """This function is used by JAX to flatten the object for JIT compilation."""
+        children = (
+            self.normalizing_matrix,
+            self.mean_vector,
+        )  # arrays / dynamic values
+        aux_data = {}  # static values
+
+        return (children, aux_data)
+
+    @classmethod
+    def _tree_unflatten(cls, aux_data, children):
+        """This function is used by JAX to unflatten the object for JIT compilation."""
+        return cls(*children, **aux_data)
+
+
+tree_util.register_pytree_node(
+    DataNormalizer,
+    DataNormalizer._tree_flatten,
+    DataNormalizer._tree_unflatten,
+)
+
+
+class DataPCA(DataTransformation):
+    """The DataPCA class can be used to transform the data using the Principal Component Analysis."""
+
+    def __init__(
+        self,
+        pca: Optional[PCA] = None,
+    ):
+        """Initialize a DataPCA object.
+
+        Args:
+            pca (Optional[PCA], optional): The PCA object to be used for the transformation. Defaults to None.
+        """
+        super().__init__()
+
+        self.pca = pca
+
+    @classmethod
+    def from_data(
+        cls, data: jnp.ndarray, n_components: Optional[int] = None
+    ) -> "DataTransformation":
+        """Initialize a DataPCA object by calculating the PCA from the given data.
+
+        Args:
+            data (jnp.ndarray): The data to be used for the PCA.
+            n_components (Optional[int], optional): The number of components to keep. If None is passed, min(n_samples,n_features) is used. Defaults to None.
+
+        Returns:
+            DataTransformation: The initialized DataPCA object.
+        """
+        instance = cls()
+        instance.pca = PCA(n_components=n_components)
+        instance.pca.fit(data)
+        return instance
+
+    def transform(self, data: jnp.ndarray) -> jnp.ndarray:
+        """Transform the given data using the PCA.
+
+        Args:
+            data (jnp.ndarray): The data to be transformed.
+
+        Returns:
+            jnp.ndarray: The data projected onto and expressed in the basis of the principal components.
+        """
+        return self.pca.transform(data.reshape(-1, data.shape[-1])).reshape(
+            -1, self.pca.n_components_
+        )
+
+    def jacobian(self, data: jnp.ndarray) -> jnp.ndarray:
+        """Return the jacobian of the pca transformation at the given data point(s).
+
+        Args:
+            data (jnp.ndarray): The data point(s) at which the jacobian should be evaluated.
+
+        Returns:
+            jnp.ndarray: The jacobian of the pca transformation at the given data point(s).
+        """
+        return self.pca.components_
diff --git a/eulerpi/core/data_transformation_types.py b/eulerpi/core/data_transformation_types.py
new file mode 100644
index 00000000..c364847a
--- /dev/null
+++ b/eulerpi/core/data_transformation_types.py
@@ -0,0 +1,10 @@
+from enum import Enum
+
+
+class DataTransformationType(Enum):
+    """Defines the allowed DataTransformation types for the inference function."""
+
+    Identity = 0  # No data transformation
+    Normalize = 1  # Normalizes the data to zero mean and no correlation
+    PCA = 2  # Apply a pca transformation
+    Custom = 3  # Apply a custom transformation to the data
diff --git a/eulerpi/core/dense_grid.py b/eulerpi/core/dense_grid.py
index 12b8681a..24135f71 100644
--- a/eulerpi/core/dense_grid.py
+++ b/eulerpi/core/dense_grid.py
@@ -6,6 +6,7 @@
 import numpy as np
 from numpy.polynomial.chebyshev import chebpts1
 
+from eulerpi.core.data_transformation import DataTransformation
 from eulerpi.core.dense_grid_types import DenseGridType
 from eulerpi.core.inference_types import InferenceType
 from eulerpi.core.model import Model
@@ -68,7 +69,14 @@ def generate_regular_grid(
 
 
 def evaluate_on_grid_chunk(
-    args: typing.Tuple[np.ndarray, Model, np.ndarray, np.ndarray, np.ndarray]
+    args: typing.Tuple[
+        np.ndarray,
+        Model,
+        np.ndarray,
+        DataTransformation,
+        np.ndarray,
+        np.ndarray,
+    ]
 ) -> np.ndarray:
     """Define a function which evaluates the density for a given grid chunk. The input args contains the grid chunk, the model, the data, the data_stdevs and the slice.
 
@@ -76,22 +84,23 @@ def evaluate_on_grid_chunk(
         grid_chunk(np.ndarray): The grid chunk contains the grid points (parameter vectors) for which the density should be evaluated.
         model(Model): The model used for the inference.
         data(np.ndarray): The data points used for the inference.
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         data_stdevs(np.ndarray): The standard deviations of the data points. (Currently the kernel width, #TODO!)
         slice(np.ndarray): The slice defines for which dimensions of the grid points / paramater vectors the marginal density should be evaluated.
 
     Returns:
         np.ndarray: The evaluation results for the given grid chunk. It is a vector, containing the parameters in the first columns, the simulation results in the second columns and the density evaluations in the last columns.
     """
-    grid_chunk, model, data, data_stdevs, slice = args
+    grid_chunk, model, data, data_transformation, data_stdevs, slice = args
 
     # Init the result array
     evaluation_results = np.zeros(
-        (grid_chunk.shape[0], data.shape[1] + slice.shape[0] + 1)
+        (grid_chunk.shape[0], model.data_dim + slice.shape[0] + 1)
     )
     # Evaluate the grid points
     for i, gridPoint in enumerate(grid_chunk):
         density, param_sim_res_density = evaluate_density(
-            gridPoint, model, data, data_stdevs, slice
+            gridPoint, model, data, data_transformation, data_stdevs, slice
         )
         evaluation_results[i] = param_sim_res_density
     return evaluation_results
@@ -100,6 +109,7 @@ def evaluate_on_grid_chunk(
 def run_dense_grid_evaluation(
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     slice: np.ndarray,
     result_manager: ResultManager,
     num_grid_points: np.ndarray,
@@ -112,6 +122,7 @@ def run_dense_grid_evaluation(
     Args:
         model(Model): The model for which the evaluation should be performed.
         data(np.ndarray): The data for which the evaluation should be performed.
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         slice(np.ndarray): The slice for which the evaluation should be performed.
         result_manager(ResultManager): The result manager that should be used to save the results.
         num_grid_points(np.ndarray): The number of grid points for each dimension.
@@ -155,6 +166,7 @@ def run_dense_grid_evaluation(
         grid_chunks,
         repeat(model),
         repeat(data),
+        repeat(data_transformation),
         repeat(data_stdevs),
         repeat(slice),
     )
@@ -163,22 +175,25 @@ def run_dense_grid_evaluation(
     pool.join()
     results = np.concatenate(results)
 
+    data_dim = model.data_dim
+
     result_manager.save_overall(
         slice,
         results[:, 0 : slice.shape[0]],
-        results[:, slice.shape[0] : slice.shape[0] + data.shape[1]],
-        results[:, slice.shape[0] + data.shape[1] :],
+        results[:, slice.shape[0] : slice.shape[0] + data_dim],
+        results[:, slice.shape[0] + data_dim :],
     )
     return (
         results[:, 0 : slice.shape[0]],
-        results[:, slice.shape[0] : slice.shape[0] + data.shape[1]],
-        results[:, slice.shape[0] + data.shape[1] :],
+        results[:, slice.shape[0] : slice.shape[0] + data_dim],
+        results[:, slice.shape[0] + data_dim :],
     )
 
 
 def inference_dense_grid(
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     result_manager: ResultManager,
     slices: list[np.ndarray],
     num_processes: int,
@@ -196,6 +211,7 @@ def inference_dense_grid(
     Args:
         model (Model): The model describing the mapping from parameters to data.
         data (np.ndarray): The data to be used for the inference.
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         result_manager (ResultManager): The result manager to be used for the inference.
         slices (np.ndarray): A list of slices to be used for the inference.
         num_processes (int): The number of processes to be used for the inference.
@@ -233,6 +249,7 @@ def inference_dense_grid(
         ) = run_dense_grid_evaluation(
             model=model,
             data=data,
+            data_transformation=data_transformation,
             slice=slice,
             result_manager=result_manager,
             num_grid_points=n_points,
diff --git a/eulerpi/core/inference.py b/eulerpi/core/inference.py
index 2c52e3b1..ea7086a4 100644
--- a/eulerpi/core/inference.py
+++ b/eulerpi/core/inference.py
@@ -4,7 +4,15 @@
 
 import jax.numpy as jnp
 import numpy as np
+import psutil
 
+from eulerpi.core.data_transformation import (
+    DataIdentity,
+    DataNormalizer,
+    DataPCA,
+    DataTransformation,
+)
+from eulerpi.core.data_transformation_types import DataTransformationType
 from eulerpi.core.dense_grid import inference_dense_grid
 from eulerpi.core.inference_types import InferenceType
 from eulerpi.core.model import Model
@@ -18,10 +26,13 @@ def inference(
     data: Union[str, os.PathLike, np.ndarray],
     inference_type: InferenceType = InferenceType.MCMC,
     slices: Optional[list[np.ndarray]] = None,
-    num_processes: int = 4,
+    num_processes: Optional[int] = None,
     run_name: str = "default_run",
     result_manager: ResultManager = None,
     continue_sampling: bool = False,
+    data_transformation: DataTransformationType = DataTransformationType.Normalize,
+    custom_data_transformation: Optional[DataTransformation] = None,
+    n_components_pca: Optional[int] = None,
     **kwargs,
 ) -> Tuple[
     Dict[str, np.ndarray],
@@ -36,10 +47,13 @@ def inference(
         data(Union[str, os.PathLike, np.ndarray]): The data to be used for the inference. If a string is given, it is assumed to be a path to a file containing the data.
         inference_type(InferenceType, optional): The type of inference to be used. (Default value = InferenceType.MCMC)
         slices(list[np.ndarray], optional): A list of slices to be used for the inference. If None, the full joint distribution is computed. (Default value = None)
-        num_processes(int, optional): The number of processes to be used for the inference. (Default value = 4)
+        num_processes(int, optional): The number of processes to be used for the inference. Per default the number of cores is used. (Default value = Non)
         run_name(str): The name of the run. (Default value = "default_run")
         result_manager(ResultManager, optional): The result manager to be used for the inference. If None, a new result manager is created with default paths and saving methods. (Default value = None)
         continue_sampling(bool, optional): If True, the inference will continue sampling from the last saved point. (Default value = False)
+        data_transformation(DataTransformationType): The type of data transformation to use. (Default value = DataTransformationType.Normalize)
+        custom_data_transformation(DataTransformation, optional): The data transformation to be used for the inference. If None, a normalization is applied to the data. (Default value = None)
+        n_components_pca(int, optional): If using the PCA as data_transformation, selects how many dimensions are kept in the pca. Per default the number of dimensions equals the dimension of the parameter space. (Default value = None)
         **kwargs: Additional keyword arguments to be passed to the inference function. The possible parameters depend on the inference type.
 
     Returns:
@@ -56,7 +70,28 @@ def inference(
             f"The data argument must be a path to a file or a numpy array. The argument passed was of type {type(data)}."
         )
 
-    # TODO Calculate Transformation to normalize data and normalize data, create transformation object
+    # Transform the data
+    if data_transformation == DataTransformationType.Identity:
+        data_transformation = DataIdentity()
+    elif data_transformation == DataTransformationType.Normalize:
+        data_transformation = DataNormalizer.from_data(data)
+    elif data_transformation == DataTransformationType.PCA:
+        n_components = n_components_pca or model.param_dim
+        data_transformation = DataPCA.from_data(
+            data=data, n_components=n_components
+        )
+    elif data_transformation == DataTransformationType.Custom:
+        data_transformation = custom_data_transformation
+        if not issubclass(custom_data_transformation, DataTransformation):
+            raise TypeError(
+                f"The custom_data_transformation must be an instance of a subclass from DataTransformation. It is of type {type(data_transformation)}."
+            )
+    else:
+        raise TypeError(
+            "The data_transformation must be one of the enum values of DataTransformationType."
+        )
+    data = data_transformation.transform(data)
+
     # TODO rename std_dev to kernel_width, adapt calculation of kernel width
 
     slices = slices or [
@@ -70,11 +105,14 @@ def inference(
         result_manager.delete_application_folder_structure()
     result_manager.create_application_folder_structure()
 
-    # TODO give transformation object to inference functions
+    if not num_processes:
+        num_processes = psutil.cpu_count(logical=False)
+
     if inference_type == InferenceType.DENSE_GRID:
         return inference_dense_grid(
             model=model,
             data=data,
+            data_transformation=data_transformation,
             result_manager=result_manager,
             slices=slices,
             num_processes=num_processes,
@@ -84,6 +122,7 @@ def inference(
         return inference_mcmc(
             model=model,
             data=data,
+            data_transformation=data_transformation,
             result_manager=result_manager,
             slices=slices,
             num_processes=num_processes,
@@ -93,6 +132,7 @@ def inference(
         return inference_sparse_grid(
             model=model,
             data=data,
+            data_transformation=data_transformation,
             result_manager=result_manager,
             slices=slices,
             num_processes=num_processes,
diff --git a/eulerpi/core/model.py b/eulerpi/core/model.py
index 7569e5c7..fbd3e934 100644
--- a/eulerpi/core/model.py
+++ b/eulerpi/core/model.py
@@ -10,6 +10,7 @@
 from jax import jacrev, jit, vmap
 
 import amici
+from eulerpi import logger
 from eulerpi.jax_extension import value_and_jacrev
 
 
@@ -307,11 +308,11 @@ def __init__(
                 self.amici_model_name, self.amici_model_dir
             )
 
+        self.param_names = param_names
+
         # Load the generated model
         self.load_amici_model_and_solver()
 
-        self.param_names = param_names or self.amici_model.getParameterNames()
-
     def load_amici_model_and_solver(self):
         """Loads the AMICI model from the previously generated model."""
         amici_model_module = amici.import_model_module(
@@ -322,7 +323,21 @@ def load_amici_model_and_solver(self):
 
         # TODO: Maybe this is redundant when using settings to hdf5
         self.amici_model.setTimepoints([self.tEnd])
-        self.amici_model.requireSensitivitiesForAllParameters()
+
+        if self.param_names is not None:
+            # We need the indices for setParameterList, not the ids or names
+            amici_param_indices = []
+            for i, param_id in enumerate(self.amici_model.getParameterIds()):
+                if param_id in self.param_names:
+                    amici_param_indices.append(i)
+                else:
+                    logger.warning(
+                        f"Parameter {param_id} is specified in the sbml file, but not in the passed parameter list. It will be ignored."
+                    )
+            self.amici_model.setParameterList(amici_param_indices)
+        else:
+            self.param_names = self.amici_model.getParameterIds()
+            self.amici_model.requireSensitivitiesForAllParameters()
         self.amici_solver.setSensitivityMethod(amici.SensitivityMethod.forward)
         self.amici_solver.setSensitivityOrder(amici.SensitivityOrder.first)
 
diff --git a/eulerpi/core/sampling.py b/eulerpi/core/sampling.py
index 007068cc..c133d691 100644
--- a/eulerpi/core/sampling.py
+++ b/eulerpi/core/sampling.py
@@ -20,6 +20,7 @@
 import numpy as np
 
 from eulerpi import logger
+from eulerpi.core.data_transformation import DataTransformation
 from eulerpi.core.inference_types import InferenceType
 from eulerpi.core.kde import calc_kernel_width
 from eulerpi.core.model import Model
@@ -35,11 +36,11 @@
 # )
 
 
-# TODO Give transformation object to eval function
 def evaluate_sample(
     param: np.ndarray,
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     data_stdevs: np.ndarray,
     slice: np.ndarray,
 ) -> typing.Tuple[float, np.ndarray]:
@@ -49,6 +50,7 @@ def evaluate_sample(
         param (np.ndarray): parameter values
         model (Model): The model which will be sampled
         data (np.ndarray): data
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         data_stdevs (np.ndarray): kernel width for the data
         slice (np.ndarray): slice of the parameter space which will be sampled
 
@@ -57,7 +59,7 @@ def evaluate_sample(
     """
 
     log_samplerresult = eval_log_transformed_density(
-        param, model, data, data_stdevs, slice
+        param, model, data, data_transformation, data_stdevs, slice
     )
     return log_samplerresult
 
@@ -65,6 +67,7 @@ def evaluate_sample(
 def run_emcee_once(
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     data_stdevs: np.ndarray,
     slice: np.ndarray,
     initial_walker_positions: np.ndarray,
@@ -77,6 +80,7 @@ def run_emcee_once(
     Args:
         model (Model): The model which will be sampled
         data (np.ndarray): data
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         data_stdevs (np.ndarray): kernel width for the data
         slice (np.ndarray): slice of the parameter space which will be sampled
         initial_walker_positions (np.ndarray): initial parameter values for the walkers
@@ -100,7 +104,7 @@ def run_emcee_once(
             sampling_dim,
             evaluate_sample,
             pool=pool,
-            args=[model, data, data_stdevs, slice],
+            args=[model, data, data_transformation, data_stdevs, slice],
         )
         # Extract the final walker position and close the pool of worker processes.
         final_walker_positions, _, _, _ = sampler.run_mcmc(
@@ -124,7 +128,7 @@ def run_emcee_once(
     # TODO: Keep as 3d array?
     # Should have shape (num_steps, num_walkers, param_dim+data_dim+1)
     sampler_results = sampler.get_blobs()
-    data_dim = data.shape[1]
+    data_dim = model.data_dim
     sampler_results = sampler_results.reshape(
         num_steps * num_walkers, sampling_dim + data_dim + 1
     )
@@ -149,6 +153,7 @@ def run_emcee_once(
 def run_emcee_sampling(
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     slice: np.ndarray,
     result_manager: ResultManager,
     num_processes: int,
@@ -164,6 +169,7 @@ def run_emcee_sampling(
     Args:
         model (Model): The model which will be sampled
         data (np.ndarray): data
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         slice (np.ndarray): slice of the parameter space which will be sampled
         result_manager (ResultManager): ResultManager which will store the results
         num_processes (int): number of parallel threads.
@@ -215,6 +221,7 @@ def run_emcee_sampling(
         sampler_results, final_walker_positions = run_emcee_once(
             model,
             data,
+            data_transformation,
             data_stdevs,
             slice,
             initial_walker_positions,
@@ -302,6 +309,7 @@ def calc_walker_acceptance(
 def inference_mcmc(
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     result_manager: ResultManager,
     slices: list[np.ndarray],
     num_processes: int,
@@ -322,6 +330,7 @@ def inference_mcmc(
     Args:
         model (Model): The model describing the mapping from parameters to data.
         data (np.ndarray): The data to be used for the inference.
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         result_manager (ResultManager): The result manager to be used for the inference.
         slices (np.ndarray): A list of slices to be used for the inference.
         num_processes (int): The number of processes to be used for the inference.
@@ -366,6 +375,7 @@ def inference_mcmc(
         ) = run_emcee_sampling(
             model=model,
             data=data,
+            data_transformation=data_transformation,
             slice=slice,
             result_manager=result_manager,
             num_runs=num_runs,
diff --git a/eulerpi/core/sparsegrid.py b/eulerpi/core/sparsegrid.py
index 3b88e1b5..4b63d6c2 100644
--- a/eulerpi/core/sparsegrid.py
+++ b/eulerpi/core/sparsegrid.py
@@ -15,6 +15,7 @@
 import numpy as np
 
 from eulerpi import logger
+from eulerpi.core.data_transformation import DataTransformation
 from eulerpi.core.inference_types import InferenceType
 from eulerpi.core.kde import calc_kernel_width
 from eulerpi.core.model import Model
@@ -369,7 +370,14 @@ def pass_contributions2higher_levels(self):
 
 
 def evaluate_on_sparse_grid(
-    args: typing.Tuple[np.ndarray, Model, np.ndarray, np.ndarray, np.ndarray]
+    args: typing.Tuple[
+        np.ndarray,
+        Model,
+        np.ndarray,
+        DataTransformation,
+        np.ndarray,
+        np.ndarray,
+    ]
 ) -> np.ndarray:
     """This function is used to evaluate a function on a sparse grid in parallel. The input args is a tuple containing the function, the sparse grid and the number of processes to be used.
 
@@ -377,6 +385,7 @@ def evaluate_on_sparse_grid(
         params (np.ndarray): The parameters to be evaluated.
         model(Model): The model used for the inference.
         data(np.ndarray): The data points used for the inference.
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         data_stdevs(np.ndarray): The standard deviations of the data points. (Currently the kernel width, #TODO!)
         slice(np.ndarray): The slice defines for which dimensions of the grid points / paramater vectors the marginal density should be evaluated.
 
@@ -384,13 +393,16 @@ def evaluate_on_sparse_grid(
         np.ndarray: The density values for the given params.
     """
 
-    params, model, data, data_stdevs, slice = args
-    return evaluate_density(params, model, data, data_stdevs, slice)[1]
+    params, model, data, data_transformation, data_stdevs, slice = args
+    return evaluate_density(
+        params, model, data, data_transformation, data_stdevs, slice
+    )[1]
 
 
 def inference_sparse_grid(
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     result_manager: ResultManager,
     slices: typing.List[np.ndarray],
     num_processes: int,
@@ -404,10 +416,11 @@ def inference_sparse_grid(
     """Evaluates the transformed parameter density over a set of points resembling a sparse grid, thereby attempting parameter inference. If a data path is given, it is used to load the data for the model. Else, the default data path of the model is used.
 
     Args:
-      model(Model): The model describing the mapping from parameters to data.
-      data(np.ndarray): The data to be used for inference.
-      num_processes(int): number of processes to use for parallel evaluation of the model.
-      num_levels(int, optional): Maximum sparse grid level depth that mainly defines the number of points. Defaults to 5.
+        model(Model): The model describing the mapping from parameters to data.
+        data(np.ndarray): The data to be used for inference.
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
+        num_processes(int): number of processes to use for parallel evaluation of the model.
+        num_levels(int, optional): Maximum sparse grid level depth that mainly defines the number of points. Defaults to 5.
 
     Returns:
         Tuple[typing.Dict[str, np.ndarray], typing.Dict[str, np.ndarray], typing.Dict[str, np.ndarray], ResultManager]: The parameter samples, the corresponding simulation results, the corresponding density
@@ -441,6 +454,7 @@ def inference_sparse_grid(
             scaledSparseGridPoints,
             repeat(model),
             repeat(data),
+            repeat(data_transformation),
             repeat(data_stdevs),
             repeat(slice),
         )
@@ -459,20 +473,22 @@ def inference_sparse_grid(
         # Take care! The results here are for single points, therefore we cant use np.concatenate
         results = np.vstack(results)
 
+        data_dim = model.data_dim
+
         # save the results
         result_manager.save_overall(
             slice,
             results[:, 0 : slice.shape[0]],
-            results[:, slice.shape[0] : slice.shape[0] + data.shape[1]],
-            results[:, slice.shape[0] + data.shape[1] :],
+            results[:, slice.shape[0] : slice.shape[0] + data_dim],
+            results[:, slice.shape[0] + data_dim :],
         )
         slice_name = result_manager.get_slice_name(slice)
         overall_params[slice_name] = results[:, 0 : slice.shape[0]]
         overall_sim_results[slice_name] = results[
-            :, slice.shape[0] : slice.shape[0] + data.shape[1]
+            :, slice.shape[0] : slice.shape[0] + data_dim
         ]
         overall_density_evals[slice_name] = results[
-            :, slice.shape[0] + data.shape[1] :
+            :, slice.shape[0] + data_dim :
         ]
         result_manager.save_inference_information(
             slice=slice,
diff --git a/eulerpi/core/transformations.py b/eulerpi/core/transformations.py
index 31a22d55..701410b0 100644
--- a/eulerpi/core/transformations.py
+++ b/eulerpi/core/transformations.py
@@ -5,6 +5,7 @@
 from jax import jit
 
 from eulerpi import logger
+from eulerpi.core.data_transformation import DataTransformation
 from eulerpi.core.kde import eval_kde_gauss
 from eulerpi.core.model import Model
 
@@ -13,6 +14,7 @@ def evaluate_density(
     param: np.ndarray,
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     data_stdevs: np.ndarray,
     slice: np.ndarray,
 ) -> Tuple[np.double, np.ndarray]:
@@ -22,6 +24,7 @@ def evaluate_density(
         param (np.ndarray): parameter for which the transformed density shall be evaluated
         model (Model): model to be evaluated
         data (np.ndarray): data for the model. 2D array with shape (#num_data_points, #data_dim)
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         data_stdevs (np.ndarray): array of suitable kernel width for each data dimension
         slice (np.ndarray): slice of the parameter vector that is to be evaluated
 
@@ -47,15 +50,24 @@ def evaluate_density(
     # If the parameter is within the valid ranges...
     else:
         # Evaluating the model and the jacobian for the specified parameter simultaneously provide a little speedup over calculating it separately in some cases.
-        sim_res, jac = model.forward_and_jacobian(fullParam)
+        sim_res, model_jac = model.forward_and_jacobian(fullParam)
 
-        # TODO transform sim res and use transformed simres below
+        # normalize sim_res
+        transformed_sim_res = data_transformation.transform(sim_res)
 
-        # Evaluate the data density in the simulation result. # TODO scale with determinant transformation matrix
-        densityEvaluation = eval_kde_gauss(data, sim_res, data_stdevs)
+        # Evaluate the data density in the simulation result.
+        densityEvaluation = eval_kde_gauss(
+            data, transformed_sim_res, data_stdevs
+        )
 
         # Calculate the simulation model's pseudo-determinant in the parameter point (also called the correction factor).
-        correction = calc_gram_determinant(jac)
+        # Scale with the determinant of the transformation matrix.
+        transformation_jac = data_transformation.jacobian(sim_res)
+        correction = calc_gram_determinant(
+            jnp.dot(
+                transformation_jac, model_jac
+            )  # We use dot, because matmul does not support scalars
+        )
 
         # Multiply data density and correction factor.
         trafo_density_evaluation = densityEvaluation * correction
@@ -72,6 +84,7 @@ def eval_log_transformed_density(
     param: np.ndarray,
     model: Model,
     data: np.ndarray,
+    data_transformation: DataTransformation,
     data_stdevs: np.ndarray,
     slice: np.ndarray,
 ) -> Tuple[np.double, np.ndarray]:
@@ -82,6 +95,7 @@ def eval_log_transformed_density(
         param (np.ndarray): parameter for which the transformed density shall be evaluated
         model (Model): model to be evaluated
         data (np.ndarray): data for the model. 2D array with shape (#num_data_points, #data_dim)
+        data_transformation (DataTransformation): The data transformation used to normalize the data.
         data_stdevs (np.ndarray): array of suitable kernel width for each data dimension
         slice (np.ndarray): slice of the parameter vector that is to be evaluated
 
@@ -91,9 +105,8 @@ def eval_log_transformed_density(
             : sampler_results (array concatenation of parameters, simulation results and evaluated density, stored as "blob" by the emcee sampler)
 
     """
-    # TODO give transformation object to evaluation function
     trafo_density_evaluation, evaluation_results = evaluate_density(
-        param, model, data, data_stdevs, slice
+        param, model, data, data_transformation, data_stdevs, slice
     )
     if trafo_density_evaluation == 0:
         return -np.inf, evaluation_results
diff --git a/eulerpi/examples/sbml/sbml_menten_model.py b/eulerpi/examples/sbml/sbml_menten_model.py
index 62b1f383..c78a69b1 100644
--- a/eulerpi/examples/sbml/sbml_menten_model.py
+++ b/eulerpi/examples/sbml/sbml_menten_model.py
@@ -7,7 +7,8 @@
 
 class MentenSBMLModel(SBMLModel, ArtificialModelInterface):
     """The MentenModel is a simple model for the Michaelis-Menten enzyme kinetics.
-    It requires a few adaptations to the SBMLModel class to work with the inference, because the first output of the model does not depend on the parameters.
+    It requires a few adaptations to the SBMLModel class to work with the inference, because
+    the first output of the model does not depend on the parameters and the second output is linear dependent on the third output.
     See data_dim, forward, jacobian, forward_and_jacobian.
     """
 
@@ -34,17 +35,17 @@ def __init__(
 
     @property
     def data_dim(self) -> int:
-        return 2
+        return 1
 
     def forward(self, params) -> np.ndarray:
-        return super().forward(params)[1:]
+        return super().forward(params)[2:]
 
     def jacobian(self, params) -> np.ndarray:
-        return super().jacobian(params)[1:, :]
+        return super().jacobian(params)[2:, :]
 
     def forward_and_jacobian(self, params) -> np.ndarray:
         val, jac = super().forward_and_jacobian(params)
-        return val[1:], jac[1:, :]
+        return val[2:], jac[2:, :]
 
     def generate_artificial_params(self, num_samples: int) -> np.ndarray:
         diff0 = 5.0
diff --git a/poetry.lock b/poetry.lock
index d19a9f1f..d903b2c5 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1214,6 +1214,18 @@ MarkupSafe = ">=2.0"
 [package.extras]
 i18n = ["Babel (>=2.7)"]
 
+[[package]]
+name = "joblib"
+version = "1.2.0"
+description = "Lightweight pipelining with Python functions"
+category = "main"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "joblib-1.2.0-py3-none-any.whl", hash = "sha256:091138ed78f800342968c523bdde947e7a305b8594b910a0fea2ab83c3c6d385"},
+    {file = "joblib-1.2.0.tar.gz", hash = "sha256:e1cee4a79e4af22881164f218d4311f60074197fb707e082e803b61f6d137018"},
+]
+
 [[package]]
 name = "jupyter-client"
 version = "8.0.3"
@@ -2586,6 +2598,49 @@ urllib3 = ">=1.21.1,<1.27"
 socks = ["PySocks (>=1.5.6,!=1.5.7)"]
 use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
 
+[[package]]
+name = "scikit-learn"
+version = "1.2.2"
+description = "A set of python modules for machine learning and data mining"
+category = "main"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "scikit-learn-1.2.2.tar.gz", hash = "sha256:8429aea30ec24e7a8c7ed8a3fa6213adf3814a6efbea09e16e0a0c71e1a1a3d7"},
+    {file = "scikit_learn-1.2.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:99cc01184e347de485bf253d19fcb3b1a3fb0ee4cea5ee3c43ec0cc429b6d29f"},
+    {file = "scikit_learn-1.2.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:e6e574db9914afcb4e11ade84fab084536a895ca60aadea3041e85b8ac963edb"},
+    {file = "scikit_learn-1.2.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6fe83b676f407f00afa388dd1fdd49e5c6612e551ed84f3b1b182858f09e987d"},
+    {file = "scikit_learn-1.2.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2e2642baa0ad1e8f8188917423dd73994bf25429f8893ddbe115be3ca3183584"},
+    {file = "scikit_learn-1.2.2-cp310-cp310-win_amd64.whl", hash = "sha256:ad66c3848c0a1ec13464b2a95d0a484fd5b02ce74268eaa7e0c697b904f31d6c"},
+    {file = "scikit_learn-1.2.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:dfeaf8be72117eb61a164ea6fc8afb6dfe08c6f90365bde2dc16456e4bc8e45f"},
+    {file = "scikit_learn-1.2.2-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:fe0aa1a7029ed3e1dcbf4a5bc675aa3b1bc468d9012ecf6c6f081251ca47f590"},
+    {file = "scikit_learn-1.2.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:065e9673e24e0dc5113e2dd2b4ca30c9d8aa2fa90f4c0597241c93b63130d233"},
+    {file = "scikit_learn-1.2.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bf036ea7ef66115e0d49655f16febfa547886deba20149555a41d28f56fd6d3c"},
+    {file = "scikit_learn-1.2.2-cp311-cp311-win_amd64.whl", hash = "sha256:8b0670d4224a3c2d596fd572fb4fa673b2a0ccfb07152688ebd2ea0b8c61025c"},
+    {file = "scikit_learn-1.2.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:9c710ff9f9936ba8a3b74a455ccf0dcf59b230caa1e9ba0223773c490cab1e51"},
+    {file = "scikit_learn-1.2.2-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:2dd3ffd3950e3d6c0c0ef9033a9b9b32d910c61bd06cb8206303fb4514b88a49"},
+    {file = "scikit_learn-1.2.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:44b47a305190c28dd8dd73fc9445f802b6ea716669cfc22ab1eb97b335d238b1"},
+    {file = "scikit_learn-1.2.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:953236889928d104c2ef14027539f5f2609a47ebf716b8cbe4437e85dce42744"},
+    {file = "scikit_learn-1.2.2-cp38-cp38-win_amd64.whl", hash = "sha256:7f69313884e8eb311460cc2f28676d5e400bd929841a2c8eb8742ae78ebf7c20"},
+    {file = "scikit_learn-1.2.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:8156db41e1c39c69aa2d8599ab7577af53e9e5e7a57b0504e116cc73c39138dd"},
+    {file = "scikit_learn-1.2.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:fe175ee1dab589d2e1033657c5b6bec92a8a3b69103e3dd361b58014729975c3"},
+    {file = "scikit_learn-1.2.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7d5312d9674bed14f73773d2acf15a3272639b981e60b72c9b190a0cffed5bad"},
+    {file = "scikit_learn-1.2.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ea061bf0283bf9a9f36ea3c5d3231ba2176221bbd430abd2603b1c3b2ed85c89"},
+    {file = "scikit_learn-1.2.2-cp39-cp39-win_amd64.whl", hash = "sha256:6477eed40dbce190f9f9e9d0d37e020815825b300121307942ec2110302b66a3"},
+]
+
+[package.dependencies]
+joblib = ">=1.1.1"
+numpy = ">=1.17.3"
+scipy = ">=1.3.2"
+threadpoolctl = ">=2.0.0"
+
+[package.extras]
+benchmark = ["matplotlib (>=3.1.3)", "memory-profiler (>=0.57.0)", "pandas (>=1.0.5)"]
+docs = ["Pillow (>=7.1.2)", "matplotlib (>=3.1.3)", "memory-profiler (>=0.57.0)", "numpydoc (>=1.2.0)", "pandas (>=1.0.5)", "plotly (>=5.10.0)", "pooch (>=1.6.0)", "scikit-image (>=0.16.2)", "seaborn (>=0.9.0)", "sphinx (>=4.0.1)", "sphinx-gallery (>=0.7.0)", "sphinx-prompt (>=1.3.0)", "sphinxext-opengraph (>=0.4.2)"]
+examples = ["matplotlib (>=3.1.3)", "pandas (>=1.0.5)", "plotly (>=5.10.0)", "pooch (>=1.6.0)", "scikit-image (>=0.16.2)", "seaborn (>=0.9.0)"]
+tests = ["black (>=22.3.0)", "flake8 (>=3.8.2)", "matplotlib (>=3.1.3)", "mypy (>=0.961)", "numpydoc (>=1.2.0)", "pandas (>=1.0.5)", "pooch (>=1.6.0)", "pyamg (>=4.0.0)", "pytest (>=5.3.1)", "pytest-cov (>=2.9.0)", "scikit-image (>=0.16.2)"]
+
 [[package]]
 name = "scipy"
 version = "1.9.3"
@@ -2958,6 +3013,18 @@ files = [
 [package.dependencies]
 mpmath = ">=0.19"
 
+[[package]]
+name = "threadpoolctl"
+version = "3.1.0"
+description = "threadpoolctl"
+category = "main"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "threadpoolctl-3.1.0-py3-none-any.whl", hash = "sha256:8b99adda265feb6773280df41eece7b2e6561b772d21ffd52e372f999024907b"},
+    {file = "threadpoolctl-3.1.0.tar.gz", hash = "sha256:a335baacfaa4400ae1f0d8e3a58d6674d2f8828e3716bb2802c44955ad391380"},
+]
+
 [[package]]
 name = "tomli"
 version = "2.0.1"
@@ -3201,4 +3268,4 @@ testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8.1"
-content-hash = "3e25adc2bb472f3f310332302a4e4f94236749728ab135ebffc33ae2cc42aac4"
+content-hash = "4e16e4712f12629b446230f6cf02a5f39acc4990015fe7085205cd3d5ce4d521"
diff --git a/pyproject.toml b/pyproject.toml
index d2310768..70f3d517 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "eulerpi"
-version = "0.2.1"
+version = "0.3.0"
 description = "The eulerian parameter inference (eulerpi) returns a parameter distribution, which is consistent with the observed data by solving the inverse problem directly. In the case of a one-to-one mapping, this is the true underlying distribution."
 authors = ["Lars Kaiser <lars.g.kaiser@gmx.de>", "Sebastian Hoepfl <sebastian.hoepfl@ist.uni-stuttgart.de>", "Vincent Wagner <vincent.wagner@ist.uni-stuttgart.de>"]
 readme = "README.md"
@@ -29,6 +29,7 @@ tqdm = "^4.65.0"
 seedir = "^0.4.2"
 yfinance = "^0.2.14"
 amici = "^0.16.1"
+scikit-learn = "^1.2.2"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.2.2"


Test patch:
diff --git a/docs/source/tests.rst b/docs/source/tests.rst
index 80aa7aed..ef07594c 100644
--- a/docs/source/tests.rst
+++ b/docs/source/tests.rst
@@ -9,6 +9,7 @@ Submodules
 
    configuration<tests.conftest>
    Profiling<tests.profiling>
+   Test Data Transformation<tests.test_data_transformation>
    Test Dense Grid<tests.test_dense_grid>
    Test Examples<tests.test_examples>
    Test Fixed Params<tests.test_fixed_params>
diff --git a/docs/source/tests.test_data_transformation.rst b/docs/source/tests.test_data_transformation.rst
new file mode 100644
index 00000000..1fa9e878
--- /dev/null
+++ b/docs/source/tests.test_data_transformation.rst
@@ -0,0 +1,7 @@
+tests.test\_data\_transformation module
+=======================================
+
+.. automodule:: tests.test_data_transformation
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/tests/test_data_transformation.py b/tests/test_data_transformation.py
new file mode 100644
index 00000000..17e22e2e
--- /dev/null
+++ b/tests/test_data_transformation.py
@@ -0,0 +1,63 @@
+import numpy as np
+
+from eulerpi.core.data_transformation import (
+    DataIdentity,
+    DataNormalizer,
+    DataPCA,
+)
+
+
+def test_DataIdentity():
+    """Test whether the DataIdentity transformation does not change the data."""
+    data0dim = np.random.rand()
+    data1dim = np.random.rand(100)
+    data2dim = np.random.rand(100, 2)
+
+    data_transformation = DataIdentity()
+
+    assert np.allclose(data_transformation.transform(data0dim), data0dim)
+    assert np.allclose(data_transformation.transform(data1dim), data1dim)
+    assert np.allclose(data_transformation.transform(data2dim), data2dim)
+
+
+def test_DataNormalizer():
+    """Test whether the DataNormalizer transformation normalizes the data to zero mean and unit variance."""
+    data1d = np.random.rand(100, 1)
+    data2d = np.random.rand(100, 2)
+    test_data = [(1, data1d), (2, data2d)]
+
+    for dim, data in test_data:
+        data_transformation = DataNormalizer.from_data(data)
+        transformed_data = data_transformation.transform(data)
+        assert np.allclose(
+            np.mean(transformed_data, axis=0),
+            np.zeros_like(transformed_data[0]),
+        )
+        assert np.allclose(np.cov(transformed_data, rowvar=False), np.eye(dim))
+
+        # Check if transform also works for single datapoints
+        transformed_datapoint = data_transformation.transform(data[0])
+        assert transformed_datapoint.shape == data[0].shape
+
+
+def test_DataPCA():
+    """Test whether the DataPCA transformation is able to run on data with different dimensions."""
+    n_samples = 100
+    data1d = np.random.rand(n_samples, 1)
+    data2d = np.random.rand(n_samples, 2)
+    test_data = [(1, 1, data1d), (2, 2, data2d), (2, 1, data2d)]
+
+    for data_dim, pca_dim, data in test_data:
+        data_transformation = DataPCA.from_data(
+            data=data, n_components=pca_dim
+        )
+        transformed_data = data_transformation.transform(data)
+
+        assert transformed_data.shape[0] == n_samples
+        assert transformed_data.shape[1] == pca_dim
+
+        transformed_datapoint = data_transformation.transform(data[0])
+        assert transformed_datapoint.shape == (1, pca_dim)
+
+        jacobian = data_transformation.jacobian(data[0])
+        assert jacobian.shape == (pca_dim, data_dim)
diff --git a/tests/test_examples.py b/tests/test_examples.py
index b622b6ab..29de01a3 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -76,6 +76,26 @@ def test_examples(example, inference_type):
     ModelClass = getattr(module, className)
     model: Model = ModelClass()
 
+    # test the shapes
+    assert model.param_dim > 0
+    assert model.data_dim > 0
+    assert model.central_param.shape == (model.param_dim,)
+    assert model.param_limits.shape == (model.param_dim, 2)
+
+    model_forward = model.forward(model.central_param)
+    assert (
+        model_forward.shape == (1, model.data_dim)
+        or model_forward.shape == (model.data_dim,)
+        or model_forward.shape == ()
+    )
+
+    model_jac = model.jacobian(model.central_param)
+    assert (
+        model_jac.shape == (model.data_dim, model.param_dim)
+        or (model.data_dim == 1 and model_jac.shape == (model.param_dim,))
+        or (model.param_dim == 1 and model_jac.shape == (model.data_dim,))
+    )
+
     # generate artificial data if necessary
     if model.is_artificial():
         num_data_points = 100
diff --git a/tests/test_simple_models.py b/tests/test_simple_models.py
index 4d3a3b50..4c2a8d6c 100644
--- a/tests/test_simple_models.py
+++ b/tests/test_simple_models.py
@@ -7,6 +7,7 @@
 from jax import vmap
 from matplotlib import cm
 
+from eulerpi.core.data_transformation import DataIdentity
 from eulerpi.core.inference import InferenceType, inference
 from eulerpi.core.kde import calc_kernel_width, eval_kde_gauss
 from eulerpi.core.transformations import evaluate_density
@@ -26,10 +27,11 @@ def test_transformationLinear():
     xMesh, yMesh = np.meshgrid(xGrid, yGrid)
 
     data = np.array([xMesh.flatten(), yMesh.flatten()]).T
+    data_transformation = DataIdentity()
 
     # define standard deviations according to silverman
     data_stdevs = calc_kernel_width(data)
-
+    data_transformation = DataIdentity()
     # Now plot the data Gaussian KDE
     KDEresolution = 25
 
@@ -75,6 +77,7 @@ def test_transformationLinear():
                 paramPoint,
                 model,
                 data,
+                data_transformation,
                 data_stdevs,
                 slice=np.arange(model.param_dim),
             )
@@ -103,6 +106,7 @@ def test_transformationExponential():
     trueParam = np.random.rand(num_data_points, 2) + 1
 
     data = vmap(model.forward, in_axes=0)(trueParam)
+    data_transformation = DataIdentity()
 
     # define standard deviations according to silverman
     data_stdevs = calc_kernel_width(data)
@@ -147,7 +151,12 @@ def test_transformationExponential():
         for j in range(paramResolution):
             paramPoint = np.array([paramxMesh[i, j], paramyMesh[i, j]])
             paramEvals[i, j], _ = evaluate_density(
-                paramPoint, model, data, data_stdevs, slice=fullSlice
+                paramPoint,
+                model,
+                data,
+                data_transformation,
+                data_stdevs,
+                slice=fullSlice,
             )
 
     fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
diff --git a/tests/test_transformations.py b/tests/test_transformations.py
index 48228f47..65345180 100644
--- a/tests/test_transformations.py
+++ b/tests/test_transformations.py
@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 
+from eulerpi.core.data_transformation import DataIdentity
 from eulerpi.core.kde import calc_kernel_width, eval_kde_gauss
 from eulerpi.core.model import ArtificialModelInterface, JaxModel
 from eulerpi.core.transformations import calc_gram_determinant
@@ -60,26 +61,29 @@ def test_evaluate_density(caplog):
 
     # KDE has its own tests, so we can use it here to test the transformations
     data = np.array([[0.0], [2.0]])
+    data_transformation = DataIdentity()
     data_stdevs = calc_kernel_width(data)
     pure_density = eval_kde_gauss(data, sim_res, data_stdevs)
 
     # Test case 1: When the slice is one dimensional
     slice = np.array([0])
-    density, _ = evaluate_density(param, x2_model, data, data_stdevs, slice)
+    density, _ = evaluate_density(
+        param, x2_model, data, data_transformation, data_stdevs, slice
+    )
     assert density == pure_density * correction
 
     # Test case 2: When the slice is empty
     slice = np.array([])
     with pytest.raises(IndexError):
         density, _ = evaluate_density(
-            param, x2_model, data, data_stdevs, slice
+            param, x2_model, data, data_transformation, data_stdevs, slice
         )
 
     # Test case 3: When the slice is two dimensional, but the model is one dimensional
     slice = np.array([0, 1])
     with pytest.raises(IndexError):
         density, _ = evaluate_density(
-            param, x2_model, data, data_stdevs, slice
+            param, x2_model, data, data_transformation, data_stdevs, slice
         )
 
     # Test case 4: When the param is out of bounds
@@ -90,6 +94,8 @@ def test_evaluate_density(caplog):
     from eulerpi import logger
 
     logger.setLevel("INFO")
-    density, _ = evaluate_density(param, x2_model, data, data_stdevs, slice)
+    density, _ = evaluate_density(
+        param, x2_model, data, data_transformation, data_stdevs, slice
+    )
     assert density == 0.0
     assert "Parameters outside of predefined range" in caplog.text