apply suggestions

Author: dweindl

petab/v1/distributions.py

@@ -33,27 +33,27 @@ class Distribution(abc.ABC):
     def __init__(self, log: bool | float = False):
         if log is True:
             log = np.exp(1)
-        self._log = log
+        self._logbase = log
 
     def _undo_log(self, x: np.ndarray | float) -> np.ndarray | float:
         """Undo the log transformation.
 
         :param x: The sample to transform.
         :return: The transformed sample
         """
-        if self._log is False:
+        if self._logbase is False:
             return x
-        return self._log**x
+        return self._logbase**x
 
     def _apply_log(self, x: np.ndarray | float) -> np.ndarray | float:
         """Apply the log transformation.
 
         :param x: The value to transform.
         :return: The transformed value.
         """
-        if self._log is False:
+        if self._logbase is False:
             return x
-        return np.log(x) / np.log(self._log)
+        return np.log(x) / np.log(self._logbase)
 
     def sample(self, shape=None) -> np.ndarray:
         """Sample from the distribution.
@@ -82,7 +82,9 @@ def pdf(self, x):
         """
         # handle the log transformation; see also:
         #  https://en.wikipedia.org/wiki/Probability_density_function#Scalar_to_scalar
-        chain_rule_factor = (1 / (x * np.log(self._log))) if self._log else 1
+        chain_rule_factor = (
+            (1 / (x * np.log(self._logbase))) if self._logbase else 1
+        )
         return self._pdf(self._apply_log(x)) * chain_rule_factor
 
     @abc.abstractmethod
@@ -94,6 +96,14 @@ def _pdf(self, x):
         """
         ...
 
+    @property
+    def logbase(self) -> bool | float:
+        """The base of the log transformation.
+
+        If ``False``, no transformation is applied.
+        """
+        return self._logbase
+
 
 class Normal(Distribution):
     """A (log-)normal distribution.
@@ -127,7 +137,7 @@ def __init__(
 
     def __repr__(self):
         trunc = f", truncation={self._truncation}" if self._truncation else ""
-        log = f", log={self._log}" if self._log else ""
+        log = f", log={self._logbase}" if self._logbase else ""
         return f"Normal(loc={self._loc}, scale={self._scale}{trunc}{log})"
 
     def _sample(self, shape=None):
@@ -172,7 +182,7 @@ def __init__(
         self._high = high
 
     def __repr__(self):
-        log = f", log={self._log}" if self._log else ""
+        log = f", log={self._logbase}" if self._logbase else ""
         return f"Uniform(low={self._low}, high={self._high}{log})"
 
     def _sample(self, shape=None):
@@ -213,7 +223,7 @@ def __init__(
 
     def __repr__(self):
         trunc = f", truncation={self._truncation}" if self._truncation else ""
-        log = f", log={self._log}" if self._log else ""
+        log = f", log={self._logbase}" if self._logbase else ""
         return f"Laplace(loc={self._loc}, scale={self._scale}{trunc}{log})"
 
     def _sample(self, shape=None):

tests/v1/test_distributions.py

@@ -1,41 +1,42 @@
-from itertools import product
-
 import numpy as np
 import pytest
+from numpy.testing import assert_allclose
 from scipy.integrate import cumulative_trapezoid
-from scipy.stats import kstest
+from scipy.stats import (
+    kstest,
+    laplace,
+    loglaplace,
+    lognorm,
+    loguniform,
+    norm,
+    uniform,
+)
 
 from petab.v1.distributions import *
 from petab.v2.C import *
 
 
 @pytest.mark.parametrize(
-    "distribution, transform",
-    list(
-        product(
-            [
-                Normal(2, 1),
-                Normal(2, 1, log=True),
-                Normal(2, 1, log=10),
-                Uniform(2, 4),
-                Uniform(-2, 4, log=True),
-                Uniform(2, 4, log=10),
-                Laplace(1, 2),
-                Laplace(1, 0.5, log=True),
-            ],
-            [LIN, LOG, LOG10],
-        )
-    ),
+    "distribution",
+    [
+        Normal(2, 1),
+        Normal(2, 1, log=True),
+        Normal(2, 1, log=10),
+        Uniform(2, 4),
+        Uniform(-2, 4, log=True),
+        Uniform(2, 4, log=10),
+        Laplace(1, 2),
+        Laplace(1, 0.5, log=True),
+    ],
 )
-def test_sample_matches_pdf(distribution, transform):
+def test_sample_matches_pdf(distribution):
     """Test that the sample matches the PDF."""
     np.random.seed(1)
     N_SAMPLES = 10_000
-    distribution.transform = transform
     sample = distribution.sample(N_SAMPLES)
 
-    # pdf -> cdf
     def cdf(x):
+        # pdf -> cdf
         return cumulative_trapezoid(distribution.pdf(x), x)
 
     # Kolmogorov-Smirnov test to check if the sample is drawn from the CDF
@@ -49,3 +50,38 @@ def cdf(x):
     #     plt.show()
 
     assert p > 0.05, (p, distribution)
+
+    # Test samples match scipy CDFs
+    reference_pdf = None
+    if isinstance(distribution, Normal) and distribution.logbase is False:
+        reference_pdf = norm.pdf(sample, distribution.loc, distribution.scale)
+    elif isinstance(distribution, Uniform) and distribution.logbase is False:
+        reference_pdf = uniform.pdf(
+            sample, distribution._low, distribution._high - distribution._low
+        )
+    elif isinstance(distribution, Laplace) and distribution.logbase is False:
+        reference_pdf = laplace.pdf(
+            sample, distribution.loc, distribution.scale
+        )
+    elif isinstance(distribution, Normal) and distribution.logbase == np.exp(
+        1
+    ):
+        reference_pdf = lognorm.pdf(
+            sample, scale=np.exp(distribution.loc), s=distribution.scale
+        )
+    elif isinstance(distribution, Uniform) and distribution.logbase == np.exp(
+        1
+    ):
+        reference_pdf = loguniform.pdf(
+            sample, np.exp(distribution._low), np.exp(distribution._high)
+        )
+    elif isinstance(distribution, Laplace) and distribution.logbase == np.exp(
+        1
+    ):
+        reference_pdf = loglaplace.pdf(
+            sample, c=1 / distribution.scale, scale=np.exp(distribution.loc)
+        )
+    if reference_pdf is not None:
+        assert_allclose(
+            distribution.pdf(sample), reference_pdf, rtol=1e-10, atol=1e-14
+        )