modelbase2 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

modelbase2/experimental/symbolic.py

@@ -0,0 +1,286 @@
+# ruff: noqa: D100, D101, D102, D103, D104, D105, D106, D107, D200, D203, D400, D401
+
+import ast
+import inspect
+import textwrap
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, cast
+
+import sympy
+
+from modelbase2.model import Model
+
+__all__ = ["Context", "SymbolicModel", "model_fn_to_sympy", "to_symbolic_model"]
+
+
+@dataclass
+class Context:
+    symbols: dict[str, sympy.Symbol | sympy.Expr]
+    caller: Callable
+
+
+@dataclass
+class SymbolicModel:
+    variables: dict[str, sympy.Symbol]
+    parameters: dict[str, sympy.Symbol]
+    eqs: list[sympy.Expr]
+
+
+def to_symbolic_model(model: Model) -> SymbolicModel:
+    cache = model._create_cache()  # noqa: SLF001
+
+    variables = dict(
+        zip(model.variables, sympy.symbols(list(model.variables)), strict=True)
+    )
+    parameters = dict(
+        zip(model.parameters, sympy.symbols(list(model.parameters)), strict=True)
+    )
+    symbols = variables | parameters
+
+    for k, v in model.derived.items():
+        symbols[k] = model_fn_to_sympy(v.fn, [symbols[i] for i in v.args])
+
+    rxns = {
+        k: model_fn_to_sympy(v.fn, [symbols[i] for i in v.args])
+        for k, v in model.reactions.items()
+    }
+
+    eqs: dict[str, sympy.Expr] = {}
+    for cpd, stoich in cache.stoich_by_cpds.items():
+        for rxn, stoich_value in stoich.items():
+            eqs[cpd] = (
+                eqs.get(cpd, sympy.Float(0.0)) + sympy.Float(stoich_value) * rxns[rxn]  # type: ignore
+            )
+
+    for cpd, dstoich in cache.dyn_stoich_by_cpds.items():
+        for rxn, der in dstoich.items():
+            eqs[cpd] = eqs.get(cpd, sympy.Float(0.0)) + model_fn_to_sympy(
+                der.fn,
+                [symbols[i] for i in der.args] * rxns[rxn],  # type: ignore
+            )  # type: ignore
+
+    return SymbolicModel(
+        variables=variables,
+        parameters=parameters,
+        eqs=[eqs[i] for i in cache.var_names],
+    )
+
+
+def model_fn_to_sympy(
+    fn: Callable, model_args: list[sympy.Symbol | sympy.Expr] | None = None
+) -> sympy.Expr:
+    source = textwrap.dedent(inspect.getsource(fn))
+
+    if not isinstance(fn_def := ast.parse(source).body[0], ast.FunctionDef):
+        msg = "Expected a function definition"
+        raise TypeError(msg)
+
+    fn_args = [str(arg.arg) for arg in fn_def.args.args]
+
+    sympy_expr = _handle_fn_body(
+        fn_def.body,
+        ctx=Context(
+            symbols={name: sympy.Symbol(name) for name in fn_args},
+            caller=fn,
+        ),
+    )
+
+    if model_args is not None:
+        sympy_expr = sympy_expr.subs(dict(zip(fn_args, model_args, strict=True)))
+
+    return cast(sympy.Expr, sympy_expr)
+
+
+def _handle_fn_body(body: list[ast.stmt], ctx: Context) -> sympy.Expr:
+    pieces = []
+    remaining_body = list(body)
+
+    while remaining_body:
+        node = remaining_body.pop(0)
+
+        if isinstance(node, ast.If):
+            condition = _handle_expr(node.test, ctx)
+            if_expr = _handle_fn_body(node.body, ctx)
+            pieces.append((if_expr, condition))
+
+            # If there's an else clause
+            if node.orelse:
+                # Check if it's an elif (an If node in orelse)
+                if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If):
+                    # Push the elif back to the beginning of remaining_body to process next
+                    remaining_body.insert(0, node.orelse[0])
+                else:
+                    # It's a regular else
+                    else_expr = _handle_fn_body(node.orelse, ctx)  # FIXME: copy here
+                    pieces.append((else_expr, True))
+                    break  # We're done with this chain
+
+            elif not remaining_body and any(
+                isinstance(n, ast.Return) for n in body[body.index(node) + 1 :]
+            ):
+                else_expr = _handle_fn_body(
+                    body[body.index(node) + 1 :], ctx
+                )  # FIXME: copy here
+                pieces.append((else_expr, True))
+
+        elif isinstance(node, ast.Return):
+            if (value := node.value) is None:
+                msg = "Return value cannot be None"
+                raise ValueError(msg)
+
+            expr = _handle_expr(value, ctx)
+            if not pieces:
+                return expr
+            pieces.append((expr, True))
+            break
+
+        elif isinstance(node, ast.Assign):
+            # Handle tuple assignments like c, d = a, b
+            if isinstance(node.targets[0], ast.Tuple):
+                # Handle tuple unpacking
+                target_elements = node.targets[0].elts
+
+                if isinstance(node.value, ast.Tuple):
+                    # Direct unpacking like c, d = a, b
+                    value_elements = node.value.elts
+                    for target, value_expr in zip(
+                        target_elements, value_elements, strict=True
+                    ):
+                        if isinstance(target, ast.Name):
+                            ctx.symbols[target.id] = _handle_expr(value_expr, ctx)
+                else:
+                    # Handle potential iterable unpacking
+                    value = _handle_expr(node.value, ctx)
+            else:
+                # Regular single assignment
+                if not isinstance(target := node.targets[0], ast.Name):
+                    msg = "Only single variable assignments are supported"
+                    raise TypeError(msg)
+                target_name = target.id
+                value = _handle_expr(node.value, ctx)
+                ctx.symbols[target_name] = value
+
+    # If we have pieces to combine into a Piecewise
+    if pieces:
+        return sympy.Piecewise(*pieces)
+
+    # If no return was found but we have assignments, return the last assigned variable
+    for node in reversed(body):
+        if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name):
+            target_name = node.targets[0].id
+            return ctx.symbols[target_name]
+
+    msg = "No return value found in function body"
+    raise ValueError(msg)
+
+
+def _handle_unaryop(node: ast.UnaryOp, ctx: Context) -> sympy.Expr:
+    left = _handle_expr(node.operand, ctx)
+    left = cast(Any, left)  # stupid sympy types don't allow ops on symbols
+
+    if isinstance(node.op, ast.UAdd):
+        return +left
+    if isinstance(node.op, ast.USub):
+        return -left
+
+    msg = f"Operation {type(node.op).__name__} not implemented"
+    raise NotImplementedError(msg)
+
+
+def _handle_binop(node: ast.BinOp, ctx: Context) -> sympy.Expr:
+    left = _handle_expr(node.left, ctx)
+    left = cast(Any, left)  # stupid sympy types don't allow ops on symbols
+
+    right = _handle_expr(node.right, ctx)
+    right = cast(Any, right)  # stupid sympy types don't allow ops on symbols
+
+    if isinstance(node.op, ast.Add):
+        return left + right
+    if isinstance(node.op, ast.Sub):
+        return left - right
+    if isinstance(node.op, ast.Mult):
+        return left * right
+    if isinstance(node.op, ast.Div):
+        return left / right
+    if isinstance(node.op, ast.Pow):
+        return left**right
+    if isinstance(node.op, ast.Mod):
+        return left % right
+    if isinstance(node.op, ast.FloorDiv):
+        return left // right
+
+    msg = f"Operation {type(node.op).__name__} not implemented"
+    raise NotImplementedError(msg)
+
+
+def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr:
+    if not isinstance(callee := node.func, ast.Name):
+        msg = "Only function calls with names are supported"
+        raise TypeError(msg)
+
+    fn_name = str(callee.id)
+    parent_module = inspect.getmodule(ctx.caller)
+    fns = dict(inspect.getmembers(parent_module, predicate=callable))
+
+    return model_fn_to_sympy(
+        fns[fn_name],
+        model_args=[_handle_expr(i, ctx) for i in node.args],
+    )
+
+
+def _handle_name(node: ast.Name, ctx: Context) -> sympy.Symbol | sympy.Expr:
+    return ctx.symbols[node.id]
+
+
+def _handle_expr(node: ast.expr, ctx: Context) -> sympy.Expr:
+    if isinstance(node, ast.UnaryOp):
+        return _handle_unaryop(node, ctx)
+    if isinstance(node, ast.BinOp):
+        return _handle_binop(node, ctx)
+    if isinstance(node, ast.Name):
+        return _handle_name(node, ctx)
+    if isinstance(node, ast.Constant):
+        return node.value
+    if isinstance(node, ast.Compare):
+        # Handle chained comparisons like 1 < a < 2
+        left = cast(Any, _handle_expr(node.left, ctx))
+        comparisons = []
+
+        # Build all individual comparisons from the chain
+        prev_value = left
+        for op, comparator in zip(node.ops, node.comparators, strict=True):
+            right = cast(Any, _handle_expr(comparator, ctx))
+
+            if isinstance(op, ast.Gt):
+                comparisons.append(prev_value > right)
+            elif isinstance(op, ast.GtE):
+                comparisons.append(prev_value >= right)
+            elif isinstance(op, ast.Lt):
+                comparisons.append(prev_value < right)
+            elif isinstance(op, ast.LtE):
+                comparisons.append(prev_value <= right)
+            elif isinstance(op, ast.Eq):
+                comparisons.append(prev_value == right)
+            elif isinstance(op, ast.NotEq):
+                comparisons.append(prev_value != right)
+
+            prev_value = right
+
+        # Combine all comparisons with logical AND
+        result = comparisons[0]
+        for comp in comparisons[1:]:
+            result = sympy.And(result, comp)
+        return cast(sympy.Expr, result)
+    if isinstance(node, ast.Call):
+        return _handle_call(node, ctx)
+
+    # Handle conditional expressions (ternary operators)
+    if isinstance(node, ast.IfExp):
+        condition = _handle_expr(node.test, ctx)
+        if_true = _handle_expr(node.body, ctx)
+        if_false = _handle_expr(node.orelse, ctx)
+        return sympy.Piecewise((if_true, condition), (if_false, True))
+
+    msg = f"Expression type {type(node).__name__} not implemented"
+    raise NotImplementedError(msg)

modelbase2/fit.py

@@ -50,7 +50,7 @@ class SteadyStateResidualFn(Protocol):
         data: pd.Series,
         model: Model,
         y0: dict[str, float],
-        integrator: type[IntegratorProtocol],
+        integrator: Callable[[Callable, ArrayLike], IntegratorProtocol],
     ) -> float:
         """Calculate residual error between model steady state and experimental data."""
         ...
@@ -67,7 +67,7 @@ class TimeSeriesResidualFn(Protocol):
         data: pd.DataFrame,
         model: Model,
         y0: dict[str, float],
-        integrator: type[IntegratorProtocol],
+        integrator: Callable[[Callable, ArrayLike], IntegratorProtocol],
     ) -> float:
         """Calculate residual error between model time course and experimental data."""
         ...
@@ -101,7 +101,7 @@ def _steady_state_residual(
     data: pd.Series,
     model: Model,
     y0: dict[str, float] | None,
-    integrator: type[IntegratorProtocol],
+    integrator: Callable[[Callable, ArrayLike], IntegratorProtocol],
 ) -> float:
     """Calculate residual error between model steady state and experimental data.
 
@@ -148,7 +148,7 @@ def _time_course_residual(
     data: pd.DataFrame,
     model: Model,
     y0: dict[str, float],
-    integrator: type[IntegratorProtocol],
+    integrator: Callable[[Callable, ArrayLike], IntegratorProtocol],
 ) -> float:
     """Calculate residual error between model time course and experimental data.
 
@@ -187,7 +187,7 @@ def steady_state(
     y0: dict[str, float] | None = None,
     minimize_fn: MinimizeFn = _default_minimize_fn,
     residual_fn: SteadyStateResidualFn = _steady_state_residual,
-    integrator: type[IntegratorProtocol] = DefaultIntegrator,
+    integrator: Callable[[Callable, ArrayLike], IntegratorProtocol] = DefaultIntegrator,
 ) -> dict[str, float]:
     """Fit model parameters to steady-state experimental data.
 
@@ -241,7 +241,7 @@ def time_course(
     y0: dict[str, float] | None = None,
     minimize_fn: MinimizeFn = _default_minimize_fn,
     residual_fn: TimeSeriesResidualFn = _time_course_residual,
-    integrator: type[IntegratorProtocol] = DefaultIntegrator,
+    integrator: Callable[[Callable, ArrayLike], IntegratorProtocol] = DefaultIntegrator,
 ) -> dict[str, float]:
     """Fit model parameters to time course of experimental data.
 

modelbase2/model.py

@@ -10,7 +10,6 @@ from __future__ import annotations
 import copy
 import inspect
 import itertools as it
-import math
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Self, cast
 

modelbase2/nnarchitectures.py

@@ -0,0 +1,128 @@
+"""Neural network architectures.
+
+This module provides implementations of neural network architectures used for mechanistic learning.
+
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+__all__ = ["DefaultDevice", "LSTMnn", "MLP"]
+
+DefaultDevice = torch.device("cpu")
+
+
+class MLP(nn.Module):
+    """Multilayer Perceptron (MLP) for surrogate modeling and neural posterior estimation.
+
+    Attributes:
+        net: Sequential neural network model.
+
+    Methods:
+        forward: Forward pass through the neural network.
+
+    """
+
+    def __init__(
+        self,
+        n_inputs: int,
+        layers: list[int],
+        activation: Callable | None = nn.ReLU(),
+        output_activation: Callable | None = None,
+    ) -> None:
+        """Initializes the MLP with the given number of inputs and list of (hidden) layers.
+
+        Args:
+            n_inputs (int): The number of input features.
+            n_outputs list(int): A list containing the number of neurons in hidden and output layer.
+            activation Callable | None (default nn.ReLU()): The activation function to be applied after each hidden layer
+            activation Callable | None (default None): The activation function to be applied after the final (output) layer
+
+        For instance, MLP(10, layers = [50, 50, 10]) initializes a neural network with the following architecture:
+        - Linear layer with `n_inputs` inputs and 50 outputs
+        - ReLU activation
+        - Linear layer with 50 inputs and 50 outputs
+        - ReLU activation
+        - Linear layer with 50 inputs and 10 outputs
+
+        The weights of the linear layers are initialized with a normal distribution
+        (mean=0, std=0.1) and the biases are initialized to 0.
+
+        """
+        super().__init__()
+        self.layers = layers
+        self.activation = activation
+        self.output_activation = output_activation
+
+        levels = []
+        previous_neurons = n_inputs
+
+        for idx, neurons in enumerate(self.layers):
+            if idx == (len(self.layers) - 1):
+                levels.append(nn.Linear(previous_neurons, neurons))
+
+                if self.output_activation:
+                    levels.append(self.output_activation)
+
+            else:
+                levels.append(nn.Linear(previous_neurons, neurons))
+
+                if self.activation:
+                    levels.append(self.activation)
+
+            previous_neurons = neurons
+
+        self.net = nn.Sequential(*levels)
+
+        for m in self.net.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, mean=0, std=0.1)
+                nn.init.constant_(m.bias, val=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the neural network.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor.
+
+        """
+        return self.net(x)
+
+
+class LSTMnn(nn.Module):
+    """Default LSTM neural network model for time-series approximation."""
+
+    def __init__(self, n_inputs: int, n_outputs: int, n_hidden: int) -> None:
+        """Initializes the neural network model.
+
+        Args:
+            n_inputs (int): Number of input features.
+            n_outputs (int): Number of output features.
+            n_hidden (int): Number of hidden units in the LSTM layer.
+
+        """
+        super().__init__()
+
+        self.n_hidden = n_hidden
+
+        self.lstm = nn.LSTM(n_inputs, n_hidden)
+        self.to_out = nn.Linear(n_hidden, n_outputs)
+
+        nn.init.normal_(self.to_out.weight, mean=0, std=0.1)
+        nn.init.constant_(self.to_out.bias, val=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the neural network."""
+        # lstm_out, (hidden_state, cell_state)
+        _, (hn, _) = self.lstm(x)
+        return cast(torch.Tensor, self.to_out(hn[-1]))  # Use last hidden state

modelbase2/npe.py

@@ -4,10 +4,6 @@ This module provides classes and functions for training neural network models to
 parameters in metabolic models. It includes functionality for both steady-state and
 time-series data.
 
-Classes:
-    DefaultSSAproximator: Default neural network model for steady-state approximation
-    DefaultTimeSeriesApproximator: Default neural network model for time-series approximation
-
 Functions:
     train_torch_surrogate: Train a PyTorch surrogate model
     train_torch_time_course_estimator: Train a PyTorch time course estimator
@@ -18,9 +14,6 @@ from __future__ import annotations
 __all__ = [
     "AbstractEstimator",
     "DefaultCache",
-    "DefaultDevice",
-    "DefaultSSAproximator",
-    "DefaultTimeSeriesApproximator",
     "TorchSSEstimator",
     "TorchTimeCourseEstimator",
     "train_torch_ss_estimator",
@@ -30,7 +23,7 @@ __all__ = [
 from abc import abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 import pandas as pd
@@ -39,75 +32,15 @@ import tqdm
 from torch import nn
 from torch.optim.adam import Adam
 
+from modelbase2.nnarchitectures import MLP, DefaultDevice, LSTMnn
 from modelbase2.parallel import Cache
 
-DefaultDevice = torch.device("cpu")
-DefaultCache = Cache(Path(".cache"))
-
-
-class DefaultSSAproximator(nn.Module):
-    """Default neural network model for steady-state approximation."""
-
-    def __init__(self, n_inputs: int, n_outputs: int, n_hidden: int = 50) -> None:
-        """Initializes the neural network with the specified number of inputs and outputs.
-
-        Args:
-            n_inputs (int): The number of input features.
-            n_outputs (int): The number of output features.
-            n_hidden (int): The number of hidden units in the fully connected layers
-
-        The network consists of three fully connected layers with ReLU activations in between.
-        The weights of the linear layers are initialized with a normal distribution (mean=0, std=0.1),
-        and the biases are initialized to zero.
-
-        """
-        super().__init__()
-
-        self.net = nn.Sequential(
-            nn.Linear(n_inputs, n_hidden),
-            nn.ReLU(),
-            nn.Linear(n_hidden, n_hidden),
-            nn.ReLU(),
-            nn.Linear(n_hidden, n_outputs),
-        )
-
-        for m in self.net.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, mean=0, std=0.1)
-                nn.init.constant_(m.bias, val=0)
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass through the neural network."""
-        return cast(torch.Tensor, self.net(x))
+    from torch.optim.optimizer import ParamsT
 
-
-class DefaultTimeSeriesApproximator(nn.Module):
-    """Default neural network model for time-series approximation."""
-
-    def __init__(self, n_inputs: int, n_outputs: int, n_hidden: int) -> None:
-        """Initializes the neural network model.
-
-        Args:
-            n_inputs (int): Number of input features.
-            n_outputs (int): Number of output features.
-            n_hidden (int): Number of hidden units in the LSTM layer.
-
-        """
-        super().__init__()
-
-        self.n_hidden = n_hidden
-
-        self.lstm = nn.LSTM(n_inputs, n_hidden)
-        self.to_out = nn.Linear(n_hidden, n_outputs)
-
-        nn.init.normal_(self.to_out.weight, mean=0, std=0.1)
-        nn.init.constant_(self.to_out.bias, val=0)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass through the neural network."""
-        # lstm_out, (hidden_state, cell_state)
-        _, (hn, _) = self.lstm(x)
-        return cast(torch.Tensor, self.to_out(hn[-1]))  # Use last hidden state
+DefaultCache = Cache(Path(".cache"))
 
 
 @dataclass(kw_only=True)
@@ -212,7 +145,7 @@ def train_torch_ss_estimator(
     epochs: int,
     batch_size: int | None = None,
     approximator: nn.Module | None = None,
-    optimimzer_cls: type[Adam] = Adam,
+    optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
     device: torch.device = DefaultDevice,
 ) -> tuple[TorchSSEstimator, pd.Series]:
     """Train a PyTorch steady state estimator.
@@ -229,7 +162,7 @@ def train_torch_ss_estimator(
         targets: DataFrame containing the target values for training
         epochs: Number of training epochs
         batch_size: Size of mini-batches for training (None for full-batch)
-        approximator: Predefined neural network model (None to use default)
+        approximator: Predefined neural network model (None to use default MLP)
         optimimzer_cls: Optimizer class to use for training (default: Adam)
         device: Device to run the training on (default: DefaultDevice)
 
@@ -238,10 +171,10 @@ def train_torch_ss_estimator(
 
     """
     if approximator is None:
-        approximator = DefaultSSAproximator(
-            n_inputs=len(features.columns),
-            n_outputs=len(targets.columns),
-            n_hidden=max(2 * len(features.columns) * len(targets.columns), 10),
+        n_hidden = max(2 * len(features.columns) * len(targets.columns), 10)
+        n_outputs = len(targets.columns)
+        approximator = MLP(
+            n_inputs=len(features.columns), layers=[n_hidden, n_hidden, n_outputs]
         ).to(device)
 
     features_ = torch.Tensor(features.to_numpy(), device=device)
@@ -278,7 +211,7 @@ def train_torch_time_course_estimator(
     epochs: int,
     batch_size: int | None = None,
     approximator: nn.Module | None = None,
-    optimimzer_cls: type[Adam] = Adam,
+    optimimzer_cls: Callable[[ParamsT], Adam] = Adam,
     device: torch.device = DefaultDevice,
 ) -> tuple[TorchTimeCourseEstimator, pd.Series]:
     """Train a PyTorch time course estimator.
@@ -295,7 +228,7 @@ def train_torch_time_course_estimator(
         targets: DataFrame containing the target values for training
         epochs: Number of training epochs
         batch_size: Size of mini-batches for training (None for full-batch)
-        approximator: Predefined neural network model (None to use default)
+        approximator: Predefined neural network model (None to use default LSTM)
         optimimzer_cls: Optimizer class to use for training (default: Adam)
         device: Device to run the training on (default: DefaultDevice)
 
@@ -304,7 +237,7 @@ def train_torch_time_course_estimator(
 
     """
     if approximator is None:
-        approximator = DefaultTimeSeriesApproximator(
+        approximator = LSTMnn(
             n_inputs=len(features.columns),
             n_outputs=len(targets.columns),
             n_hidden=1,

modelbase2/plot.py

@@ -399,7 +399,10 @@ def lines(
     fig, ax = _default_fig_ax(ax=ax, grid=grid)
     ax.plot(x.index, x)
     _default_labels(ax, xlabel=x.index.name, ylabel=None)
-    ax.legend(x.columns)
+    if isinstance(x, pd.Series):
+        ax.legend([str(x.name)])
+    else:
+        ax.legend(x.columns)
     return fig, ax
 
 

modelbase2/simulator.py

@@ -21,6 +21,8 @@ from modelbase2.integrators import DefaultIntegrator
 __all__ = ["Simulator"]
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from modelbase2.model import Model
     from modelbase2.types import ArrayLike, IntegratorProtocol
 
@@ -83,7 +85,9 @@ class Simulator:
         self,
         model: Model,
         y0: dict[str, float] | None = None,
-        integrator: type[IntegratorProtocol] = DefaultIntegrator,
+        integrator: Callable[
+            [Callable, ArrayLike], IntegratorProtocol
+        ] = DefaultIntegrator,
         *,
         test_run: bool = True,
     ) -> None:
@@ -93,7 +97,7 @@ class Simulator:
             model (Model): The model to be simulated.
             y0 (dict[str, float] | None, optional): Initial conditions for the model variables.
                 If None, the initial conditions are obtained from the model. Defaults to None.
-            integrator (type[IntegratorProtocol], optional): The integrator to use for the simulation.
+            integrator (Callable[[Callable, ArrayLike], IntegratorProtocol], optional): The integrator to use for the simulation.
                 Defaults to DefaultIntegrator.
             test_run (bool, optional): If True, performs a test run to ensure the model's methods
                 (get_full_concs, get_fluxes, get_right_hand_side) work correctly with the initial conditions.
@@ -104,7 +108,7 @@ class Simulator:
         y0 = model.get_initial_conditions() if y0 is None else y0
         self.y0 = [y0[k] for k in model.get_variable_names()]
 
-        self.integrator = integrator(self.model, y0=self.y0)
+        self.integrator = integrator(self.model, self.y0)
         self.concs = None
         self.args = None
         self.simulation_parameters = None