mxlpy 0.25.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

mxlpy/fuzzy.py

@@ -0,0 +1,139 @@
+"""Fuzzy / bayesian fitting methods."""
+
+from __future__ import annotations
+
+import multiprocessing
+import sys
+from dataclasses import dataclass, field
+from functools import partial
+from math import ceil
+from typing import TYPE_CHECKING, Self
+
+import numpy as np
+import pandas as pd
+import pebble
+from tqdm import tqdm, trange
+
+from mxlpy.simulator import Simulator
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from mxlpy import Model
+
+__all__ = ["ThompsonState", "thompson_sampling"]
+
+
+@dataclass
+class ThompsonState:
+    """State of thompson sampling."""
+
+    rng: np.random.Generator = field(default_factory=np.random.default_rng)
+    state: dict[str, pd.DataFrame] = field(default_factory=dict)
+
+    @classmethod
+    def from_parameter_values(cls, parameters: dict[str, Iterable[float]]) -> Self:
+        """Create state from parameter values."""
+        return cls(
+            state={
+                k: pd.DataFrame(
+                    {
+                        "x": v,
+                        "success": np.ones_like(v, dtype=int),
+                        "fail": np.ones_like(v, dtype=int),
+                    }
+                )
+                for k, v in parameters.items()
+            },
+        )
+
+    def sample(self) -> tuple[dict[str, int], dict[str, float]]:
+        """Sample idxs and parameters."""
+        idxs = {
+            k: int(np.argmax(self.rng.beta(v["success"], v["fail"])))
+            for k, v in self.state.items()
+        }
+        parameters = {k: v["x"][idxs[k]] for k, v in self.state.items()}
+        return idxs, parameters
+
+    def update(
+        self,
+        idxs: dict[str, int],
+        pred: pd.DataFrame | None,
+        data: pd.DataFrame,
+        rtol: float,
+    ) -> None:
+        """Sample state."""
+        accept: bool = (
+            False if pred is None else np.sqrt(np.mean(np.square(pred - data))) < rtol
+        )
+        for k, v in self.state.items():
+            v.loc[idxs[k], "success" if accept else "fail"] += 1  # type: ignore
+
+
+def _thompson_worker(
+    inp: tuple[dict[str, int], dict[str, float]],
+    model: Model,
+    data: pd.DataFrame,
+) -> tuple[dict[str, int], pd.DataFrame | None]:
+    idxs, parameters = inp
+    if (
+        res := (
+            Simulator(model)
+            .update_parameters(parameters)
+            .simulate_time_course(data.index)
+            .get_result()
+        )
+    ) is None:
+        return idxs, None
+    return idxs, res.get_variables()
+
+
+def thompson_sampling(
+    model: Model,
+    data: pd.DataFrame,
+    state: ThompsonState,
+    rtol: float,
+    n: int,
+    *,
+    max_workers: int | None = None,
+    disable_tqdm: bool = False,
+    timeout: float | None = None,
+    parallel: bool = True,
+) -> ThompsonState:
+    """Perform thompson sampling."""
+    if sys.platform in ["win32", "cygwin"]:
+        parallel = False
+
+    max_workers = multiprocessing.cpu_count() if max_workers is None else max_workers
+    worker = partial(_thompson_worker, model=model, data=data)
+
+    if not parallel:
+        for _ in trange(n):
+            idxs, pred = worker(state.sample())
+            state.update(idxs, pred, data=data, rtol=rtol)
+    else:
+        # FIXME: think about whether this is ok to do. Thompson sampling is state-
+        # dependent. We are breaking up that state a bit by chunking the approach
+        # Is that fine to do?
+        with (
+            tqdm(total=n, disable=disable_tqdm) as pbar,
+            pebble.ProcessPool(max_workers=max_workers) as pool,
+        ):
+            for _ in range(ceil(n / max_workers)):
+                future = pool.map(
+                    worker,
+                    [state.sample() for _ in range(max_workers)],
+                    timeout=timeout,
+                )
+                it = future.result()
+                while True:
+                    try:
+                        idxs, pred = next(it)
+                        state.update(idxs, pred, data=data, rtol=rtol)
+                        pbar.update(1)
+                    except StopIteration:
+                        break
+                    except TimeoutError:
+                        pbar.update(1)
+    return state

mxlpy/identify.py

@@ -9,9 +9,8 @@ import numpy as np
 import pandas as pd
 from tqdm import tqdm
 
-from mxlpy import fit_local
+from mxlpy import fit
 from mxlpy.distributions import LogNormal, sample
-from mxlpy.fit.common import LossFn, rmse
 from mxlpy.parallel import parallelise
 
 if TYPE_CHECKING:
@@ -27,13 +26,14 @@ def _mc_fit_time_course_worker(
     p0: pd.Series,
     model: Model,
     data: pd.DataFrame,
-    loss_fn: fit_local.LossFn,
+    loss_fn: fit.LossFn,
 ) -> float:
-    fit_result = fit_local.time_course(
+    fit_result = fit.time_course(
         model=model,
         p0=p0.to_dict(),
         data=data,
         loss_fn=loss_fn,
+        minimizer=fit.LocalScipyMinimizer(),
     )
     if fit_result is None:
         return np.inf
@@ -46,7 +46,7 @@ def profile_likelihood(
     parameter_name: str,
     parameter_values: Array,
     n_random: int = 10,
-    loss_fn: LossFn = rmse,
+    loss_fn: fit.LossFn = fit.rmse,
 ) -> pd.Series:
     """Estimate the profile likelihood of model parameters given data.
 

mxlpy/integrators/int_scipy.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import copy
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Literal, cast
 
 import numpy as np
 import scipy.integrate as spi
@@ -48,6 +48,7 @@ class Scipy:
     atol: float = 1e-8
     rtol: float = 1e-8
     t0: float = 0.0
+    method: Literal["RK45", "RK23", "DOP853", "Radau", "BDF", "LSODA"] = "LSODA"
     _y0_orig: tuple[float, ...] = field(default_factory=tuple)
 
     def __post_init__(self) -> None:
@@ -111,7 +112,7 @@ class Scipy:
             jac=self.jacobian,
             atol=self.atol,
             rtol=self.rtol,
-            method="LSODA",
+            method=self.method,
         )
 
         if res.success:
@@ -149,7 +150,7 @@ class Scipy:
         # If rhs returns a tuple, we get weird errors, so we need
         # to wrap this in a list for some reason
         integ = spi.ode(lambda t, x: list(self.rhs(t, x)), jac=self.jacobian)
-        integ.set_integrator(name="lsoda")
+        integ.set_integrator(name=self.method)
         integ.set_initial_value(self.y0)
 
         t = self.t0 + step_size

mxlpy/meta/codegen_latex.py

@@ -795,6 +795,7 @@ def to_tex_export(model: Model) -> TexExport:
     for rxn_name, rxn in model.get_raw_reactions().items():
         for var_name, factor in rxn.stoichiometry.items():
             diff_eqs.setdefault(var_name, {})[rxn_name] = factor
+    # FIXME: think about surrogates here
 
     return TexExport(
         parameters=model.get_parameter_values(),

mxlpy/meta/source_tools.py

@@ -635,7 +635,7 @@ def _handle_attribute(node: ast.Attribute, ctx: Context) -> sympy.Expr | None:
                     )
                 elif (var := variables.get(level)) is not None:
                     _LOGGER.debug("var %s", var)
-                    return _get_inner_object(var, levels[(idx + 1) :] + [node.attr])
+                    return _get_inner_object(var, [*levels[idx + 1 :], node.attr])
 
                 else:
                     _LOGGER.debug("No target found")

mxlpy/model.py

@@ -928,14 +928,14 @@ class Model:
         if stoichiometries is not None:
             for rxn_name, value in stoichiometries.items():
                 target = False
-                if rxn_name in self._reactions:
+                if (rxn := self._reactions.get(rxn_name)) is not None:
                     target = True
-                    cast(dict, self._reactions[name].stoichiometry)[name] = value
+                    cast(dict, rxn.stoichiometry)[name] = value
                 else:
                     for surrogate in self._surrogates.values():
-                        if rxn_name in surrogate.stoichiometries:
+                        if stoich := surrogate.stoichiometries.get(rxn_name):
                             target = True
-                            surrogate.stoichiometries[rxn_name][name] = value
+                            stoich[name] = value
                 if not target:
                     msg = f"Reaction '{rxn_name}' not found in reactions or surrogates"
                     raise KeyError(msg)
@@ -1099,7 +1099,12 @@ class Model:
         return self
 
     @_invalidate_cache
-    def remove_variable(self, name: str) -> Self:
+    def remove_variable(
+        self,
+        name: str,
+        *,
+        remove_stoichiometries: bool = True,
+    ) -> Self:
         """Remove a variable from the model.
 
         Examples:
@@ -1107,16 +1112,31 @@ class Model:
 
         Args:
             name: The name of the variable to remove.
+            remove_stoichiometries: whether to remove the variable from all reactions
 
         Returns:
             Self: The instance of the model with the variable removed.
 
         """
+        if remove_stoichiometries:
+            for rxn in self._reactions.values():
+                if name in rxn.stoichiometry:
+                    cast(dict, rxn.stoichiometry).pop(name)
+            for surrogate in self._surrogates.values():
+                for stoich in surrogate.stoichiometries.values():
+                    if name in stoich:
+                        cast(dict, stoich).pop(name)
+
         self._remove_id(name=name)
         del self._variables[name]
         return self
 
-    def remove_variables(self, variables: Iterable[str]) -> Self:
+    def remove_variables(
+        self,
+        variables: Iterable[str],
+        *,
+        remove_stoichiometries: bool = True,
+    ) -> Self:
         """Remove multiple variables from the model.
 
         Examples:
@@ -1124,13 +1144,16 @@ class Model:
 
         Args:
             variables: An iterable of variable names to be removed.
+            remove_stoichiometries: whether to remove the variables from all reactions
 
         Returns:
             Self: The instance of the model with the specified variables removed.
 
         """
         for variable in variables:
-            self.remove_variable(name=variable)
+            self.remove_variable(
+                name=variable, remove_stoichiometries=remove_stoichiometries
+            )
         return self
 
     @_invalidate_cache
@@ -1219,24 +1242,17 @@ class Model:
         value_or_derived = (
             self._variables[name].initial_value if value is None else value
         )
-        self.remove_variable(name)
-
-        # FIXME: better handling of unit
-        if isinstance(value_or_derived, Derived):
-            self.add_derived(name, value_or_derived.fn, args=value_or_derived.args)
+        self.remove_variable(name, remove_stoichiometries=True)
+
+        if isinstance(der := value_or_derived, Derived):
+            self.add_derived(
+                name,
+                der.fn,
+                args=der.args,
+                unit=der.unit,
+            )
         else:
             self.add_parameter(name, value_or_derived)
-
-        # Remove from stoichiometries
-        for reaction in self._reactions.values():
-            if name in reaction.stoichiometry:
-                cast(dict, reaction.stoichiometry).pop(name)
-        for surrogate in self._surrogates.values():
-            surrogate.stoichiometries = {
-                k: {k2: v2 for k2, v2 in v.items() if k2 != name}
-                for k, v in surrogate.stoichiometries.items()
-                if k != name
-            }
         return self
 
     ##########################################################################
@@ -1680,7 +1696,8 @@ class Model:
 
     ##########################################################################
     # Readouts
-    # They are like derived variables, but only calculated on demand
+    # They are like derived variables, but only calculated on demand, e.g. after
+    # a simulation
     # Think of something like NADPH / (NADP + NADPH) as a proxy for energy state
     ##########################################################################
 

mxlpy/nn/__init__.py

@@ -8,11 +8,16 @@ if TYPE_CHECKING:
     import contextlib
 
     with contextlib.suppress(ImportError):
+        from . import _equinox as equinox
         from . import _keras as keras
         from . import _torch as torch
 else:
     from lazy_import import lazy_module
 
+    equinox = lazy_module(
+        "mxlpy.nn._equinox",
+        error_strings={"module": "equinox", "install_name": "mxlpy[equinox]"},
+    )
     keras = lazy_module(
         "mxlpy.nn._keras",
         error_strings={"module": "keras", "install_name": "mxlpy[tf]"},

mxlpy/nn/_equinox.py

@@ -0,0 +1,293 @@
+"""Neural network architectures.
+
+This module provides implementations of neural network architectures used for mechanistic learning.
+
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import torch
+import tqdm
+from jaxtyping import Array, PyTree
+from torch.utils.data import DataLoader, TensorDataset
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import optax
+
+
+__all__ = [
+    "LSTM",
+    "LossFn",
+    "MLP",
+    "cosine_similarity",
+    "mean_abs_error",
+    "mean_absolute_percentage",
+    "mean_error",
+    "mean_squared_error",
+    "mean_squared_logarithmic",
+    "rms_error",
+    "train",
+]
+
+
+###############################################################################
+# Loss functions
+###############################################################################
+
+type LossFn = Callable[[eqx.Module, Array, Array], Array]
+
+
+@eqx.filter_jit
+def mean_error(model: eqx.Module, inp: Array, true: Array) -> Array:
+    """Calculate mean error."""
+    pred = jax.vmap(model)(inp)  # type: ignore
+    return jnp.mean(pred - true)
+
+
+@eqx.filter_jit
+def mean_squared_error(model: eqx.Module, inp: Array, true: Array) -> Array:
+    """Calculate mean squared error."""
+    pred = jax.vmap(model)(inp)  # type: ignore
+    return jnp.mean(jnp.square(pred - true))
+
+
+@eqx.filter_jit
+def rms_error(model: eqx.Module, inp: Array, true: Array) -> Array:
+    """Calculate root mean square error."""
+    pred = jax.vmap(model)(inp)  # type: ignore
+    return jnp.sqrt(jnp.mean(jnp.square(pred - true)))
+
+
+@eqx.filter_jit
+def mean_abs_error(model: eqx.Module, inp: Array, true: Array) -> Array:
+    """Calculate mean absolute error."""
+    pred = jax.vmap(model)(inp)  # type: ignore
+    return jnp.mean(jnp.abs(pred - true))
+
+
+@eqx.filter_jit
+def mean_absolute_percentage(model: eqx.Module, inp: Array, true: Array) -> Array:
+    """Calculate mean absolute percentag error."""
+    pred = jax.vmap(model)(inp)  # type: ignore
+    return 100 * jnp.mean(jnp.abs((true - pred) / pred))
+
+
+@eqx.filter_jit
+def mean_squared_logarithmic(model: eqx.Module, inp: Array, true: Array) -> Array:
+    """Calculate root mean square error between model and data."""
+    pred = jax.vmap(model)(inp)  # type: ignore
+    return jnp.mean(jnp.square(jnp.log(pred + 1) - jnp.log(true + 1)))
+
+
+@eqx.filter_jit
+def cosine_similarity(model: eqx.Module, inp: Array, true: Array) -> Array:
+    """Calculate root mean square error between model and data."""
+    pred = jax.vmap(model)(inp)  # type: ignore
+    return -jnp.sum(jnp.linalg.norm(pred, 2) * jnp.linalg.norm(true, 2))
+
+
+###############################################################################
+# Training routines
+###############################################################################
+
+
+def train(
+    model: eqx.Module,
+    features: Array,
+    targets: Array,
+    epochs: int,
+    optimizer: optax.GradientTransformation,
+    batch_size: int | None,
+    loss_fn: LossFn,
+) -> pd.Series:
+    """Train the neural network using mini-batch gradient descent.
+
+    Args:
+        model: Neural network model to train.
+        features: Input features as a tensor.
+        targets: Target values as a tensor.
+        epochs: Number of training epochs.
+        optimizer: Optimizer for training.
+        device: torch device
+        batch_size: Size of mini-batches for training.
+        loss_fn: Loss function
+
+    Returns:
+        pd.Series: Series containing the training loss history.
+
+    """
+    losses = {}
+
+    data = TensorDataset(
+        torch.tensor(features.astype(np.float32), dtype=torch.float32),
+        torch.tensor(targets.astype(np.float32), dtype=torch.float32),
+    )
+    data_loader = DataLoader(
+        data,
+        batch_size=len(features) if batch_size is None else batch_size,
+        shuffle=True,
+    )
+
+    opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
+
+    @eqx.filter_jit
+    def make_step(
+        model: eqx.Module,
+        opt_state: PyTree,
+        x: Array,
+        y: Array,
+    ) -> tuple[eqx.Module, Array, Array]:
+        loss_value, grads = eqx.filter_value_and_grad(loss_fn)(model, x, y)
+        updates, opt_state = optimizer.update(
+            grads, opt_state, eqx.filter(model, eqx.is_array)
+        )
+        model = eqx.apply_updates(model, updates)
+        return model, opt_state, loss_value
+
+    for i in tqdm.trange(epochs):
+        epoch_loss = 0
+        for xb, yb in data_loader:
+            model, opt_state, train_loss = make_step(
+                model,
+                opt_state,
+                xb.numpy(),
+                yb.numpy(),
+            )
+            epoch_loss += train_loss * xb.size(0)
+        losses[i] = epoch_loss / len(data_loader.dataset)  # type: ignore
+    return pd.Series(losses, dtype=float)
+
+
+###############################################################################
+# Actual models
+###############################################################################
+
+
+class MLP(eqx.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.
+
+    """
+
+    layers: list
+
+    def __init__(
+        self,
+        n_inputs: int,
+        neurons_per_layer: list[int],
+        key: Array,
+    ) -> None:
+        """Initializes the MLP with the given number of inputs and list of (hidden) layers.
+
+        Args:
+            n_inputs: The number of input features.
+            neurons_per_layer: Number of neurons per layer
+            n_outputs: A list containing the number of neurons in hidden and output layer.
+            key: jax.random.PRNGKey(SEED) for initial parameters
+
+        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.
+
+        """
+        keys = iter(jax.random.split(key, len(neurons_per_layer)))
+        previous_neurons = n_inputs
+        layers = []
+        for neurons in neurons_per_layer:
+            layers.append(eqx.nn.Linear(previous_neurons, neurons, key=next(keys)))
+            previous_neurons = neurons
+        self.layers = layers
+
+    def __call__(self, x: Array) -> Array:
+        """Forward pass through the neural network.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            Output tensor.
+
+        """
+        for layer in self.layers[:-1]:
+            x = jax.nn.relu(layer(x))
+        return self.layers[-1](x)
+
+
+class LSTM(eqx.Module):
+    """Default LSTM neural network model for time-series approximation."""
+
+    lstm_cell: eqx.nn.LSTMCell
+    n_hidden: int
+    linear: eqx.nn.Linear
+
+    def __init__(
+        self,
+        n_inputs: int,
+        n_outputs: int,
+        n_hidden: int,
+        key: Array,
+    ) -> None:
+        """Initializes the LSTM 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.
+            key (Array): JAX random key for initialization.
+
+        """
+        k1, k2 = jax.random.split(key, 2)
+        self.lstm_cell = eqx.nn.LSTMCell(n_inputs, n_hidden, key=k1)
+        self.n_hidden = n_hidden
+        self.linear = eqx.nn.Linear(n_hidden, n_outputs, key=k2)
+
+    def __call__(
+        self,
+        x: Array,
+        *,
+        h: Array | None = None,
+        c: Array | None = None,
+    ) -> Array:
+        """Forward pass through the LSTM network.
+
+        Args:
+            x: Input tensor of shape (seq_len, batch_size, n_inputs).
+            h: Optional initial hidden state (batch_size, n_hidden).
+            c: Optional initial cell state (batch_size, n_hidden).
+
+        Returns:
+            Output tensor of shape (seq_len, batch_size, n_outputs).
+
+        """
+        seq_len, batch_size, _ = x.shape
+        if h is None:
+            h = jnp.zeros((batch_size, self.n_hidden))
+        if c is None:
+            c = jnp.zeros((batch_size, self.n_hidden))
+
+        outputs = []
+        for t in range(seq_len):
+            h, c = self.lstm_cell(x[t], (h, c))
+            outputs.append(h)
+        outputs = jnp.stack(outputs, axis=0)
+        return jax.vmap(self.linear)(outputs)