Refactor priors

Introduce `Distribution` classes for handling prior distributions and sampling from them.
Later on this can be extended to noise models for measurements.

Author: dweindl

petab/v1/distributions.py

@@ -0,0 +1,359 @@
+"""Probability distributions used by PEtab."""
+from __future__ import annotations
+
+import abc
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+from scipy.stats import laplace, lognorm, norm, uniform
+
+from . import C
+
+
+class Distribution(abc.ABC):
+    """A univariate probability distribution.
+
+    :param type_: The type of the distribution.
+    :param transformation: The transformation to be applied to the sample.
+        Ignored if `parameter_scale` is `True`.
+    :param parameters: The parameters of the distribution.
+    :param bounds: The bounds of the distribution (lower, upper).
+    :param parameter_scale: Whether the parameters are already on the correct
+        scale. If `False`, the parameters are transformed to the correct scale.
+        If `True`, the parameters are assumed to be on the correct scale and
+        no transformation is applied.
+    :param transformation: The transformation of the distribution.
+
+    """
+
+    _type_to_cls: dict[str, type[Distribution]] = {}
+
+    def __init__(
+        self,
+        type_: str,
+        transformation: str,
+        parameters: tuple,
+        bounds: tuple = None,
+        parameter_scale: bool = False,
+    ):
+        self.type = type_
+        self.parameters = parameters
+        self.bounds = bounds
+        self.transformation = transformation
+        self.parameter_scale = parameter_scale
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.__dict__})"
+
+    @abc.abstractmethod
+    def sample(self, shape=None) -> np.ndarray:
+        """Sample from the distribution.
+
+        :param shape: The shape of the sample.
+        :return: A sample from the distribution.
+        """
+        ...
+
+    def _scale_sample(self, sample):
+        """Scale the sample to the parameter space"""
+        if self.parameter_scale:
+            return sample
+        if self.transformation == C.LIN:
+            return sample
+        if self.transformation == C.LOG:
+            return np.log(sample)
+        if self.transformation == C.LOG10:
+            return np.log10(sample)
+        raise NotImplementedError(
+            f"Transformation {self.transformation} not implemented."
+        )
+
+    def _clip_to_bounds(self, x):
+        """Clip values in array x to bounds."""
+        # TODO: replace this by proper truncation
+        if self.bounds is None:
+            return x
+
+        return np.maximum(
+            np.minimum(self._scale_sample(self.bounds[1]), x),
+            self._scale_sample(self.bounds[0]),
+        )
+
+    @abc.abstractmethod
+    def pdf(self, x):
+        """Probability density function at x.
+
+        ``x`` is assumed to be on the parameter scale.
+        """
+        ...
+
+    @staticmethod
+    def from_par_dict(
+        d, type_=Literal["initialization", "objective"]
+    ) -> Distribution:
+        """Create a distribution from a row of the parameter table.
+
+        :param d: A dictionary representing a row of the parameter table.
+        :return: A distribution object.
+        """
+        dist_type = d.get(f"{type_}PriorType", C.NORMAL)
+        if not isinstance(dist_type, str) and np.isnan(dist_type):
+            dist_type = C.PARAMETER_SCALE_UNIFORM
+
+        cls = Distribution._type_to_cls[dist_type]
+        if (
+            pd.isna(d[f"{type_}PriorParameters"])
+            and dist_type == C.PARAMETER_SCALE_UNIFORM
+        ):
+            params = d[C.LOWER_BOUND], d[C.UPPER_BOUND]
+        else:
+            params = tuple(
+                map(
+                    float,
+                    d[f"{type_}PriorParameters"].split(C.PARAMETER_SEPARATOR),
+                )
+            )
+        return cls(
+            *params,
+            bounds=(d[C.LOWER_BOUND], d[C.UPPER_BOUND]),
+            transformation=d.get(C.PARAMETER_SCALE, C.LIN),
+        )
+
+
+class Normal(Distribution):
+    """A normal distribution."""
+
+    def __init__(
+        self,
+        mean: float,
+        std: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.NORMAL,
+            transformation=transformation,
+            parameters=(mean, std),
+            bounds=bounds,
+            parameter_scale=False,
+        )
+
+    def sample(self, shape=None):
+        sample = np.random.normal(
+            loc=self.parameters[0], scale=self.parameters[1], size=shape
+        )
+        return self._clip_to_bounds(self._scale_sample(sample))
+
+    def pdf(self, x):
+        return norm.pdf(x, loc=self.parameters[0], scale=self.parameters[1])
+
+
+class LogNormal(Distribution):
+    """A log-normal distribution."""
+
+    def __init__(
+        self,
+        mean: float,
+        std: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.LOG_NORMAL,
+            transformation=transformation,
+            parameters=(mean, std),
+            bounds=bounds,
+            parameter_scale=False,
+        )
+
+    def sample(self, shape=None):
+        sample = np.random.lognormal(
+            mean=self.parameters[0], sigma=self.parameters[1], size=shape
+        )
+        return self._clip_to_bounds(self._scale_sample(sample))
+
+    def pdf(self, x):
+        return lognorm.pdf(
+            x, loc=self.parameters[0], scale=np.exp(self.parameters[1])
+        )
+
+
+class Uniform(Distribution):
+    """A uniform distribution."""
+
+    def __init__(
+        self,
+        lower_bound: float,
+        upper_bound: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.UNIFORM,
+            transformation=transformation,
+            parameters=(lower_bound, upper_bound),
+            bounds=bounds,
+            parameter_scale=False,
+        )
+
+    def sample(self, shape=None):
+        sample = np.random.uniform(
+            low=self.parameters[0], high=self.parameters[1], size=shape
+        )
+        return self._clip_to_bounds(self._scale_sample(sample))
+
+    def pdf(self, x):
+        return uniform.pdf(
+            x,
+            loc=self.parameters[0],
+            scale=self.parameters[1] - self.parameters[0],
+        )
+
+
+class Laplace(Distribution):
+    """A Laplace distribution."""
+
+    def __init__(
+        self,
+        mean: float,
+        scale: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.LAPLACE,
+            transformation=transformation,
+            parameters=(mean, scale),
+            bounds=bounds,
+            parameter_scale=False,
+        )
+
+    def sample(self, shape=None):
+        sample = np.random.laplace(
+            loc=self.parameters[0], scale=self.parameters[1], size=shape
+        )
+        return self._clip_to_bounds(self._scale_sample(sample))
+
+    def pdf(self, x):
+        return laplace.pdf(x, loc=self.parameters[0], scale=self.parameters[1])
+
+
+class LogLaplace(Distribution):
+    """A log-Laplace distribution."""
+
+    def __init__(
+        self,
+        mean: float,
+        scale: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.LOG_LAPLACE,
+            transformation=transformation,
+            parameters=(mean, scale),
+            bounds=bounds,
+            parameter_scale=False,
+        )
+
+    @property
+    def mean(self):
+        """The mean of the underlying Laplace distribution."""
+        return self.parameters[0]
+
+    @property
+    def scale(self):
+        """The scale of the underlying Laplace distribution."""
+        return self.parameters[1]
+
+    def sample(self, shape=None):
+        sample = np.exp(
+            np.random.laplace(loc=self.mean, scale=self.scale, size=shape)
+        )
+        return self._clip_to_bounds(self._scale_sample(sample))
+
+    def pdf(self, x):
+        return (
+            1
+            / (2 * self.scale * x)
+            * np.exp(-np.abs(np.log(x) - self.mean) / self.scale)
+        )
+
+
+class ParameterScaleNormal(Distribution):
+    """A normal distribution with parameters on the parameter scale."""
+
+    def __init__(
+        self,
+        mean: float,
+        std: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.PARAMETER_SCALE_NORMAL,
+            transformation=transformation,
+            parameters=(mean, std),
+            bounds=bounds,
+            parameter_scale=True,
+        )
+
+    sample = Normal.sample
+    pdf = Normal.pdf
+
+
+class ParameterScaleUniform(Distribution):
+    """A uniform distribution with parameters on the parameter scale."""
+
+    def __init__(
+        self,
+        lower_bound: float,
+        upper_bound: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.PARAMETER_SCALE_UNIFORM,
+            transformation=transformation,
+            parameters=(lower_bound, upper_bound),
+            bounds=bounds,
+            parameter_scale=True,
+        )
+
+    sample = Uniform.sample
+    pdf = Uniform.pdf
+
+
+class ParameterScaleLaplace(Distribution):
+    """A Laplace distribution with parameters on the parameter scale."""
+
+    def __init__(
+        self,
+        mean: float,
+        scale: float,
+        bounds: tuple = None,
+        transformation: str = C.LIN,
+    ):
+        super().__init__(
+            C.PARAMETER_SCALE_LAPLACE,
+            transformation=transformation,
+            parameters=(mean, scale),
+            bounds=bounds,
+            parameter_scale=True,
+        )
+
+    sample = Laplace.sample
+    pdf = Laplace.pdf
+
+
+Distribution._type_to_cls = {
+    C.NORMAL: Normal,
+    C.LOG_NORMAL: LogNormal,
+    C.UNIFORM: Uniform,
+    C.LAPLACE: Laplace,
+    C.LOG_LAPLACE: LogLaplace,
+    C.PARAMETER_SCALE_NORMAL: ParameterScaleNormal,
+    C.PARAMETER_SCALE_UNIFORM: ParameterScaleUniform,
+    C.PARAMETER_SCALE_LAPLACE: ParameterScaleLaplace,
+}