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

Issue:
How to (properly) vectorize the forward function
@kaiserls raised a good point in one of his recent reviews concerning the vectorized forward pass:

I think the problem was that for some models a 1-d array is a single parameter and for some, a 1-d array is a list of parameters. I see several options:

    * Make another function: forward_vectorized(...)
    * Make another parameter: multiple_parameters:bool
    * Try to use the model param dim to understand whether multiple evaluations are needed
    * "Force" the user to avoid scalars. then a 1d-array is always one param, and a 2d-array is always multiple params. I am actually not sure how our sampler emcee is using the forward function and how this would work.
    * Is there another, better option? I am not really a fan of any of my options at the moment

From the proposed solutions, I prefer forward_vectorized. 
However, I find it even more elegant to overload the forward pass.
In principle, we should have the information about the parameter dimension of the model, shouldn't we? 

Or didn't I get your point?



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6034d297..f1751516 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,6 +23,8 @@ All notable changes to this project will be documented in this file.
 
 - Added the `AffineTransformation` class to perform affine transformations on the data using jax
 
+- Added the method function `forward_vectorized` to the (Base)Model class and a specialized jax version in the `JaxModel` class.
+
 ### Changed
 
 - Refactored the module `eulerpi.core.data_transformation` into the subpackage `eulerpi.core.data_transformations`. The data transformation classes can be import from this subpackage directly
@@ -31,6 +33,10 @@ All notable changes to this project will be documented in this file.
 - The inference function no longer has the parameter `custom_data_transformation`. You can pass an instance of the provided `DataTransformation` subclasses or pass a custom subclass.
 - The inference function no longer has the parameter `n_components_pca`. Instantiate a `PCATransformation` with the data and the number of components to keep. Pass it as as argument for the parameter `data_transformation` in the inference function.
 
+- Renamed the abstract base `Model` class to `BaseModel`. The usage of `Model` is still supported, but deprecated, and possibly removed in a future release.
+- Splitted the file model.py into multiple modules, one for each class, located in the subpackage `eulerpi.core.models`. The imports change from `eulerpi.core.model` to `eulerpi.core.models`.
+- Removed the function `is_artificial` from the (Base)Model class. Use `isinstance(custom_model, ArtificialModelInterface)` or `issubclass(CustomModel, ArtificialModelInterface)` instead.
+
 ### Fixed
 
 ## [0.9.1] - 2024-09-08
diff --git a/README.md b/README.md
index b8531aa1..cf61c960 100644
--- a/README.md
+++ b/README.md
@@ -55,16 +55,16 @@ You can also build the library from the latest source code by following the [Dev
 
 ## Using the library
 
-To use EPI, derive your model from the `Model` class and implement the abstract functions. Here's an example code snippet:
+To use EPI, derive your model from the `BaseModel` class and implement the abstract functions. Here's an example code snippet:
 
 ```python
 # my_model.py
 
 import jax.numpy as jnp
 
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 
-class MyModel(Model):
+class MyModel(BaseModel):
 
     param_dim = N # The dimension of a parameter point
     data_dim = M # The dimension of a data point
diff --git a/docs/source/eulerpi.core.model.rst b/docs/source/eulerpi.core.model.rst
deleted file mode 100644
index 6e0e30ef..00000000
--- a/docs/source/eulerpi.core.model.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-eulerpi.core.model module
-=========================
-
-.. automodule:: eulerpi.core.model
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/docs/source/eulerpi.core.models.artificial_model_interface.rst b/docs/source/eulerpi.core.models.artificial_model_interface.rst
new file mode 100644
index 00000000..394a369d
--- /dev/null
+++ b/docs/source/eulerpi.core.models.artificial_model_interface.rst
@@ -0,0 +1,7 @@
+eulerpi.core.models.artificial\_model\_interface module
+=======================================================
+
+.. automodule:: eulerpi.core.models.artificial_model_interface
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.models.base_model.rst b/docs/source/eulerpi.core.models.base_model.rst
new file mode 100644
index 00000000..b13a8a84
--- /dev/null
+++ b/docs/source/eulerpi.core.models.base_model.rst
@@ -0,0 +1,7 @@
+eulerpi.core.models.base_model module
+================================
+
+.. automodule:: eulerpi.core.models.base_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.models.jax_extension.rst b/docs/source/eulerpi.core.models.jax_extension.rst
new file mode 100644
index 00000000..d14125f4
--- /dev/null
+++ b/docs/source/eulerpi.core.models.jax_extension.rst
@@ -0,0 +1,7 @@
+eulerpi.core.models.jax\_extension module
+=========================================
+
+.. automodule:: eulerpi.core.models.jax_extension
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.models.jax_model.rst b/docs/source/eulerpi.core.models.jax_model.rst
new file mode 100644
index 00000000..c3c28ac9
--- /dev/null
+++ b/docs/source/eulerpi.core.models.jax_model.rst
@@ -0,0 +1,7 @@
+eulerpi.core.models.jax\_model module
+=====================================
+
+.. automodule:: eulerpi.core.models.jax_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.models.rst b/docs/source/eulerpi.core.models.rst
new file mode 100644
index 00000000..971eba26
--- /dev/null
+++ b/docs/source/eulerpi.core.models.rst
@@ -0,0 +1,11 @@
+eulerpi.core.models package
+===========================
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Classes
+
+.. automodule:: eulerpi.core.models
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.models.sbml_model.rst b/docs/source/eulerpi.core.models.sbml_model.rst
new file mode 100644
index 00000000..eeb3fa22
--- /dev/null
+++ b/docs/source/eulerpi.core.models.sbml_model.rst
@@ -0,0 +1,7 @@
+eulerpi.core.models.sbml\_model module
+======================================
+
+.. automodule:: eulerpi.core.models.sbml_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.core.rst b/docs/source/eulerpi.core.rst
index 37f5c0a4..356f6dbc 100644
--- a/docs/source/eulerpi.core.rst
+++ b/docs/source/eulerpi.core.rst
@@ -11,7 +11,7 @@ eulerpi.core package
    Inference<eulerpi.core.inference>
    Inference Types<eulerpi.core.inference_types>
    KDE<eulerpi.core.kde>
-   Model<eulerpi.core.model>
+   Models<eulerpi.core.models>
    Model Check<eulerpi.core.model_check>
    Plotting<eulerpi.core.plotting>
    Result Manager<eulerpi.core.result_manager>
diff --git a/docs/source/eulerpi.examples.cpp.cpp_model.rst b/docs/source/eulerpi.examples.cpp.cpp_model.rst
new file mode 100644
index 00000000..cc897167
--- /dev/null
+++ b/docs/source/eulerpi.examples.cpp.cpp_model.rst
@@ -0,0 +1,7 @@
+eulerpi.examples.cpp.cpp\_model module
+======================================
+
+.. automodule:: eulerpi.examples.cpp.cpp_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/eulerpi.examples.cpp.rst b/docs/source/eulerpi.examples.cpp.rst
index 656d6464..21ee230d 100644
--- a/docs/source/eulerpi.examples.cpp.rst
+++ b/docs/source/eulerpi.examples.cpp.rst
@@ -3,7 +3,7 @@ External C++ Models
 
 Many libraries in the field of scientific computing are written in C++
 to achieve fast code execution and do not have python bindings. The C++ model example shows how you can
-call C++ code from your python :py:class:`~eulerpi.core.model.Model` class.
+call C++ code from your python :py:class:`~eulerpi.core.model.BaseModel` class.
 It is primarily intended for fast implementations of the `forward` and `jacobian` method.
 
 This example model uses the following features:
diff --git a/docs/source/eulerpi.jax_extension.rst b/docs/source/eulerpi.jax_extension.rst
deleted file mode 100644
index 68484d1c..00000000
--- a/docs/source/eulerpi.jax_extension.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-eulerpi.jax\_extension module
-=============================
-
-.. automodule:: eulerpi.jax_extension
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/docs/source/eulerpi.rst b/docs/source/eulerpi.rst
index 84bf4308..8a050096 100644
--- a/docs/source/eulerpi.rst
+++ b/docs/source/eulerpi.rst
@@ -1,26 +1,8 @@
 eulerpi package
 ===============
 
-Subpackages
------------
-
 .. toctree::
    :maxdepth: 4
 
    Core<eulerpi.core>
-
-Submodules
-----------
-
-.. toctree::
-   :maxdepth: 4
-
-   Jax Extension<eulerpi.jax_extension>
-
-Module contents
----------------
-
-.. automodule:: eulerpi
-   :members:
-   :undoc-members:
-   :show-inheritance:
+   Examples<eulerpi.examples>
diff --git a/docs/source/tutorial_material/temperature.py b/docs/source/tutorial_material/temperature.py
index 7bccce73..89ff3482 100644
--- a/docs/source/tutorial_material/temperature.py
+++ b/docs/source/tutorial_material/temperature.py
@@ -3,10 +3,10 @@
 import jax.numpy as jnp
 import numpy as np
 
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 
 
-class Temperature(Model):
+class Temperature(BaseModel):
 
     param_dim = 1
     data_dim = 1
diff --git a/docs/source/tutorial_material/tutorial.ipynb b/docs/source/tutorial_material/tutorial.ipynb
index a5087c0f..2467e8e3 100644
--- a/docs/source/tutorial_material/tutorial.ipynb
+++ b/docs/source/tutorial_material/tutorial.ipynb
@@ -87,7 +87,7 @@
     "## Define your model\n",
     "\n",
     "Next you need to define your model. The most basic way is to derive from\n",
-    "the `eulerpi.core.model.Model` base class."
+    "the `eulerpi.core.models.BaseModel` class."
    ]
   },
   {
@@ -99,7 +99,7 @@
     "import importlib\n",
     "import jax.numpy as jnp\n",
     "import numpy as np\n",
-    "from eulerpi.core.model import Model"
+    "from eulerpi.core.models import BaseModel"
    ]
   },
   {
@@ -107,7 +107,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "A model inhereting from `Model` must implement the methods\n",
+    "A model inhereting from `BaseModel` must implement the methods\n",
     "- `forward`\n",
     "- `jacobian`\n",
     "\n",
@@ -133,10 +133,10 @@
     "import jax.numpy as jnp\n",
     "import numpy as np\n",
     "\n",
-    "from eulerpi.core.model import Model\n",
+    "from eulerpi.core.models import BaseModel\n",
     "\n",
     "\n",
-    "class Temperature(Model):\n",
+    "class Temperature(BaseModel):\n",
     "\n",
     "    param_dim = 1\n",
     "    data_dim = 1\n",
@@ -218,7 +218,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.10.15"
   },
   "orig_nbformat": 4,
   "vscode": {
diff --git a/docs/source/tutorial_material/tutorial.rst b/docs/source/tutorial_material/tutorial.rst
index 332567a9..353edb1d 100644
--- a/docs/source/tutorial_material/tutorial.rst
+++ b/docs/source/tutorial_material/tutorial.rst
@@ -74,7 +74,7 @@ It has 455 data points with two dimensions each. Nonuniform data is not supporte
 Define your model
 -----------------
 
-Next, you need to define your model. The most basic way is to derive from the :py:class:`eulerpi.core.model.Model` base class.
+Next, you need to define your model. The most basic way is to derive from the :py:class:`eulerpi.core.models.BaseModel` class.
 
 .. literalinclude:: ../../../eulerpi/examples/temperature/temperature.py
   :language: python
@@ -87,10 +87,10 @@ Of course, you also need the imports:
     import importlib
     import jax.numpy as jnp
     import numpy as np
-    from eulerpi.core.model import Model
+    from eulerpi.core.models import BaseModel
 
-A model inheriting from :py:class:`~eulerpi.core.model.Model` must implement the methods :py:meth:`~eulerpi.core.model.Model.forward` and :py:meth:`~eulerpi.core.model.Model.jacobian`.
-In addition, it must provide the methods :py:meth:`~eulerpi.core.model.Model.getcentral_param` and :py:meth:`~eulerpi.core.model.Model.getParamSamplingLimits` to provide the sampling algorithm with sensible starting values and boundary values.
+A model inheriting from :py:class:`~eulerpi.core.models.BaseModel` must implement the methods :py:meth:`~eulerpi.core.models.BaseModel.forward` and :py:meth:`~eulerpi.core.models.BaseModel.jacobian`.
+In addition, it must provide the methods :py:meth:`~eulerpi.core.models.BaseModel.getcentral_param` and :py:meth:`~eulerpi.core.models.BaseModel.getParamSamplingLimits` to provide the sampling algorithm with sensible starting values and boundary values.
 The jacobian is derived analytically here and implemented explicitly.
 
 .. important::
diff --git a/docs/source/welcome.rst b/docs/source/welcome.rst
index 13fb44f1..44a2003a 100644
--- a/docs/source/welcome.rst
+++ b/docs/source/welcome.rst
@@ -67,18 +67,18 @@ You can also build the library from the latest source code by following the :doc
 How to start
 ------------
 
-.. To use EPI, derive your model from :py:class:`eulerpi.core.model.Model` and implement the abstract functions :py:meth:`~eulerpi.core.model.Model.forward` and :py:meth:`~eulerpi.core.model.Model.jacobian`. You also need to define the data and parameter dimension, :py:attr:`~eulerpi.core.model.Model.data_dim` and :py:attr:`~eulerpi.core.model.Model.param_dim` of your model.
+.. To use EPI, derive your model from :py:class:`eulerpi.core.model.BaseModel` and implement the abstract functions :py:meth:`~eulerpi.core.model.BaseModel.forward` and :py:meth:`~eulerpi.core.model.BaseModel.jacobian`. You also need to define the data and parameter dimension, :py:attr:`~eulerpi.core.model.BaseModel.data_dim` and :py:attr:`~eulerpi.core.model.BaseModel.param_dim` of your model.
 
-To use EPI, derive your model from the Model class and implement the abstract functions. Here's an example code snippet:
+To use EPI, derive your model from the BaseModel class and implement the abstract functions. Here's an example code snippet:
 
 .. code-block:: python
     
     # my_model.py
     import jax.numpy as jnp
 
-    from eulerpi.core.model import Model
+    from eulerpi.core.model import BaseModel
 
-    class MyModel(Model):
+    class MyModel(BaseModel):
 
         param_dim = N # The dimension of a parameter point
         data_dim = M # The dimension of a data point
diff --git a/eulerpi/core/__init__.py b/eulerpi/core/__init__.py
index c960a234..450ac56d 100644
--- a/eulerpi/core/__init__.py
+++ b/eulerpi/core/__init__.py
@@ -1,6 +1,11 @@
 import jax
 from jax import config
 
+from .models import ArtificialModelInterface as ArtificialModelInterface
+from .models import BaseModel as BaseModel
+from .models import JaxModel as JaxModel
+from .models import SBMLModel as SBMLModel
+
 config.update("jax_enable_x64", True)
 
 # TODO: Remove restriction to cpu
diff --git a/eulerpi/core/dense_grid.py b/eulerpi/core/dense_grid.py
index 2d009738..4dc33e66 100644
--- a/eulerpi/core/dense_grid.py
+++ b/eulerpi/core/dense_grid.py
@@ -20,7 +20,7 @@
 from eulerpi.core.data_transformations import DataTransformation
 from eulerpi.core.dense_grid_types import DenseGridType
 from eulerpi.core.inference_types import InferenceType
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.core.result_manager import ResultManager
 from eulerpi.core.sampling import calc_kernel_width
 from eulerpi.core.transformations import evaluate_density
@@ -86,7 +86,7 @@ def generate_regular_grid(
 def evaluate_on_grid_chunk(
     args: typing.Tuple[
         np.ndarray,
-        Model,
+        BaseModel,
         np.ndarray,
         DataTransformation,
         np.ndarray,
@@ -97,7 +97,7 @@ def evaluate_on_grid_chunk(
 
     Args:
         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.
+        model(BaseModel): 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!)
@@ -122,7 +122,7 @@ def evaluate_on_grid_chunk(
 
 
 def run_dense_grid_evaluation(
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     slice: np.ndarray,
@@ -135,7 +135,7 @@ def run_dense_grid_evaluation(
     """This function runs a dense grid evaluation for the given model and data.
 
     Args:
-        model(Model): The model for which the evaluation should be performed.
+        model(BaseModel): 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.
@@ -209,7 +209,7 @@ def run_dense_grid_evaluation(
 
 
 def inference_dense_grid(
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     result_manager: ResultManager,
@@ -227,7 +227,7 @@ def inference_dense_grid(
     """This function runs a dense grid inference for the given model and data.
 
     Args:
-        model (Model): The model describing the mapping from parameters to data.
+        model (BaseModel): 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.
diff --git a/eulerpi/core/inference.py b/eulerpi/core/inference.py
index a63f558a..c04fea7e 100644
--- a/eulerpi/core/inference.py
+++ b/eulerpi/core/inference.py
@@ -14,14 +14,14 @@
 )
 from eulerpi.core.dense_grid import inference_dense_grid
 from eulerpi.core.inference_types import InferenceType
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.core.result_manager import ResultManager
 from eulerpi.core.sampling import inference_mcmc
 from eulerpi.core.sparsegrid import inference_sparse_grid
 
 
 def inference(
-    model: Model,
+    model: BaseModel,
     data: Union[str, os.PathLike, np.ndarray],
     inference_type: InferenceType = InferenceType.MCMC,
     slices: Optional[list[np.ndarray]] = None,
@@ -40,7 +40,7 @@ def inference(
     """Starts the parameter inference for the given model and data.
 
     Args:
-        model(Model): The model describing the mapping from parameters to data.
+        model(BaseModel): The model describing the mapping from parameters to data.
         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)
diff --git a/eulerpi/core/model.py b/eulerpi/core/model.py
deleted file mode 100644
index cc33d2ad..00000000
--- a/eulerpi/core/model.py
+++ /dev/null
@@ -1,571 +0,0 @@
-"""Module containing an abstract base class and specialized subclasses for use with the EPI algorithm."""
-
-import inspect
-import os
-import tempfile
-from abc import ABC, abstractmethod
-from functools import partial
-from typing import Optional, Tuple, Union
-
-import jax.numpy as jnp
-import numpy as np
-from jax import jacrev, jit, vmap
-
-from eulerpi.jax_extension import value_and_jacrev
-
-amici_available = False
-try:
-    import amici
-
-    amici_available = True
-except ImportError:
-    pass
-
-
-def is_amici_available():
-    return amici_available
-
-
-class Model(ABC):
-    """The base class for all models using the EPI algorithm.
-
-    Args:
-        central_param(np.ndarray): The central parameter for the model. (Default value = None)
-        param_limits(np.ndarray): Box limits for the parameters. The limits are given as a 2D array with shape (param_dim, 2). The parameter limits are used as limits as well as for the movement policy for MCMC sampling, and as boundaries for the grid when using grid-based inference. Overwrite the function param_is_within_domain if the domain is more complex than a box - the grid will still be build based on param_limits, but actual model evaluations only take place within the limits specified in param_is_within_domain. (Default value = None)
-        name(str): The name of the model. The class name is used if no name is given. (Default value = None)
-
-    .. note::
-        Examples of model implementations can be found in the :doc:`Example Models </examples>`.
-    """
-
-    param_dim: Optional[int] = (
-        None  #: The dimension of the parameter space of the model. It must be defined in the subclass.
-    )
-    data_dim: Optional[int] = (
-        None  #: The dimension of the data space of the model. It must be defined in the subclass.
-    )
-
-    def __init_subclass__(cls, **kwargs):
-        """Check if the required attributes are set."""
-        if not inspect.isabstract(cls):
-            for required in (
-                "param_dim",
-                "data_dim",
-            ):
-                if not getattr(cls, required):
-                    raise AttributeError(
-                        f"Can't instantiate abstract class {cls.__name__} without {required} attribute defined"
-                    )
-        return cls
-
-    def __init__(
-        self,
-        central_param: np.ndarray,
-        param_limits: np.ndarray,
-        name: Optional[str] = None,
-    ) -> None:
-        self.central_param = central_param
-        self.param_limits = param_limits
-
-        self.name = name or self.__class__.__name__
-
-    @abstractmethod
-    def forward(self, param: np.ndarray) -> np.ndarray:
-        """Executed the forward pass of the model to obtain data from a parameter.
-
-        Args:
-            param(np.ndarray): The parameter for which the data should be generated.
-
-        Returns:
-            np.ndarray: The data generated from the parameter.
-
-        Examples:
-
-        .. code-block:: python
-
-            import numpy as np
-            from eulerpi.examples.heat import Heat
-            from eulerpi.core.model import JaxModel
-            from jax import vmap
-
-            # instantiate the heat model
-            model = Heat()
-
-            # define a 3D example parameter for the heat model
-            example_param = np.array([1.4, 1.6, 0.5])
-
-            # the forward simulation is achieved by using the forward method of the model
-            sim_result = model.forward(example_param)
-
-            # in a more realistic scenario, we would like to perform the forward pass on multiple parameters at once
-            multiple_params = np.array([[1.5, 1.5, 0.5],
-                                        [1.4, 1.4, 0.6],
-                                        [1.6, 1.6, 0.4],
-                                        model.central_param,
-                                        [1.5, 1.4, 0.4]])
-
-            # try to use jax vmap to perform the forward pass on multiple parameters at once
-            if isinstance(model, JaxModel):
-                multiple_sim_results = vmap(model.forward, in_axes=0)(multiple_params)
-
-            # if the model is not a jax model, we can use numpy vectorize to perform the forward pass
-            else:
-                multiple_sim_results = np.vectorize(model.forward, signature="(n)->(m)")(multiple_params)
-
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def jacobian(self, param: np.ndarray) -> np.ndarray:
-        """Evaluates the jacobian of the :func:`~eulerpi.core.model.Model.forward` method.
-
-        Args:
-            param(np.ndarray): The parameter for which the jacobian should be evaluated.
-
-        Returns:
-            np.ndarray: The jacobian for the variables returned by the :func:`~eulerpi.core.model.Model.forward` method with respect to the parameters.
-
-        Examples:
-
-        .. code-block:: python
-
-            import numpy as np
-            from eulerpi.examples.heat import Heat
-            from eulerpi.core.model import JaxModel
-            from jax import vmap
-
-            # instantiate the heat model
-            model = Heat()
-
-            # define a 3D example parameter for the heat model
-            example_param = np.array([1.4, 1.6, 0.5])
-
-            sim_jacobian = model.jacobian(example_param)
-
-            # Similar to the forward pass, also the evaluation of the jacobian can be vectorized.
-            # This yields a 3D array of shape (num_params, data_dim, param_dim) = (4,5,3) in this example.
-
-            multiple_params = np.array([[1.5, 1.5, 0.5],
-                                        [1.4, 1.4, 0.6],
-                                        model.central_param,
-                                        [1.5, 1.4, 0.4]])
-
-            # try to use jax vmap for vectorization if possible
-            if isinstance(model, JaxModel):
-                multiple_sim_jacobians = vmap(model.jacobian, in_axes=0)(multiple_params)
-
-            # if the model is not a jax model, we can use numpy vectorize to vectorize
-            else:
-                multiple_sim_jacobians = np.vectorize(model.jacobian, signature="(n)->(m)")(multiple_params)
-
-        """
-        raise NotImplementedError
-
-    def forward_and_jacobian(
-        self, param: np.ndarray
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        """Evaluates the jacobian and the forward pass of the model at the same time. If the method is not overwritten in a subclass it,
-        it simply calls :func:`~eulerpi.core.model.Model.forward` and :func:`~eulerpi.core.model.Model.jacobian`.
-        It can be vectorized in the same way as the forward and jacobian methods.
-
-        Args:
-            param(np.ndarray): The parameter for which the jacobian should be evaluated.
-
-        Returns:
-            typing.Tuple[np.ndarray, np.ndarray]: The data generated from the parameter and the jacobian for the variables returned by the :func:`~eulerpi.core.model.Model.forward` method with respect to the parameters.
-
-        """
-
-        return self.forward(param), self.jacobian(param)
-
-    def param_is_within_domain(self, param: np.ndarray) -> bool:
-        """Checks whether a parameter is within the parameter domain of the model.
-        Overwrite this function if your model has a more complex parameter domain than a box. The param_limits are checked automatically.
-
-        Args:
-            param(np.ndarray): The parameter to check.
-
-        Returns:
-            bool: True if the parameter is within the limits.
-
-        """
-        return True
-
-    def is_artificial(self) -> bool:
-        """Determines whether the model provides artificial parameter and data sets.
-
-        Returns:
-            bool: True if the model inherits from the ArtificialModelInterface
-
-        """
-        return issubclass(self.__class__, ArtificialModelInterface)
-
-
-class ArtificialModelInterface(ABC):
-    """By inheriting from this interface you indicate that you are providing an artificial parameter dataset,
-    and the corresponding artificial data dataset, which can be used to compare the results from eulerpi with the ground truth.
-
-    """
-
-    @abstractmethod
-    def generate_artificial_params(self, num_samples: int) -> np.ndarray:
-        """This method must be overwritten an return an numpy array of num_samples parameters.
-
-        Args:
-            num_samples(int): The number of parameters to generate.
-
-        Returns:
-            np.ndarray: The generated parameters.
-
-        Raises:
-            NotImplementedError: If the method is not overwritten in a subclass.
-
-        """
-        raise NotImplementedError
-
-    def generate_artificial_data(
-        self,
-        params: Union[os.PathLike, str, np.ndarray],
-    ) -> np.ndarray:
-        """This method is called when the user wants to generate artificial data from the model.
-
-        Args:
-            params: typing.Union[os.PathLike, str, np.ndarray]: The parameters for which the data should be generated. Can be either a path to a file, a numpy array or a string.
-
-        Returns:
-            np.ndarray: The data generated from the parameters.
-
-        Raises:
-            TypeError: If the params argument is not a path to a file, a numpy array or a string.
-
-        """
-        if isinstance(params, str) or isinstance(params, os.PathLike):
-            params = np.loadtxt(params, delimiter=",", ndmin=2)
-        elif isinstance(params, np.ndarray) or isinstance(params, jnp.ndarray):
-            pass
-        else:
-            raise TypeError(
-                f"The params argument has to be either a path to a file or a numpy array. The passed argument was of type {type(params)}"
-            )
-
-        # try to use jax vmap to perform the forward pass on multiple parameters at once
-        if isinstance(self, JaxModel):
-            return vmap(self.forward, in_axes=0)(params)
-        else:
-            return np.vectorize(self.forward, signature="(n)->(m)")(params)
-
-
-def add_autodiff(_cls):
-    """
-    Decorator to automatically create the jacobian method based on the forward method.
-    Additionally it jit compiles the forward and jacobian method with jax.
-
-    Args:
-        _cls: The class to decorate.
-
-    Returns:
-        The decorated class with the jacobian method and the forward and jacobian method jit compiled with jax.
-
-    """
-    _cls.init_fw_and_bw()
-    return _cls
-
-
-class JaxModel(Model):
-    """The JaxModel is a base class for models using the JAX library.
-
-    It automatically creates the jacobian method based on the forward method.
-    Additionally it jit compiles the forward and jacobian method with jax for faster execution.
-
-    .. note::
-
-        To use this class you have to implement your forward method using jax, e. g. jax.numpy.
-        Dont overwrite the __init__ method of JaxModel without calling the super constructor.
-        Else your forward method wont be jitted.
-    """
-
-    def __init__(
-        self,
-        central_param: np.ndarray,
-        param_limits: np.ndarray,
-        name: Optional[str] = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            central_param=central_param,
-            param_limits=param_limits,
-            name=name,
-            **kwargs,
-        )
-        # TODO: Check performance implications of not setting this at the class level but for each instance.
-        type(self).forward = partial(JaxModel.forward_method, self)
-
-    def __init_subclass__(cls, **kwargs):
-        """Automatically create the jacobian method based on the forward method for the subclass."""
-        return add_autodiff(super().__init_subclass__(**kwargs))
-
-    @classmethod
-    def init_fw_and_bw(cls):
-        """Calculates the jitted methods for the subclass(es).
-        It is an unintended sideeffect that this happens for all intermediate classes also.
-        E.g. for: class CoronaArtificial(Corona)
-        """
-        cls.fw = jit(cls.forward)
-        cls.bw = jit(jacrev(cls.forward))
-        cls.vj = jit(value_and_jacrev(cls.forward))
-
-    @staticmethod
-    def forward_method(self, param: np.ndarray) -> np.ndarray:
-        """This method is called by the jitted forward method. It is not intended to be called directly.
-
-        Args:
-            param(np.ndarray): The parameter for which the data should be generated.
-
-        Returns:
-            np.ndarray: The data generated from the parameter.
-
-        """
-        return type(self).fw(param)
-
-    def jacobian(self, param: np.ndarray) -> np.ndarray:
-        """Jacobian of the forward pass with respect to the parameters.
-
-        Args:
-            param(np.ndarray): The parameter for which the jacobian should be evaluated.
-
-        Returns:
-            np.ndarray: The jacobian for the variables returned by the :func:`~eulerpi.core.model.Model.forward` method with respect to the parameters.
-
-        """
-        return type(self).bw(param)
-
-    def forward_and_jacobian(
-        self, param: np.ndarray
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        """Evaluates the jacobian and the forward pass of the model at the same time. This can be more efficient than calling the :func:`~eulerpi.core.model.Model.forward` and :func:`~eulerpi.core.model.Model.jacobian` methods separately.
-
-        Args:
-            param(np.ndarray): The parameter for which the jacobian should be evaluated.
-
-        Returns:
-            typing.Tuple[np.ndarray, np.ndarray]: The data and the jacobian for a given parameter.
-
-        """
-        return type(self).vj(param)
-
-
-class SBMLModel(Model):
-    """The SBMLModel class is a wrapper for the AMICI python interface to simulate SBML models using this package.
-
-    Args:
-        sbml_file(str): The path to the SBML model file.
-        param_ids(list): A list of ids of parameter, which will be estimated during the inference. If None all parameter ids are extracted from the SBML model.
-        state_ids(list): A list of state ids, for which data will be given during the inference. If None all state ids are extracted from the SBML model.
-        timepoints(list): List of measurement time points, this is where the sbml model is evaluated and compared to the data
-        skip_creation(bool): If True the model is not created againg based on the SBML file. Instead the model is loaded from a previously created model. (Default value = False)
-        central_param(np.ndarray): The central parameter for the model
-        param_limits(np.ndarray): The parameter limits for the model
-    """
-
-    @staticmethod
-    def indices_from_ids(ids: list, all_ids: list) -> list:
-        """Returns the indices of the ids in the all_ids list.
-
-        Args:
-            ids(list): The ids for which the indices should be returned.
-            all_ids(list): The list of all ids.
-
-        Returns:
-            list: The indices of the ids in the all_ids list.
-
-        Throws:
-            ValueError: If one of the ids is not in the all_ids list.
-
-        """
-        indices = []
-        for id in ids:
-            try:
-                indices.append(all_ids.index(id))
-            except ValueError:
-                raise ValueError(
-                    f"Parameter / State id '{id}' is not in the list of the relevant ids {all_ids}"
-                )
-        return indices
-
-    @property
-    def param_dim(self):
-        """The number of parameters of the model."""
-        return len(self.param_ids)
-
-    @property
-    def data_dim(self):
-        """The dimension of a data point returned by the model."""
-        return len(self.state_ids) * len(self.timepoints)
-
-    def __init__(
-        self,
-        sbml_file: str,
-        central_param: np.ndarray,
-        param_limits: np.ndarray,
-        timepoints: list,
-        param_ids: Optional[list] = None,
-        state_ids: Optional[list] = None,
-        skip_creation: bool = False,
-        name: Optional[str] = None,
-        **kwargs,
-    ) -> None:
-        if not is_amici_available():
-            raise ImportError(
-                "The SBMLModel class requires an AMICI installation. "
-                "Please install AMICI first."
-                "You can do this manually or install eulerpi with the [amici] extra: pip install eulerpi[amici]"
-            )
-
-        super().__init__(central_param, param_limits, name, **kwargs)
-
-        self.amici_model_name = self.name
-        self.amici_model_dir = (
-            "./generated_sbml_models/" + self.amici_model_name
-        )
-
-        # Generate python code
-        if not skip_creation:
-            sbml_importer = amici.SbmlImporter(sbml_file)
-            sbml_importer.sbml2amici(
-                self.amici_model_name,
-                self.amici_model_dir,
-            )
-
-        # Load the generated model
-        self.timepoints = timepoints
-        self.load_amici_model_and_solver()
-
-        self.param_ids = param_ids or self.amici_model.getParametersIds()
-        self.state_ids = state_ids or self.amici_model.getStateIds()
-        self.param_indices = self.indices_from_ids(
-            self.param_ids, self.amici_model.getParameterIds()
-        )
-        self.state_indices = self.indices_from_ids(
-            self.state_ids, self.amici_model.getStateIds()
-        )
-        self.setSensitivities()
-
-    def load_amici_model_and_solver(self):
-        """Loads the AMICI model from the previously generated model."""
-        amici_model_module = amici.import_model_module(
-            self.amici_model_name, self.amici_model_dir
-        )
-        self.amici_model = amici_model_module.getModel()
-        self.amici_solver = self.amici_model.getSolver()
-
-        self.amici_model.setTimepoints(self.timepoints)
-        self.amici_solver.setAbsoluteTolerance(1e-10)
-
-    def setSensitivities(self):
-        """Tell the underlying amici solver to calculate sensitivities based on the attribute `self.param_ids`"""
-        if self.param_ids == self.amici_model.getParameterIds():
-            self.amici_model.requireSensitivitiesForAllParameters()
-        else:
-            self.amici_model.setParameterList(self.param_indices)
-
-        self.amici_solver.setSensitivityMethod(amici.SensitivityMethod.forward)
-        self.amici_solver.setSensitivityOrder(amici.SensitivityOrder.first)
-
-    def forward(self, params):
-        for i, param in enumerate(params):
-            self.amici_model.setParameterById(self.param_ids[i], param)
-        rdata = amici.runAmiciSimulation(self.amici_model, self.amici_solver)
-        return rdata.x[:, self.state_indices].reshape(self.data_dim)
-
-    def jacobian(self, params):
-        for i, param in enumerate(params):
-            self.amici_model.setParameterById(self.param_ids[i], param)
-        rdata = amici.runAmiciSimulation(self.amici_model, self.amici_solver)
-        return (
-            rdata.sx[:, :, self.state_indices]
-            .transpose(1, 0, 2)
-            .reshape(self.param_dim, self.data_dim)
-            .T
-        )
-
-    def forward_and_jacobian(
-        self, params: np.ndarray
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        for i, param in enumerate(params):
-            self.amici_model.setParameterById(self.param_ids[i], param)
-        rdata = amici.runAmiciSimulation(self.amici_model, self.amici_solver)
-        return (
-            rdata.x[:, self.state_indices].reshape(self.data_dim),
-            rdata.sx[:, :, self.state_indices]
-            .transpose(1, 0, 2)
-            .reshape(self.param_dim, self.data_dim)
-            .T,
-        )
-
-    # Allow the model to be pickled
-    def __getstate__(self):
-        # Create a copy of the object's state
-        state = self.__dict__.copy()
-
-        # Save the amici solver settings to
-        _fd, _file = tempfile.mkstemp()
-
-        try:
-            # write amici solver settings to file
-            try:
-                amici.writeSolverSettingsToHDF5(self.amici_solver, _file)
-            except AttributeError as e:
-                e.args += (
-                    "Pickling the SBMLModel requires an AMICI "
-                    "installation with HDF5 support.",
-                )
-                raise
-            # read in byte stream
-            with open(_fd, "rb", closefd=False) as f:
-                state["amici_solver_settings"] = f.read()
-        finally:
-            # close file descriptor and remove temporary file
-            os.close(_fd)
-            os.remove(_file)
-
-        state["amici_model_settings"] = amici.get_model_settings(
-            self.amici_model
-        )
-
-        # Remove the unpicklable entries.
-        del state["amici_model"]
-        del state["amici_solver"]
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-
-        # Restore amici model and solver
-        self.load_amici_model_and_solver()
-        self.setSensitivities()
-
-        _fd, _file = tempfile.mkstemp()
-        try:
-            # write solver settings to temporary file
-            with open(_fd, "wb", closefd=False) as f:
-                f.write(state["amici_solver_settings"])
-            # read in solver settings
-            try:
-                amici.readSolverSettingsFromHDF5(_file, self.amici_solver)
-            except AttributeError as err:
-                if not err.args:
-                    err.args = ("",)
-                err.args += (
-                    "Unpickling an AmiciObjective requires an AMICI "
-                    "installation with HDF5 support.",
-                )
-                raise
-        finally:
-            # close file descriptor and remove temporary file
-            os.close(_fd)
-            os.remove(_file)
-
-        amici.set_model_settings(
-            self.amici_model,
-            state["amici_model_settings"],
-        )
diff --git a/eulerpi/core/model_check.py b/eulerpi/core/model_check.py
index 75cb2e6c..f6db3db6 100644
--- a/eulerpi/core/model_check.py
+++ b/eulerpi/core/model_check.py
@@ -1,15 +1,14 @@
-"""Check a custom :py:class:`Model <eulerpi.core.model.Model>` for implementation errors or test them in a quick inference run on an artificially created dataset."""
+"""Check a custom :py:class:`BaseModel <eulerpi.core.models.BaseModel>` for implementation errors or test them in a quick inference run on an artificially created dataset."""
 
 import jax.numpy as jnp
 import numpy as np
-from jax import vmap
 
 from eulerpi.core.inference import InferenceType, inference
-from eulerpi.core.model import JaxModel, Model
+from eulerpi.core.models import BaseModel
 from eulerpi.core.plotting import sample_violin_plot
 
 
-def basic_model_check(model: Model) -> None:
+def basic_model_check(model: BaseModel) -> None:
     """Perform a simple sanity check on the model.
 
     It tests the following:
@@ -26,7 +25,7 @@ def basic_model_check(model: Model) -> None:
     - The jacobi matrix has full rank
 
     Args:
-        model(Model): The model describing the mapping from parameters to data.
+        model(BaseModel): The model describing the mapping from parameters to data.
 
     Raises:
         AssertionError: Raised if any of the tests fails.
@@ -124,7 +123,7 @@ def basic_model_check(model: Model) -> None:
 
 
 def inference_model_check(
-    model: Model,
+    model: BaseModel,
     num_data_points: int = 1000,
     num_model_evaluations: int = 11000,
 ) -> None:
@@ -133,7 +132,7 @@ def inference_model_check(
     It produces a violin plot comparing the artificially created parameters and data to the respectively inferred samples.
 
     Args:
-        model(Model): The model describing the mapping from parameters to data.
+        model(BaseModel): The model describing the mapping from parameters to data.
         num_data_points (int, optional): The number of data data points to artificially generate (Default value = 1000)
         num_model_evaluations (int, optional): The number of model evaluations to perform in the inference. (Default value = 11000)
 
@@ -168,13 +167,7 @@ def inference_model_check(
         (np.random.rand(num_data_points, model.param_dim) - 0.5) / 3.0
     )
 
-    # try to use jax vmap to perform the forward pass on multiple parameters at once
-    if isinstance(model, JaxModel):
-        data_sample = vmap(model.forward, in_axes=0)(param_sample)
-    else:
-        data_sample = np.vectorize(model.forward, signature="(n)->(m)")(
-            param_sample
-        )
+    data_sample = model.forward_vectorized(param_sample)
 
     print(
         f"Successfully created an artificial data set of size {num_data_points}.\n"
@@ -234,7 +227,7 @@ def inference_model_check(
 
 
 def full_model_check(
-    model: Model,
+    model: BaseModel,
     num_data_points: int = 1000,
     num_model_evaluations: int = 11000,
 ) -> None:
@@ -245,7 +238,7 @@ def full_model_check(
     It runs the functions :py:func:`basic_model_check <basic_model_check>` and :py:func:`inference_model_check <inference_model_check>` to perform the checks.
 
     Args:
-        model(Model): The model describing the mapping from parameters to data.
+        model(BaseModel): The model describing the mapping from parameters to data.
         num_data_points (int, optional): The number of data data points to artificially generate (Default value = 1000)
         num_model_evaluations (int, optional): The number of model evaluations to perform in the inference. (Default value = 11000)
 
diff --git a/eulerpi/core/models/__init__.py b/eulerpi/core/models/__init__.py
new file mode 100644
index 00000000..45e14935
--- /dev/null
+++ b/eulerpi/core/models/__init__.py
@@ -0,0 +1,17 @@
+"""Package containing an abstract model base class and specialized subclasses for use with the EPI algorithm."""
+
+from .artificial_model_interface import (
+    ArtificialModelInterface as ArtificialModelInterface,
+)
+from .base_model import BaseModel as BaseModel  # Backwards Compatibility
+from .base_model import BaseModel as Model
+from .jax_model import JaxModel as JaxModel
+from .sbml_model import SBMLModel as SBMLModel
+
+__all__ = [
+    "ArtificialModelInterface",
+    "BaseModel",
+    "JaxModel",
+    "SBMLModel",
+    "Model",
+]
diff --git a/eulerpi/core/models/artificial_model_interface.py b/eulerpi/core/models/artificial_model_interface.py
new file mode 100644
index 00000000..37b8c836
--- /dev/null
+++ b/eulerpi/core/models/artificial_model_interface.py
@@ -0,0 +1,56 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Union
+
+import jax.numpy as jnp
+import numpy as np
+
+
+class ArtificialModelInterface(ABC):
+    """By inheriting from this interface you indicate that you are providing an artificial parameter dataset,
+    and the corresponding artificial data dataset, which can be used to compare the results from eulerpi with the ground truth.
+
+    """
+
+    @abstractmethod
+    def generate_artificial_params(self, num_samples: int) -> np.ndarray:
+        """This method must be overwritten an return an numpy array of num_samples parameters.
+
+        Args:
+            num_samples(int): The number of parameters to generate.
+
+        Returns:
+            np.ndarray: The generated parameters.
+
+        Raises:
+            NotImplementedError: If the method is not overwritten in a subclass.
+
+        """
+        raise NotImplementedError
+
+    def generate_artificial_data(
+        self,
+        params: Union[os.PathLike, str, np.ndarray],
+    ) -> np.ndarray:
+        """This method is called when the user wants to generate artificial data from the model.
+
+        Args:
+            params: typing.Union[os.PathLike, str, np.ndarray]: The parameters for which the data should be generated. Can be either a path to a file, a numpy array or a string.
+
+        Returns:
+            np.ndarray: The data generated from the parameters.
+
+        Raises:
+            TypeError: If the params argument is not a path to a file, a numpy array or a string.
+
+        """
+        if isinstance(params, str) or isinstance(params, os.PathLike):
+            params = np.loadtxt(params, delimiter=",", ndmin=2)
+        elif isinstance(params, np.ndarray) or isinstance(params, jnp.ndarray):
+            pass
+        else:
+            raise TypeError(
+                f"The params argument has to be either a path to a file or a numpy array. The passed argument was of type {type(params)}"
+            )
+
+        return self.forward_vectorized(params)
diff --git a/eulerpi/core/models/base_model.py b/eulerpi/core/models/base_model.py
new file mode 100644
index 00000000..66bf3bde
--- /dev/null
+++ b/eulerpi/core/models/base_model.py
@@ -0,0 +1,175 @@
+import inspect
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple
+
+import numpy as np
+
+
+class BaseModel(ABC):
+    """The base class for all models using the EPI algorithm.
+
+    Args:
+        central_param(np.ndarray): The central parameter for the model. (Default value = None)
+        param_limits(np.ndarray): Box limits for the parameters. The limits are given as a 2D array with shape (param_dim, 2). The parameter limits are used as limits as well as for the movement policy for MCMC sampling, and as boundaries for the grid when using grid-based inference. Overwrite the function param_is_within_domain if the domain is more complex than a box - the grid will still be build based on param_limits, but actual model evaluations only take place within the limits specified in param_is_within_domain. (Default value = None)
+        name(str): The name of the model. The class name is used if no name is given. (Default value = None)
+
+    .. note::
+        Examples of model implementations can be found in the :doc:`Example Models </examples>`.
+    """
+
+    param_dim: Optional[int] = (
+        None  #: The dimension of the parameter space of the model. It must be defined in the subclass.
+    )
+    data_dim: Optional[int] = (
+        None  #: The dimension of the data space of the model. It must be defined in the subclass.
+    )
+
+    def __init_subclass__(cls, **kwargs):
+        """Check if the required attributes are set."""
+        if not inspect.isabstract(cls):
+            for required in (
+                "param_dim",
+                "data_dim",
+            ):
+                if not getattr(cls, required):
+                    raise AttributeError(
+                        f"Can't instantiate abstract class {cls.__name__} without {required} attribute defined"
+                    )
+        return cls
+
+    def __init__(
+        self,
+        central_param: np.ndarray,
+        param_limits: np.ndarray,
+        name: Optional[str] = None,
+    ) -> None:
+        self.central_param = central_param
+        self.param_limits = param_limits
+
+        self.name = name or self.__class__.__name__
+
+    @abstractmethod
+    def forward(self, param: np.ndarray) -> np.ndarray:
+        """Executed the forward pass of the model to obtain data from a parameter.
+
+        Args:
+            param(np.ndarray): The parameter for which the data should be generated.
+
+        Returns:
+            np.ndarray: The data generated from the parameter.
+
+        Examples:
+
+        .. code-block:: python
+
+            import numpy as np
+            from eulerpi.examples.heat import Heat
+            from eulerpi.core.models import JaxModel
+            from jax import vmap
+
+            # instantiate the heat model
+            model = Heat()
+
+            # define a 3D example parameter for the heat model
+            example_param = np.array([1.4, 1.6, 0.5])
+
+            # the forward simulation is achieved by using the forward method of the model
+            sim_result = model.forward(example_param)
+
+            # in a more realistic scenario, we would like to perform the forward pass on multiple parameters at once
+            multiple_params = np.array([[1.5, 1.5, 0.5],
+                                        [1.4, 1.4, 0.6],
+                                        [1.6, 1.6, 0.4],
+                                        model.central_param,
+                                        [1.5, 1.4, 0.4]])
+            multiple_sim_results = model.forward_vectorized(multiple_params)
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def jacobian(self, param: np.ndarray) -> np.ndarray:
+        """Evaluates the jacobian of the :func:`~eulerpi.core.models.BaseModel.forward` method.
+
+        Args:
+            param(np.ndarray): The parameter for which the jacobian should be evaluated.
+
+        Returns:
+            np.ndarray: The jacobian for the variables returned by the :func:`~eulerpi.core.models.BaseModel.forward` method with respect to the parameters.
+
+        Examples:
+
+        .. code-block:: python
+
+            import numpy as np
+            from eulerpi.examples.heat import Heat
+            from eulerpi.core.models import JaxModel
+            from jax import vmap
+
+            # instantiate the heat model
+            model = Heat()
+
+            # define a 3D example parameter for the heat model
+            example_param = np.array([1.4, 1.6, 0.5])
+
+            sim_jacobian = model.jacobian(example_param)
+
+            # Similar to the forward pass, also the evaluation of the jacobian can be vectorized.
+            # This yields a 3D array of shape (num_params, data_dim, param_dim) = (4,5,3) in this example.
+
+            multiple_params = np.array([[1.5, 1.5, 0.5],
+                                        [1.4, 1.4, 0.6],
+                                        model.central_param,
+                                        [1.5, 1.4, 0.4]])
+
+            # try to use jax vmap for vectorization if possible
+            if isinstance(model, JaxModel):
+                multiple_sim_jacobians = vmap(model.jacobian, in_axes=0)(multiple_params)
+
+            # if the model is not a jax model, we can use numpy vectorize to vectorize
+            else:
+                multiple_sim_jacobians = np.vectorize(model.jacobian, signature="(n)->(m)")(multiple_params)
+
+        """
+        raise NotImplementedError
+
+    def forward_and_jacobian(
+        self, param: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """Evaluates the jacobian and the forward pass of the model at the same time. If the method is not overwritten in a subclass it,
+        it simply calls :func:`~eulerpi.core.models.BaseModel.forward` and :func:`~eulerpi.core.models.BaseModel.jacobian`.
+        It can be vectorized in the same way as the forward and jacobian methods.
+
+        Args:
+            param(np.ndarray): The parameter for which the jacobian should be evaluated.
+
+        Returns:
+            typing.Tuple[np.ndarray, np.ndarray]: The data generated from the parameter and the jacobian for the variables returned by the :func:`~eulerpi.core.models.BaseModel.forward` method with respect to the parameters.
+
+        """
+
+        return self.forward(param), self.jacobian(param)
+
+    def param_is_within_domain(self, param: np.ndarray) -> bool:
+        """Checks whether a parameter is within the parameter domain of the model.
+        Overwrite this function if your model has a more complex parameter domain than a box. The param_limits are checked automatically.
+
+        Args:
+            param(np.ndarray): The parameter to check.
+
+        Returns:
+            bool: True if the parameter is within the limits.
+
+        """
+        return True
+
+    def forward_vectorized(self, params: np.ndarray) -> np.ndarray:
+        """A vectorized version of the forward function
+
+        Args:
+            params (np.ndarray): an array of parameters, shape (n, self.param_dim)
+
+        Returns:
+            np.ndarray: The data vector generated from the vector of parameters, shape (n, self.data_dim)
+        """
+        return np.vectorize(self.forward, signature="(n)->(m)")(params)
diff --git a/eulerpi/jax_extension.py b/eulerpi/core/models/jax_extension.py
similarity index 100%
rename from eulerpi/jax_extension.py
rename to eulerpi/core/models/jax_extension.py
diff --git a/eulerpi/core/models/jax_model.py b/eulerpi/core/models/jax_model.py
new file mode 100644
index 00000000..dae869af
--- /dev/null
+++ b/eulerpi/core/models/jax_model.py
@@ -0,0 +1,110 @@
+from functools import partial
+from typing import Optional, Tuple
+
+import numpy as np
+from jax import jacrev, jit, vmap
+
+from .base_model import BaseModel
+from .jax_extension import value_and_jacrev
+
+
+def add_autodiff(_cls):
+    """
+    Decorator to automatically create the jacobian method based on the forward method.
+    Additionally it jit compiles the forward and jacobian method with jax.
+
+    Args:
+        _cls: The class to decorate.
+
+    Returns:
+        The decorated class with the jacobian method and the forward and jacobian method jit compiled with jax.
+
+    """
+    _cls.init_fw_and_bw()
+    return _cls
+
+
+class JaxModel(BaseModel):
+    """The JaxModel is a base class for models using the JAX library.
+
+    It automatically creates the jacobian method based on the forward method.
+    Additionally it jit compiles the forward and jacobian method with jax for faster execution.
+
+    .. note::
+
+        To use this class you have to implement your forward method using jax, e. g. jax.numpy.
+        Dont overwrite the __init__ method of JaxModel without calling the super constructor.
+        Else your forward method wont be jitted.
+    """
+
+    def __init__(
+        self,
+        central_param: np.ndarray,
+        param_limits: np.ndarray,
+        name: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            central_param=central_param,
+            param_limits=param_limits,
+            name=name,
+            **kwargs,
+        )
+        # TODO: Check performance implications of not setting this at the class level but for each instance.
+        type(self).forward = partial(JaxModel.forward_method, self)
+
+    def __init_subclass__(cls, **kwargs):
+        """Automatically create the jacobian method based on the forward method for the subclass."""
+        return add_autodiff(super().__init_subclass__(**kwargs))
+
+    @classmethod
+    def init_fw_and_bw(cls):
+        """Calculates the jitted methods for the subclass(es).
+        It is an unintended sideeffect that this happens for all intermediate classes also.
+        E.g. for: class CoronaArtificial(Corona)
+        """
+        cls.fw = jit(cls.forward)
+        cls.bw = jit(jacrev(cls.forward))
+        cls.vj = jit(value_and_jacrev(cls.forward))
+
+    @staticmethod
+    def forward_method(self, param: np.ndarray) -> np.ndarray:
+        """This method is called by the jitted forward method. It is not intended to be called directly.
+
+        Args:
+            param(np.ndarray): The parameter for which the data should be generated.
+
+        Returns:
+            np.ndarray: The data generated from the parameter.
+
+        """
+        return type(self).fw(param)
+
+    def jacobian(self, param: np.ndarray) -> np.ndarray:
+        """Jacobian of the forward pass with respect to the parameters.
+
+        Args:
+            param(np.ndarray): The parameter for which the jacobian should be evaluated.
+
+        Returns:
+            np.ndarray: The jacobian for the variables returned by the :func:`~eulerpi.core.models.BaseModel.forward` method with respect to the parameters.
+
+        """
+        return type(self).bw(param)
+
+    def forward_and_jacobian(
+        self, param: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """Evaluates the jacobian and the forward pass of the model at the same time. This can be more efficient than calling the :func:`~eulerpi.core.models.BaseModel.forward` and :func:`~eulerpi.core.models.BaseModel.jacobian` methods separately.
+
+        Args:
+            param(np.ndarray): The parameter for which the jacobian should be evaluated.
+
+        Returns:
+            typing.Tuple[np.ndarray, np.ndarray]: The data and the jacobian for a given parameter.
+
+        """
+        return type(self).vj(param)
+
+    def forward_vectorized(self, params: np.ndarray) -> np.ndarray:
+        return vmap(self.forward, in_axes=0)(params)
diff --git a/eulerpi/core/models/sbml_model.py b/eulerpi/core/models/sbml_model.py
new file mode 100644
index 00000000..a6242865
--- /dev/null
+++ b/eulerpi/core/models/sbml_model.py
@@ -0,0 +1,236 @@
+import os
+import tempfile
+from typing import Optional, Tuple
+
+import numpy as np
+
+from .base_model import BaseModel
+
+amici_available = False
+try:
+    import amici
+
+    amici_available = True
+except ImportError:
+    pass
+
+
+def is_amici_available():
+    return amici_available
+
+
+class SBMLModel(BaseModel):
+    """The SBMLModel class is a wrapper for the AMICI python interface to simulate SBML models using this package.
+
+    Args:
+        sbml_file(str): The path to the SBML model file.
+        param_ids(list): A list of ids of parameter, which will be estimated during the inference. If None all parameter ids are extracted from the SBML model.
+        state_ids(list): A list of state ids, for which data will be given during the inference. If None all state ids are extracted from the SBML model.
+        timepoints(list): List of measurement time points, this is where the sbml model is evaluated and compared to the data
+        skip_creation(bool): If True the model is not created againg based on the SBML file. Instead the model is loaded from a previously created model. (Default value = False)
+        central_param(np.ndarray): The central parameter for the model
+        param_limits(np.ndarray): The parameter limits for the model
+    """
+
+    @staticmethod
+    def indices_from_ids(ids: list, all_ids: list) -> list:
+        """Returns the indices of the ids in the all_ids list.
+
+        Args:
+            ids(list): The ids for which the indices should be returned.
+            all_ids(list): The list of all ids.
+
+        Returns:
+            list: The indices of the ids in the all_ids list.
+
+        Throws:
+            ValueError: If one of the ids is not in the all_ids list.
+
+        """
+        indices = []
+        for id in ids:
+            try:
+                indices.append(all_ids.index(id))
+            except ValueError:
+                raise ValueError(
+                    f"Parameter / State id '{id}' is not in the list of the relevant ids {all_ids}"
+                )
+        return indices
+
+    @property
+    def param_dim(self):
+        """The number of parameters of the model."""
+        return len(self.param_ids)
+
+    @property
+    def data_dim(self):
+        """The dimension of a data point returned by the model."""
+        return len(self.state_ids) * len(self.timepoints)
+
+    def __init__(
+        self,
+        sbml_file: str,
+        central_param: np.ndarray,
+        param_limits: np.ndarray,
+        timepoints: list,
+        param_ids: Optional[list] = None,
+        state_ids: Optional[list] = None,
+        skip_creation: bool = False,
+        name: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        if not is_amici_available():
+            raise ImportError(
+                "The SBMLModel class requires an AMICI installation. "
+                "Please install AMICI first."
+                "You can do this manually or install eulerpi with the [amici] extra: pip install eulerpi[amici]"
+            )
+
+        super().__init__(central_param, param_limits, name, **kwargs)
+
+        self.amici_model_name = self.name
+        self.amici_model_dir = (
+            "./generated_sbml_models/" + self.amici_model_name
+        )
+
+        # Generate python code
+        if not skip_creation:
+            sbml_importer = amici.SbmlImporter(sbml_file)
+            sbml_importer.sbml2amici(
+                self.amici_model_name,
+                self.amici_model_dir,
+            )
+
+        # Load the generated model
+        self.timepoints = timepoints
+        self.load_amici_model_and_solver()
+
+        self.param_ids = param_ids or self.amici_model.getParametersIds()
+        self.state_ids = state_ids or self.amici_model.getStateIds()
+        self.param_indices = self.indices_from_ids(
+            self.param_ids, self.amici_model.getParameterIds()
+        )
+        self.state_indices = self.indices_from_ids(
+            self.state_ids, self.amici_model.getStateIds()
+        )
+        self.setSensitivities()
+
+    def load_amici_model_and_solver(self):
+        """Loads the AMICI model from the previously generated model."""
+        amici_model_module = amici.import_model_module(
+            self.amici_model_name, self.amici_model_dir
+        )
+        self.amici_model = amici_model_module.getModel()
+        self.amici_solver = self.amici_model.getSolver()
+
+        self.amici_model.setTimepoints(self.timepoints)
+        self.amici_solver.setAbsoluteTolerance(1e-10)
+
+    def setSensitivities(self):
+        """Tell the underlying amici solver to calculate sensitivities based on the attribute `self.param_ids`"""
+        if self.param_ids == self.amici_model.getParameterIds():
+            self.amici_model.requireSensitivitiesForAllParameters()
+        else:
+            self.amici_model.setParameterList(self.param_indices)
+
+        self.amici_solver.setSensitivityMethod(amici.SensitivityMethod.forward)
+        self.amici_solver.setSensitivityOrder(amici.SensitivityOrder.first)
+
+    def forward(self, params):
+        for i, param in enumerate(params):
+            self.amici_model.setParameterById(self.param_ids[i], param)
+        rdata = amici.runAmiciSimulation(self.amici_model, self.amici_solver)
+        return rdata.x[:, self.state_indices].reshape(self.data_dim)
+
+    def jacobian(self, params):
+        for i, param in enumerate(params):
+            self.amici_model.setParameterById(self.param_ids[i], param)
+        rdata = amici.runAmiciSimulation(self.amici_model, self.amici_solver)
+        return (
+            rdata.sx[:, :, self.state_indices]
+            .transpose(1, 0, 2)
+            .reshape(self.param_dim, self.data_dim)
+            .T
+        )
+
+    def forward_and_jacobian(
+        self, params: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        for i, param in enumerate(params):
+            self.amici_model.setParameterById(self.param_ids[i], param)
+        rdata = amici.runAmiciSimulation(self.amici_model, self.amici_solver)
+        return (
+            rdata.x[:, self.state_indices].reshape(self.data_dim),
+            rdata.sx[:, :, self.state_indices]
+            .transpose(1, 0, 2)
+            .reshape(self.param_dim, self.data_dim)
+            .T,
+        )
+
+    # Allow the model to be pickled
+    def __getstate__(self):
+        # Create a copy of the object's state
+        state = self.__dict__.copy()
+
+        # Save the amici solver settings to
+        _fd, _file = tempfile.mkstemp()
+
+        try:
+            # write amici solver settings to file
+            try:
+                amici.writeSolverSettingsToHDF5(self.amici_solver, _file)
+            except AttributeError as e:
+                e.args += (
+                    "Pickling the SBMLModel requires an AMICI "
+                    "installation with HDF5 support.",
+                )
+                raise
+            # read in byte stream
+            with open(_fd, "rb", closefd=False) as f:
+                state["amici_solver_settings"] = f.read()
+        finally:
+            # close file descriptor and remove temporary file
+            os.close(_fd)
+            os.remove(_file)
+
+        state["amici_model_settings"] = amici.get_model_settings(
+            self.amici_model
+        )
+
+        # Remove the unpicklable entries.
+        del state["amici_model"]
+        del state["amici_solver"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+        # Restore amici model and solver
+        self.load_amici_model_and_solver()
+        self.setSensitivities()
+
+        _fd, _file = tempfile.mkstemp()
+        try:
+            # write solver settings to temporary file
+            with open(_fd, "wb", closefd=False) as f:
+                f.write(state["amici_solver_settings"])
+            # read in solver settings
+            try:
+                amici.readSolverSettingsFromHDF5(_file, self.amici_solver)
+            except AttributeError as err:
+                if not err.args:
+                    err.args = ("",)
+                err.args += (
+                    "Unpickling an AmiciObjective requires an AMICI "
+                    "installation with HDF5 support.",
+                )
+                raise
+        finally:
+            # close file descriptor and remove temporary file
+            os.close(_fd)
+            os.remove(_file)
+
+        amici.set_model_settings(
+            self.amici_model,
+            state["amici_model_settings"],
+        )
diff --git a/eulerpi/core/plotting.py b/eulerpi/core/plotting.py
index cc2c01a1..14bef0a4 100644
--- a/eulerpi/core/plotting.py
+++ b/eulerpi/core/plotting.py
@@ -14,14 +14,14 @@
 from matplotlib import axes
 
 from eulerpi.core.kde import calc_kernel_width, eval_kde_gauss
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.core.result_manager import ResultManager
 
 # general plotting function for joint runs
 
 
 def sample_violin_plot(
-    model: Model,
+    model: BaseModel,
     reference_sample: Union[str, os.PathLike, np.ndarray, None] = None,
     run_name: str = "default_run",
     what_to_plot: str = "param",
@@ -33,7 +33,7 @@ def sample_violin_plot(
        Can be used for parameters and simulation results and compares reference (or true underlying) and inferred values when possible.
 
     Args:
-        model(Model): The model describing the mapping from parameters to data.
+        model(BaseModel): The model describing the mapping from parameters to data.
         reference_sample(Union[str, os.PathLike, np.ndarray]): Depending on what_to_plot, this is either the data sample used for the inference, or a reference sample of "true" parameter samples. If a string is given, it is assumed to be a path to a file containing the respective sample. (Default value = None)
         run_name(str): The name of the inference run. (Default value = "default_run")
         what_to_plot(str): Choose between "param" and "data" to respectively visualize either the model parameters or output. (Default value = "param")
diff --git a/eulerpi/core/result_manager.py b/eulerpi/core/result_manager.py
index 53290999..341c0d0a 100644
--- a/eulerpi/core/result_manager.py
+++ b/eulerpi/core/result_manager.py
@@ -11,7 +11,7 @@
 
 from eulerpi import logger
 from eulerpi.core.inference_types import InferenceType
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 
 
 class ResultManager:
@@ -180,7 +180,7 @@ def get_application_path(self) -> str:
 
     def save_run(
         self,
-        model: Model,
+        model: BaseModel,
         slice: np.ndarray,
         run,
         sampler_results: np.ndarray,
@@ -191,7 +191,7 @@ def save_run(
         as seperate files in the folders 'Params' and'SimResults' and 'DensityEvals'.
 
         Args:
-            model(Model): The model for which the results will be saved
+            model(BaseModel): The model for which the results will be saved
             slice(np.ndarray): The slice for which the results will be saved
             run(int): The run for which the results will be saved
             sampler_results(np.ndarray): The results of the sampler, expects an np.array with shape (num_walkers * num_steps, sampling_dim + data_dim + 1)
@@ -263,7 +263,7 @@ def save_overall(
     def save_inference_information(
         self,
         slice: np.ndarray,
-        model: Model,
+        model: BaseModel,
         inference_type: InferenceType,
         num_processes: int,
         **kwargs,
@@ -272,7 +272,7 @@ def save_inference_information(
 
         Args:
             slice(np.ndarray): The slice for which the results will be saved.
-            model(Model): The model for which the results will be saved.
+            model(BaseModel): The model for which the results will be saved.
             inference_type(InferenceType): The type of inference that was performed.
             num_processes(int): The number of processes that were used for the inference.
             **kwargs: Additional information about the inference run.
diff --git a/eulerpi/core/sampling.py b/eulerpi/core/sampling.py
index 9d6e894a..5864374a 100644
--- a/eulerpi/core/sampling.py
+++ b/eulerpi/core/sampling.py
@@ -23,13 +23,13 @@
 from eulerpi.core.data_transformations import DataTransformation
 from eulerpi.core.inference_types import InferenceType
 from eulerpi.core.kde import calc_kernel_width
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.core.result_manager import ResultManager
 from eulerpi.core.transformations import eval_log_transformed_density
 
 
 def run_emcee_once(
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     data_stdevs: np.ndarray,
@@ -42,7 +42,7 @@ def run_emcee_once(
     """Run the emcee particle swarm sampler once.
 
     Args:
-        model (Model): The model which will be sampled
+        model (BaseModel): 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
@@ -112,7 +112,7 @@ def run_emcee_once(
 
 
 def run_emcee_sampling(
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     slice: np.ndarray,
@@ -128,7 +128,7 @@ def run_emcee_sampling(
        Inital values are not stored in the chain and each file contains <num_steps> blocks of size num_walkers.
 
     Args:
-        model (Model): The model which will be sampled
+        model (BaseModel): 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
@@ -229,7 +229,7 @@ def run_emcee_sampling(
 
 
 def calc_walker_acceptance(
-    model: Model,
+    model: BaseModel,
     slice: np.ndarray,
     num_walkers: int,
     num_burn_in_samples: int,
@@ -239,7 +239,7 @@ def calc_walker_acceptance(
     This is especially important to find "zombie" walkers, that are never moving.
 
     Args:
-        model (Model): The model for which the acceptance ratio should be calculated
+        model (BaseModel): The model for which the acceptance ratio should be calculated
         slice (np.ndarray): slice for which the acceptance ratio should be calculated
         num_walkers (int): number of walkers in the emcee chain
         num_burn_in_samples(int): Number of samples that will be deleted (burned) per chain (i.e. walker). Only for mcmc inference.
@@ -281,7 +281,7 @@ def calc_walker_acceptance(
 
 
 def inference_mcmc(
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     result_manager: ResultManager,
@@ -302,7 +302,7 @@ def inference_mcmc(
     """This function runs a MCMC sampling for the given model and data.
 
     Args:
-        model (Model): The model describing the mapping from parameters to data.
+        model (BaseModel): 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.
diff --git a/eulerpi/core/sparsegrid.py b/eulerpi/core/sparsegrid.py
index 56b8711e..dd2a80f2 100644
--- a/eulerpi/core/sparsegrid.py
+++ b/eulerpi/core/sparsegrid.py
@@ -18,7 +18,7 @@
 from eulerpi.core.data_transformations import DataTransformation
 from eulerpi.core.inference_types import InferenceType
 from eulerpi.core.kde import calc_kernel_width
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.core.result_manager import ResultManager
 from eulerpi.core.transformations import evaluate_density
 
@@ -372,7 +372,7 @@ def pass_contributions2higher_levels(self):
 def evaluate_on_sparse_grid(
     args: typing.Tuple[
         np.ndarray,
-        Model,
+        BaseModel,
         np.ndarray,
         DataTransformation,
         np.ndarray,
@@ -383,7 +383,7 @@ def evaluate_on_sparse_grid(
 
     Args:
         params (np.ndarray): The parameters to be evaluated.
-        model(Model): The model used for the inference.
+        model(BaseModel): 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!)
@@ -400,7 +400,7 @@ def evaluate_on_sparse_grid(
 
 
 def inference_sparse_grid(
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     result_manager: ResultManager,
@@ -416,7 +416,7 @@ 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.
+        model(BaseModel): 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.
diff --git a/eulerpi/core/transformations.py b/eulerpi/core/transformations.py
index 26558176..f1f7d743 100644
--- a/eulerpi/core/transformations.py
+++ b/eulerpi/core/transformations.py
@@ -10,12 +10,12 @@
 from eulerpi import logger
 from eulerpi.core.data_transformations import DataTransformation
 from eulerpi.core.kde import eval_kde_gauss
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 
 
 def evaluate_density(
     param: np.ndarray,
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     data_stdevs: np.ndarray,
@@ -29,7 +29,7 @@ def evaluate_density(
 
     Args:
         param (np.ndarray): parameter for which the transformed density shall be evaluated
-        model (Model): model to be evaluated
+        model (BaseModel): 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
@@ -139,7 +139,7 @@ def evaluate_density(
 
 def eval_log_transformed_density(
     param: np.ndarray,
-    model: Model,
+    model: BaseModel,
     data: np.ndarray,
     data_transformation: DataTransformation,
     data_stdevs: np.ndarray,
@@ -157,7 +157,7 @@ def eval_log_transformed_density(
 
     Args:
         param (np.ndarray): parameter for which the transformed density shall be evaluated
-        model (Model): model to be evaluated
+        model (BaseModel): 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
diff --git a/eulerpi/examples/corona/corona.py b/eulerpi/examples/corona/corona.py
index edd9ad79..84d53a0c 100644
--- a/eulerpi/examples/corona/corona.py
+++ b/eulerpi/examples/corona/corona.py
@@ -5,7 +5,7 @@
 import numpy as np
 
 from eulerpi import logger
-from eulerpi.core.model import ArtificialModelInterface, JaxModel
+from eulerpi.core.models import ArtificialModelInterface, JaxModel
 
 
 class Corona(JaxModel):
@@ -30,9 +30,9 @@ class Corona(JaxModel):
     .. note::
 
         * ODE Solver: To solve the ODE problem the jax based ode solver library :code:`diffrax` is used: https://github.com/patrick-kidger/diffrax.
-        * Automatic Differentiation: The derivatives are calculated automatically with jax by deriving from the class :py:class:`~eulerpi.core.model.JaxModel`,
-          which automatically calculates sets :py:meth:`~eulerpi.core.model.Model.jacobian`.
-        * JIT compilation: Inheriting from :py:class:`~eulerpi.core.model.JaxModel` also enables jit compilation / optimization for the forward and jacobian method.
+        * Automatic Differentiation: The derivatives are calculated automatically with jax by deriving from the class :py:class:`~eulerpi.core.models.JaxModel`,
+          which automatically calculates sets :py:meth:`~eulerpi.core.models.BaseModel.jacobian`.
+        * JIT compilation: Inheriting from :py:class:`~eulerpi.core.models.JaxModel` also enables jit compilation / optimization for the forward and jacobian method.
           This usually results in a significant execution speedup. It also allows to run your model on the gpu.
     """
 
diff --git a/eulerpi/examples/cpp/cpp_plant.py b/eulerpi/examples/cpp/cpp_plant.py
index 885b4e13..b5a6b81c 100644
--- a/eulerpi/examples/cpp/cpp_plant.py
+++ b/eulerpi/examples/cpp/cpp_plant.py
@@ -1,10 +1,10 @@
 import numpy as np
 
-from eulerpi.core.model import ArtificialModelInterface, Model
+from eulerpi.core.models import ArtificialModelInterface, BaseModel
 from eulerpi.examples.cpp import cpp_model
 
 
-class CppPlant(Model, ArtificialModelInterface):
+class CppPlant(BaseModel, ArtificialModelInterface):
     """A plant model which uses a c++ library with eigen3 to evaluate the forward pass and the gradient
     Param0: Water [0,1]
     Param1: Sun   [0,1]
diff --git a/eulerpi/examples/cpp/python_reference_plants.py b/eulerpi/examples/cpp/python_reference_plants.py
index 01c9f3c9..4a0bf659 100644
--- a/eulerpi/examples/cpp/python_reference_plants.py
+++ b/eulerpi/examples/cpp/python_reference_plants.py
@@ -2,7 +2,7 @@
 import numpy as np
 from jax import jacrev, jit
 
-from eulerpi.core.model import ArtificialModelInterface, JaxModel, Model
+from eulerpi.core.models import ArtificialModelInterface, BaseModel, JaxModel
 
 
 class JaxPlant(JaxModel, ArtificialModelInterface):
@@ -68,7 +68,7 @@ def bw(param):
     return fwJac(param)
 
 
-class ExternalPlant(Model, ArtificialModelInterface):
+class ExternalPlant(BaseModel, ArtificialModelInterface):
     """A plant model which uses functions defined outside the class to evaluate the forward pass and the jacobian
     Param0: Water [0,1]
     Param1: Sun   [0,1]
diff --git a/eulerpi/examples/heat/heat.py b/eulerpi/examples/heat/heat.py
index 89693b78..5ae6359c 100644
--- a/eulerpi/examples/heat/heat.py
+++ b/eulerpi/examples/heat/heat.py
@@ -4,7 +4,7 @@
 import jax.numpy as jnp
 import numpy as np
 
-from eulerpi.core.model import ArtificialModelInterface, JaxModel
+from eulerpi.core.models import ArtificialModelInterface, JaxModel
 
 
 def heat_rhs(t: float, u: jnp.ndarray, args: tuple | list) -> jnp.ndarray:
diff --git a/eulerpi/examples/sbml/sbml_caffeine_model.py b/eulerpi/examples/sbml/sbml_caffeine_model.py
index e3be8bc1..53dffc19 100644
--- a/eulerpi/examples/sbml/sbml_caffeine_model.py
+++ b/eulerpi/examples/sbml/sbml_caffeine_model.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from eulerpi.core.model import ArtificialModelInterface, SBMLModel
+from eulerpi.core.models import ArtificialModelInterface, SBMLModel
 
 
 class CaffeineSBMLModel(SBMLModel, ArtificialModelInterface):
diff --git a/eulerpi/examples/sbml/sbml_menten_model.py b/eulerpi/examples/sbml/sbml_menten_model.py
index 758d882f..da477626 100644
--- a/eulerpi/examples/sbml/sbml_menten_model.py
+++ b/eulerpi/examples/sbml/sbml_menten_model.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from eulerpi.core.model import ArtificialModelInterface, SBMLModel
+from eulerpi.core.models import ArtificialModelInterface, SBMLModel
 
 
 class MentenSBMLModel(SBMLModel, ArtificialModelInterface):
diff --git a/eulerpi/examples/simple_models.py b/eulerpi/examples/simple_models.py
index 6ffffcbc..b44bab01 100644
--- a/eulerpi/examples/simple_models.py
+++ b/eulerpi/examples/simple_models.py
@@ -3,7 +3,7 @@
 import jax.numpy as jnp
 import numpy as np
 
-from eulerpi.core.model import ArtificialModelInterface, JaxModel
+from eulerpi.core.models import ArtificialModelInterface, JaxModel
 
 
 class Linear(JaxModel, ArtificialModelInterface):
diff --git a/eulerpi/examples/temperature/temperature.py b/eulerpi/examples/temperature/temperature.py
index e2068ba9..cb9e1372 100644
--- a/eulerpi/examples/temperature/temperature.py
+++ b/eulerpi/examples/temperature/temperature.py
@@ -4,13 +4,13 @@
 import jax.numpy as jnp
 import numpy as np
 
-from eulerpi.core.model import ArtificialModelInterface, Model
+from eulerpi.core.models import ArtificialModelInterface, BaseModel
 
 # from functools import partial
 # from jax import jit
 
 
-class Temperature(Model):
+class Temperature(BaseModel):
     """The model describes the temperature :math:`y` in degree celsius at a given latitude :math:`q` in degree.
 
 


Test patch:
diff --git a/tests/profiling.py b/tests/profiling.py
index e075e495..e1551cbd 100644
--- a/tests/profiling.py
+++ b/tests/profiling.py
@@ -9,13 +9,13 @@
 import numpy as np
 
 from eulerpi.core.inference import InferenceType, inference
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.examples.corona import CoronaArtificial
 
 
 def profiling_with_slices(inference_type):
     """ """
-    model: Model = CoronaArtificial()
+    model: BaseModel = CoronaArtificial()
 
     # generate artificial data
     num_data_points = 1000
diff --git a/tests/test_dense_grid.py b/tests/test_dense_grid.py
index 1424fbbe..81a1a567 100644
--- a/tests/test_dense_grid.py
+++ b/tests/test_dense_grid.py
@@ -8,7 +8,7 @@
 from eulerpi.core.dense_grid_types import DenseGridType
 from eulerpi.core.inference import inference
 from eulerpi.core.inference_types import InferenceType
-from eulerpi.core.model import Model
+from eulerpi.core.models import ArtificialModelInterface, BaseModel
 from eulerpi.examples.corona import CoronaArtificial
 
 
@@ -20,10 +20,10 @@
 )
 def test_dense_grid(dense_grid_type):
     """ """
-    model: Model = CoronaArtificial()
+    model: BaseModel = CoronaArtificial()
 
     # generate artificial data
-    if model.is_artificial():
+    if isinstance(model, ArtificialModelInterface):
         num_data_points = 100
         params = model.generate_artificial_params(num_data_points)
         data = model.generate_artificial_data(params)
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 138bf5e1..6c89b016 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -7,8 +7,9 @@
 import pytest
 
 from eulerpi.core.inference import InferenceType, inference
-from eulerpi.core.model import Model, is_amici_available
 from eulerpi.core.model_check import basic_model_check
+from eulerpi.core.models import ArtificialModelInterface, BaseModel
+from eulerpi.core.models.sbml_model import is_amici_available
 
 cpp_plant_example = pytest.param(
     ("eulerpi.examples.cpp.cpp_plant", "CppPlant"),
@@ -88,7 +89,7 @@ def test_model_requirements(example):
     # Import class dynamically to avoid error on imports at the top which cant be tracked back to a specific test
     module = importlib.import_module(module_location)
     ModelClass = getattr(module, className)
-    model: Model = ModelClass()
+    model: BaseModel = ModelClass()
 
     # Use the basic model check to assert all model attribute dimension requirements as well as forward simulations and the corresponding jacobian.
     basic_model_check(model)
@@ -115,7 +116,7 @@ def test_examples(example, inference_type):
     # Import class dynamically to avoid error on imports at the top which cant be tracked back to a specific test
     module = importlib.import_module(module_location)
     ModelClass = getattr(module, className)
-    model: Model = ModelClass()
+    model: BaseModel = ModelClass()
 
     # Define kwargs for inference
     kwargs = {}
@@ -129,7 +130,7 @@ def test_examples(example, inference_type):
         kwargs["num_levels"] = 3
 
     # generate artificial data if necessary
-    if model.is_artificial():
+    if isinstance(model, ArtificialModelInterface):
         num_data_points = 100
         params = model.generate_artificial_params(num_data_points)
         data = model.generate_artificial_data(params)
diff --git a/tests/test_fixed_params.py b/tests/test_fixed_params.py
index b516bf84..42609394 100644
--- a/tests/test_fixed_params.py
+++ b/tests/test_fixed_params.py
@@ -7,7 +7,7 @@
 import pytest
 
 from eulerpi.core.inference import inference
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.examples.temperature import TemperatureWithFixedParams
 
 
@@ -35,7 +35,7 @@ def temp_name(low_T, high_T):
 )
 def test_fixed_params(temperatures):
     """ """
-    model: Model = TemperatureWithFixedParams(*temperatures)
+    model: BaseModel = TemperatureWithFixedParams(*temperatures)
     run_name = temp_name(*temperatures)
     data_resource = importlib.resources.files(
         "eulerpi.examples.temperature"
diff --git a/tests/test_logging.py b/tests/test_logging.py
index 6e34c9ff..7feff556 100644
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -2,11 +2,11 @@
 
 from eulerpi.core.data_transformations import DataIdentity
 from eulerpi.core.kde import calc_kernel_width
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from eulerpi.core.transformations import evaluate_density
 
 
-class CrashModel(Model):
+class CrashModel(BaseModel):
     data_dim = 1
     param_dim = 1
 
diff --git a/tests/test_pickling.py b/tests/test_pickling.py
index 66a763e8..5b6bb9ee 100644
--- a/tests/test_pickling.py
+++ b/tests/test_pickling.py
@@ -6,7 +6,7 @@
 import numpy as np
 import pytest
 
-from eulerpi.core.model import Model
+from eulerpi.core.models import BaseModel
 from tests.test_examples import Examples, get_example_name
 
 
@@ -21,7 +21,7 @@ def test_pickling_model(example):
     # Import class dynamically to avoid error on imports at the top which cant be tracked back to a specific test
     module = importlib.import_module(module_location)
     ModelClass = getattr(module, className)
-    model: Model = ModelClass()
+    model: BaseModel = ModelClass()
 
     # Test pickling
     dumped = pickle.dumps(model)
diff --git a/tests/test_simple_models.py b/tests/test_simple_models.py
index 263c9725..a24ccee5 100644
--- a/tests/test_simple_models.py
+++ b/tests/test_simple_models.py
@@ -4,7 +4,6 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from jax import vmap
 from matplotlib import cm
 
 from eulerpi.core.data_transformations import DataIdentity
@@ -105,7 +104,7 @@ def test_transformationExponential():
     num_data_points = 10000
     trueParam = np.random.rand(num_data_points, 2) + 1
 
-    data = vmap(model.forward, in_axes=0)(trueParam)
+    data = model.forward_vectorized(trueParam)
     data_transformation = DataIdentity()
 
     # define standard deviations according to silverman
diff --git a/tests/test_slices.py b/tests/test_slices.py
index 3a045841..e657aea5 100644
--- a/tests/test_slices.py
+++ b/tests/test_slices.py
@@ -6,7 +6,7 @@
 import pytest
 
 from eulerpi.core.inference import InferenceType, inference
-from eulerpi.core.model import Model
+from eulerpi.core.models import ArtificialModelInterface, BaseModel
 from eulerpi.examples.corona import CoronaArtificial
 
 # Parametrize the test to run for each inference type
@@ -19,10 +19,10 @@
 )
 def test_slices(inference_type):
     """ """
-    model: Model = CoronaArtificial()
+    model: BaseModel = CoronaArtificial()
 
     # generate artificial data
-    if model.is_artificial():
+    if isinstance(model, ArtificialModelInterface):
         num_data_points = 100
         params = model.generate_artificial_params(num_data_points)
         data = model.generate_artificial_data(params)
diff --git a/tests/test_transformations.py b/tests/test_transformations.py
index 40319b9b..96dbc2a6 100644
--- a/tests/test_transformations.py
+++ b/tests/test_transformations.py
@@ -4,7 +4,7 @@
 
 from eulerpi.core.data_transformations import DataIdentity
 from eulerpi.core.kde import calc_kernel_width, eval_kde_gauss
-from eulerpi.core.model import ArtificialModelInterface, JaxModel
+from eulerpi.core.models import ArtificialModelInterface, JaxModel
 from eulerpi.core.transformations import calc_gram_determinant