From 4904212eb34a1aba2c69476d455ed61783b662f0 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Mon, 2 Jun 2025 11:52:28 -0400 Subject: [PATCH 01/73] Add `gempyor.distributions` module Added a module for representing distributions that can be used throughout gempyor. Started with `DistributionABC` as and abstract base for all distributions. Implemented `FixedDistribution` and `NormalDistribution` using that ABC. Finally, exposed a field discriminated type to easily create distributions from config. --- .../gempyor_pkg/src/gempyor/distributions.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 flepimop/gempyor_pkg/src/gempyor/distributions.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py new file mode 100644 index 000000000..380b302a9 --- /dev/null +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -0,0 +1,92 @@ +"""Representations of distributions used for modifiers, likelihoods, etc.""" + +__all__: tuple[str, ...] = ( + "DistributionABC", + "FixedDistribution", + "NormalDistribution", + "Distribution", +) + + +from abc import ABC, abstractmethod +from typing import Annotated, Literal + +import numpy as np +from numpy.random import Generator +import numpy.typing as npt +from pydantic import BaseModel, Field + + +class DistributionABC(ABC, BaseModel): + """Base class for distributions used in modifiers, likelihoods, etc.""" + + distribution: str + + @abstractmethod + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64 | np.int64]: + """Sample from the distribution.""" + raise NotImplementedError + + +class FixedDistribution(DistributionABC): + """ + Represents a fixed distribution that always returns the same value. + + Examples: + >>> from gempyor.distributions import FixedDistribution + >>> dist = FixedDistribution(value=1.23) + >>> dist.sample() + array([1.23]) + >>> dist.sample(size=(3, 5)) + array([[1.23, 1.23, 1.23, 1.23, 1.23], + [1.23, 1.23, 1.23, 1.23, 1.23], + [1.23, 1.23, 1.23, 1.23, 1.23]]) + """ + + distribution: Literal["fixed"] = "fixed" + value: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: + """Sample from the fixed distribution.""" + return np.full(size, self.value) + + +class NormalDistribution(DistributionABC): + """ + Represents a normal distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import NormalDistribution + >>> rng = np.random.default_rng(42) + >>> dist = NormalDistribution(mu=2.3, sigma=4.5) + >>> dist + NormalDistribution(distribution='norm', mu=2.3, sigma=4.5) + >>> dist.sample(rng=rng) + array([3.67122686]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[-2.37992848, 5.67703038, 6.53254122, -6.47965835, -3.55980778], + [ 2.87528181, 0.87690833, 2.22439479, -1.53869767, 6.25729089], + [ 5.80006371, 2.59713814, 7.37258543, 4.40379204, -1.56681608]]) + """ + + distribution: Literal["norm"] = "norm" + mu: float + sigma: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: + """Sample from the normal distribution.""" + if rng is None: + rng = np.random.default_rng() + return rng.normal(loc=self.mu, scale=self.sigma, size=size) + + +Distribution = Annotated[ + FixedDistribution | NormalDistribution, Field(discriminator="distribution") +] From 3413d855a573cffe84fbf332dd88929767e2e162 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Mon, 2 Jun 2025 12:10:45 -0400 Subject: [PATCH 02/73] Add `UniformDistribution` Added a representation for uniform distributions to `gempyor.distributions`. --- .../gempyor_pkg/src/gempyor/distributions.py | 51 +++++++++++++++++-- 1 file changed, 48 insertions(+), 3 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 380b302a9..094c3d78f 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -1,20 +1,22 @@ """Representations of distributions used for modifiers, likelihoods, etc.""" __all__: tuple[str, ...] = ( + "Distribution", "DistributionABC", "FixedDistribution", "NormalDistribution", - "Distribution", + "UniformDistribution", ) from abc import ABC, abstractmethod +from math import isclose from typing import Annotated, Literal import numpy as np from numpy.random import Generator import numpy.typing as npt -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator class DistributionABC(ABC, BaseModel): @@ -87,6 +89,49 @@ def sample( return rng.normal(loc=self.mu, scale=self.sigma, size=size) +class UniformDistribution(DistributionABC): + """ + Represents a uniform distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import UniformDistribution + >>> rng = np.random.default_rng(42) + >>> dist = UniformDistribution(low=-0.5, high=1.5) + >>> dist + UniformDistribution(distribution='uniform', low=-0.5, high=1.5) + >>> dist.sample(rng=rng) + array([1.0479121]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[ 0.37775688, 1.21719584, 0.89473606, -0.3116453 , 1.4512447 ], + [ 1.0222794 , 1.07212861, -0.24377273, 0.40077188, 0.24159605], + [ 1.35352998, 0.78773024, 1.14552323, 0.3868284 , -0.04552256]]) + """ + + distribution: Literal["uniform"] = "uniform" + low: float + high: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: + """Sample from the uniform distribution.""" + if rng is None: + rng = np.random.default_rng() + return rng.uniform(low=self.low, high=self.high, size=size) + + @model_validator(mode="after") + def _ensure_high_greater_than_low(self) -> "UniformDistribution": + """Ensure that high is greater than low.""" + if self.high < self.low or isclose(self.high, self.low): + raise ValueError( + f"The 'high' value, {self.high}, must be " + f"greater than the 'low' value, {self.low}." + ) + return self + + Distribution = Annotated[ - FixedDistribution | NormalDistribution, Field(discriminator="distribution") + FixedDistribution | NormalDistribution | UniformDistribution, + Field(discriminator="distribution"), ] From 33bc45662747ef6e0d110a65dd9b79b61e45e36b Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Mon, 2 Jun 2025 12:16:22 -0400 Subject: [PATCH 03/73] Add `PoissonDistribution` Added a representation of a poisson distribution to `gempyor.distributions` --- .../gempyor_pkg/src/gempyor/distributions.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 094c3d78f..387ef4a4c 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -5,6 +5,7 @@ "DistributionABC", "FixedDistribution", "NormalDistribution", + "PoissonDistribution", "UniformDistribution", ) @@ -131,7 +132,38 @@ def _ensure_high_greater_than_low(self) -> "UniformDistribution": return self +class PoissonDistribution(DistributionABC): + """ + Represents a Poisson distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import PoissonDistribution + >>> rng = np.random.default_rng(42) + >>> dist = PoissonDistribution(lam=3.0) + >>> dist + PoissonDistribution(distribution='poisson', lam=3.0) + >>> dist.sample(rng=rng) + array([4]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[4, 5, 1, 7, 1], + [4, 2, 2, 5, 4], + [1, 6, 2, 5, 0]]) + """ + + distribution: Literal["poisson"] = "poisson" + lam: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.int64]: + """Sample from the Poisson distribution.""" + if rng is None: + rng = np.random.default_rng() + return rng.poisson(lam=self.lam, size=size) + + Distribution = Annotated[ - FixedDistribution | NormalDistribution | UniformDistribution, + FixedDistribution | NormalDistribution | PoissonDistribution | UniformDistribution, Field(discriminator="distribution"), ] From bf06e4a7082c35aae01a417f19d73a1888c57510 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Mon, 2 Jun 2025 13:04:04 -0400 Subject: [PATCH 04/73] Add `BinomialDistribution` Added a representation for binomial distribution to `gempyor.distributions`. --- .../gempyor_pkg/src/gempyor/distributions.py | 39 ++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 387ef4a4c..7043f0efd 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -1,6 +1,7 @@ """Representations of distributions used for modifiers, likelihoods, etc.""" __all__: tuple[str, ...] = ( + "BinomialDistribution", "Distribution", "DistributionABC", "FixedDistribution", @@ -163,7 +164,43 @@ def sample( return rng.poisson(lam=self.lam, size=size) +class BinomialDistribution(DistributionABC): + """ + Represents a binomial distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import BinomialDistribution + >>> rng = np.random.default_rng(42) + >>> dist = BinomialDistribution(n=10, p=0.5) + >>> dist + BinomialDistribution(distribution='binomial', n=10, p=0.5) + >>> dist.sample(rng=rng) + array([6]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[5, 7, 6, 3, 8], + [6, 6, 3, 5, 4], + [7, 6, 6, 5, 4]]) + """ + + distribution: Literal["binomial"] = "binomial" + n: int + p: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.int64]: + """Sample from the binomial distribution.""" + if rng is None: + rng = np.random.default_rng() + return rng.binomial(n=self.n, p=self.p, size=size) + + Distribution = Annotated[ - FixedDistribution | NormalDistribution | PoissonDistribution | UniformDistribution, + BinomialDistribution + | FixedDistribution + | NormalDistribution + | PoissonDistribution + | UniformDistribution, Field(discriminator="distribution"), ] From 44d02c528c1d5bacb4029700cdb02e39b747209e Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Mon, 2 Jun 2025 14:27:41 -0400 Subject: [PATCH 05/73] Add `LognormalDistribution` Added a lognormal distribution to `gempyor.distributions`. --- .../gempyor_pkg/src/gempyor/distributions.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 7043f0efd..29166d41d 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -5,6 +5,7 @@ "Distribution", "DistributionABC", "FixedDistribution", + "LognormalDistribution", "NormalDistribution", "PoissonDistribution", "UniformDistribution", @@ -133,6 +134,38 @@ def _ensure_high_greater_than_low(self) -> "UniformDistribution": return self +class LognormalDistribution(DistributionABC): + """ + Represents a Lognormal distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import LognormalDistribution + >>> rng = np.random.default_rng(42) + >>> dist = LognormalDistribution(meanlog=0.0, sdlog=1.0) + >>> dist + LognormalDistribution(distribution='lognorm', meanlog=0.0, sdlog=1.0) + >>> dist.sample(rng=rng) + array([1.35624124]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[0.3534603 , 2.11795541, 2.56142749, 0.14212687, 0.27193845], + [1.13637163, 0.72888261, 0.98333919, 0.42611589, 2.40944872], + [2.17666075, 1.06825951, 3.08712799, 1.59601411, 0.42346159]]) + """ + + distribution: Literal["lognorm"] = "lognorm" + meanlog: float + sdlog: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.int64]: + """Sample from the Lognormal distribution.""" + if rng is None: + rng = np.random.default_rng() + return rng.lognormal(mean=self.meanlog, sigma=self.sdlog, size=size) + + class PoissonDistribution(DistributionABC): """ Represents a Poisson distribution. @@ -199,6 +232,7 @@ def sample( Distribution = Annotated[ BinomialDistribution | FixedDistribution + | LognormalDistribution | NormalDistribution | PoissonDistribution | UniformDistribution, From 1ebc8ac60a4eadba15114a1128c5c1f003978a5f Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Mon, 2 Jun 2025 14:47:45 -0400 Subject: [PATCH 06/73] Add `TruncatedNormalDistribution` Added a truncated normal distribution to `gempyor.distributions`. --- .../gempyor_pkg/src/gempyor/distributions.py | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 29166d41d..c985d84c3 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -8,6 +8,7 @@ "LognormalDistribution", "NormalDistribution", "PoissonDistribution", + "TruncatedNormalDistribution", "UniformDistribution", ) @@ -20,6 +21,7 @@ from numpy.random import Generator import numpy.typing as npt from pydantic import BaseModel, Field, model_validator +from scipy.stats import truncnorm class DistributionABC(ABC, BaseModel): @@ -166,6 +168,51 @@ def sample( return rng.lognormal(mean=self.meanlog, sigma=self.sdlog, size=size) +class TruncatedNormalDistribution(DistributionABC): + # pylint: disable=line-too-long + """ + Represents a truncated normal distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import TruncatedNormalDistribution + >>> rng = np.random.default_rng(42) + >>> dist = TruncatedNormalDistribution(mean=1.0, sd=1.0, a=0.0, b=10.0) + >>> dist + TruncatedNormalDistribution(distribution='truncnorm', mean=1.0, sd=1.0, a=0.0, b=10.0) + >>> dist.sample(rng=rng) + array([1.87722989]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[1.07000038, 2.18016199, 1.66002835, 0.28689654, 3.04332767], + [1.83818339, 1.9153892 , 0.37639339, 1.09435167, 0.92629918], + [2.54134935, 1.52545861, 2.04022114, 1.07959285, 0.6142512 ]]) + """ + # pylint: enable=line-too-long + + distribution: Literal["truncnorm"] = "truncnorm" + mean: float + sd: float + a: float + b: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: + """Sample from the truncated normal distribution.""" + if rng is None: + rng = np.random.default_rng() + lower = (self.a - self.mean) / self.sd + upper = (self.b - self.mean) / self.sd + return truncnorm.rvs( + a=lower, + b=upper, + loc=self.mean, + scale=self.sd, + size=size, + random_state=rng, + ) + + class PoissonDistribution(DistributionABC): """ Represents a Poisson distribution. @@ -235,6 +282,7 @@ def sample( | LognormalDistribution | NormalDistribution | PoissonDistribution + | TruncatedNormalDistribution | UniformDistribution, Field(discriminator="distribution"), ] From d5739544bebf29e91352a73176a56f3be41cfff6 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 20 Jun 2025 13:11:40 -0400 Subject: [PATCH 07/73] Add GammaDistribution Added a representation for gamma distributions to `gempyor.distributions` --- .../gempyor_pkg/src/gempyor/distributions.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index c985d84c3..dc6f1392d 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -5,6 +5,7 @@ "Distribution", "DistributionABC", "FixedDistribution", + "GammaDistribution", "LognormalDistribution", "NormalDistribution", "PoissonDistribution", @@ -274,11 +275,44 @@ def sample( if rng is None: rng = np.random.default_rng() return rng.binomial(n=self.n, p=self.p, size=size) + + +class GammaDistribution(DistributionABC): + """ + Represents a gamma distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import GammaDistribution + >>> rng = np.random.default_rng(42) + >>> dist = GammaDistribution(shape=2.0, scale=1.5) + >>> dist + GammaDistribution(distribution='gamma', shape=2.0, scale=1.5) + >>> dist.sample(rng=rng) + array([2.78453141]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[1.0116844 , 5.1632733 , 6.51682397, 0.49053229, 0.82522731], + [2.45524838, 1.4870423 , 1.95460596, 0.94191942, 5.92679803], + [5.08819234, 2.21976694, 7.82570081, 3.42761858, 1.1070266 ]]) + """ + + distribution: Literal["gamma"] = "gamma" + shape: float + scale: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: + """Sample from the gamma distribution.""" + if rng is None: + rng = np.random.default_rng() + return rng.gamma(shape=self.shape, scale=self.scale, size=size) Distribution = Annotated[ BinomialDistribution | FixedDistribution + | GammaDistribution | LognormalDistribution | NormalDistribution | PoissonDistribution From 0ce1aa05d4963ed64d2a78e9d18cc9fc51838806 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 20 Jun 2025 13:13:33 -0400 Subject: [PATCH 08/73] Add `WeibullDistribution` Added a representation for Weibull distributions to `gempyor.distributions`. Also, fixed a return type hint error. --- .../gempyor_pkg/src/gempyor/distributions.py | 39 ++++++++++++++++++- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index dc6f1392d..e8082ba9a 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -11,6 +11,7 @@ "PoissonDistribution", "TruncatedNormalDistribution", "UniformDistribution", + "WeibullDistribution" ) @@ -162,7 +163,7 @@ class LognormalDistribution(DistributionABC): def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None - ) -> npt.NDArray[np.int64]: + ) -> npt.NDArray[np.float64]: """Sample from the Lognormal distribution.""" if rng is None: rng = np.random.default_rng() @@ -308,6 +309,39 @@ def sample( rng = np.random.default_rng() return rng.gamma(shape=self.shape, scale=self.scale, size=size) +class WeibullDistribution(DistributionABC): + """ + Represents a weibull distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import WeibullDistribution + >>> rng = np.random.default_rng(42) + >>> dist = WeibullDistribution(shape=2.5, scale=5.0) + >>> dist + WeibullDistribution(distribution='weibull', shape=2.5, scale=5.0) + >>> dist.sample(rng=rng) + array([5.18664161]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[4.09267151, 6.25764353, 6.57560416, 2.72348545, 3.73147743], + [5.08637841, 4.63044996, 4.96639599, 3.99658252, 6.42564251], + [6.29525049, 5.03450379, 6.88371131, 5.65602432, 3.99307222]]) + """ + + + distribution: Literal["weibull"] = "weibull" + shape: float + scale: float + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: + """Sample from the Weibull distribution.""" + if rng is None: + rng = np.random.default_rng() + # Multiply by scale b/c rng.weibull assumes standard weibull dist (scale of 1) + return self.scale * rng.weibull(a=self.shape, size=size) + Distribution = Annotated[ BinomialDistribution @@ -317,6 +351,7 @@ def sample( | NormalDistribution | PoissonDistribution | TruncatedNormalDistribution - | UniformDistribution, + | UniformDistribution + | WeibullDistribution, Field(discriminator="distribution"), ] From 565899343f19c1c7f711afc9982a8cf605ce6800 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 20 Jun 2025 15:45:02 -0400 Subject: [PATCH 09/73] Add tests for `NormalDistribution` Unit tests and validation for the `NormalDistribution` class in `gempyor_pkg/tests/distributions/` --- .../gempyor_pkg/src/gempyor/distributions.py | 18 ++--- .../distributions/test_NormalDistribution.py | 78 +++++++++++++++++++ 2 files changed, 87 insertions(+), 9 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index e8082ba9a..8995a5a83 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -11,7 +11,7 @@ "PoissonDistribution", "TruncatedNormalDistribution", "UniformDistribution", - "WeibullDistribution" + "WeibullDistribution", ) @@ -85,7 +85,7 @@ class NormalDistribution(DistributionABC): distribution: Literal["norm"] = "norm" mu: float - sigma: float + sigma: float = Field(..., gt=0) def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None @@ -276,11 +276,11 @@ def sample( if rng is None: rng = np.random.default_rng() return rng.binomial(n=self.n, p=self.p, size=size) - + class GammaDistribution(DistributionABC): """ - Represents a gamma distribution. + Represents a gamma distribution. Examples: >>> import numpy as np @@ -302,16 +302,17 @@ class GammaDistribution(DistributionABC): scale: float def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the gamma distribution.""" if rng is None: rng = np.random.default_rng() return rng.gamma(shape=self.shape, scale=self.scale, size=size) + class WeibullDistribution(DistributionABC): """ - Represents a weibull distribution. + Represents a weibull distribution. Examples: >>> import numpy as np @@ -328,20 +329,19 @@ class WeibullDistribution(DistributionABC): [6.29525049, 5.03450379, 6.88371131, 5.65602432, 3.99307222]]) """ - distribution: Literal["weibull"] = "weibull" shape: float scale: float def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the Weibull distribution.""" if rng is None: rng = np.random.default_rng() # Multiply by scale b/c rng.weibull assumes standard weibull dist (scale of 1) return self.scale * rng.weibull(a=self.shape, size=size) - + Distribution = Annotated[ BinomialDistribution diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py new file mode 100644 index 000000000..f0637137d --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py @@ -0,0 +1,78 @@ +import numpy as np +from pydantic import ValidationError + +import pytest + +from gempyor.distributions import NormalDistribution + + +@pytest.mark.parametrize("mu", [-100.0, 0.0, 50.5], ids=["mu_neg", "mu_zero", "mu_pos"]) +@pytest.mark.parametrize( + "sigma", [10.0, 1.0, 0.5], ids=["sigma_ten", "sigma_one", "sigma_pointfive"] +) +def test_normal_distribution_init(mu: float, sigma: float) -> None: + """ + Test that a NormalDistribution object can be initialized with given parameters. + """ + + dist = NormalDistribution(mu=mu, sigma=sigma) + assert dist.mu == mu + assert dist.sigma == sigma + assert dist.distribution == "norm" + + +@pytest.mark.parametrize( + "invalid_sigma", [0.0, -0.5, -50.5], ids=["sigma_zero", "sigma_neg", "sigma_larger_neg"] +) +def test_normal_distribution_sample_raises_error_for_invalid_sigma( + invalid_sigma: float, +) -> None: + """ + Test that initializing a NormalDistribution object with an + invalid sigma raises a ValidationError. + """ + + with pytest.raises(ValidationError, match="Input should be greater than 0") as e: + NormalDistribution(mu=10.0, sigma=invalid_sigma) + + +def test_normal_distribution_sample_rng_reproducibility() -> None: + """ + Test that sampling is reproducible when a seed is given to RNG. + """ + + dist = NormalDistribution(mu=10.0, sigma=1.5) + + rng1 = np.random.default_rng(seed=100) + sample1 = dist.sample(size=(2, 2), rng=rng1) + + rng2 = np.random.default_rng(seed=100) + sample2 = dist.sample(size=(2, 2), rng=rng2) + + assert np.array_equal(sample1, sample2) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((3, 2), (3, 2)), # Case 1: size is not provided, should default to a single value + (10, (10,)), # Case 2: size is an integer + ((3, 4), (3, 4)), # Case 3: size is a tuple + ], + ids=["tuple_size1", "integer_size", "tuple_size2"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_normal_distribution_sample_properties(size, expected_shape, use_rng) -> None: + """ + Test that .sample() method produces a np.array of expected size wtih float64s. + """ + + dist = NormalDistribution(mu=0.0, sigma=1.0) + kwargs = {} + kwargs["size"] = size + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 From 65093e397b6fa71a03f7da3f229a472e0185f4d1 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 20 Jun 2025 16:16:02 -0400 Subject: [PATCH 10/73] Add tests for `FixedDistribution` Unit tests and validation for the `FixedDistribution` class in `gmepyor_pkg/tests/distributions/`. Also removal of an unused parameter in the `FixedDistribution` `.sample()` method. --- .../gempyor_pkg/src/gempyor/distributions.py | 2 +- .../distributions/test_FixedDistribution.py | 46 +++++++++++++++++++ .../distributions/test_NormalDistribution.py | 26 ++--------- 3 files changed, 51 insertions(+), 23 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 8995a5a83..69e820f1e 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -58,7 +58,7 @@ class FixedDistribution(DistributionABC): value: float def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + self, size: int | tuple[int, ...] = 1 ) -> npt.NDArray[np.float64]: """Sample from the fixed distribution.""" return np.full(size, self.value) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py new file mode 100644 index 000000000..951831ae6 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py @@ -0,0 +1,46 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import FixedDistribution + + +@pytest.mark.parametrize("value", [123.45, -100.0, 0.0], ids=["positive", "negative", "zero"]) +def test_fixed_distribution_ini(value: float ) -> None: + dist = FixedDistribution(value=value) + assert dist.value == value + assert dist.distribution == "fixed" + + +@pytest.mark.parametrize( + "value", + [5.5, -1.2], + ids=["value_positive", "value_negative"] +) +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((2, 5), (2, 5)), + (5, (5,)), + ((2, 3), (2, 3)), + ], + ids=["default_size", "integer_size", "tuple_size"] +) +@pytest.mark.parametrize( + "use_rng", + [True, False], + ids=["with_rng", "without_rng"] +) +def test_fixed_distribution_sample_properties( + value: float, size, expected_shape, use_rng +) -> None: + dist = FixedDistribution(value=value) + kwargs = {} + kwargs['size'] = size + if use_rng: + kwargs['rng'] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample == value) \ No newline at end of file diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py index f0637137d..b4902b6ed 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py @@ -1,7 +1,6 @@ import numpy as np -from pydantic import ValidationError - import pytest +from pydantic import ValidationError from gempyor.distributions import NormalDistribution @@ -11,10 +10,6 @@ "sigma", [10.0, 1.0, 0.5], ids=["sigma_ten", "sigma_one", "sigma_pointfive"] ) def test_normal_distribution_init(mu: float, sigma: float) -> None: - """ - Test that a NormalDistribution object can be initialized with given parameters. - """ - dist = NormalDistribution(mu=mu, sigma=sigma) assert dist.mu == mu assert dist.sigma == sigma @@ -27,20 +22,11 @@ def test_normal_distribution_init(mu: float, sigma: float) -> None: def test_normal_distribution_sample_raises_error_for_invalid_sigma( invalid_sigma: float, ) -> None: - """ - Test that initializing a NormalDistribution object with an - invalid sigma raises a ValidationError. - """ - with pytest.raises(ValidationError, match="Input should be greater than 0") as e: NormalDistribution(mu=10.0, sigma=invalid_sigma) def test_normal_distribution_sample_rng_reproducibility() -> None: - """ - Test that sampling is reproducible when a seed is given to RNG. - """ - dist = NormalDistribution(mu=10.0, sigma=1.5) rng1 = np.random.default_rng(seed=100) @@ -55,18 +41,14 @@ def test_normal_distribution_sample_rng_reproducibility() -> None: @pytest.mark.parametrize( "size, expected_shape", [ - ((3, 2), (3, 2)), # Case 1: size is not provided, should default to a single value - (10, (10,)), # Case 2: size is an integer - ((3, 4), (3, 4)), # Case 3: size is a tuple + ((3, 2), (3, 2)), + (10, (10,)), + ((3, 4), (3, 4)), ], ids=["tuple_size1", "integer_size", "tuple_size2"], ) @pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) def test_normal_distribution_sample_properties(size, expected_shape, use_rng) -> None: - """ - Test that .sample() method produces a np.array of expected size wtih float64s. - """ - dist = NormalDistribution(mu=0.0, sigma=1.0) kwargs = {} kwargs["size"] = size From c530c97a9689edb09ab53f278ededc2f53a1c6bc Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 20 Jun 2025 16:33:54 -0400 Subject: [PATCH 11/73] Add tests for `UniformDistribution` Unit tests and validation for the `UniformDistribution` class in `gempyor_pkg/tests/distributions/`. Also linting from a previous commit I forgot to add. --- .../gempyor_pkg/src/gempyor/distributions.py | 4 +- .../distributions/test_FixedDistribution.py | 33 ++++---- .../distributions/test_NormalDistribution.py | 6 +- .../distributions/test_UniformDistribution.py | 77 +++++++++++++++++++ 4 files changed, 94 insertions(+), 26 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 69e820f1e..9cbe0ea18 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -57,9 +57,7 @@ class FixedDistribution(DistributionABC): distribution: Literal["fixed"] = "fixed" value: float - def sample( - self, size: int | tuple[int, ...] = 1 - ) -> npt.NDArray[np.float64]: + def sample(self, size: int | tuple[int, ...] = 1) -> npt.NDArray[np.float64]: """Sample from the fixed distribution.""" return np.full(size, self.value) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py index 951831ae6..5dce35533 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py @@ -1,46 +1,39 @@ import numpy as np import pytest -from pydantic import ValidationError from gempyor.distributions import FixedDistribution -@pytest.mark.parametrize("value", [123.45, -100.0, 0.0], ids=["positive", "negative", "zero"]) -def test_fixed_distribution_ini(value: float ) -> None: +@pytest.mark.parametrize( + "value", [123.45, -100.0, 0.0], ids=["positive", "negative", "zero"] +) +def test_fixed_distribution_ini(value: float) -> None: dist = FixedDistribution(value=value) assert dist.value == value assert dist.distribution == "fixed" -@pytest.mark.parametrize( - "value", - [5.5, -1.2], - ids=["value_positive", "value_negative"] -) +@pytest.mark.parametrize("value", [5.5, -1.2], ids=["value_positive", "value_negative"]) @pytest.mark.parametrize( "size, expected_shape", [ - ((2, 5), (2, 5)), - (5, (5,)), - ((2, 3), (2, 3)), + ((2, 5), (2, 5)), + (5, (5,)), + ((2, 3), (2, 3)), ], - ids=["default_size", "integer_size", "tuple_size"] -) -@pytest.mark.parametrize( - "use_rng", - [True, False], - ids=["with_rng", "without_rng"] + ids=["default_size", "integer_size", "tuple_size"], ) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) def test_fixed_distribution_sample_properties( value: float, size, expected_shape, use_rng ) -> None: dist = FixedDistribution(value=value) kwargs = {} - kwargs['size'] = size + kwargs["size"] = size if use_rng: - kwargs['rng'] = np.random.default_rng() + kwargs["rng"] = np.random.default_rng() sample = dist.sample(**kwargs) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 - assert np.all(sample == value) \ No newline at end of file + assert np.all(sample == value) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py index b4902b6ed..98af7723a 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py @@ -41,9 +41,9 @@ def test_normal_distribution_sample_rng_reproducibility() -> None: @pytest.mark.parametrize( "size, expected_shape", [ - ((3, 2), (3, 2)), - (10, (10,)), - ((3, 4), (3, 4)), + ((3, 2), (3, 2)), + (10, (10,)), + ((3, 4), (3, 4)), ], ids=["tuple_size1", "integer_size", "tuple_size2"], ) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py new file mode 100644 index 000000000..c03614bcf --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py @@ -0,0 +1,77 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import UniformDistribution + + +@pytest.mark.parametrize( + "low, high", + [ + (0.0, 10.0), + (-10.0, -5.0), + (-5.0, 5.0), + ], + ids=["positive_range", "negative_range", "spanning_range"], +) +def test_uniform_distribution_init_valid(low: float, high: float) -> None: + dist = UniformDistribution(low=low, high=high) + assert dist.low == low + assert dist.high == high + assert dist.distribution == "uniform" + + +@pytest.mark.parametrize( + "low, high", + [ + (10.0, 0.0), + (5.0, 5.0), + ], + ids=["high_less_than_low", "high_equals_low"], +) +def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> None: + """ + Tests that creating a UniformDistribution with invalid bounds (high <= low) + raises a Pydantic ValidationError, triggering the custom validator. + """ + # The custom validator raises a ValueError, which Pydantic wraps. + # We can match the custom error message. + with pytest.raises(ValidationError, match="must be greater than the 'low' value"): + UniformDistribution(low=low, high=high) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((6, 4), (6, 4)), # Case for the default size + (20, (20,)), # Case for an integer size + ((4, 5), (4, 5)), # Case for a tuple size + ], + ids=["default_size", "integer_size", "tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_uniform_distribution_sample_properties(size, expected_shape, use_rng) -> None: + low, high = 10.0, 20.0 + dist = UniformDistribution(low=low, high=high) + kwargs = {} + kwargs["size"] = size + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample >= low) + assert np.all(sample < high) + + +def test_uniform_distribution_sample_rng_reproducibility() -> None: + dist = UniformDistribution(low=-10.0, high=10.0) + + rng1 = np.random.default_rng(seed=123) + sample1 = dist.sample(size=(3, 3), rng=rng1) + + rng2 = np.random.default_rng(seed=123) + sample2 = dist.sample(size=(3, 3), rng=rng2) + + assert np.array_equal(sample1, sample2) From dd21f0495777954b3a3a794b7ddd71dc1c666a51 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 20 Jun 2025 16:52:35 -0400 Subject: [PATCH 12/73] Add tests for `LognormalDistribution` Unit tests and validation for the `LognormalDistribution` class in `gempyor_pkg/tests/distributions/` + a fix for an issue that I created in `2c5fce5` --- .../gempyor_pkg/src/gempyor/distributions.py | 6 +- .../distributions/test_FixedDistribution.py | 2 +- .../test_LognormalDistribution.py | 63 +++++++++++++++++++ .../distributions/test_NormalDistribution.py | 2 +- .../distributions/test_UniformDistribution.py | 12 +--- 5 files changed, 72 insertions(+), 13 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 9cbe0ea18..8b2335d59 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -57,7 +57,9 @@ class FixedDistribution(DistributionABC): distribution: Literal["fixed"] = "fixed" value: float - def sample(self, size: int | tuple[int, ...] = 1) -> npt.NDArray[np.float64]: + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: """Sample from the fixed distribution.""" return np.full(size, self.value) @@ -157,7 +159,7 @@ class LognormalDistribution(DistributionABC): distribution: Literal["lognorm"] = "lognorm" meanlog: float - sdlog: float + sdlog: float = Field(..., gt=0) def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None diff --git a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py index 5dce35533..de4247359 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py @@ -21,7 +21,7 @@ def test_fixed_distribution_ini(value: float) -> None: (5, (5,)), ((2, 3), (2, 3)), ], - ids=["default_size", "integer_size", "tuple_size"], + ids=["tuple_size1", "integer_size", "tuple_size2"], ) @pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) def test_fixed_distribution_sample_properties( diff --git a/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py new file mode 100644 index 000000000..39371c5da --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py @@ -0,0 +1,63 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import LognormalDistribution + + +@pytest.mark.parametrize( + "meanlog", [-10.0, 0.0, 5.0], ids=["meanlog_neg", "meanlog_zero", "meanlog_pos"] +) +@pytest.mark.parametrize( + "sdlog", [0.1, 1.0, 10.0], ids=["sdlog_small", "sdlog_one", "sdlog_large"] +) +def test_lognormal_distribution_init_valid(meanlog: float, sdlog: float) -> None: + dist = LognormalDistribution(meanlog=meanlog, sdlog=sdlog) + assert dist.meanlog == meanlog + assert dist.sdlog == sdlog + assert dist.distribution == "lognorm" + + +@pytest.mark.parametrize( + "invalid_sdlog", [0.0, -0.1, -10.0], ids=["zero", "small_negative", "large_negative"] +) +def test_lognormal_distribution_init_raises_error_for_invalid_sdlog( + invalid_sdlog: float, +) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + LognormalDistribution(meanlog=0.0, sdlog=invalid_sdlog) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((2, 3), (2, 3)), + (25, (25,)), + ((5, 5), (5, 5)), + ], + ids=["tuple_size1", "integer_size", "tuple_size2"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_lognormal_distribution_sample_properties(size, expected_shape, use_rng) -> None: + dist = LognormalDistribution(meanlog=0.0, sdlog=1.0) + kwargs = {} + kwargs["size"] = size + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample > 0) + + +def test_lognormal_distribution_sample_rng_reproducibility() -> None: + dist = LognormalDistribution(meanlog=2.0, sdlog=0.5) + + rng1 = np.random.default_rng(seed=42) + sample1 = dist.sample(size=(3, 3), rng=rng1) + + rng2 = np.random.default_rng(seed=42) + sample2 = dist.sample(size=(3, 3), rng=rng2) + + np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py index 98af7723a..adafa359e 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py @@ -22,7 +22,7 @@ def test_normal_distribution_init(mu: float, sigma: float) -> None: def test_normal_distribution_sample_raises_error_for_invalid_sigma( invalid_sigma: float, ) -> None: - with pytest.raises(ValidationError, match="Input should be greater than 0") as e: + with pytest.raises(ValidationError, match="Input should be greater than 0"): NormalDistribution(mu=10.0, sigma=invalid_sigma) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py index c03614bcf..9ea07194f 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py @@ -30,12 +30,6 @@ def test_uniform_distribution_init_valid(low: float, high: float) -> None: ids=["high_less_than_low", "high_equals_low"], ) def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> None: - """ - Tests that creating a UniformDistribution with invalid bounds (high <= low) - raises a Pydantic ValidationError, triggering the custom validator. - """ - # The custom validator raises a ValueError, which Pydantic wraps. - # We can match the custom error message. with pytest.raises(ValidationError, match="must be greater than the 'low' value"): UniformDistribution(low=low, high=high) @@ -43,9 +37,9 @@ def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> No @pytest.mark.parametrize( "size, expected_shape", [ - ((6, 4), (6, 4)), # Case for the default size - (20, (20,)), # Case for an integer size - ((4, 5), (4, 5)), # Case for a tuple size + ((6, 4), (6, 4)), + (20, (20,)), + ((4, 5), (4, 5)), ], ids=["default_size", "integer_size", "tuple_size"], ) From 0f0c8ebd9c9172e02b5d8e07c8cd984b0b1141dc Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 23 Jun 2025 10:53:57 -0400 Subject: [PATCH 13/73] Add tests for `TruncatedNormalDistribution` Add tests for the `TruncatedNormalDistribution` class from `gempyor.distributions`. Also, formatting fixes and bound fixes. --- .../gempyor_pkg/src/gempyor/distributions.py | 12 +++- .../distributions/test_NormalDistribution.py | 3 - .../test_TruncatedNormalDistribution.py | 69 +++++++++++++++++++ .../distributions/test_UniformDistribution.py | 3 - 4 files changed, 80 insertions(+), 7 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_TruncatedNormalDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 8b2335d59..65166cb50 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -193,7 +193,7 @@ class TruncatedNormalDistribution(DistributionABC): distribution: Literal["truncnorm"] = "truncnorm" mean: float - sd: float + sd: float = Field(..., gt=0) a: float b: float @@ -214,6 +214,16 @@ def sample( random_state=rng, ) + @model_validator(mode="after") + def _ensure_b_greater_than_a(self) -> "TruncatedNormalDistribution": + """Ensure that the upper bound 'b' is strictly greater than the lower bound 'a'.""" + if self.b < self.a or isclose(self.b, self.a): + raise ValueError( + f"The upper bound 'b' ({self.b}) must be strictly " + f"greater than the lower bound 'a' ({self.a})." + ) + return self + class PoissonDistribution(DistributionABC): """ diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py index adafa359e..4ec871025 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py @@ -28,13 +28,10 @@ def test_normal_distribution_sample_raises_error_for_invalid_sigma( def test_normal_distribution_sample_rng_reproducibility() -> None: dist = NormalDistribution(mu=10.0, sigma=1.5) - rng1 = np.random.default_rng(seed=100) sample1 = dist.sample(size=(2, 2), rng=rng1) - rng2 = np.random.default_rng(seed=100) sample2 = dist.sample(size=(2, 2), rng=rng2) - assert np.array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_TruncatedNormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_TruncatedNormalDistribution.py new file mode 100644 index 000000000..107def8c3 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_TruncatedNormalDistribution.py @@ -0,0 +1,69 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import TruncatedNormalDistribution + + +@pytest.mark.parametrize( + "mean, sd, a, b", + [ + (0.0, 1.0, -2.0, 2.0), + (10.0, 2.0, 5.0, 15.0), + (-5.0, 0.5, -6.0, -4.0), + (0.0, 10.0, -1.0, 1.0), + ], + ids=["standard", "shifted_positive", "shifted_negative", "wide_sd"], +) +def test_truncated_normal_distribution_init_valid( + mean: float, sd: float, a: float, b: float +) -> None: + dist = TruncatedNormalDistribution(mean=mean, sd=sd, a=a, b=b) + assert dist.mean == mean + assert dist.sd == sd + assert dist.a == a + assert dist.b == b + assert dist.distribution == "truncnorm" + + +@pytest.mark.parametrize( + "invalid_sd", [0.0, -0.1, -10.0], ids=["zero", "small_negative", "large_negative"] +) +def test_truncated_normal_distribution_init_invalid_sd(invalid_sd: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + TruncatedNormalDistribution(mean=0.0, sd=invalid_sd, a=0.0, b=1.0) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((3, 5), (3, 5)), + (10, (10,)), + ((2, 3, 4), (2, 3, 4)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_truncated_normal_distribution_sample_properties( + size, expected_shape, use_rng +) -> None: + a, b = 0.0, 10.0 + dist = TruncatedNormalDistribution(mean=5.0, sd=2.0, a=a, b=b) + kwargs = {"size": size} + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample >= a) + assert np.all(sample <= b) + + +def test_truncated_normal_distribution_sample_rng_reproducibility() -> None: + dist = TruncatedNormalDistribution(mean=10.0, sd=3.0, a=5.0, b=15.0) + rng1 = np.random.default_rng(seed=100) + sample1 = dist.sample(size=(4, 4), rng=rng1) + rng2 = np.random.default_rng(seed=100) + sample2 = dist.sample(size=(4, 4), rng=rng2) + np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py index 9ea07194f..ef47b2306 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py @@ -61,11 +61,8 @@ def test_uniform_distribution_sample_properties(size, expected_shape, use_rng) - def test_uniform_distribution_sample_rng_reproducibility() -> None: dist = UniformDistribution(low=-10.0, high=10.0) - rng1 = np.random.default_rng(seed=123) sample1 = dist.sample(size=(3, 3), rng=rng1) - rng2 = np.random.default_rng(seed=123) sample2 = dist.sample(size=(3, 3), rng=rng2) - assert np.array_equal(sample1, sample2) From fa715621ccef47a9f823eeb5b2f9aabeaa1da3a8 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 23 Jun 2025 11:06:05 -0400 Subject: [PATCH 14/73] Add tests for `PoissonDistribution` Unit tests and validation for `PoissonDIstribution` class in `gempyor.distributions`. Also, formatting and bounds fixes. --- .../gempyor_pkg/src/gempyor/distributions.py | 2 +- .../distributions/test_FixedDistribution.py | 2 +- .../test_LognormalDistribution.py | 9 +-- .../distributions/test_NormalDistribution.py | 20 +++---- .../distributions/test_PoissonDistribution.py | 60 +++++++++++++++++++ .../distributions/test_UniformDistribution.py | 6 +- 6 files changed, 78 insertions(+), 21 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 65166cb50..06c28c2a9 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -245,7 +245,7 @@ class PoissonDistribution(DistributionABC): """ distribution: Literal["poisson"] = "poisson" - lam: float + lam: float = Field(..., ge=0) def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None diff --git a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py index de4247359..f46553323 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py @@ -19,7 +19,7 @@ def test_fixed_distribution_ini(value: float) -> None: [ ((2, 5), (2, 5)), (5, (5,)), - ((2, 3), (2, 3)), + ((2, 3, 4), (2, 3, 4)), ], ids=["tuple_size1", "integer_size", "tuple_size2"], ) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py index 39371c5da..3163c524c 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py @@ -33,7 +33,7 @@ def test_lognormal_distribution_init_raises_error_for_invalid_sdlog( [ ((2, 3), (2, 3)), (25, (25,)), - ((5, 5), (5, 5)), + ((2, 3, 4), (2, 3, 4)), ], ids=["tuple_size1", "integer_size", "tuple_size2"], ) @@ -53,11 +53,8 @@ def test_lognormal_distribution_sample_properties(size, expected_shape, use_rng) def test_lognormal_distribution_sample_rng_reproducibility() -> None: dist = LognormalDistribution(meanlog=2.0, sdlog=0.5) - - rng1 = np.random.default_rng(seed=42) + rng1 = np.random.default_rng(seed=100) sample1 = dist.sample(size=(3, 3), rng=rng1) - - rng2 = np.random.default_rng(seed=42) + rng2 = np.random.default_rng(seed=100) sample2 = dist.sample(size=(3, 3), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py index 4ec871025..169c7ac5f 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py @@ -26,21 +26,12 @@ def test_normal_distribution_sample_raises_error_for_invalid_sigma( NormalDistribution(mu=10.0, sigma=invalid_sigma) -def test_normal_distribution_sample_rng_reproducibility() -> None: - dist = NormalDistribution(mu=10.0, sigma=1.5) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(2, 2), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(2, 2), rng=rng2) - assert np.array_equal(sample1, sample2) - - @pytest.mark.parametrize( "size, expected_shape", [ ((3, 2), (3, 2)), (10, (10,)), - ((3, 4), (3, 4)), + ((2, 3, 4), (2, 3, 4)), ], ids=["tuple_size1", "integer_size", "tuple_size2"], ) @@ -55,3 +46,12 @@ def test_normal_distribution_sample_properties(size, expected_shape, use_rng) -> assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 + + +def test_normal_distribution_sample_rng_reproducibility() -> None: + dist = NormalDistribution(mu=10.0, sigma=1.5) + rng1 = np.random.default_rng(seed=100) + sample1 = dist.sample(size=(2, 2), rng=rng1) + rng2 = np.random.default_rng(seed=100) + sample2 = dist.sample(size=(2, 2), rng=rng2) + assert np.array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py new file mode 100644 index 000000000..4429c05c8 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py @@ -0,0 +1,60 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import PoissonDistribution + +@pytest.mark.parametrize( + "valid_lam", + [ + (0.0), + (1.0), + (25.5), + ], + ids=["zero_rate", "unit_rate", "positive_rate"], +) +def test_poisson_distribution_init_valid(valid_lam: float) -> None: + dist = PoissonDistribution(lam=valid_lam) + assert dist.lam == valid_lam + assert dist.distribution == "poisson" + + +@pytest.mark.parametrize( + "invalid_lam", [-0.1, -10.0], ids=["small_negative", "large_negative"] +) +def test_poisson_distribution_init_invalid_lam(invalid_lam: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): + PoissonDistribution(lam=invalid_lam) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((2, 5), (2, 5)), + (15, (15,)), + ((2, 3, 2), (2, 3, 2)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_poisson_distribution_sample_properties( + size, expected_shape, use_rng +) -> None: + dist = PoissonDistribution(lam=5.0) + kwargs = {"size": size} + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.int64 + assert np.all(sample >= 0) + + +def test_poisson_distribution_sample_rng_reproducibility() -> None: + dist = PoissonDistribution(lam=10.0) + rng1 = np.random.default_rng(seed=100) + sample1 = dist.sample(size=(5, 5), rng=rng1) + rng2 = np.random.default_rng(seed=100) + sample2 = dist.sample(size=(5, 5), rng=rng2) + np.testing.assert_array_equal(sample1, sample2) \ No newline at end of file diff --git a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py index ef47b2306..6e52def51 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py @@ -39,7 +39,7 @@ def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> No [ ((6, 4), (6, 4)), (20, (20,)), - ((4, 5), (4, 5)), + ((2, 3, 4), (2, 3, 4)), ], ids=["default_size", "integer_size", "tuple_size"], ) @@ -61,8 +61,8 @@ def test_uniform_distribution_sample_properties(size, expected_shape, use_rng) - def test_uniform_distribution_sample_rng_reproducibility() -> None: dist = UniformDistribution(low=-10.0, high=10.0) - rng1 = np.random.default_rng(seed=123) + rng1 = np.random.default_rng(seed=100) sample1 = dist.sample(size=(3, 3), rng=rng1) - rng2 = np.random.default_rng(seed=123) + rng2 = np.random.default_rng(seed=100) sample2 = dist.sample(size=(3, 3), rng=rng2) assert np.array_equal(sample1, sample2) From bc1edeebed14c97a0831d5837b463ca7d19ec95a Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 23 Jun 2025 11:06:32 -0400 Subject: [PATCH 15/73] Reformatting with `black` --- .../tests/distributions/test_PoissonDistribution.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py index 4429c05c8..df32b06dd 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py @@ -4,6 +4,7 @@ from gempyor.distributions import PoissonDistribution + @pytest.mark.parametrize( "valid_lam", [ @@ -37,9 +38,7 @@ def test_poisson_distribution_init_invalid_lam(invalid_lam: float) -> None: ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) @pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_poisson_distribution_sample_properties( - size, expected_shape, use_rng -) -> None: +def test_poisson_distribution_sample_properties(size, expected_shape, use_rng) -> None: dist = PoissonDistribution(lam=5.0) kwargs = {"size": size} if use_rng: @@ -57,4 +56,4 @@ def test_poisson_distribution_sample_rng_reproducibility() -> None: sample1 = dist.sample(size=(5, 5), rng=rng1) rng2 = np.random.default_rng(seed=100) sample2 = dist.sample(size=(5, 5), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) \ No newline at end of file + np.testing.assert_array_equal(sample1, sample2) From 48c3481d4c96bc50413ad24633f33804a16aeffb Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 23 Jun 2025 11:23:33 -0400 Subject: [PATCH 16/73] Add tests for `BinomialDistribution` Unit tests and validation for the `BinomialDistribution` class in `gempyor.distributions`. Also a bounds fix in `distributions.py`. --- .../gempyor_pkg/src/gempyor/distributions.py | 4 +- .../test_BinomialDistribution.py | 93 +++++++++++++++++++ 2 files changed, 95 insertions(+), 2 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_BinomialDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 06c28c2a9..ff3adf021 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -276,8 +276,8 @@ class BinomialDistribution(DistributionABC): """ distribution: Literal["binomial"] = "binomial" - n: int - p: float + n: int = Field(..., ge=0) + p: float = Field(..., ge=0, le=1) def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None diff --git a/flepimop/gempyor_pkg/tests/distributions/test_BinomialDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_BinomialDistribution.py new file mode 100644 index 000000000..29825fe9f --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_BinomialDistribution.py @@ -0,0 +1,93 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import BinomialDistribution + + +@pytest.mark.parametrize( + "n, p", + [ + (10, 0.5), + (0, 0.5), + (100, 0.0), + (100, 1.0), + (50, 0.25), + ], + ids=[ + "standard_case", + "zero_trials", + "zero_probability", + "full_probability", + "standard_case_2", + ], +) +def test_binomial_distribution_init_valid(n: int, p: float) -> None: + dist = BinomialDistribution(n=n, p=p) + assert dist.n == n + assert dist.p == p + assert dist.distribution == "binomial" + + +@pytest.mark.parametrize( + "invalid_n", + [-1, -10], + ids=["-1", "-10"], +) +def test_binomial_distribution_init_invalid_n(invalid_n: int) -> None: + with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): + BinomialDistribution(n=invalid_n, p=0.5) + + +@pytest.mark.parametrize( + "invalid_p", + [-0.1, -10.0], + ids=["small_negative", "large_negative"], +) +def test_binomial_distribution_init_invalid_p_below_zero(invalid_p: float) -> None: + """Test that initialization fails for probability 'p' less than 0.""" + with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): + BinomialDistribution(n=10, p=invalid_p) + + +@pytest.mark.parametrize( + "invalid_p", + [1.1, 25.0], + ids=["small_positive", "large_positive"], +) +def test_binomial_distribution_init_invalid_p_above_one(invalid_p: float) -> None: + with pytest.raises(ValidationError, match="Input should be less than or equal to 1"): + BinomialDistribution(n=10, p=invalid_p) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((4, 4), (4, 4)), + (30, (30,)), + ((2, 3, 4), (2, 3, 4)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_binomial_distribution_sample_properties(size, expected_shape, use_rng) -> None: + n = 20 + dist = BinomialDistribution(n=n, p=0.5) + kwargs = {"size": size} + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.int64 + assert np.all(sample >= 0) + assert np.all(sample <= n) + + +def test_binomial_distribution_sample_rng_reproducibility() -> None: + dist = BinomialDistribution(n=50, p=0.75) + rng1 = np.random.default_rng(seed=100) + sample1 = dist.sample(size=(3, 3), rng=rng1) + rng2 = np.random.default_rng(seed=100) + sample2 = dist.sample(size=(3, 3), rng=rng2) + np.testing.assert_array_equal(sample1, sample2) From ed2e81d0d851cbff32170ea5e2907c667cf1e05e Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 23 Jun 2025 11:42:31 -0400 Subject: [PATCH 17/73] Add tests for `GammaDistribution` Unit tests and validation for the `GammaDistribution` class of `gempyor.distributions`. Also, a bounds fix in `distributions.py`. --- .../gempyor_pkg/src/gempyor/distributions.py | 4 +- .../distributions/test_GammaDistribution.py | 78 +++++++++++++++++++ 2 files changed, 80 insertions(+), 2 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_GammaDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index ff3adf021..8aa009149 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -308,8 +308,8 @@ class GammaDistribution(DistributionABC): """ distribution: Literal["gamma"] = "gamma" - shape: float - scale: float + shape: float = Field(..., gt=0) + scale: float = Field(..., gt=0) def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None diff --git a/flepimop/gempyor_pkg/tests/distributions/test_GammaDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_GammaDistribution.py new file mode 100644 index 000000000..50877ba19 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_GammaDistribution.py @@ -0,0 +1,78 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import GammaDistribution + + +@pytest.mark.parametrize( + "shape, scale", + [ + (2.0, 1.5), + (0.1, 0.1), + (9.0, 2.0), + (100.0, 50.0), + ], + ids=[ + "standard_case", + "small_values", + "numpy_docs_example", + "large_values", + ], +) +def test_gamma_distribution_init_valid(shape: float, scale: float) -> None: + dist = GammaDistribution(shape=shape, scale=scale) + assert dist.shape == shape + assert dist.scale == scale + assert dist.distribution == "gamma" + + +@pytest.mark.parametrize( + "invalid_shape", + [0.0, -0.1, -10.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_gamma_distribution_init_invalid_shape(invalid_shape: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + GammaDistribution(shape=invalid_shape, scale=1.0) + + +@pytest.mark.parametrize( + "invalid_scale", + [0.0, -1.0, -50.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_gamma_distribution_init_invalid_scale(invalid_scale: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + GammaDistribution(shape=1.0, scale=invalid_scale) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((3, 3), (3, 3)), + (5, (5,)), + ((2, 3, 4), (2, 3, 4)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_gamma_distribution_sample_properties(size, expected_shape, use_rng) -> None: + dist = GammaDistribution(shape=2.0, scale=2.0) + kwargs = {"size": size} + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample > 0) + + +def test_gamma_distribution_sample_rng_reproducibility() -> None: + dist = GammaDistribution(shape=9.0, scale=0.5) + rng1 = np.random.default_rng(seed=100) + sample1 = dist.sample(size=(4, 4), rng=rng1) + rng2 = np.random.default_rng(seed=100) + sample2 = dist.sample(size=(4, 4), rng=rng2) + np.testing.assert_array_equal(sample1, sample2) From e660b0d8dd39c3d979024ccf4f5e93d95b917354 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 23 Jun 2025 11:50:40 -0400 Subject: [PATCH 18/73] Add tests for `WeibullDistribution` Unit tests and validation for the `WeibullDistribution` class of `gempyor.distributions`. Also, bounds and formatting fixes. --- .../gempyor_pkg/src/gempyor/distributions.py | 4 +- .../distributions/test_FixedDistribution.py | 2 +- .../test_LognormalDistribution.py | 2 +- .../distributions/test_NormalDistribution.py | 2 +- .../distributions/test_UniformDistribution.py | 2 +- .../distributions/test_WeibullDistribution.py | 76 +++++++++++++++++++ 6 files changed, 82 insertions(+), 6 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_WeibullDistribution.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 8aa009149..c19791b2e 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -340,8 +340,8 @@ class WeibullDistribution(DistributionABC): """ distribution: Literal["weibull"] = "weibull" - shape: float - scale: float + shape: float = Field(..., gt=0) + scale: float = Field(..., gt=0) def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None diff --git a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py index f46553323..5a991efd8 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py @@ -21,7 +21,7 @@ def test_fixed_distribution_ini(value: float) -> None: (5, (5,)), ((2, 3, 4), (2, 3, 4)), ], - ids=["tuple_size1", "integer_size", "tuple_size2"], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) @pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) def test_fixed_distribution_sample_properties( diff --git a/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py index 3163c524c..4614c1804 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py @@ -35,7 +35,7 @@ def test_lognormal_distribution_init_raises_error_for_invalid_sdlog( (25, (25,)), ((2, 3, 4), (2, 3, 4)), ], - ids=["tuple_size1", "integer_size", "tuple_size2"], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) @pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) def test_lognormal_distribution_sample_properties(size, expected_shape, use_rng) -> None: diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py index 169c7ac5f..4a7dadddc 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py @@ -33,7 +33,7 @@ def test_normal_distribution_sample_raises_error_for_invalid_sigma( (10, (10,)), ((2, 3, 4), (2, 3, 4)), ], - ids=["tuple_size1", "integer_size", "tuple_size2"], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) @pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) def test_normal_distribution_sample_properties(size, expected_shape, use_rng) -> None: diff --git a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py index 6e52def51..c831b38b6 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py @@ -41,7 +41,7 @@ def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> No (20, (20,)), ((2, 3, 4), (2, 3, 4)), ], - ids=["default_size", "integer_size", "tuple_size"], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) @pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) def test_uniform_distribution_sample_properties(size, expected_shape, use_rng) -> None: diff --git a/flepimop/gempyor_pkg/tests/distributions/test_WeibullDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_WeibullDistribution.py new file mode 100644 index 000000000..a06a96c04 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_WeibullDistribution.py @@ -0,0 +1,76 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import WeibullDistribution + + +@pytest.mark.parametrize( + "shape, scale", + [ + (2.5, 5.0), + (0.5, 1.0), + (5.0, 20.0), + ], + ids=[ + "standard_case", + "shape_less_than_one", + "large_values", + ], +) +def test_weibull_distribution_init_valid(shape: float, scale: float) -> None: + dist = WeibullDistribution(shape=shape, scale=scale) + assert dist.shape == shape + assert dist.scale == scale + assert dist.distribution == "weibull" + + +@pytest.mark.parametrize( + "invalid_shape", + [0.0, -1.5, -20.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_weibull_distribution_init_invalid_shape(invalid_shape: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + WeibullDistribution(shape=invalid_shape, scale=1.0) + + +@pytest.mark.parametrize( + "invalid_scale", + [0.0, -2.0, -100.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_weibull_distribution_init_invalid_scale(invalid_scale: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + WeibullDistribution(shape=1.0, scale=invalid_scale) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((5, 5), (5, 5)), + (5, (5,)), + ((2, 3, 4), (2, 3, 4)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_weibull_distribution_sample_properties(size, expected_shape, use_rng) -> None: + dist = WeibullDistribution(shape=2.0, scale=10.0) + kwargs = {"size": size} + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample > 0) + + +def test_weibull_distribution_sample_rng_reproducibility() -> None: + dist = WeibullDistribution(shape=2.5, scale=5.0) + rng1 = np.random.default_rng(seed=100) + sample1 = dist.sample(size=(3, 3), rng=rng1) + rng2 = np.random.default_rng(seed=100) + sample2 = dist.sample(size=(3, 3), rng=rng2) + np.testing.assert_array_equal(sample1, sample2) From 8c6ad61ee0170ca07b5a8c583e34bd2d82a5f805 Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Mon, 23 Jun 2025 16:01:13 -0400 Subject: [PATCH 19/73] Update flepimop/gempyor_pkg/src/gempyor/distributions.py Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- flepimop/gempyor_pkg/src/gempyor/distributions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index c19791b2e..4baad4748 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -335,8 +335,8 @@ class WeibullDistribution(DistributionABC): array([5.18664161]) >>> dist.sample(size=(3, 5), rng=rng) array([[4.09267151, 6.25764353, 6.57560416, 2.72348545, 3.73147743], - [5.08637841, 4.63044996, 4.96639599, 3.99658252, 6.42564251], - [6.29525049, 5.03450379, 6.88371131, 5.65602432, 3.99307222]]) + [5.08637841, 4.63044996, 4.96639599, 3.99658252, 6.42564251], + [6.29525049, 5.03450379, 6.88371131, 5.65602432, 3.99307222]]) """ distribution: Literal["weibull"] = "weibull" From 5829da15c5e6851e9f002e49180ad99a86e58c7d Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 24 Jun 2025 14:39:27 -0400 Subject: [PATCH 20/73] Rename test files and update a docstring in `distributions.py` --- flepimop/gempyor_pkg/src/gempyor/distributions.py | 8 +++++++- ...istribution.py => test_binomial_distribution_class.py} | 0 ...edDistribution.py => test_fixed_distribution_class.py} | 0 ...maDistribution.py => test_gamma_distribution_class.py} | 0 ...stribution.py => test_lognormal_distribution_class.py} | 0 ...lDistribution.py => test_normal_distribution_class.py} | 0 ...Distribution.py => test_poisson_distribution_class.py} | 0 ...tion.py => test_truncatednormal_distribution_class.py} | 0 ...Distribution.py => test_uniform_distribution_class.py} | 0 ...Distribution.py => test_weibull_distribution_class.py} | 0 10 files changed, 7 insertions(+), 1 deletion(-) rename flepimop/gempyor_pkg/tests/distributions/{test_BinomialDistribution.py => test_binomial_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_FixedDistribution.py => test_fixed_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_GammaDistribution.py => test_gamma_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_LognormalDistribution.py => test_lognormal_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_NormalDistribution.py => test_normal_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_PoissonDistribution.py => test_poisson_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_TruncatedNormalDistribution.py => test_truncatednormal_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_UniformDistribution.py => test_uniform_distribution_class.py} (100%) rename flepimop/gempyor_pkg/tests/distributions/{test_WeibullDistribution.py => test_weibull_distribution_class.py} (100%) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 4baad4748..eccd2695e 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -35,7 +35,13 @@ class DistributionABC(ABC, BaseModel): def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64 | np.int64]: - """Sample from the distribution.""" + """ + Draw random sample(s) from the distribution. + + Args: + size: The desired output size of samples to be drawn. + rng: A NumPy random number generator instance used for sampling to ensure reproducibility. + """ raise NotImplementedError diff --git a/flepimop/gempyor_pkg/tests/distributions/test_BinomialDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_BinomialDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_FixedDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_GammaDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_GammaDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_LognormalDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_NormalDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_PoissonDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_TruncatedNormalDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_TruncatedNormalDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_UniformDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py diff --git a/flepimop/gempyor_pkg/tests/distributions/test_WeibullDistribution.py b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py similarity index 100% rename from flepimop/gempyor_pkg/tests/distributions/test_WeibullDistribution.py rename to flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py From ffb1d3c16ba5e7399c18179d6d17b9496fc73bf9 Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Wed, 25 Jun 2025 10:21:23 -0400 Subject: [PATCH 21/73] Update flepimop/gempyor_pkg/src/gempyor/distributions.py Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- flepimop/gempyor_pkg/src/gempyor/distributions.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index eccd2695e..c2d1cff63 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -41,6 +41,10 @@ def sample( Args: size: The desired output size of samples to be drawn. rng: A NumPy random number generator instance used for sampling to ensure reproducibility. + + Returns: + A numpy array of either floats/ints (depending on distribution) + with shape `size` drawn from the distribution. """ raise NotImplementedError From 346d7e11506b5c3d3a4e00955f035ee0a0f237aa Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Wed, 2 Jul 2025 13:17:53 -0400 Subject: [PATCH 22/73] Re-implement recent changes after force push by @TimothyWillard Add `BetaDistribution`, de-duplicate `rng` code, remove unnecessary `rng` tests. --- .../gempyor_pkg/src/gempyor/distributions.py | 63 +++++--- .../teset_beta_distribution_class.py | 144 ++++++++++++++++++ .../test_binomial_distribution_class.py | 9 -- .../test_gamma_distribution_class.py | 9 -- .../test_lognormal_distribution_class.py | 9 -- .../test_normal_distribution_class.py | 9 -- .../test_poisson_distribution_class.py | 9 -- ...test_truncatednormal_distribution_class.py | 9 -- .../test_uniform_distribution_class.py | 9 -- .../test_weibull_distribution_class.py | 9 -- 10 files changed, 190 insertions(+), 89 deletions(-) create mode 100644 flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index c2d1cff63..388515719 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -1,6 +1,7 @@ """Representations of distributions used for modifiers, likelihoods, etc.""" __all__: tuple[str, ...] = ( + "BetaDistribution", "BinomialDistribution", "Distribution", "DistributionABC", @@ -31,6 +32,10 @@ class DistributionABC(ABC, BaseModel): distribution: str + def _rng(self, rng: Generator | None = None) -> Generator: + """Return the provided rng or a default NumPy Generator""" + return rng if rng is not None else np.random.default_rng() + @abstractmethod def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None @@ -101,8 +106,7 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the normal distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) return rng.normal(loc=self.mu, scale=self.sigma, size=size) @@ -133,8 +137,7 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the uniform distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) return rng.uniform(low=self.low, high=self.high, size=size) @model_validator(mode="after") @@ -175,8 +178,7 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the Lognormal distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) return rng.lognormal(mean=self.meanlog, sigma=self.sdlog, size=size) @@ -211,8 +213,7 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the truncated normal distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) lower = (self.a - self.mean) / self.sd upper = (self.b - self.mean) / self.sd return truncnorm.rvs( @@ -261,8 +262,7 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.int64]: """Sample from the Poisson distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) return rng.poisson(lam=self.lam, size=size) @@ -293,8 +293,7 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.int64]: """Sample from the binomial distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) return rng.binomial(n=self.n, p=self.p, size=size) @@ -325,8 +324,7 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the gamma distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) return rng.gamma(shape=self.shape, scale=self.scale, size=size) @@ -357,14 +355,45 @@ def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64]: """Sample from the Weibull distribution.""" - if rng is None: - rng = np.random.default_rng() + rng = self._rng(rng) # Multiply by scale b/c rng.weibull assumes standard weibull dist (scale of 1) return self.scale * rng.weibull(a=self.shape, size=size) +class BetaDistribution(DistributionABC): + """ + Represents a beta distribution. + + Examples: + >>> import numpy as np + >>> from gempyor.distributions import BetaDistribution + >>> rng = np.random.default_rng(42) + >>> dist = BetaDistribution(alpha=2.0, beta=5.0) + >>> dist + BetaDistribution(distribution='beta', alpha=2.0, beta=5.0) + >>> dist.sample(rng=rng) + array([0.31153835]) + >>> dist.sample(size=(3, 5), rng=rng) + array([[0.18343469, 0.43545995, 0.48512838, 0.08625927, 0.22223849], + [0.34020943, 0.28014294, 0.3069152 , 0.201726 , 0.4729195 ], + [0.44331908, 0.32357912, 0.5282496 , 0.38318282, 0.20141687]]) + """ + + distribution: Literal["beta"] = "beta" + alpha: float = Field(..., gt=0) + beta: float = Field(..., gt=0) + + def sample( + self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + ) -> npt.NDArray[np.float64]: + """Sample from the beta distribution.""" + rng = self._rng(rng) + return rng.beta(a=self.alpha, b=self.beta, size=size) + + Distribution = Annotated[ - BinomialDistribution + BetaDistribution + | BinomialDistribution | FixedDistribution | GammaDistribution | LognormalDistribution diff --git a/flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py new file mode 100644 index 000000000..78ceb9503 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py @@ -0,0 +1,144 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import BetaDistribution + + +@pytest.mark.parametrize( + "alpha, beta", + [ + (2.0, 5.0), + (0.5, 0.5), + (1.0, 1.0), + (100.0, 25.0), + ], + ids=[ + "standard_case", + "small_values", + "uniform_case", + "large_values", + ], +) +def test_beta_distribution_init_valid(alpha: float, beta: float) -> None: + dist = BetaDistribution(alpha=alpha, beta=beta) + assert dist.alpha == alpha + assert dist.beta == beta + assert dist.distribution == "beta" + + +@pytest.mark.parametrize( + "invalid_alpha", + [0.0, -0.1, -10.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_beta_distribution_init_invalid_alpha(invalid_alpha: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + BetaDistribution(alpha=invalid_alpha, beta=1.0) + + +@pytest.mark.parametrize( + "invalid_beta", + [0.0, -1.0, -50.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_beta_distribution_init_invalid_beta(invalid_beta: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + BetaDistribution(alpha=1.0, beta=invalid_beta) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((3, 3), (3, 3)), + (5, (5,)), + ((2, 3, 4), (2, 3, 4)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_beta_distribution_sample_properties( + size: int | tuple[int, ...], expected_shape: tuple[int, ...], use_rng: bool +) -> None: + dist = BetaDistribution(alpha=2.0, beta=5.0) + kwargs: dict[str, int | tuple[int, ...] | np.random.Generator] = {"size": size} + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample >= 0) and np.all(sample <= 1) + + +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import BetaDistribution + + +@pytest.mark.parametrize( + "alpha, beta", + [ + (2.0, 5.0), + (0.5, 0.5), + (1.0, 1.0), + (100.0, 25.0), + ], + ids=[ + "standard_case", + "small_values", + "uniform_case", + "large_values", + ], +) +def test_beta_distribution_init_valid(alpha: float, beta: float) -> None: + dist = BetaDistribution(alpha=alpha, beta=beta) + assert dist.alpha == alpha + assert dist.beta == beta + assert dist.distribution == "beta" + + +@pytest.mark.parametrize( + "invalid_alpha", + [0.0, -0.1, -10.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_beta_distribution_init_invalid_alpha(invalid_alpha: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + BetaDistribution(alpha=invalid_alpha, beta=1.0) + + +@pytest.mark.parametrize( + "invalid_beta", + [0.0, -1.0, -50.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_beta_distribution_init_invalid_beta(invalid_beta: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + BetaDistribution(alpha=1.0, beta=invalid_beta) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((3, 3), (3, 3)), + (5, (5,)), + ((2, 3, 4), (2, 3, 4)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) +def test_beta_distribution_sample_properties( + size: int | tuple[int, ...], expected_shape: tuple[int, ...], use_rng: bool +) -> None: + dist = BetaDistribution(alpha=2.0, beta=5.0) + kwargs: dict[str, int | tuple[int, ...] | np.random.Generator] = {"size": size} + if use_rng: + kwargs["rng"] = np.random.default_rng() + sample = dist.sample(**kwargs) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample >= 0) and np.all(sample <= 1) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py index 29825fe9f..379d65e4e 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py @@ -82,12 +82,3 @@ def test_binomial_distribution_sample_properties(size, expected_shape, use_rng) assert sample.dtype == np.int64 assert np.all(sample >= 0) assert np.all(sample <= n) - - -def test_binomial_distribution_sample_rng_reproducibility() -> None: - dist = BinomialDistribution(n=50, p=0.75) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(3, 3), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(3, 3), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py index 50877ba19..3d9c0ff9f 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py @@ -67,12 +67,3 @@ def test_gamma_distribution_sample_properties(size, expected_shape, use_rng) -> assert sample.shape == expected_shape assert sample.dtype == np.float64 assert np.all(sample > 0) - - -def test_gamma_distribution_sample_rng_reproducibility() -> None: - dist = GammaDistribution(shape=9.0, scale=0.5) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(4, 4), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(4, 4), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py index 4614c1804..87d8691f0 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py @@ -49,12 +49,3 @@ def test_lognormal_distribution_sample_properties(size, expected_shape, use_rng) assert sample.shape == expected_shape assert sample.dtype == np.float64 assert np.all(sample > 0) - - -def test_lognormal_distribution_sample_rng_reproducibility() -> None: - dist = LognormalDistribution(meanlog=2.0, sdlog=0.5) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(3, 3), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(3, 3), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py index 4a7dadddc..4b31ff14e 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py @@ -46,12 +46,3 @@ def test_normal_distribution_sample_properties(size, expected_shape, use_rng) -> assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 - - -def test_normal_distribution_sample_rng_reproducibility() -> None: - dist = NormalDistribution(mu=10.0, sigma=1.5) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(2, 2), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(2, 2), rng=rng2) - assert np.array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py index df32b06dd..607a1208a 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py @@ -48,12 +48,3 @@ def test_poisson_distribution_sample_properties(size, expected_shape, use_rng) - assert sample.shape == expected_shape assert sample.dtype == np.int64 assert np.all(sample >= 0) - - -def test_poisson_distribution_sample_rng_reproducibility() -> None: - dist = PoissonDistribution(lam=10.0) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(5, 5), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(5, 5), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py index 107def8c3..857141abe 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py @@ -58,12 +58,3 @@ def test_truncated_normal_distribution_sample_properties( assert sample.dtype == np.float64 assert np.all(sample >= a) assert np.all(sample <= b) - - -def test_truncated_normal_distribution_sample_rng_reproducibility() -> None: - dist = TruncatedNormalDistribution(mean=10.0, sd=3.0, a=5.0, b=15.0) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(4, 4), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(4, 4), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py index c831b38b6..92df57749 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py @@ -57,12 +57,3 @@ def test_uniform_distribution_sample_properties(size, expected_shape, use_rng) - assert sample.dtype == np.float64 assert np.all(sample >= low) assert np.all(sample < high) - - -def test_uniform_distribution_sample_rng_reproducibility() -> None: - dist = UniformDistribution(low=-10.0, high=10.0) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(3, 3), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(3, 3), rng=rng2) - assert np.array_equal(sample1, sample2) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py index a06a96c04..c7a56c286 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py @@ -65,12 +65,3 @@ def test_weibull_distribution_sample_properties(size, expected_shape, use_rng) - assert sample.shape == expected_shape assert sample.dtype == np.float64 assert np.all(sample > 0) - - -def test_weibull_distribution_sample_rng_reproducibility() -> None: - dist = WeibullDistribution(shape=2.5, scale=5.0) - rng1 = np.random.default_rng(seed=100) - sample1 = dist.sample(size=(3, 3), rng=rng1) - rng2 = np.random.default_rng(seed=100) - sample2 = dist.sample(size=(3, 3), rng=rng2) - np.testing.assert_array_equal(sample1, sample2) From febe7cbf2805a521c39ee2d0a1714c3c5349583e Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Wed, 9 Jul 2025 12:45:08 -0400 Subject: [PATCH 23/73] Split `gempyor` test suite into slow and not slow tests --- .github/workflows/gempyor-ci.yml | 10 ++++++++-- flepimop/gempyor_pkg/pyproject.toml | 5 +++++ flepimop/gempyor_pkg/tests/inference/test_inference.py | 1 + flepimop/gempyor_pkg/tests/npi/test_npis.py | 1 + flepimop/gempyor_pkg/tests/seir/test_seir.py | 5 +++++ flepimop/gempyor_pkg/tests/shared_cli/test_cli.py | 5 +++++ 6 files changed, 25 insertions(+), 2 deletions(-) diff --git a/.github/workflows/gempyor-ci.yml b/.github/workflows/gempyor-ci.yml index 0624fdee1..1af373e62 100644 --- a/.github/workflows/gempyor-ci.yml +++ b/.github/workflows/gempyor-ci.yml @@ -47,11 +47,17 @@ jobs: python -m pip install --upgrade pip python -m pip install "flepimop/gempyor_pkg[dev]" shell: bash - - name: Run gempyor tests + - name: Run not slow tests run: | export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg - pytest --exitfirst + pytest -m "not slow" --exitfirst + shell: bash + - name: Run slow tests + run: | + export FLEPI_PATH=$(pwd) + cd flepimop/gempyor_pkg + pytest -m "slow" --exitfirst shell: bash - name: Run gempyor-cli integration tests from examples run: | diff --git a/flepimop/gempyor_pkg/pyproject.toml b/flepimop/gempyor_pkg/pyproject.toml index 8bd5c55e7..6936d3edc 100644 --- a/flepimop/gempyor_pkg/pyproject.toml +++ b/flepimop/gempyor_pkg/pyproject.toml @@ -59,3 +59,8 @@ package-data = { "gempyor" = ["templates/*.j2"] } [tool.setuptools.packages.find] where = ["src"] namespaces = false + +[tool.pytest.ini_options] +markers = [ + "slow: marks tests as slow to run" +] diff --git a/flepimop/gempyor_pkg/tests/inference/test_inference.py b/flepimop/gempyor_pkg/tests/inference/test_inference.py index b911c7643..bb35043ea 100644 --- a/flepimop/gempyor_pkg/tests/inference/test_inference.py +++ b/flepimop/gempyor_pkg/tests/inference/test_inference.py @@ -20,6 +20,7 @@ class TestGempyorInference: + @pytest.mark.slow def test_GempyorInference_success(self): os.chdir(os.path.dirname(__file__)) diff --git a/flepimop/gempyor_pkg/tests/npi/test_npis.py b/flepimop/gempyor_pkg/tests/npi/test_npis.py index beacd33e0..abadd8cde 100644 --- a/flepimop/gempyor_pkg/tests/npi/test_npis.py +++ b/flepimop/gempyor_pkg/tests/npi/test_npis.py @@ -23,6 +23,7 @@ os.chdir(os.path.dirname(__file__)) +@pytest.mark.slow def test_full_npis_read_write(): os.chdir(os.path.dirname(__file__)) diff --git a/flepimop/gempyor_pkg/tests/seir/test_seir.py b/flepimop/gempyor_pkg/tests/seir/test_seir.py index b6c58a5c5..afc6ac38b 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_seir.py +++ b/flepimop/gempyor_pkg/tests/seir/test_seir.py @@ -214,6 +214,7 @@ def test_constant_population_rk4jit_integration(): @ignore_non_csv_mobility_warning +@pytest.mark.slow def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): os.chdir(os.path.dirname(__file__)) config.clear() @@ -323,6 +324,7 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): @ignore_non_csv_mobility_warning +@pytest.mark.slow def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): os.chdir(os.path.dirname(__file__)) config.clear() @@ -404,6 +406,7 @@ def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): @ignore_non_csv_mobility_warning +@pytest.mark.slow def test_steps_SEIR_no_spread(): os.chdir(os.path.dirname(__file__)) print("test mobility with no spread") @@ -675,6 +678,7 @@ def test_inference_resume(): ## Clean up after ourselves +@pytest.mark.slow def test_parallel_compartments_with_vacc(): os.chdir(os.path.dirname(__file__)) config.clear() @@ -769,6 +773,7 @@ def test_parallel_compartments_with_vacc(): ) +@pytest.mark.slow def test_parallel_compartments_no_vacc(): config.clear() config.read(user=False) diff --git a/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py b/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py index d6461411f..7b31e1391 100644 --- a/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py +++ b/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py @@ -1,4 +1,5 @@ import os +import pytest import subprocess from pathlib import Path @@ -15,6 +16,7 @@ tutorialpath = os.path.dirname(__file__) + "/../../../../examples/tutorials" +@pytest.mark.slow def test_config_sample_2pop(): os.chdir(tutorialpath) runner = CliRunner() @@ -22,6 +24,7 @@ def test_config_sample_2pop(): assert result.exit_code == 0 +@pytest.mark.slow def test_config_sample_2pop_deprecated(): os.chdir(tutorialpath) runner = CliRunner() @@ -29,6 +32,7 @@ def test_config_sample_2pop_deprecated(): assert result.exit_code == 0 +@pytest.mark.slow def test_sample_2pop_modifiers(): os.chdir(tutorialpath) runner = CliRunner() @@ -74,6 +78,7 @@ def test_sample_2pop_modifiers_combined(tmp_path: Path): } +@pytest.mark.slow def test_simple_usa_statelevel_more_deprecated(): os.chdir(tutorialpath + "/../simple_usa_statelevel") result = subprocess.run( From bb68fce4d7b7721a6e507dc08f6ecf1b221aeafc Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Wed, 9 Jul 2025 13:02:43 -0400 Subject: [PATCH 24/73] Add `test_npis::test_spatial_groups()` to slow test group to resolve dependency. `test_spatial_groups()` relies on an inference config file used in the "slow" groups. --- flepimop/gempyor_pkg/tests/npi/test_npis.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flepimop/gempyor_pkg/tests/npi/test_npis.py b/flepimop/gempyor_pkg/tests/npi/test_npis.py index abadd8cde..ad4565de8 100644 --- a/flepimop/gempyor_pkg/tests/npi/test_npis.py +++ b/flepimop/gempyor_pkg/tests/npi/test_npis.py @@ -121,6 +121,7 @@ def test_full_npis_read_write(): assert (hnpi_read == hnpi_wrote).all().all() +@pytest.mark.slow def test_spatial_groups(): inference_simulator = gempyor.GempyorInference( config_filepath=f"{config_filepath_prefix}config_test_spatial_group_npi.yml", From be0664aa66b04308a9d115461dc6b4ef9fce0904 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Wed, 9 Jul 2025 13:14:00 -0400 Subject: [PATCH 25/73] Add second `test_spatial_groups()` test to slow group These tests are duplicates, which gets addressed in #570 , but for now we need to add both to the slow group to allow workflow success. --- flepimop/gempyor_pkg/tests/npi/test_npis.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flepimop/gempyor_pkg/tests/npi/test_npis.py b/flepimop/gempyor_pkg/tests/npi/test_npis.py index ad4565de8..5d6e51d9a 100644 --- a/flepimop/gempyor_pkg/tests/npi/test_npis.py +++ b/flepimop/gempyor_pkg/tests/npi/test_npis.py @@ -211,6 +211,7 @@ def test_spatial_groups(): ) +@pytest.mark.slow def test_spatial_groups(): inference_simulator = gempyor.GempyorInference( config_filepath=f"{config_filepath_prefix}config_test_spatial_group_npi.yml", From 2e22442df4c328eba53414574b7403255db024d3 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Wed, 9 Jul 2025 17:42:13 -0400 Subject: [PATCH 26/73] Change "slow" test group (fast group takes ~1/7 of total time) Remove old "slow" group, add timing feature to workflow for testing purposes (to be removed later). --- .github/workflows/gempyor-ci.yml | 4 ++-- flepimop/gempyor_pkg/tests/inference/test_inference.py | 1 - flepimop/gempyor_pkg/tests/npi/test_npis.py | 3 --- .../gempyor_pkg/tests/outcomes/test_run_parallel_outcomes.py | 1 + flepimop/gempyor_pkg/tests/seir/test_run_parallel_SEIR.py | 1 + flepimop/gempyor_pkg/tests/seir/test_seir.py | 4 ---- flepimop/gempyor_pkg/tests/shared_cli/test_cli.py | 3 --- .../gempyor_pkg/tests/simulate/test_random_parameter_draws.py | 2 ++ 8 files changed, 6 insertions(+), 13 deletions(-) diff --git a/.github/workflows/gempyor-ci.yml b/.github/workflows/gempyor-ci.yml index 1af373e62..e7bc5dc61 100644 --- a/.github/workflows/gempyor-ci.yml +++ b/.github/workflows/gempyor-ci.yml @@ -51,13 +51,13 @@ jobs: run: | export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg - pytest -m "not slow" --exitfirst + time pytest -m "not slow" --exitfirst shell: bash - name: Run slow tests run: | export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg - pytest -m "slow" --exitfirst + time pytest -m "slow" --exitfirst shell: bash - name: Run gempyor-cli integration tests from examples run: | diff --git a/flepimop/gempyor_pkg/tests/inference/test_inference.py b/flepimop/gempyor_pkg/tests/inference/test_inference.py index bb35043ea..b911c7643 100644 --- a/flepimop/gempyor_pkg/tests/inference/test_inference.py +++ b/flepimop/gempyor_pkg/tests/inference/test_inference.py @@ -20,7 +20,6 @@ class TestGempyorInference: - @pytest.mark.slow def test_GempyorInference_success(self): os.chdir(os.path.dirname(__file__)) diff --git a/flepimop/gempyor_pkg/tests/npi/test_npis.py b/flepimop/gempyor_pkg/tests/npi/test_npis.py index 5d6e51d9a..beacd33e0 100644 --- a/flepimop/gempyor_pkg/tests/npi/test_npis.py +++ b/flepimop/gempyor_pkg/tests/npi/test_npis.py @@ -23,7 +23,6 @@ os.chdir(os.path.dirname(__file__)) -@pytest.mark.slow def test_full_npis_read_write(): os.chdir(os.path.dirname(__file__)) @@ -121,7 +120,6 @@ def test_full_npis_read_write(): assert (hnpi_read == hnpi_wrote).all().all() -@pytest.mark.slow def test_spatial_groups(): inference_simulator = gempyor.GempyorInference( config_filepath=f"{config_filepath_prefix}config_test_spatial_group_npi.yml", @@ -211,7 +209,6 @@ def test_spatial_groups(): ) -@pytest.mark.slow def test_spatial_groups(): inference_simulator = gempyor.GempyorInference( config_filepath=f"{config_filepath_prefix}config_test_spatial_group_npi.yml", diff --git a/flepimop/gempyor_pkg/tests/outcomes/test_run_parallel_outcomes.py b/flepimop/gempyor_pkg/tests/outcomes/test_run_parallel_outcomes.py index 1269723bb..5225a16f9 100644 --- a/flepimop/gempyor_pkg/tests/outcomes/test_run_parallel_outcomes.py +++ b/flepimop/gempyor_pkg/tests/outcomes/test_run_parallel_outcomes.py @@ -16,6 +16,7 @@ @pytest.mark.parametrize("n_jobs", (1, 2)) @pytest.mark.parametrize("config_file", ("config_sample_2pop_vaccine_scenarios.yml",)) @pytest.mark.parametrize("start_method", mp.get_all_start_methods()) +@pytest.mark.slow def test_run_parallel_outcomes_by_multiprocessing_start_method( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/flepimop/gempyor_pkg/tests/seir/test_run_parallel_SEIR.py b/flepimop/gempyor_pkg/tests/seir/test_run_parallel_SEIR.py index 75a2ed33f..de6a7199d 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_run_parallel_SEIR.py +++ b/flepimop/gempyor_pkg/tests/seir/test_run_parallel_SEIR.py @@ -16,6 +16,7 @@ @pytest.mark.parametrize("config_file", ("config_sample_2pop_vaccine_scenarios.yml",)) @pytest.mark.parametrize("n_jobs", (1, 2)) @pytest.mark.parametrize("start_method", mp.get_all_start_methods()) +@pytest.mark.slow def test_run_parallel_SEIR_by_multiprocessing_start_method( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, diff --git a/flepimop/gempyor_pkg/tests/seir/test_seir.py b/flepimop/gempyor_pkg/tests/seir/test_seir.py index afc6ac38b..421cffcda 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_seir.py +++ b/flepimop/gempyor_pkg/tests/seir/test_seir.py @@ -214,7 +214,6 @@ def test_constant_population_rk4jit_integration(): @ignore_non_csv_mobility_warning -@pytest.mark.slow def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): os.chdir(os.path.dirname(__file__)) config.clear() @@ -324,7 +323,6 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices(): @ignore_non_csv_mobility_warning -@pytest.mark.slow def test_steps_SEIR_nb_simple_spread_with_csv_matrices(): os.chdir(os.path.dirname(__file__)) config.clear() @@ -678,7 +676,6 @@ def test_inference_resume(): ## Clean up after ourselves -@pytest.mark.slow def test_parallel_compartments_with_vacc(): os.chdir(os.path.dirname(__file__)) config.clear() @@ -773,7 +770,6 @@ def test_parallel_compartments_with_vacc(): ) -@pytest.mark.slow def test_parallel_compartments_no_vacc(): config.clear() config.read(user=False) diff --git a/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py b/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py index 7b31e1391..4157ad68d 100644 --- a/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py +++ b/flepimop/gempyor_pkg/tests/shared_cli/test_cli.py @@ -16,7 +16,6 @@ tutorialpath = os.path.dirname(__file__) + "/../../../../examples/tutorials" -@pytest.mark.slow def test_config_sample_2pop(): os.chdir(tutorialpath) runner = CliRunner() @@ -24,7 +23,6 @@ def test_config_sample_2pop(): assert result.exit_code == 0 -@pytest.mark.slow def test_config_sample_2pop_deprecated(): os.chdir(tutorialpath) runner = CliRunner() @@ -78,7 +76,6 @@ def test_sample_2pop_modifiers_combined(tmp_path: Path): } -@pytest.mark.slow def test_simple_usa_statelevel_more_deprecated(): os.chdir(tutorialpath + "/../simple_usa_statelevel") result = subprocess.run( diff --git a/flepimop/gempyor_pkg/tests/simulate/test_random_parameter_draws.py b/flepimop/gempyor_pkg/tests/simulate/test_random_parameter_draws.py index bc992fdc6..1662c02e7 100644 --- a/flepimop/gempyor_pkg/tests/simulate/test_random_parameter_draws.py +++ b/flepimop/gempyor_pkg/tests/simulate/test_random_parameter_draws.py @@ -387,6 +387,7 @@ def assert_df_passes( ), ), ) +@pytest.mark.slow def test_parameter_draw_per_slot( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -428,6 +429,7 @@ def test_parameter_draw_per_slot( reason="The $FLEPI_PATH environment variable is not set.", ) @pytest.mark.parametrize("n_jobs", [1, 2]) +@pytest.mark.slow def test_parameter_draws_per_slot_across_scenarios( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, n_jobs: int ) -> None: From 365e196e2b59dfc039be2deb0642c5436e806561 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Wed, 9 Jul 2025 21:20:25 -0400 Subject: [PATCH 27/73] Remove unecessary `time` command in `gempyor CI`. Remove temporary time marker to ensure the test timing is appropriate. --- .github/workflows/gempyor-ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/gempyor-ci.yml b/.github/workflows/gempyor-ci.yml index e7bc5dc61..1af373e62 100644 --- a/.github/workflows/gempyor-ci.yml +++ b/.github/workflows/gempyor-ci.yml @@ -51,13 +51,13 @@ jobs: run: | export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg - time pytest -m "not slow" --exitfirst + pytest -m "not slow" --exitfirst shell: bash - name: Run slow tests run: | export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg - time pytest -m "slow" --exitfirst + pytest -m "slow" --exitfirst shell: bash - name: Run gempyor-cli integration tests from examples run: | From 8f7daf848aca7a809d0e1eaa9b92c8cd952dbf8b Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 11 Jul 2025 11:47:31 -0400 Subject: [PATCH 28/73] Refactor tests for new ABC `rng` property. --- .../gempyor_pkg/src/gempyor/distributions.py | 104 +++++++------ .../teset_beta_distribution_class.py | 144 ------------------ .../test_beta_distribution_class.py | 65 ++++++++ .../test_binomial_distribution_class.py | 8 +- .../distributions/test_distributions_abc.py | 51 +++++++ .../test_fixed_distribution_class.py | 11 +- .../test_gamma_distribution_class.py | 8 +- .../test_lognormal_distribution_class.py | 9 +- .../test_normal_distribution_class.py | 9 +- .../test_poisson_distribution_class.py | 8 +- ...test_truncatednormal_distribution_class.py | 10 +- .../test_uniform_distribution_class.py | 9 +- .../test_weibull_distribution_class.py | 8 +- 13 files changed, 191 insertions(+), 253 deletions(-) delete mode 100644 flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py create mode 100644 flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 388515719..546e6d4a1 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -23,7 +23,7 @@ import numpy as np from numpy.random import Generator import numpy.typing as npt -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, PrivateAttr, Field, model_validator from scipy.stats import truncnorm @@ -31,12 +31,17 @@ class DistributionABC(ABC, BaseModel): """Base class for distributions used in modifiers, likelihoods, etc.""" distribution: str + _rng: Generator | None = PrivateAttr(default=None) - def _rng(self, rng: Generator | None = None) -> Generator: - """Return the provided rng or a default NumPy Generator""" - return rng if rng is not None else np.random.default_rng() + @property + def rng(self) -> Generator: + """ + Inheritable property for RNG. + """ + if self._rng is None: + self._rng = np.random.default_rng() + return self._rng - @abstractmethod def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None ) -> npt.NDArray[np.float64 | np.int64]: @@ -45,11 +50,25 @@ def sample( Args: size: The desired output size of samples to be drawn. - rng: A NumPy random number generator instance used for sampling to ensure reproducibility. + rng: A NumPy random number generator instance used for sampling. Returns: - A numpy array of either floats/ints (depending on distribution) - with shape `size` drawn from the distribution. + A NumPy array of either floats/ints (depending on distribution) + drawn from the distribution with shape `size`. + """ + rng = rng if rng is not None else self.rng + return self._sample_from_generator(size=size, rng=rng) + + @abstractmethod + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator + ) -> npt.NDArray[np.float64 | np.int64]: + """ + Distribution-specific sampling logic. + + Args: + size: The desired output size of samples to be drawn. + rng: A NumPy random number generator instance used for sampling. """ raise NotImplementedError @@ -72,10 +91,10 @@ class FixedDistribution(DistributionABC): distribution: Literal["fixed"] = "fixed" value: float - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the fixed distribution.""" + """Sampling logic for fixed distributions.""" return np.full(size, self.value) @@ -102,11 +121,10 @@ class NormalDistribution(DistributionABC): mu: float sigma: float = Field(..., gt=0) - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the normal distribution.""" - rng = self._rng(rng) + """Sampling logic for normal distributions.""" return rng.normal(loc=self.mu, scale=self.sigma, size=size) @@ -133,11 +151,10 @@ class UniformDistribution(DistributionABC): low: float high: float - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the uniform distribution.""" - rng = self._rng(rng) + """Sampling logic for uniform distributions.""" return rng.uniform(low=self.low, high=self.high, size=size) @model_validator(mode="after") @@ -174,11 +191,10 @@ class LognormalDistribution(DistributionABC): meanlog: float sdlog: float = Field(..., gt=0) - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the Lognormal distribution.""" - rng = self._rng(rng) + """Sampling logic for lognormal distributions.""" return rng.lognormal(mean=self.meanlog, sigma=self.sdlog, size=size) @@ -209,11 +225,10 @@ class TruncatedNormalDistribution(DistributionABC): a: float b: float - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the truncated normal distribution.""" - rng = self._rng(rng) + """Sampling logic for truncated normal distributions.""" lower = (self.a - self.mean) / self.sd upper = (self.b - self.mean) / self.sd return truncnorm.rvs( @@ -258,11 +273,10 @@ class PoissonDistribution(DistributionABC): distribution: Literal["poisson"] = "poisson" lam: float = Field(..., ge=0) - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.int64]: - """Sample from the Poisson distribution.""" - rng = self._rng(rng) + """Sampling logic for Poisson distributions.""" return rng.poisson(lam=self.lam, size=size) @@ -289,11 +303,10 @@ class BinomialDistribution(DistributionABC): n: int = Field(..., ge=0) p: float = Field(..., ge=0, le=1) - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.int64]: - """Sample from the binomial distribution.""" - rng = self._rng(rng) + """Sampling logic for binomial distributions.""" return rng.binomial(n=self.n, p=self.p, size=size) @@ -320,11 +333,10 @@ class GammaDistribution(DistributionABC): shape: float = Field(..., gt=0) scale: float = Field(..., gt=0) - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the gamma distribution.""" - rng = self._rng(rng) + """Sampling logic for Gamma distributions.""" return rng.gamma(shape=self.shape, scale=self.scale, size=size) @@ -351,11 +363,10 @@ class WeibullDistribution(DistributionABC): shape: float = Field(..., gt=0) scale: float = Field(..., gt=0) - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the Weibull distribution.""" - rng = self._rng(rng) + """Sampling logic for Weibull distributions.""" # Multiply by scale b/c rng.weibull assumes standard weibull dist (scale of 1) return self.scale * rng.weibull(a=self.shape, size=size) @@ -383,11 +394,10 @@ class BetaDistribution(DistributionABC): alpha: float = Field(..., gt=0) beta: float = Field(..., gt=0) - def sample( - self, size: int | tuple[int, ...] = 1, rng: Generator | None = None + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: - """Sample from the beta distribution.""" - rng = self._rng(rng) + """Sampling logic for beta distributions.""" return rng.beta(a=self.alpha, b=self.beta, size=size) diff --git a/flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py deleted file mode 100644 index 78ceb9503..000000000 --- a/flepimop/gempyor_pkg/tests/distributions/teset_beta_distribution_class.py +++ /dev/null @@ -1,144 +0,0 @@ -import numpy as np -import pytest -from pydantic import ValidationError - -from gempyor.distributions import BetaDistribution - - -@pytest.mark.parametrize( - "alpha, beta", - [ - (2.0, 5.0), - (0.5, 0.5), - (1.0, 1.0), - (100.0, 25.0), - ], - ids=[ - "standard_case", - "small_values", - "uniform_case", - "large_values", - ], -) -def test_beta_distribution_init_valid(alpha: float, beta: float) -> None: - dist = BetaDistribution(alpha=alpha, beta=beta) - assert dist.alpha == alpha - assert dist.beta == beta - assert dist.distribution == "beta" - - -@pytest.mark.parametrize( - "invalid_alpha", - [0.0, -0.1, -10.0], - ids=["zero", "small_negative", "large_negative"], -) -def test_beta_distribution_init_invalid_alpha(invalid_alpha: float) -> None: - with pytest.raises(ValidationError, match="Input should be greater than 0"): - BetaDistribution(alpha=invalid_alpha, beta=1.0) - - -@pytest.mark.parametrize( - "invalid_beta", - [0.0, -1.0, -50.0], - ids=["zero", "small_negative", "large_negative"], -) -def test_beta_distribution_init_invalid_beta(invalid_beta: float) -> None: - with pytest.raises(ValidationError, match="Input should be greater than 0"): - BetaDistribution(alpha=1.0, beta=invalid_beta) - - -@pytest.mark.parametrize( - "size, expected_shape", - [ - ((3, 3), (3, 3)), - (5, (5,)), - ((2, 3, 4), (2, 3, 4)), - ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], -) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_beta_distribution_sample_properties( - size: int | tuple[int, ...], expected_shape: tuple[int, ...], use_rng: bool -) -> None: - dist = BetaDistribution(alpha=2.0, beta=5.0) - kwargs: dict[str, int | tuple[int, ...] | np.random.Generator] = {"size": size} - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample >= 0) and np.all(sample <= 1) - - -import numpy as np -import pytest -from pydantic import ValidationError - -from gempyor.distributions import BetaDistribution - - -@pytest.mark.parametrize( - "alpha, beta", - [ - (2.0, 5.0), - (0.5, 0.5), - (1.0, 1.0), - (100.0, 25.0), - ], - ids=[ - "standard_case", - "small_values", - "uniform_case", - "large_values", - ], -) -def test_beta_distribution_init_valid(alpha: float, beta: float) -> None: - dist = BetaDistribution(alpha=alpha, beta=beta) - assert dist.alpha == alpha - assert dist.beta == beta - assert dist.distribution == "beta" - - -@pytest.mark.parametrize( - "invalid_alpha", - [0.0, -0.1, -10.0], - ids=["zero", "small_negative", "large_negative"], -) -def test_beta_distribution_init_invalid_alpha(invalid_alpha: float) -> None: - with pytest.raises(ValidationError, match="Input should be greater than 0"): - BetaDistribution(alpha=invalid_alpha, beta=1.0) - - -@pytest.mark.parametrize( - "invalid_beta", - [0.0, -1.0, -50.0], - ids=["zero", "small_negative", "large_negative"], -) -def test_beta_distribution_init_invalid_beta(invalid_beta: float) -> None: - with pytest.raises(ValidationError, match="Input should be greater than 0"): - BetaDistribution(alpha=1.0, beta=invalid_beta) - - -@pytest.mark.parametrize( - "size, expected_shape", - [ - ((3, 3), (3, 3)), - (5, (5,)), - ((2, 3, 4), (2, 3, 4)), - ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], -) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_beta_distribution_sample_properties( - size: int | tuple[int, ...], expected_shape: tuple[int, ...], use_rng: bool -) -> None: - dist = BetaDistribution(alpha=2.0, beta=5.0) - kwargs: dict[str, int | tuple[int, ...] | np.random.Generator] = {"size": size} - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample >= 0) and np.all(sample <= 1) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py new file mode 100644 index 000000000..0e724b870 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py @@ -0,0 +1,65 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from gempyor.distributions import BetaDistribution + + +@pytest.mark.parametrize( + "alpha, beta", + [ + (2.0, 5.0), + (0.5, 0.5), + (1.0, 1.0), + (100.0, 25.0), + ], + ids=[ + "standard_case", + "small_values", + "uniform_case", + "large_values", + ], +) +def test_beta_distribution_init_valid(alpha: float, beta: float) -> None: + dist = BetaDistribution(alpha=alpha, beta=beta) + assert dist.alpha == alpha + assert dist.beta == beta + assert dist.distribution == "beta" + + +@pytest.mark.parametrize( + "invalid_alpha", + [0.0, -0.1, -10.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_beta_distribution_init_invalid_alpha(invalid_alpha: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + BetaDistribution(alpha=invalid_alpha, beta=1.0) + + +@pytest.mark.parametrize( + "invalid_beta", + [0.0, -1.0, -50.0], + ids=["zero", "small_negative", "large_negative"], +) +def test_beta_distribution_init_invalid_beta(invalid_beta: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + BetaDistribution(alpha=1.0, beta=invalid_beta) + + +@pytest.mark.parametrize( + "size, expected_shape", + [ + ((3, 3), (3, 3)), + (5, (5,)), + ((2, 3, 4), (2, 3, 4)), + ], + ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], +) +def test_beta_distribution_sample_properties(size, expected_shape) -> None: + dist = BetaDistribution(alpha=2.0, beta=5.0) + sample = dist.sample(size=size) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert sample.dtype == np.float64 + assert np.all(sample >= 0) and np.all(sample <= 1) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py index 379d65e4e..9317521a1 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py @@ -69,14 +69,10 @@ def test_binomial_distribution_init_invalid_p_above_one(invalid_p: float) -> Non ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_binomial_distribution_sample_properties(size, expected_shape, use_rng) -> None: +def test_binomial_distribution_sample_properties(size, expected_shape) -> None: n = 20 dist = BinomialDistribution(n=n, p=0.5) - kwargs = {"size": size} - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.int64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py b/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py new file mode 100644 index 000000000..f769a1064 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py @@ -0,0 +1,51 @@ +import numpy as np +import pytest +from gempyor.distributions import ( + BetaDistribution, + BinomialDistribution, + # No FixedDistribution; does not use RNGs + GammaDistribution, + LognormalDistribution, + NormalDistribution, + PoissonDistribution, + TruncatedNormalDistribution, + UniformDistribution, + WeibullDistribution, +) + + +@pytest.mark.parametrize( + "dist", + [ + BetaDistribution(alpha=2.0, beta=5.0), + BinomialDistribution(n=10, p=0.5), + GammaDistribution(shape=2.0, scale=1.5), + LognormalDistribution(meanlog=0.0, sdlog=1.0), + NormalDistribution(mu=2.3, sigma=4.5), + PoissonDistribution(lam=3.0), + TruncatedNormalDistribution(mean=5.0, sd=2.0, a=0.0, b=10.0), + UniformDistribution(low=-0.5, high=1.5), + WeibullDistribution(shape=2.5, scale=5.0), + ], + ids=[ + "Beta", + "Binomial", + "Gamma", + "Lognormal", + "Normal", + "Poisson", + "TruncatedNormal", + "Uniform", + "Weibull", + ], +) +def test_distribution_abc_rng_behavior(dist) -> None: + rng1 = np.random.default_rng(seed=42) + sample1 = dist.sample(rng=rng1) + rng2 = np.random.default_rng(seed=42) + sample2 = dist.sample(rng=rng2) + np.testing.assert_array_equal(sample1, sample2) + + sample3_default = dist.sample() + sample4_default = dist.sample() + assert not np.array_equal(sample3_default, sample4_default) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py index 5a991efd8..bacc63404 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py @@ -23,16 +23,9 @@ def test_fixed_distribution_ini(value: float) -> None: ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_fixed_distribution_sample_properties( - value: float, size, expected_shape, use_rng -) -> None: +def test_fixed_distribution_sample_properties(value: float, size, expected_shape) -> None: dist = FixedDistribution(value=value) - kwargs = {} - kwargs["size"] = size - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py index 3d9c0ff9f..7e23479be 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py @@ -56,13 +56,9 @@ def test_gamma_distribution_init_invalid_scale(invalid_scale: float) -> None: ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_gamma_distribution_sample_properties(size, expected_shape, use_rng) -> None: +def test_gamma_distribution_sample_properties(size, expected_shape) -> None: dist = GammaDistribution(shape=2.0, scale=2.0) - kwargs = {"size": size} - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py index 87d8691f0..2687f1239 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py @@ -37,14 +37,9 @@ def test_lognormal_distribution_init_raises_error_for_invalid_sdlog( ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_lognormal_distribution_sample_properties(size, expected_shape, use_rng) -> None: +def test_lognormal_distribution_sample_properties(size, expected_shape) -> None: dist = LognormalDistribution(meanlog=0.0, sdlog=1.0) - kwargs = {} - kwargs["size"] = size - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py index 4b31ff14e..0d42e73ea 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py @@ -35,14 +35,9 @@ def test_normal_distribution_sample_raises_error_for_invalid_sigma( ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_normal_distribution_sample_properties(size, expected_shape, use_rng) -> None: +def test_normal_distribution_sample_properties(size, expected_shape) -> None: dist = NormalDistribution(mu=0.0, sigma=1.0) - kwargs = {} - kwargs["size"] = size - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py index 607a1208a..0dbb7a501 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py @@ -37,13 +37,9 @@ def test_poisson_distribution_init_invalid_lam(invalid_lam: float) -> None: ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_poisson_distribution_sample_properties(size, expected_shape, use_rng) -> None: +def test_poisson_distribution_sample_properties(size, expected_shape) -> None: dist = PoissonDistribution(lam=5.0) - kwargs = {"size": size} - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.int64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py index 857141abe..4a69797a3 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py @@ -43,16 +43,10 @@ def test_truncated_normal_distribution_init_invalid_sd(invalid_sd: float) -> Non ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_truncated_normal_distribution_sample_properties( - size, expected_shape, use_rng -) -> None: +def test_truncated_normal_distribution_sample_properties(size, expected_shape) -> None: a, b = 0.0, 10.0 dist = TruncatedNormalDistribution(mean=5.0, sd=2.0, a=a, b=b) - kwargs = {"size": size} - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py index 92df57749..ca03a2ec1 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py @@ -43,15 +43,10 @@ def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> No ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_uniform_distribution_sample_properties(size, expected_shape, use_rng) -> None: +def test_uniform_distribution_sample_properties(size, expected_shape) -> None: low, high = 10.0, 20.0 dist = UniformDistribution(low=low, high=high) - kwargs = {} - kwargs["size"] = size - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py index c7a56c286..a5679f5de 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py @@ -54,13 +54,9 @@ def test_weibull_distribution_init_invalid_scale(invalid_scale: float) -> None: ], ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], ) -@pytest.mark.parametrize("use_rng", [True, False], ids=["with_rng", "without_rng"]) -def test_weibull_distribution_sample_properties(size, expected_shape, use_rng) -> None: +def test_weibull_distribution_sample_properties(size, expected_shape) -> None: dist = WeibullDistribution(shape=2.0, scale=10.0) - kwargs = {"size": size} - if use_rng: - kwargs["rng"] = np.random.default_rng() - sample = dist.sample(**kwargs) + sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert sample.dtype == np.float64 From 3b75ca50a6427fec54cb260b452a76e05a811ea1 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 11 Jul 2025 12:00:39 -0400 Subject: [PATCH 29/73] Reduce chance transient sampling error in `test_distributions_abc.py` Take 10 samples rather than 1 to reduce the risk of an identical series of samples being drawn (causing a test failure when arrays are compared). --- .../gempyor_pkg/tests/distributions/test_distributions_abc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py b/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py index f769a1064..6080ee409 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py @@ -46,6 +46,6 @@ def test_distribution_abc_rng_behavior(dist) -> None: sample2 = dist.sample(rng=rng2) np.testing.assert_array_equal(sample1, sample2) - sample3_default = dist.sample() - sample4_default = dist.sample() + sample3_default = dist.sample(size=10) + sample4_default = dist.sample(size=10) assert not np.array_equal(sample3_default, sample4_default) From 711d25fe188b8a1f4eda1cd589bfa25d4bf1d53c Mon Sep 17 00:00:00 2001 From: "Carl A. B. Pearson" Date: Mon, 14 Jul 2025 09:29:46 -0400 Subject: [PATCH 30/73] Apply suggestions from code review use pydantic capabilities for default_factory Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- flepimop/gempyor_pkg/src/gempyor/distributions.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 546e6d4a1..10af5d46b 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -31,16 +31,7 @@ class DistributionABC(ABC, BaseModel): """Base class for distributions used in modifiers, likelihoods, etc.""" distribution: str - _rng: Generator | None = PrivateAttr(default=None) - - @property - def rng(self) -> Generator: - """ - Inheritable property for RNG. - """ - if self._rng is None: - self._rng = np.random.default_rng() - return self._rng + _rng: Generator = PrivateAttr(default_factory=np.random.default_rng) def sample( self, size: int | tuple[int, ...] = 1, rng: Generator | None = None @@ -56,7 +47,7 @@ def sample( A NumPy array of either floats/ints (depending on distribution) drawn from the distribution with shape `size`. """ - rng = rng if rng is not None else self.rng + rng = rng if rng is not None else self._rng return self._sample_from_generator(size=size, rng=rng) @abstractmethod From 69756ebd96722fd706c9ffaf5dfeef690b04eb08 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 14 Jul 2025 17:46:43 -0400 Subject: [PATCH 31/73] Abstract shape tests into `test_distributions_common.py` Move shape tests to something that is tested at the `DistributionsABC` level, small verbiage changes in parametrization, re-naming of `DistributionsABC` testing file. --- .../test_beta_distribution_class.py | 21 ++++----- .../test_binomial_distribution_class.py | 23 ++++------ ...ns_abc.py => test_distributions_common.py} | 45 ++++++++++++++++++- .../test_fixed_distribution_class.py | 20 +++------ .../test_gamma_distribution_class.py | 21 ++++----- .../test_lognormal_distribution_class.py | 21 ++++----- .../test_normal_distribution_class.py | 17 ------- .../test_poisson_distribution_class.py | 20 +++------ ...test_truncatednormal_distribution_class.py | 25 +++++------ .../test_uniform_distribution_class.py | 21 ++++----- .../test_weibull_distribution_class.py | 21 ++++----- 11 files changed, 121 insertions(+), 134 deletions(-) rename flepimop/gempyor_pkg/tests/distributions/{test_distributions_abc.py => test_distributions_common.py} (53%) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py index 0e724b870..f95cf6582 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py @@ -48,18 +48,15 @@ def test_beta_distribution_init_invalid_beta(invalid_beta: float) -> None: @pytest.mark.parametrize( - "size, expected_shape", + "alpha, beta", [ - ((3, 3), (3, 3)), - (5, (5,)), - ((2, 3, 4), (2, 3, 4)), + (2.0, 5.0), + (1.0, 1.0), + (0.5, 0.5), ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + ids=["alpha_gt_1", "uniform_case", "u_shaped_case"], ) -def test_beta_distribution_sample_properties(size, expected_shape) -> None: - dist = BetaDistribution(alpha=2.0, beta=5.0) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample >= 0) and np.all(sample <= 1) +def test_beta_distribution_sample_range(alpha: float, beta: float) -> None: + dist = BetaDistribution(alpha=alpha, beta=beta) + sample = dist.sample(size=(10, 10)) + assert np.all((sample >= 0.0) & (sample <= 1.0)) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py index 9317521a1..c1be91e59 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py @@ -45,7 +45,6 @@ def test_binomial_distribution_init_invalid_n(invalid_n: int) -> None: ids=["small_negative", "large_negative"], ) def test_binomial_distribution_init_invalid_p_below_zero(invalid_p: float) -> None: - """Test that initialization fails for probability 'p' less than 0.""" with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): BinomialDistribution(n=10, p=invalid_p) @@ -53,7 +52,7 @@ def test_binomial_distribution_init_invalid_p_below_zero(invalid_p: float) -> No @pytest.mark.parametrize( "invalid_p", [1.1, 25.0], - ids=["small_positive", "large_positive"], + ids=["small_gt_one", "large_gt_one"], ) def test_binomial_distribution_init_invalid_p_above_one(invalid_p: float) -> None: with pytest.raises(ValidationError, match="Input should be less than or equal to 1"): @@ -61,20 +60,16 @@ def test_binomial_distribution_init_invalid_p_above_one(invalid_p: float) -> Non @pytest.mark.parametrize( - "size, expected_shape", + "n, p", [ - ((4, 4), (4, 4)), - (30, (30,)), - ((2, 3, 4), (2, 3, 4)), + (10, 0.5), + (100, 0.1), + (20, 0.99), ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + ids=["n_10_p_0.5", "n_100_p_0.1", "n_20_p_0.99"], ) -def test_binomial_distribution_sample_properties(size, expected_shape) -> None: - n = 20 - dist = BinomialDistribution(n=n, p=0.5) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.int64 +def test_binomial_sample_properties(n: int, p: float) -> None: + dist = BinomialDistribution(n=n, p=p) + sample = dist.sample(size=(10, 10)) assert np.all(sample >= 0) assert np.all(sample <= n) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py similarity index 53% rename from flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py rename to flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py index 6080ee409..411be984f 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_distributions_abc.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py @@ -3,7 +3,7 @@ from gempyor.distributions import ( BetaDistribution, BinomialDistribution, - # No FixedDistribution; does not use RNGs + FixedDistribution, GammaDistribution, LognormalDistribution, NormalDistribution, @@ -49,3 +49,46 @@ def test_distribution_abc_rng_behavior(dist) -> None: sample3_default = dist.sample(size=10) sample4_default = dist.sample(size=10) assert not np.array_equal(sample3_default, sample4_default) + + +@pytest.mark.parametrize( + "dist", + [ + BetaDistribution(alpha=2.0, beta=5.0), + BinomialDistribution(n=10, p=0.5), + FixedDistribution(value=1.23), + GammaDistribution(shape=2.0, scale=1.5), + LognormalDistribution(meanlog=0.0, sdlog=1.0), + NormalDistribution(mu=2.3, sigma=4.5), + PoissonDistribution(lam=3.0), + TruncatedNormalDistribution(mean=5.0, sd=2.0, a=0.0, b=10.0), + UniformDistribution(low=-0.5, high=1.5), + WeibullDistribution(shape=2.5, scale=5.0), + ], + ids=[ + "Beta", + "Binomial", + "Fixed" + "Gamma", + "Lognormal", + "Normal", + "Poisson", + "TruncatedNormal", + "Uniform", + "Weibull", + ], +) +@pytest.mark.parametrize( + "size, expected_shape", + [ + (15, (15,)), + ((2, 5), (2, 5)), + ((2, 3, 2), (2, 3, 2)), + ], + ids=["integer_size", "2d_tuple_size", "3d_tuple_size"], +) +def test_distribution_abc_sample_properties(dist, size, expected_shape) -> None: + sample = dist.sample(size=size) + assert isinstance(sample, np.ndarray) + assert sample.shape == expected_shape + assert np.issubdtype(sample.dtype, np.number) \ No newline at end of file diff --git a/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py index bacc63404..6330a8579 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py @@ -7,26 +7,18 @@ @pytest.mark.parametrize( "value", [123.45, -100.0, 0.0], ids=["positive", "negative", "zero"] ) -def test_fixed_distribution_ini(value: float) -> None: +def test_fixed_distribution_init(value: float) -> None: dist = FixedDistribution(value=value) assert dist.value == value assert dist.distribution == "fixed" -@pytest.mark.parametrize("value", [5.5, -1.2], ids=["value_positive", "value_negative"]) @pytest.mark.parametrize( - "size, expected_shape", - [ - ((2, 5), (2, 5)), - (5, (5,)), - ((2, 3, 4), (2, 3, 4)), - ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + "value", + [0.0, 5.5, -1.2], + ids=["zero", "positive", "negative"] ) -def test_fixed_distribution_sample_properties(value: float, size, expected_shape) -> None: +def test_fixed_distribution_samples_values(value: float) -> None: dist = FixedDistribution(value=value) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 + sample = dist.sample(size=(5, 5)) assert np.all(sample == value) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py index 7e23479be..4e2c981b7 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py @@ -48,18 +48,15 @@ def test_gamma_distribution_init_invalid_scale(invalid_scale: float) -> None: @pytest.mark.parametrize( - "size, expected_shape", + "shape, scale", [ - ((3, 3), (3, 3)), - (5, (5,)), - ((2, 3, 4), (2, 3, 4)), + (2.0, 2.0), + (1.0, 1.0), + (5.0, 0.5), ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + ids=["shape_2_scale_2", "standard_gamma", "shape_5_scale_0.5"], ) -def test_gamma_distribution_sample_properties(size, expected_shape) -> None: - dist = GammaDistribution(shape=2.0, scale=2.0) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample > 0) +def test_gamma_sample_is_non_negative(shape: float, scale: float) -> None: + dist = GammaDistribution(shape=shape, scale=scale) + sample = dist.sample(size=(10, 10)) + assert np.all(sample > 0.0) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py index 2687f1239..e2f395abd 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_lognormal_distribution_class.py @@ -29,18 +29,15 @@ def test_lognormal_distribution_init_raises_error_for_invalid_sdlog( @pytest.mark.parametrize( - "size, expected_shape", + "meanlog, sdlog", [ - ((2, 3), (2, 3)), - (25, (25,)), - ((2, 3, 4), (2, 3, 4)), + (0.0, 1.0), + (-2.0, 0.5), + (5.0, 3.0), ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + ids=["standard", "low_variance", "high_variance"], ) -def test_lognormal_distribution_sample_properties(size, expected_shape) -> None: - dist = LognormalDistribution(meanlog=0.0, sdlog=1.0) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample > 0) +def test_lognormal_sample_is_non_negative(meanlog: float, sdlog: float) -> None: + dist = LognormalDistribution(meanlog=meanlog, sdlog=sdlog) + sample = dist.sample(size=(10, 10)) + assert np.all(sample > 0.0) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py index 0d42e73ea..1bb024339 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_normal_distribution_class.py @@ -24,20 +24,3 @@ def test_normal_distribution_sample_raises_error_for_invalid_sigma( ) -> None: with pytest.raises(ValidationError, match="Input should be greater than 0"): NormalDistribution(mu=10.0, sigma=invalid_sigma) - - -@pytest.mark.parametrize( - "size, expected_shape", - [ - ((3, 2), (3, 2)), - (10, (10,)), - ((2, 3, 4), (2, 3, 4)), - ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], -) -def test_normal_distribution_sample_properties(size, expected_shape) -> None: - dist = NormalDistribution(mu=0.0, sigma=1.0) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 diff --git a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py index 0dbb7a501..1abc9a175 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py @@ -29,18 +29,12 @@ def test_poisson_distribution_init_invalid_lam(invalid_lam: float) -> None: @pytest.mark.parametrize( - "size, expected_shape", - [ - ((2, 5), (2, 5)), - (15, (15,)), - ((2, 3, 2), (2, 3, 2)), - ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + "lam", + [0.0, 1.0, 50.5], + ids=["zero_rate", "one_rate", "high_float_rate"], ) -def test_poisson_distribution_sample_properties(size, expected_shape) -> None: - dist = PoissonDistribution(lam=5.0) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.int64 +def test_poisson_samples_are_non_negative(lam: float) -> None: + dist = PoissonDistribution(lam=lam) + sample = dist.sample(size=(10, 10)) assert np.all(sample >= 0) + diff --git a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py index 4a69797a3..5a0fba0a4 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py @@ -35,20 +35,17 @@ def test_truncated_normal_distribution_init_invalid_sd(invalid_sd: float) -> Non @pytest.mark.parametrize( - "size, expected_shape", + "mean, sd, a, b", [ - ((3, 5), (3, 5)), - (10, (10,)), - ((2, 3, 4), (2, 3, 4)), + (5.0, 2.0, 0.0, 10.0), + (0.0, 1.0, -1.0, 1.0), + (100.0, 10.0, 90.0, 110.0), ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + ids=["mean_in_range", "std_normal_truncated", "high_mean"], ) -def test_truncated_normal_distribution_sample_properties(size, expected_shape) -> None: - a, b = 0.0, 10.0 - dist = TruncatedNormalDistribution(mean=5.0, sd=2.0, a=a, b=b) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample >= a) - assert np.all(sample <= b) +def test_truncated_normal_distribution_sample_range( + mean: float, sd: float, a: float, b: float +) -> None: + dist = TruncatedNormalDistribution(mean=mean, sd=sd, a=a, b=b) + sample = dist.sample(size=(10, 10)) + assert np.all((sample >= a) & (sample <= b)) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py index ca03a2ec1..510504c02 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py @@ -35,20 +35,15 @@ def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> No @pytest.mark.parametrize( - "size, expected_shape", + "low, high", [ - ((6, 4), (6, 4)), - (20, (20,)), - ((2, 3, 4), (2, 3, 4)), + (10.0, 20.0), + (-5.0, 5.0), + (0.0, 1.0), ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + ids=["positive_range", "centered_range", "unit_range"], ) -def test_uniform_distribution_sample_properties(size, expected_shape) -> None: - low, high = 10.0, 20.0 +def test_uniform_distribution_sample_range(low: float, high: float) -> None: dist = UniformDistribution(low=low, high=high) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample >= low) - assert np.all(sample < high) + sample = dist.sample(size=(10, 10)) + assert np.all((sample >= low) & (sample < high)) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py index a5679f5de..14a949ae5 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py @@ -46,18 +46,15 @@ def test_weibull_distribution_init_invalid_scale(invalid_scale: float) -> None: @pytest.mark.parametrize( - "size, expected_shape", + "shape, scale", [ - ((5, 5), (5, 5)), - (5, (5,)), - ((2, 3, 4), (2, 3, 4)), + (2.0, 10.0), + (1.0, 1.0), + (5.0, 1.0), ], - ids=["2d_tuple_size", "integer_size", "3d_tuple_size"], + ids=["shape_2_scale_10", "exponential_case", "shape_5_scale_1"], ) -def test_weibull_distribution_sample_properties(size, expected_shape) -> None: - dist = WeibullDistribution(shape=2.0, scale=10.0) - sample = dist.sample(size=size) - assert isinstance(sample, np.ndarray) - assert sample.shape == expected_shape - assert sample.dtype == np.float64 - assert np.all(sample > 0) +def test_weibull_sample_is_non_negative(shape: float, scale: float) -> None: + dist = WeibullDistribution(shape=shape, scale=scale) + sample = dist.sample(size=(10, 10)) + assert np.all(sample > 0.0) From 1e32b6d318e50777f040eaa2ee00219885bab2eb Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 15 Jul 2025 11:24:24 -0400 Subject: [PATCH 32/73] Typo fix and additional bounds test for `TruncatedNormalDistribution` --- .../distributions/test_distributions_common.py | 2 +- .../test_truncatednormal_distribution_class.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py index 411be984f..5e3642979 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py @@ -68,7 +68,7 @@ def test_distribution_abc_rng_behavior(dist) -> None: ids=[ "Beta", "Binomial", - "Fixed" + "Fixed", "Gamma", "Lognormal", "Normal", diff --git a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py index 5a0fba0a4..5af4fe6f1 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py @@ -34,6 +34,19 @@ def test_truncated_normal_distribution_init_invalid_sd(invalid_sd: float) -> Non TruncatedNormalDistribution(mean=0.0, sd=invalid_sd, a=0.0, b=1.0) +@pytest.mark.parametrize( + "a, b", + [ + (10.0, 5.0), + (5.0, 5.0), + ], + ids=["b_less_than_a", "b_equal_to_a"], +) +def test_truncated_normal_distribution_init_invalid_bounds(a: float, b: float) -> None: + with pytest.raises(ValidationError, match="must be strictly greater than"): + TruncatedNormalDistribution(mean=5.0, sd=2.0, a=a, b=b) + + @pytest.mark.parametrize( "mean, sd, a, b", [ From d5a5e983da315e45351b27821d90543fdd6a0970 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 15 Jul 2025 15:11:27 -0400 Subject: [PATCH 33/73] Build parameter `allow_edge_cases: bool = False` to distribution classes --- .../gempyor_pkg/src/gempyor/distributions.py | 57 ++++++++++++++----- .../test_beta_distribution_class.py | 2 +- .../test_distributions_common.py | 2 +- .../test_fixed_distribution_class.py | 8 +-- .../test_gamma_distribution_class.py | 2 +- .../test_poisson_distribution_class.py | 41 ++++++++----- ...test_truncatednormal_distribution_class.py | 39 +++++++++---- .../test_uniform_distribution_class.py | 37 ++++++++---- .../test_weibull_distribution_class.py | 2 +- 9 files changed, 129 insertions(+), 61 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 10af5d46b..a8551068b 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -31,6 +31,8 @@ class DistributionABC(ABC, BaseModel): """Base class for distributions used in modifiers, likelihoods, etc.""" distribution: str + allow_edge_cases: bool = False + _rng: Generator = PrivateAttr(default_factory=np.random.default_rng) def sample( @@ -149,13 +151,18 @@ def _sample_from_generator( return rng.uniform(low=self.low, high=self.high, size=size) @model_validator(mode="after") - def _ensure_high_greater_than_low(self) -> "UniformDistribution": - """Ensure that high is greater than low.""" - if self.high < self.low or isclose(self.high, self.low): - raise ValueError( - f"The 'high' value, {self.high}, must be " - f"greater than the 'low' value, {self.low}." - ) + def _validate_bounds(self) -> "UniformDistribution": + """Validate bounds based on whether or not edge cases are allowed.""" + if self.allow_edge_cases: + if self.high < self.low: + raise ValueError( + f"'high' value ({self.high}) must be ≥ to the 'low' value ({self.low})." + ) + else: + if self.high < self.low or isclose(self.high, self.low): + raise ValueError( + f"'high' value ({self.high}) must be > 'low' value ({self.low})." + ) return self @@ -220,6 +227,11 @@ def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator ) -> npt.NDArray[np.float64]: """Sampling logic for truncated normal distributions.""" + if ( + isclose(self.a, self.b) and self.allow_edge_cases + ): # use this logic b/c scipy.truncnorm doesn't support equal bounds + return np.full(size, self.a) + lower = (self.a - self.mean) / self.sd upper = (self.b - self.mean) / self.sd return truncnorm.rvs( @@ -232,13 +244,18 @@ def _sample_from_generator( ) @model_validator(mode="after") - def _ensure_b_greater_than_a(self) -> "TruncatedNormalDistribution": - """Ensure that the upper bound 'b' is strictly greater than the lower bound 'a'.""" - if self.b < self.a or isclose(self.b, self.a): - raise ValueError( - f"The upper bound 'b' ({self.b}) must be strictly " - f"greater than the lower bound 'a' ({self.a})." - ) + def _validate_bounds(self) -> "TruncatedNormalDistribution": + """Validate bounds based on whether or not edge cases are allowed..""" + if self.allow_edge_cases: + if self.b < self.a: + raise ValueError( + f"Upper bound 'b' ({self.b}) must be ≥ to lower bound 'a' ({self.a})." + ) + else: + if self.b < self.a or isclose(self.b, self.a): + raise ValueError( + f"Upper bound 'b' ({self.b}) must be > lower bound 'a' ({self.a})." + ) return self @@ -262,7 +279,7 @@ class PoissonDistribution(DistributionABC): """ distribution: Literal["poisson"] = "poisson" - lam: float = Field(..., ge=0) + lam: float def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -270,6 +287,16 @@ def _sample_from_generator( """Sampling logic for Poisson distributions.""" return rng.poisson(lam=self.lam, size=size) + @model_validator(mode="after") + def _validate_lambda(self) -> "PoissonDistribution": + if self.allow_edge_cases: + if self.lam < 0: + raise ValueError("Input for 'lam' must be ≥ 0.") + else: + if self.lam <= 0: + raise ValueError("Input for 'lam' must be > 0.") + return self + class BinomialDistribution(DistributionABC): """ diff --git a/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py index f95cf6582..f15d76ef1 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_beta_distribution_class.py @@ -52,7 +52,7 @@ def test_beta_distribution_init_invalid_beta(invalid_beta: float) -> None: [ (2.0, 5.0), (1.0, 1.0), - (0.5, 0.5), + (0.5, 0.5), ], ids=["alpha_gt_1", "uniform_case", "u_shaped_case"], ) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py index 5e3642979..e918b5708 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py @@ -91,4 +91,4 @@ def test_distribution_abc_sample_properties(dist, size, expected_shape) -> None: sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape - assert np.issubdtype(sample.dtype, np.number) \ No newline at end of file + assert np.issubdtype(sample.dtype, np.number) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py index 6330a8579..be45cb340 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_fixed_distribution_class.py @@ -13,12 +13,8 @@ def test_fixed_distribution_init(value: float) -> None: assert dist.distribution == "fixed" -@pytest.mark.parametrize( - "value", - [0.0, 5.5, -1.2], - ids=["zero", "positive", "negative"] -) +@pytest.mark.parametrize("value", [0.0, 5.5, -1.2], ids=["zero", "positive", "negative"]) def test_fixed_distribution_samples_values(value: float) -> None: dist = FixedDistribution(value=value) - sample = dist.sample(size=(5, 5)) + sample = dist.sample(size=(5, 5)) assert np.all(sample == value) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py index 4e2c981b7..126f23238 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_gamma_distribution_class.py @@ -51,7 +51,7 @@ def test_gamma_distribution_init_invalid_scale(invalid_scale: float) -> None: "shape, scale", [ (2.0, 2.0), - (1.0, 1.0), + (1.0, 1.0), (5.0, 0.5), ], ids=["shape_2_scale_2", "standard_gamma", "shape_5_scale_0.5"], diff --git a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py index 1abc9a175..1fda6441f 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py @@ -8,33 +8,46 @@ @pytest.mark.parametrize( "valid_lam", [ - (0.0), (1.0), (25.5), ], - ids=["zero_rate", "unit_rate", "positive_rate"], + ids=["unit_rate", "positive_rate"], ) def test_poisson_distribution_init_valid(valid_lam: float) -> None: dist = PoissonDistribution(lam=valid_lam) assert dist.lam == valid_lam assert dist.distribution == "poisson" + assert dist.allow_edge_cases is False -@pytest.mark.parametrize( - "invalid_lam", [-0.1, -10.0], ids=["small_negative", "large_negative"] -) -def test_poisson_distribution_init_invalid_lam(invalid_lam: float) -> None: - with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): - PoissonDistribution(lam=invalid_lam) +def test_poisson_distribution_init_fails_on_edge_case_false_by_default() -> None: + with pytest.raises(ValidationError, match="Input for 'lam' must be > 0."): + PoissonDistribution(lam=0.0) + + +def test_poisson_distribution_init_succeeds_on_edge_case_true() -> None: + dist = PoissonDistribution(lam=0.0, allow_edge_cases=True) + assert dist.lam == 0.0 + assert dist.allow_edge_cases is True + + +def test_poisson_distribution_init_fails_on_negative_lam() -> None: + with pytest.raises(ValidationError, match="Input for 'lam' must be > 0."): + PoissonDistribution(lam=-5.0, allow_edge_cases=False) + with pytest.raises(ValidationError, match="Input for 'lam' must be ≥ 0."): + PoissonDistribution(lam=-5.0, allow_edge_cases=True) @pytest.mark.parametrize( - "lam", - [0.0, 1.0, 50.5], - ids=["zero_rate", "one_rate", "high_float_rate"], + "lam, allow_edges", + [ + (1.0, False), + (50.5, False), + (0.0, True), + ], + ids=["unit_rate", "high_float_rate", "zero_rate_edge_case"], ) -def test_poisson_samples_are_non_negative(lam: float) -> None: - dist = PoissonDistribution(lam=lam) +def test_poisson_samples_are_non_negative(lam: float, allow_edges: bool) -> None: + dist = PoissonDistribution(lam=lam, allow_edge_cases=allow_edges) sample = dist.sample(size=(10, 10)) assert np.all(sample >= 0) - diff --git a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py index 5af4fe6f1..01435488a 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py @@ -24,6 +24,7 @@ def test_truncated_normal_distribution_init_valid( assert dist.a == a assert dist.b == b assert dist.distribution == "truncnorm" + assert dist.allow_edge_cases is False @pytest.mark.parametrize( @@ -34,17 +35,25 @@ def test_truncated_normal_distribution_init_invalid_sd(invalid_sd: float) -> Non TruncatedNormalDistribution(mean=0.0, sd=invalid_sd, a=0.0, b=1.0) -@pytest.mark.parametrize( - "a, b", - [ - (10.0, 5.0), - (5.0, 5.0), - ], - ids=["b_less_than_a", "b_equal_to_a"], -) -def test_truncated_normal_distribution_init_invalid_bounds(a: float, b: float) -> None: - with pytest.raises(ValidationError, match="must be strictly greater than"): - TruncatedNormalDistribution(mean=5.0, sd=2.0, a=a, b=b) +def test_truncated_normal_init_fails_on_edge_case_false_by_default() -> None: + with pytest.raises(ValidationError, match="must be > lower bound 'a'"): + TruncatedNormalDistribution(mean=5.0, sd=2.0, a=5.0, b=5.0) + + +def test_truncated_normal_init_succeeds_on_edge_case_true() -> None: + dist = TruncatedNormalDistribution( + mean=5.0, sd=2.0, a=5.0, b=5.0, allow_edge_cases=True + ) + assert dist.a == 5.0 + assert dist.b == 5.0 + assert dist.allow_edge_cases is True + + +def test_truncated_normal_init_fails_on_b_less_than_a() -> None: + with pytest.raises(ValidationError, match="must be > lower bound 'a'"): + TruncatedNormalDistribution(mean=5.0, sd=2.0, a=10.0, b=5.0) + with pytest.raises(ValidationError, match="must be ≥ to lower bound 'a'"): + TruncatedNormalDistribution(mean=5.0, sd=2.0, a=10.0, b=5.0, allow_edge_cases=True) @pytest.mark.parametrize( @@ -62,3 +71,11 @@ def test_truncated_normal_distribution_sample_range( dist = TruncatedNormalDistribution(mean=mean, sd=sd, a=a, b=b) sample = dist.sample(size=(10, 10)) assert np.all((sample >= a) & (sample <= b)) + + +def test_truncated_normal_distribution_sample_range_edge_case() -> None: + dist = TruncatedNormalDistribution( + mean=5.0, sd=2.0, a=7.0, b=7.0, allow_edge_cases=True + ) + sample = dist.sample(size=(10, 10)) + assert np.all(sample == 7.0) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py index 510504c02..5e978ac9d 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py @@ -19,19 +19,28 @@ def test_uniform_distribution_init_valid(low: float, high: float) -> None: assert dist.low == low assert dist.high == high assert dist.distribution == "uniform" + assert dist.allow_edge_cases is False -@pytest.mark.parametrize( - "low, high", - [ - (10.0, 0.0), - (5.0, 5.0), - ], - ids=["high_less_than_low", "high_equals_low"], -) -def test_uniform_distribution_init_invalid_bounds(low: float, high: float) -> None: - with pytest.raises(ValidationError, match="must be greater than the 'low' value"): - UniformDistribution(low=low, high=high) +def test_uniform_init_fails_on_edge_case_false_by_default() -> None: + with pytest.raises(ValidationError, match="'high' value .* must be > 'low' value"): + UniformDistribution(low=5.0, high=5.0) + + +def test_uniform_init_succeeds_on_edge_case_true() -> None: + dist = UniformDistribution(low=5.0, high=5.0, allow_edge_cases=True) + assert dist.low == 5.0 + assert dist.high == 5.0 + assert dist.allow_edge_cases is True + + +def test_uniform_init_fails_on_high_less_than_low() -> None: + with pytest.raises(ValidationError, match="'high' value .* must be > 'low' value"): + UniformDistribution(low=10.0, high=5.0, allow_edge_cases=False) + with pytest.raises( + ValidationError, match="'high' value .* must be ≥ to the 'low' value" + ): + UniformDistribution(low=10.0, high=5.0, allow_edge_cases=True) @pytest.mark.parametrize( @@ -47,3 +56,9 @@ def test_uniform_distribution_sample_range(low: float, high: float) -> None: dist = UniformDistribution(low=low, high=high) sample = dist.sample(size=(10, 10)) assert np.all((sample >= low) & (sample < high)) + + +def test_uniform_distribution_sample_range_edge_case() -> None: + dist = UniformDistribution(low=10.0, high=10.0, allow_edge_cases=True) + sample = dist.sample(size=(10, 10)) + assert np.all(sample == 10.0) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py index 14a949ae5..6f619346b 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_weibull_distribution_class.py @@ -49,7 +49,7 @@ def test_weibull_distribution_init_invalid_scale(invalid_scale: float) -> None: "shape, scale", [ (2.0, 10.0), - (1.0, 1.0), + (1.0, 1.0), (5.0, 1.0), ], ids=["shape_2_scale_10", "exponential_case", "shape_5_scale_1"], From 8185916e29ce1af08928b92585a0bc4484eb0dd7 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Wed, 16 Jul 2025 13:12:47 -0400 Subject: [PATCH 34/73] Apply suggestions from @TimothyWillard 's review --- .../gempyor_pkg/src/gempyor/distributions.py | 102 +++++++++++++----- .../test_binomial_distribution_class.py | 33 +++--- .../test_distributions_common.py | 95 +++++----------- .../test_poisson_distribution_class.py | 21 ++-- ...test_truncatednormal_distribution_class.py | 30 +++--- .../test_uniform_distribution_class.py | 22 ++-- 6 files changed, 157 insertions(+), 146 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index a8551068b..fb9a9da8c 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -138,6 +138,16 @@ class UniformDistribution(DistributionABC): array([[ 0.37775688, 1.21719584, 0.89473606, -0.3116453 , 1.4512447 ], [ 1.0222794 , 1.07212861, -0.24377273, 0.40077188, 0.24159605], [ 1.35352998, 0.78773024, 1.14552323, 0.3868284 , -0.04552256]]) + >>> # With `low == high` and `allow_edge_cases=True`, all samples == `low`. + >>> dist_edge = UniformDistribution(low=5.0, high=5.0, allow_edge_cases=True) + >>> dist_edge.sample(size=5) + array([5., 5., 5., 5., 5.]) + >>> # Without `allow_edge_cases` set to True, it fails by default when `low == high`. + >>> UniformDistribution(low=5.0, high=5.0) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for UniformDistribution + Value error, Upper bound `high`, 5.0, must be > to lower bound `low`, 5.0. [type=value_error, ... """ distribution: Literal["uniform"] = "uniform" @@ -153,16 +163,13 @@ def _sample_from_generator( @model_validator(mode="after") def _validate_bounds(self) -> "UniformDistribution": """Validate bounds based on whether or not edge cases are allowed.""" - if self.allow_edge_cases: - if self.high < self.low: - raise ValueError( - f"'high' value ({self.high}) must be ≥ to the 'low' value ({self.low})." - ) - else: - if self.high < self.low or isclose(self.high, self.low): - raise ValueError( - f"'high' value ({self.high}) must be > 'low' value ({self.low})." - ) + if self.high < self.low or ( + not self.allow_edge_cases and isclose(self.high, self.low) + ): + op = ">=" if self.allow_edge_cases else ">" + raise ValueError( + f"Upper bound `high`, {self.high}, must be {op} to lower bound `low`, {self.low}." + ) return self @@ -214,6 +221,16 @@ class TruncatedNormalDistribution(DistributionABC): array([[1.07000038, 2.18016199, 1.66002835, 0.28689654, 3.04332767], [1.83818339, 1.9153892 , 0.37639339, 1.09435167, 0.92629918], [2.54134935, 1.52545861, 2.04022114, 1.07959285, 0.6142512 ]]) + >>> # With `a == b` and `allow_edge_cases=True`, all samples == `a`. + >>> dist_edge = TruncatedNormalDistribution(mean=5.0, sd=2.0, a=7.0, b=7.0, allow_edge_cases=True) + >>> dist_edge.sample(size=5) + array([7., 7., 7., 7., 7.]) + >>> # Withoug `allow_edge_cases` set to True, it fails by default when `a == b`. + >>> TruncatedNormalDistribution(mean=5.0, sd=2.0, a=7.0, b=7.0) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for TruncatedNormalDistribution + Value error, Upper bound `b`, 7.0, must be > to lower bound `a`, 7.0. [type=value_error, ... """ # pylint: enable=line-too-long @@ -246,16 +263,11 @@ def _sample_from_generator( @model_validator(mode="after") def _validate_bounds(self) -> "TruncatedNormalDistribution": """Validate bounds based on whether or not edge cases are allowed..""" - if self.allow_edge_cases: - if self.b < self.a: - raise ValueError( - f"Upper bound 'b' ({self.b}) must be ≥ to lower bound 'a' ({self.a})." - ) - else: - if self.b < self.a or isclose(self.b, self.a): - raise ValueError( - f"Upper bound 'b' ({self.b}) must be > lower bound 'a' ({self.a})." - ) + if self.b < self.a or (not self.allow_edge_cases and isclose(self.a, self.b)): + op = ">=" if self.allow_edge_cases else ">" + raise ValueError( + f"Upper bound `b`, {self.b}, must be {op} to lower bound `a`, {self.a}." + ) return self @@ -276,10 +288,20 @@ class PoissonDistribution(DistributionABC): array([[4, 5, 1, 7, 1], [4, 2, 2, 5, 4], [1, 6, 2, 5, 0]]) + >>> # With `lam=0` and `allow_edge_cases=True`, all samples will be 0. + >>> dist_edge = PoissonDistribution(lam=0.0, allow_edge_cases=True) + >>> dist_edge.sample(size=5) + array([0, 0, 0, 0, 0]) + >>> # Without `allow_edge_cases` explicitly set to True, it fails by default. + >>> PoissonDistribution(lam=0.0) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for PoissonDistribution + Value error, Input for `lam` cannot be zero when `allow_edge_cases` is `False`. [type=value_error, ... """ distribution: Literal["poisson"] = "poisson" - lam: float + lam: float = Field(..., ge=0.0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -289,12 +311,10 @@ def _sample_from_generator( @model_validator(mode="after") def _validate_lambda(self) -> "PoissonDistribution": - if self.allow_edge_cases: - if self.lam < 0: - raise ValueError("Input for 'lam' must be ≥ 0.") - else: - if self.lam <= 0: - raise ValueError("Input for 'lam' must be > 0.") + if not self.allow_edge_cases and isclose(self.lam, 0.0): + raise ValueError( + "Input for `lam` cannot be zero when `allow_edge_cases` is `False`." + ) return self @@ -315,11 +335,21 @@ class BinomialDistribution(DistributionABC): array([[5, 7, 6, 3, 8], [6, 6, 3, 5, 4], [7, 6, 6, 5, 4]]) + >>> # It succeeds with `p=0` or `p=1` when `allow_edge_cases=True`. + >>> dist_edge = BinomialDistribution(n=10, p=1.0, allow_edge_cases=True) + >>> dist_edge.sample(size=5) + array([10, 10, 10, 10, 10]) + >>> # Without `allow_edge_cases` set to True, it fails by default when `p=0` or `p=1`. + >>> BinomialDistribution(n=10, p=0.0) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for BinomialDistribution + Value error, `p` cannot be 0 or 1 when `allow_edge_cases` is `False`. [type=value_error, ... """ distribution: Literal["binomial"] = "binomial" - n: int = Field(..., ge=0) - p: float = Field(..., ge=0, le=1) + n: int = Field(..., gt=0) + p: float = Field(..., gt=0.0, lt=1.0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -327,6 +357,20 @@ def _sample_from_generator( """Sampling logic for binomial distributions.""" return rng.binomial(n=self.n, p=self.p, size=size) + @model_validator(mode="after") + def _validate_params(self) -> "BinomialDistribution": + """Validate params based on whether or not edge cases are allowed.""" + if not self.allow_edge_cases: + if self.n == 0: + raise ValueError( + "Input for `n` cannot be zero when `allow_edge_cases` is `False`." + ) + if isclose(self.p, 0.0) or isclose(self.p, 1.0): + raise ValueError( + "Input for `p` cannot be 0 or 1 when `allow_edge_cases` is `False`." + ) + return self + class GammaDistribution(DistributionABC): """ diff --git a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py index c1be91e59..e3e7f37c6 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py @@ -9,17 +9,13 @@ "n, p", [ (10, 0.5), - (0, 0.5), - (100, 0.0), - (100, 1.0), (50, 0.25), + (1, 0.99), ], ids=[ - "standard_case", - "zero_trials", - "zero_probability", - "full_probability", + "standard_case_1", "standard_case_2", + "min_valid_n", ], ) def test_binomial_distribution_init_valid(n: int, p: float) -> None: @@ -27,35 +23,36 @@ def test_binomial_distribution_init_valid(n: int, p: float) -> None: assert dist.n == n assert dist.p == p assert dist.distribution == "binomial" + assert dist.allow_edge_cases is False @pytest.mark.parametrize( "invalid_n", - [-1, -10], - ids=["-1", "-10"], + [0, -1, -10], + ids=["zero", "-1", "-10"], ) def test_binomial_distribution_init_invalid_n(invalid_n: int) -> None: - with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): + with pytest.raises(ValidationError, match="Input should be greater than 0"): BinomialDistribution(n=invalid_n, p=0.5) @pytest.mark.parametrize( "invalid_p", - [-0.1, -10.0], - ids=["small_negative", "large_negative"], + [0.0, -0.1, -10.0], + ids=["zero", "small_negative", "large_negative"], ) -def test_binomial_distribution_init_invalid_p_below_zero(invalid_p: float) -> None: - with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): +def test_binomial_distribution_init_invalid_p_below_range(invalid_p: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): BinomialDistribution(n=10, p=invalid_p) @pytest.mark.parametrize( "invalid_p", - [1.1, 25.0], - ids=["small_gt_one", "large_gt_one"], + [1.0, 1.1, 25.0], + ids=["one", "small_gt_one", "large_gt_one"], ) -def test_binomial_distribution_init_invalid_p_above_one(invalid_p: float) -> None: - with pytest.raises(ValidationError, match="Input should be less than or equal to 1"): +def test_binomial_distribution_init_invalid_p_above_range(invalid_p: float) -> None: + with pytest.raises(ValidationError, match="Input should be less than 1"): BinomialDistribution(n=10, p=invalid_p) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py index e918b5708..e95a0bcf8 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py @@ -1,83 +1,37 @@ import numpy as np import pytest -from gempyor.distributions import ( - BetaDistribution, - BinomialDistribution, - FixedDistribution, - GammaDistribution, - LognormalDistribution, - NormalDistribution, - PoissonDistribution, - TruncatedNormalDistribution, - UniformDistribution, - WeibullDistribution, -) +from gempyor.distributions import DistributionABC -@pytest.mark.parametrize( - "dist", - [ - BetaDistribution(alpha=2.0, beta=5.0), - BinomialDistribution(n=10, p=0.5), - GammaDistribution(shape=2.0, scale=1.5), - LognormalDistribution(meanlog=0.0, sdlog=1.0), - NormalDistribution(mu=2.3, sigma=4.5), - PoissonDistribution(lam=3.0), - TruncatedNormalDistribution(mean=5.0, sd=2.0, a=0.0, b=10.0), - UniformDistribution(low=-0.5, high=1.5), - WeibullDistribution(shape=2.5, scale=5.0), - ], - ids=[ - "Beta", - "Binomial", - "Gamma", - "Lognormal", - "Normal", - "Poisson", - "TruncatedNormal", - "Uniform", - "Weibull", - ], -) -def test_distribution_abc_rng_behavior(dist) -> None: +class DummyDistribution(DistributionABC): + """A simple dummy implementation for testing DistributionABC logic.""" + + distribution: str = "dummy" + + def _sample_from_generator( + self, size: int | tuple[int, ...], rng: np.random.Generator + ) -> np.ndarray: + """A fake sampling implementation.""" + return rng.random(size) + + +def test_reproducible_sampling_with_seeded_rng() -> None: + dist = DummyDistribution() rng1 = np.random.default_rng(seed=42) sample1 = dist.sample(rng=rng1) rng2 = np.random.default_rng(seed=42) sample2 = dist.sample(rng=rng2) - np.testing.assert_array_equal(sample1, sample2) + assert np.array_equal(sample1, sample2) - sample3_default = dist.sample(size=10) - sample4_default = dist.sample(size=10) - assert not np.array_equal(sample3_default, sample4_default) + +def test_stochastic_sampling_with_default_rng() -> None: + dist = DummyDistribution() + sample1 = dist.sample(size=10) + sample2 = dist.sample(size=10) + + assert not np.array_equal(sample1, sample2) -@pytest.mark.parametrize( - "dist", - [ - BetaDistribution(alpha=2.0, beta=5.0), - BinomialDistribution(n=10, p=0.5), - FixedDistribution(value=1.23), - GammaDistribution(shape=2.0, scale=1.5), - LognormalDistribution(meanlog=0.0, sdlog=1.0), - NormalDistribution(mu=2.3, sigma=4.5), - PoissonDistribution(lam=3.0), - TruncatedNormalDistribution(mean=5.0, sd=2.0, a=0.0, b=10.0), - UniformDistribution(low=-0.5, high=1.5), - WeibullDistribution(shape=2.5, scale=5.0), - ], - ids=[ - "Beta", - "Binomial", - "Fixed", - "Gamma", - "Lognormal", - "Normal", - "Poisson", - "TruncatedNormal", - "Uniform", - "Weibull", - ], -) @pytest.mark.parametrize( "size, expected_shape", [ @@ -87,7 +41,8 @@ def test_distribution_abc_rng_behavior(dist) -> None: ], ids=["integer_size", "2d_tuple_size", "3d_tuple_size"], ) -def test_distribution_abc_sample_properties(dist, size, expected_shape) -> None: +def test_distribution_abc_sample_properties(size, expected_shape) -> None: + dist = DummyDistribution() sample = dist.sample(size=size) assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape diff --git a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py index 1fda6441f..3a1e31bac 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_poisson_distribution_class.py @@ -21,7 +21,10 @@ def test_poisson_distribution_init_valid(valid_lam: float) -> None: def test_poisson_distribution_init_fails_on_edge_case_false_by_default() -> None: - with pytest.raises(ValidationError, match="Input for 'lam' must be > 0."): + with pytest.raises( + ValidationError, + match="Input for `lam` cannot be zero when `allow_edge_cases` is `False`.", + ): PoissonDistribution(lam=0.0) @@ -31,11 +34,17 @@ def test_poisson_distribution_init_succeeds_on_edge_case_true() -> None: assert dist.allow_edge_cases is True -def test_poisson_distribution_init_fails_on_negative_lam() -> None: - with pytest.raises(ValidationError, match="Input for 'lam' must be > 0."): - PoissonDistribution(lam=-5.0, allow_edge_cases=False) - with pytest.raises(ValidationError, match="Input for 'lam' must be ≥ 0."): - PoissonDistribution(lam=-5.0, allow_edge_cases=True) +@pytest.mark.parametrize( + "invalid_lam", + [ + -5.0, + -50.5, + ], + ids=["small_negative_lam", "large_negative_lam"], +) +def test_poisson_distribution_init_fails_on_negative_lam(invalid_lam: float) -> None: + with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): + PoissonDistribution(lam=invalid_lam) @pytest.mark.parametrize( diff --git a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py index 01435488a..a3efb3bcc 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_truncatednormal_distribution_class.py @@ -35,9 +35,17 @@ def test_truncated_normal_distribution_init_invalid_sd(invalid_sd: float) -> Non TruncatedNormalDistribution(mean=0.0, sd=invalid_sd, a=0.0, b=1.0) -def test_truncated_normal_init_fails_on_edge_case_false_by_default() -> None: - with pytest.raises(ValidationError, match="must be > lower bound 'a'"): - TruncatedNormalDistribution(mean=5.0, sd=2.0, a=5.0, b=5.0) +@pytest.mark.parametrize( + "a, b", + [ + (5.0, 5.0), + (10.0, 5.0), + ], + ids=["b_equals_a", "b_less_than_a"], +) +def test_truncated_normal_init_fails_when_b_not_gt_a(a: float, b: float) -> None: + with pytest.raises(ValidationError, match="must be > to lower bound `a`"): + TruncatedNormalDistribution(mean=5.0, sd=2.0, a=a, b=b, allow_edge_cases=False) def test_truncated_normal_init_succeeds_on_edge_case_true() -> None: @@ -49,21 +57,15 @@ def test_truncated_normal_init_succeeds_on_edge_case_true() -> None: assert dist.allow_edge_cases is True -def test_truncated_normal_init_fails_on_b_less_than_a() -> None: - with pytest.raises(ValidationError, match="must be > lower bound 'a'"): - TruncatedNormalDistribution(mean=5.0, sd=2.0, a=10.0, b=5.0) - with pytest.raises(ValidationError, match="must be ≥ to lower bound 'a'"): +def test_truncated_normal_init_fails_when_b_lt_a_with_edge_cases() -> None: + with pytest.raises(ValidationError, match="must be >= to lower bound `a`"): TruncatedNormalDistribution(mean=5.0, sd=2.0, a=10.0, b=5.0, allow_edge_cases=True) +@pytest.mark.parametrize("mean", [0.0, 100.0], ids=["mean_zero", "mean_high"]) +@pytest.mark.parametrize("sd", [1.0, 10.0], ids=["sd_one", "sd_high"]) @pytest.mark.parametrize( - "mean, sd, a, b", - [ - (5.0, 2.0, 0.0, 10.0), - (0.0, 1.0, -1.0, 1.0), - (100.0, 10.0, 90.0, 110.0), - ], - ids=["mean_in_range", "std_normal_truncated", "high_mean"], + ("a", "b"), [(-2.0, 2.0), (90.0, 110.0)], ids=["bounds_centered", "bounds_high"] ) def test_truncated_normal_distribution_sample_range( mean: float, sd: float, a: float, b: float diff --git a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py index 5e978ac9d..4adf42fed 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_uniform_distribution_class.py @@ -22,9 +22,17 @@ def test_uniform_distribution_init_valid(low: float, high: float) -> None: assert dist.allow_edge_cases is False -def test_uniform_init_fails_on_edge_case_false_by_default() -> None: - with pytest.raises(ValidationError, match="'high' value .* must be > 'low' value"): - UniformDistribution(low=5.0, high=5.0) +@pytest.mark.parametrize( + "low, high", + [ + (5.0, 5.0), + (10.0, 5.0), + ], + ids=["high_equals_low", "high_less_than_low"], +) +def test_uniform_init_fails_when_high_not_gt_low(low: float, high: float) -> None: + with pytest.raises(ValidationError, match="must be > to lower bound `low`"): + UniformDistribution(low=low, high=high, allow_edge_cases=False) def test_uniform_init_succeeds_on_edge_case_true() -> None: @@ -34,12 +42,8 @@ def test_uniform_init_succeeds_on_edge_case_true() -> None: assert dist.allow_edge_cases is True -def test_uniform_init_fails_on_high_less_than_low() -> None: - with pytest.raises(ValidationError, match="'high' value .* must be > 'low' value"): - UniformDistribution(low=10.0, high=5.0, allow_edge_cases=False) - with pytest.raises( - ValidationError, match="'high' value .* must be ≥ to the 'low' value" - ): +def test_uniform_init_fails_when_high_lt_low_with_edge_cases() -> None: + with pytest.raises(ValidationError, match="must be >= to lower bound `low`"): UniformDistribution(low=10.0, high=5.0, allow_edge_cases=True) From 921f41a36ff2b286bba9b1129856bac1d7c71793 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Thu, 17 Jul 2025 11:18:37 -0400 Subject: [PATCH 35/73] Change `Field` constraints for `BinomialDistribution` Change to correct operator (`ge` and `le`), update tests accordingly --- .../gempyor_pkg/src/gempyor/distributions.py | 4 +- .../test_binomial_distribution_class.py | 70 ++++++++++++++++--- 2 files changed, 61 insertions(+), 13 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index fb9a9da8c..a99fa2500 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -348,8 +348,8 @@ class BinomialDistribution(DistributionABC): """ distribution: Literal["binomial"] = "binomial" - n: int = Field(..., gt=0) - p: float = Field(..., gt=0.0, lt=1.0) + n: int = Field(..., ge=0) + p: float = Field(..., ge=0.0, le=1.0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator diff --git a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py index e3e7f37c6..7d06fe896 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_binomial_distribution_class.py @@ -15,7 +15,7 @@ ids=[ "standard_case_1", "standard_case_2", - "min_valid_n", + "min_valid_n_for_non_edge_case", ], ) def test_binomial_distribution_init_valid(n: int, p: float) -> None: @@ -26,33 +26,66 @@ def test_binomial_distribution_init_valid(n: int, p: float) -> None: assert dist.allow_edge_cases is False +@pytest.mark.parametrize( + "n, p", + [(0, 0.5), (10, 0.0), (10, 1.0)], + ids=["n_is_zero", "p_is_zero", "p_is_one"], +) +def test_binomial_distribution_init_valid_edge_cases(n: int, p: float) -> None: + dist = BinomialDistribution(n=n, p=p, allow_edge_cases=True) + assert dist.n == n + assert dist.p == p + assert dist.allow_edge_cases is True + + +def test_binomial_distribution_init_invalid_n_default() -> None: + with pytest.raises( + ValidationError, + match="Input for `n` cannot be zero when `allow_edge_cases` is `False`.", + ): + BinomialDistribution(n=0, p=0.5) + + @pytest.mark.parametrize( "invalid_n", - [0, -1, -10], - ids=["zero", "-1", "-10"], + [-1, -10], + ids=["-1", "-10"], ) -def test_binomial_distribution_init_invalid_n(invalid_n: int) -> None: - with pytest.raises(ValidationError, match="Input should be greater than 0"): +def test_binomial_distribution_init_invalid_n_negative(invalid_n: int) -> None: + with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): BinomialDistribution(n=invalid_n, p=0.5) @pytest.mark.parametrize( "invalid_p", - [0.0, -0.1, -10.0], - ids=["zero", "small_negative", "large_negative"], + [0.0, 1.0], + ids=["zero", "one"], +) +def test_binomial_distribution_init_invalid_p_default(invalid_p: float) -> None: + with pytest.raises( + ValidationError, + match="Input for `p` cannot be 0 or 1 when `allow_edge_cases` is `False`.", + ): + BinomialDistribution(n=10, p=invalid_p) + + +@pytest.mark.parametrize( + "invalid_p", + [-0.1, -10.0], + ids=["small_negative", "large_negative"], ) def test_binomial_distribution_init_invalid_p_below_range(invalid_p: float) -> None: - with pytest.raises(ValidationError, match="Input should be greater than 0"): + with pytest.raises(ValidationError, match="Input should be greater than or equal to 0"): BinomialDistribution(n=10, p=invalid_p) @pytest.mark.parametrize( "invalid_p", - [1.0, 1.1, 25.0], - ids=["one", "small_gt_one", "large_gt_one"], + [1.1, 25.0], + ids=["small_gt_one", "large_gt_one"], ) def test_binomial_distribution_init_invalid_p_above_range(invalid_p: float) -> None: - with pytest.raises(ValidationError, match="Input should be less than 1"): + with pytest.raises(ValidationError, match="Input should be less than or equal to 1"): BinomialDistribution(n=10, p=invalid_p) @@ -70,3 +103,18 @@ def test_binomial_sample_properties(n: int, p: float) -> None: sample = dist.sample(size=(10, 10)) assert np.all(sample >= 0) assert np.all(sample <= n) + + +@pytest.mark.parametrize( + "n, p, expected_val", + [ + (10, 1.0, 10), + (10, 0.0, 0), + (0, 0.5, 0), + ], + ids=["p_is_one", "p_is_zero", "n_is_zero"], +) +def test_binomial_sample_edge_cases(n: int, p: float, expected_val: int) -> None: + dist = BinomialDistribution(n=n, p=p, allow_edge_cases=True) + sample = dist.sample(size=10) + assert np.all(sample == expected_val) From 6036cf3625ffb8a7f0ef99710c284aeb90df2495 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Thu, 17 Jul 2025 16:45:43 -0400 Subject: [PATCH 36/73] Incorporate `gempyor.distributions` into `gempyor.parameters` --- flepimop/gempyor_pkg/src/gempyor/parameters.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index 0b95834ae..655c0f7e2 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -19,12 +19,16 @@ import numpy as np from numpy import ndarray import pandas as pd +import pydantic from . import NPI, utils +from .distributions import Distribution logger = logging.getLogger(__name__) +DISTRIBUTION_ADAPTER = pydantic.TypeAdapter(Distribution) + class Parameters: """ @@ -97,9 +101,10 @@ def __init__( self.pdata[pn] = {} self.pdata[pn]["idx"] = idx - # Parameter characterized by it's distribution + # Parameter characterized by its distribution if self.pconfig[pn]["value"].exists(): - self.pdata[pn]["dist"] = self.pconfig[pn]["value"].as_random_distribution() + dist_dict = self.pconfig[pn]["value"].get() + self.pdata[pn]["dist"] = DISTRIBUTION_ADAPTER.validate_python(dist_dict) # Parameter given as a file elif self.pconfig[pn]["timeseries"].exists(): @@ -182,7 +187,8 @@ def reinitialize_distributions(self) -> None: """ for pn in self.pnames: if "dist" in self.pdata[pn]: - self.pdata[pn]["dist"] = self.pconfig[pn]["value"].as_random_distribution() + dist_dict = self.pconfig[pn]["value"].get() + self.pdata[pn]["dist"] = DISTRIBUTION_ADAPTER.validate_python(dist_dict) def get_pnames2pindex(self) -> dict: """ @@ -224,7 +230,7 @@ class via `subpop_names`. for idx, pn in enumerate(self.pnames): if "dist" in self.pdata[pn]: - param_arr[idx] = np.full((n_days, nsubpops), self.pdata[pn]["dist"]()) + param_arr[idx] = self.pdata[pn]["dist"].sample(size=(n_days, nsubpops)) else: param_arr[idx] = self.pdata[pn]["ts"].values @@ -272,8 +278,7 @@ def parameters_load( f"PARAM: parameter {pn} NOT found in loadID file. " "Drawing from config distribution" ) - pval = self.pdata[pn]["dist"]() - param_arr[idx] = np.full((n_days, nsubpops), pval) + param_arr[idx] = self.pdata[pn]["dist"].sample(size=(n_days, nsubpops)) return param_arr From 536e8475a8db2204294d410cfb440f8a56ef310a Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 18 Jul 2025 12:10:58 -0400 Subject: [PATCH 37/73] Update `Parameters` class testing Update the `test_parameters_class.py` file to be up to date with new `gempyor.distributions` incorporation. --- .../tests/parameters/test_parameters_class.py | 189 +++++++++++++----- 1 file changed, 136 insertions(+), 53 deletions(-) diff --git a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py index 8bbd29076..8042cbd65 100644 --- a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py +++ b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py @@ -13,6 +13,18 @@ import pytest from gempyor.parameters import Parameters +from gempyor.distributions import ( + BetaDistribution, + BinomialDistribution, + FixedDistribution, + GammaDistribution, + LognormalDistribution, + NormalDistribution, + PoissonDistribution, + TruncatedNormalDistribution, + UniformDistribution, + WeibullDistribution, +) from gempyor.testing import ( create_confuse_configview_from_dict, partials_are_similar, @@ -99,7 +111,11 @@ def get_timeseries_df( def fixed_three_valid_parameter_factory(tmp_path: pathlib.Path) -> MockParametersInput: return MockParametersInput( - config={"sigma": {"value": 0.1}, "eta": {"value": 0.2}, "nu": {"value": 0.3}}, + config={ + "sigma": {"value": {"distribution": "fixed", "value": 0.1}}, + "eta": {"value": {"distribution": "fixed", "value": 0.2}}, + "nu": {"value": {"distribution": "fixed", "value": 0.3}}, + }, ti=date(2024, 1, 1), tf=date(2024, 1, 31), subpop_names=["1", "2", "3"], @@ -134,7 +150,10 @@ def valid_parameters_factory(tmp_path: pathlib.Path) -> MockParametersInput: return MockParametersInput( config={ "sigma": {"timeseries": str(tmp_file.absolute())}, - "gamma": {"value": 0.1234, "stacked_modifier_method": "sum"}, + "gamma": { + "value": {"distribution": "fixed", "value": 0.1234}, + "stacked_modifier_method": "sum", + }, "Ro": {"value": {"distribution": "uniform", "low": 1.0, "high": 2.0}}, }, ti=df["date"].dt.date.min(), @@ -343,37 +362,45 @@ def test_parameters_instance_attributes( assert params.pdata[param_name]["ts"].equals( mock_inputs.get_timeseries_df(param_name) ) - elif isinstance(params.pdata[param_name]["dist"], partial): - if isinstance(param_conf.get("value"), float): - expected = random_distribution_sampler( - "fixed", value=param_conf.get("value") - ) + # This is the new, corrected logic for checking distributions. + elif "dist" in params.pdata[param_name]: + dist_obj = params.pdata[param_name]["dist"] + value_conf = param_conf.get("value") + + # TODO, does this just test pydantic? + # Check for the correct Pydantic model type and assert its attributes. + if isinstance(dist_obj, FixedDistribution): + assert dist_obj.value == value_conf.get("value") + elif isinstance(dist_obj, NormalDistribution): + assert dist_obj.mu == value_conf.get("mu") + assert dist_obj.sigma == value_conf.get("sigma") + elif isinstance(dist_obj, UniformDistribution): + assert dist_obj.low == value_conf.get("low") + assert dist_obj.high == value_conf.get("high") + elif isinstance(dist_obj, LognormalDistribution): + assert dist_obj.meanlog == value_conf.get("meanlog") + assert dist_obj.sdlog == value_conf.get("sdlog") + elif isinstance(dist_obj, TruncatedNormalDistribution): + assert dist_obj.mean == value_conf.get("mean") + assert dist_obj.sd == value_conf.get("sd") + assert dist_obj.a == value_conf.get("a") + assert dist_obj.b == value_conf.get("b") + elif isinstance(dist_obj, PoissonDistribution): + assert dist_obj.lam == value_conf.get("lam") + elif isinstance(dist_obj, BinomialDistribution): + assert dist_obj.n == value_conf.get("n") + assert dist_obj.p == value_conf.get("p") + elif isinstance(dist_obj, GammaDistribution): + assert dist_obj.shape == value_conf.get("shape") + assert dist_obj.scale == value_conf.get("scale") + elif isinstance(dist_obj, WeibullDistribution): + assert dist_obj.shape == value_conf.get("shape") + assert dist_obj.scale == value_conf.get("scale") + elif isinstance(dist_obj, BetaDistribution): + assert dist_obj.alpha == value_conf.get("alpha") + assert dist_obj.beta == value_conf.get("beta") else: - expected = random_distribution_sampler( - param_conf.get("value").get("distribution"), - **{ - k: v - for k, v in param_conf.get("value").items() - if k != "distribution" - }, - ) - assert partials_are_similar(params.pdata[param_name]["dist"], expected) - else: - expected = random_distribution_sampler( - param_conf.get("value").get("distribution"), - **{ - k: v - for k, v in param_conf.get("value").items() - if k != "distribution" - }, - ) - assert ( - params.pdata[param_name]["dist"].__self__.kwds == expected.__self__.kwds - ) - assert ( - params.pdata[param_name]["dist"].__self__.support() - == expected.__self__.support() - ) + raise TypeError(f"Unhandled distribution type: {type(dist_obj)}") # The `pnames` attribute assert set(params.pnames) == set(mock_inputs.config.keys()) @@ -467,20 +494,48 @@ def test_parameters_quick_draw( # Loop over each param and check it individually for param_name, conf in mock_inputs.config.items(): i = params.pnames.index(param_name) + if "timeseries" in conf: - # Check if the values in p_draw[i, :, :] match timeseries timeseries_df = mock_inputs.get_timeseries_df(param_name) assert np.allclose(p_draw[i, :, :], timeseries_df.values) - elif isinstance((fixed_value := conf.get("value")), float): - # Check if all the values in p_draw[i, :, :] match a const - assert np.allclose(p_draw[i, :, :], fixed_value) + else: - # Check if the values in p_draw[i, :, :] match the distribution - assert np.allclose(p_draw[i, :, :], p_draw[i, 0, 0]) - value = float(p_draw[i, 0, 0]) - assert sample_fits_distribution( - value, **{k: v for k, v in conf.get("value").items()} - ) + dist_obj = params.pdata[param_name]["dist"] + drawn_values = p_draw[i, :, :] + + # For a fixed distribution, all values must be equal to the specified value + if isinstance(dist_obj, FixedDistribution): + assert np.allclose(drawn_values, dist_obj.value) + + # For bounded distributions, check that all values are within the domain + elif isinstance(dist_obj, UniformDistribution): + assert np.all(drawn_values >= dist_obj.low) + assert np.all(drawn_values <= dist_obj.high) + elif isinstance(dist_obj, TruncatedNormalDistribution): + assert np.all(drawn_values >= dist_obj.a) + assert np.all(drawn_values <= dist_obj.b) + elif isinstance(dist_obj, BetaDistribution): + assert np.all(drawn_values >= 0) + assert np.all(drawn_values <= 1) + + # For discrete distributions, check for integer type and domain + elif isinstance(dist_obj, BinomialDistribution): + assert np.all(drawn_values >= 0) + assert np.all(drawn_values <= dist_obj.n) + elif isinstance(dist_obj, PoissonDistribution): + assert np.all(drawn_values >= 0) + + # For distributions with a non-negative domain + elif isinstance( + dist_obj, + (LognormalDistribution, GammaDistribution, WeibullDistribution), + ): + assert np.all(drawn_values >= 0) + + # For distributions without a simple domain to check (Normal) + else: + assert isinstance(dist_obj, NormalDistribution) + assert not np.isnan(drawn_values).any() @pytest.mark.parametrize( "factory,param_df,n_days,nsubpops", @@ -599,28 +654,56 @@ def test_parameters_load( # Loop over each param and check it individually for param_name, conf in mock_inputs.config.items(): i = params.pnames.index(param_name) + if param_name in param_df["parameter"].values: - # Check that the values in p_draw[i, :, :] match override assert np.allclose( p_draw[i, :, :], param_df[param_df["parameter"] == param_name] .iloc[0]["value"] .item(), ) + elif "timeseries" in conf: - # Check if the values in p_draw[i, :, :] match timeseries timeseries_df = mock_inputs.get_timeseries_df(param_name) assert np.allclose(p_draw[i, :, :], timeseries_df.values) - elif isinstance((fixed_value := conf.get("value")), float): - # Check if all the values in p_draw[i, :, :] match a const - assert np.allclose(p_draw[i, :, :], fixed_value) + else: - # Check if the values in p_draw[i, :, :] match the distribution - assert np.allclose(p_draw[i, :, :], p_draw[i, 0, 0]) - value = float(p_draw[i, 0, 0]) - assert sample_fits_distribution( - value, **{k: v for k, v in conf.get("value").items()} - ) + dist_obj = params.pdata[param_name]["dist"] + drawn_values = p_draw[i, :, :] + + # For a fixed distribution, all values must be equal to the specified value + if isinstance(dist_obj, FixedDistribution): + assert np.allclose(drawn_values, dist_obj.value) + + # For bounded distributions, check that all values are within the domain + elif isinstance(dist_obj, UniformDistribution): + assert np.all(drawn_values >= dist_obj.low) + assert np.all(drawn_values <= dist_obj.high) + elif isinstance(dist_obj, TruncatedNormalDistribution): + assert np.all(drawn_values >= dist_obj.a) + assert np.all(drawn_values <= dist_obj.b) + elif isinstance(dist_obj, BetaDistribution): + assert np.all(drawn_values >= 0) + assert np.all(drawn_values <= 1) + + # For discrete distributions, check for integer type and domain + elif isinstance(dist_obj, BinomialDistribution): + assert np.all(drawn_values >= 0) + assert np.all(drawn_values <= dist_obj.n) + elif isinstance(dist_obj, PoissonDistribution): + assert np.all(drawn_values >= 0) + + # For distributions with a non-negative domain + elif isinstance( + dist_obj, + (LognormalDistribution, GammaDistribution, WeibullDistribution), + ): + assert np.all(drawn_values >= 0) + + # For distributions without a simple domain to check (Normal) + else: + assert isinstance(dist_obj, NormalDistribution) + assert not np.isnan(drawn_values).any() @pytest.mark.parametrize( "factory,n_days,nsubpops", From 89453b2a2e9e9dbff53eb3236e0fe054a312f831 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 18 Jul 2025 15:58:54 -0400 Subject: [PATCH 38/73] Fix bugs created in last push --- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 34 +++++++++++++- .../gempyor_pkg/src/gempyor/distributions.py | 44 ++++++++++--------- .../gempyor_pkg/src/gempyor/parameters.py | 16 ++++--- 3 files changed, 67 insertions(+), 27 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index fd8ab45bd..f669cacb5 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -8,12 +8,14 @@ __all__ = () +import numbers from pathlib import Path -from typing import Any, TypeVar, overload +from typing import Any, TypeVar, overload, Annotated import pyarrow.parquet as pq import pandas as pd -from pydantic import BaseModel, RootModel +from pydantic import BaseModel, RootModel, BeforeValidator +import sympy T = TypeVar("T") @@ -180,3 +182,31 @@ def _read_and_validate_dataframe( model = RootModel[list[model]] data = [r.model_dump() for r in model.model_validate(data).root] return pd.DataFrame.from_records(data) + + +def _evaled_expression(val: float | str | Any) -> float | Any: + """ + Evaluates an expression and attempts to convert it to a float. + + Args: + val: Expression to be evaluated: + + Returns: + Either the expression coerced into a float, or the expression. + + Raises: + ValueError: on parsing errors + """ + + if isinstance(val, float): + return val + elif isinstance(val, str): + try: + return float(sympy.parsing.sympy_parser.parse_expr(val)) + except TypeError as e: + raise ValueError(e) from e + return val + + +# TODO: add EvaluatedInt = Annotated[int, BeforeValidator(_evaled_expression_int)] +EvaledFloat = Annotated[float, BeforeValidator(_evaled_expression)] diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index a99fa2500..bb7b81e62 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -23,9 +23,11 @@ import numpy as np from numpy.random import Generator import numpy.typing as npt -from pydantic import BaseModel, PrivateAttr, Field, model_validator +from pydantic import BaseModel, PrivateAttr, Field, TypeAdapter, model_validator from scipy.stats import truncnorm +from ._pydantic_ext import EvaledFloat + class DistributionABC(ABC, BaseModel): """Base class for distributions used in modifiers, likelihoods, etc.""" @@ -82,7 +84,7 @@ class FixedDistribution(DistributionABC): """ distribution: Literal["fixed"] = "fixed" - value: float + value: EvaledFloat def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -111,8 +113,8 @@ class NormalDistribution(DistributionABC): """ distribution: Literal["norm"] = "norm" - mu: float - sigma: float = Field(..., gt=0) + mu: EvaledFloat + sigma: EvaledFloat = Field(..., gt=0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -151,8 +153,8 @@ class UniformDistribution(DistributionABC): """ distribution: Literal["uniform"] = "uniform" - low: float - high: float + low: EvaledFloat + high: EvaledFloat def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -193,8 +195,8 @@ class LognormalDistribution(DistributionABC): """ distribution: Literal["lognorm"] = "lognorm" - meanlog: float - sdlog: float = Field(..., gt=0) + meanlog: EvaledFloat + sdlog: EvaledFloat = Field(..., gt=0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -235,10 +237,10 @@ class TruncatedNormalDistribution(DistributionABC): # pylint: enable=line-too-long distribution: Literal["truncnorm"] = "truncnorm" - mean: float - sd: float = Field(..., gt=0) - a: float - b: float + mean: EvaledFloat + sd: EvaledFloat = Field(..., gt=0) + a: EvaledFloat + b: EvaledFloat def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -301,7 +303,7 @@ class PoissonDistribution(DistributionABC): """ distribution: Literal["poisson"] = "poisson" - lam: float = Field(..., ge=0.0) + lam: EvaledFloat = Field(..., ge=0.0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -349,7 +351,7 @@ class BinomialDistribution(DistributionABC): distribution: Literal["binomial"] = "binomial" n: int = Field(..., ge=0) - p: float = Field(..., ge=0.0, le=1.0) + p: EvaledFloat = Field(..., ge=0.0, le=1.0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -392,8 +394,8 @@ class GammaDistribution(DistributionABC): """ distribution: Literal["gamma"] = "gamma" - shape: float = Field(..., gt=0) - scale: float = Field(..., gt=0) + shape: EvaledFloat = Field(..., gt=0) + scale: EvaledFloat = Field(..., gt=0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -422,8 +424,8 @@ class WeibullDistribution(DistributionABC): """ distribution: Literal["weibull"] = "weibull" - shape: float = Field(..., gt=0) - scale: float = Field(..., gt=0) + shape: EvaledFloat = Field(..., gt=0) + scale: EvaledFloat = Field(..., gt=0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -453,8 +455,8 @@ class BetaDistribution(DistributionABC): """ distribution: Literal["beta"] = "beta" - alpha: float = Field(..., gt=0) - beta: float = Field(..., gt=0) + alpha: EvaledFloat = Field(..., gt=0) + beta: EvaledFloat = Field(..., gt=0) def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -476,3 +478,5 @@ def _sample_from_generator( | WeibullDistribution, Field(discriminator="distribution"), ] + +DISTRIBUTION_ADAPTER = TypeAdapter(Distribution) diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index 655c0f7e2..b8003e844 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -22,13 +22,11 @@ import pydantic from . import NPI, utils -from .distributions import Distribution +from .distributions import Distribution, DISTRIBUTION_ADAPTER logger = logging.getLogger(__name__) -DISTRIBUTION_ADAPTER = pydantic.TypeAdapter(Distribution) - class Parameters: """ @@ -104,6 +102,8 @@ def __init__( # Parameter characterized by its distribution if self.pconfig[pn]["value"].exists(): dist_dict = self.pconfig[pn]["value"].get() + if isinstance(dist_dict, float | int | str): + dist_dict = {"distribution": "fixed", "value": dist_dict} self.pdata[pn]["dist"] = DISTRIBUTION_ADAPTER.validate_python(dist_dict) # Parameter given as a file @@ -188,6 +188,8 @@ def reinitialize_distributions(self) -> None: for pn in self.pnames: if "dist" in self.pdata[pn]: dist_dict = self.pconfig[pn]["value"].get() + if isinstance(dist_dict, float | int | str): + dist_dict = {"distribution": "fixed", "value": dist_dict} self.pdata[pn]["dist"] = DISTRIBUTION_ADAPTER.validate_python(dist_dict) def get_pnames2pindex(self) -> dict: @@ -230,7 +232,9 @@ class via `subpop_names`. for idx, pn in enumerate(self.pnames): if "dist" in self.pdata[pn]: - param_arr[idx] = self.pdata[pn]["dist"].sample(size=(n_days, nsubpops)) + param_arr[idx] = np.full( + (n_days, nsubpops), self.pdata[pn]["dist"].sample().item() + ) else: param_arr[idx] = self.pdata[pn]["ts"].values @@ -278,7 +282,9 @@ def parameters_load( f"PARAM: parameter {pn} NOT found in loadID file. " "Drawing from config distribution" ) - param_arr[idx] = self.pdata[pn]["dist"].sample(size=(n_days, nsubpops)) + param_arr[idx] = np.full( + (n_days, nsubpops), self.pdata[pn]["dist"].sample().item() + ) return param_arr From 0c74575b6b95ed90caf681f1a52f34be9f2646a5 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 18 Jul 2025 16:40:24 -0400 Subject: [PATCH 39/73] Remove unnecessary assertions in `test_parameters_class.py` We don't need to test vendor libraries! I need to write that 1,000 times on a piece of paper or something haha --- .../tests/parameters/test_parameters_class.py | 38 +------------------ 1 file changed, 2 insertions(+), 36 deletions(-) diff --git a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py index 8042cbd65..5f89bc649 100644 --- a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py +++ b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py @@ -14,6 +14,7 @@ from gempyor.parameters import Parameters from gempyor.distributions import ( + DistributionABC, BetaDistribution, BinomialDistribution, FixedDistribution, @@ -362,45 +363,10 @@ def test_parameters_instance_attributes( assert params.pdata[param_name]["ts"].equals( mock_inputs.get_timeseries_df(param_name) ) - # This is the new, corrected logic for checking distributions. elif "dist" in params.pdata[param_name]: dist_obj = params.pdata[param_name]["dist"] value_conf = param_conf.get("value") - - # TODO, does this just test pydantic? - # Check for the correct Pydantic model type and assert its attributes. - if isinstance(dist_obj, FixedDistribution): - assert dist_obj.value == value_conf.get("value") - elif isinstance(dist_obj, NormalDistribution): - assert dist_obj.mu == value_conf.get("mu") - assert dist_obj.sigma == value_conf.get("sigma") - elif isinstance(dist_obj, UniformDistribution): - assert dist_obj.low == value_conf.get("low") - assert dist_obj.high == value_conf.get("high") - elif isinstance(dist_obj, LognormalDistribution): - assert dist_obj.meanlog == value_conf.get("meanlog") - assert dist_obj.sdlog == value_conf.get("sdlog") - elif isinstance(dist_obj, TruncatedNormalDistribution): - assert dist_obj.mean == value_conf.get("mean") - assert dist_obj.sd == value_conf.get("sd") - assert dist_obj.a == value_conf.get("a") - assert dist_obj.b == value_conf.get("b") - elif isinstance(dist_obj, PoissonDistribution): - assert dist_obj.lam == value_conf.get("lam") - elif isinstance(dist_obj, BinomialDistribution): - assert dist_obj.n == value_conf.get("n") - assert dist_obj.p == value_conf.get("p") - elif isinstance(dist_obj, GammaDistribution): - assert dist_obj.shape == value_conf.get("shape") - assert dist_obj.scale == value_conf.get("scale") - elif isinstance(dist_obj, WeibullDistribution): - assert dist_obj.shape == value_conf.get("shape") - assert dist_obj.scale == value_conf.get("scale") - elif isinstance(dist_obj, BetaDistribution): - assert dist_obj.alpha == value_conf.get("alpha") - assert dist_obj.beta == value_conf.get("beta") - else: - raise TypeError(f"Unhandled distribution type: {type(dist_obj)}") + assert isinstance(dist_obj, DistributionABC) # The `pnames` attribute assert set(params.pnames) == set(mock_inputs.config.keys()) From bf1fc2df020c6fbae30461464911b69181b28c48 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Tue, 22 Jul 2025 10:31:09 -0400 Subject: [PATCH 40/73] Fix `UnboundLocalError` error in `onerun_SEIR` Could reach an `UnboundedLocalError` if `modinf.write_csv` and `modinf.write_parquet` were both `False` as `out_df` would not be set. Solution is to move the call to `states2DF` from `write_seir` to `onerun_SEIR` and return the `out_df` regardless of the `modinf.write_*` attributes. --- flepimop/gempyor_pkg/src/gempyor/seir.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/seir.py b/flepimop/gempyor_pkg/src/gempyor/seir.py index 492b85fc7..a31fc7a18 100644 --- a/flepimop/gempyor_pkg/src/gempyor/seir.py +++ b/flepimop/gempyor_pkg/src/gempyor/seir.py @@ -381,10 +381,11 @@ def onerun_SEIR( ) with Timer("onerun_SEIR.postprocess"): + out_df = states2Df(modinf, states) if modinf.write_csv or modinf.write_parquet: write_spar_snpi(sim_id2write, modinf, p_draw, npi) - out_df = write_seir(sim_id2write, modinf, states) - return out_df + write_seir(sim_id2write, modinf, out_df) + return out_df def _onerun_SEIR_with_random_seed( @@ -420,7 +421,7 @@ def _onerun_SEIR_with_random_seed( ) -def run_parallel_SEIR(modinf: ModelInfo, config, *, n_jobs=1): +def run_parallel_SEIR(modinf: ModelInfo, config, *, n_jobs=1) -> None: """ Run SEIR simulations in parallel. @@ -538,11 +539,8 @@ def write_spar_snpi(sim_id: int, modinf: ModelInfo, p_draw, npi): ) -def write_seir(sim_id, modinf: ModelInfo, states): - # print_disk_diagnosis() - out_df = states2Df(modinf, states) +def write_seir(sim_id, modinf: ModelInfo, out_df): modinf.write_simID(ftype="seir", sim_id=sim_id, df=out_df) - return out_df From 113a8bbe9986626e29c4e0e4980643051a331734 Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Thu, 24 Jul 2025 10:37:24 -0400 Subject: [PATCH 41/73] Update flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- .../tests/parameters/test_parameters_class.py | 33 ------------------- 1 file changed, 33 deletions(-) diff --git a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py index 5f89bc649..0ca75a98e 100644 --- a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py +++ b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py @@ -469,39 +469,6 @@ def test_parameters_quick_draw( dist_obj = params.pdata[param_name]["dist"] drawn_values = p_draw[i, :, :] - # For a fixed distribution, all values must be equal to the specified value - if isinstance(dist_obj, FixedDistribution): - assert np.allclose(drawn_values, dist_obj.value) - - # For bounded distributions, check that all values are within the domain - elif isinstance(dist_obj, UniformDistribution): - assert np.all(drawn_values >= dist_obj.low) - assert np.all(drawn_values <= dist_obj.high) - elif isinstance(dist_obj, TruncatedNormalDistribution): - assert np.all(drawn_values >= dist_obj.a) - assert np.all(drawn_values <= dist_obj.b) - elif isinstance(dist_obj, BetaDistribution): - assert np.all(drawn_values >= 0) - assert np.all(drawn_values <= 1) - - # For discrete distributions, check for integer type and domain - elif isinstance(dist_obj, BinomialDistribution): - assert np.all(drawn_values >= 0) - assert np.all(drawn_values <= dist_obj.n) - elif isinstance(dist_obj, PoissonDistribution): - assert np.all(drawn_values >= 0) - - # For distributions with a non-negative domain - elif isinstance( - dist_obj, - (LognormalDistribution, GammaDistribution, WeibullDistribution), - ): - assert np.all(drawn_values >= 0) - - # For distributions without a simple domain to check (Normal) - else: - assert isinstance(dist_obj, NormalDistribution) - assert not np.isnan(drawn_values).any() @pytest.mark.parametrize( "factory,param_df,n_days,nsubpops", From 8edee45412444706f13f1b6ec7e43611749acd46 Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Thu, 24 Jul 2025 10:37:48 -0400 Subject: [PATCH 42/73] Update flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index f669cacb5..508165362 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -200,7 +200,7 @@ def _evaled_expression(val: float | str | Any) -> float | Any: if isinstance(val, float): return val - elif isinstance(val, str): + if isinstance(val, str): try: return float(sympy.parsing.sympy_parser.parse_expr(val)) except TypeError as e: From 8d66efa2527e726129944f0184ad510dd35f52a8 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Thu, 24 Jul 2025 12:06:39 -0400 Subject: [PATCH 43/73] Improvements surrounding new `_evaled_expression()` function Add `@overload` ext for `_evaled_expression`, improve documentation and add unit testing --- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 26 ++++++++++++++ .../_pydantic_ext/test__evaled_expression.py | 36 +++++++++++++++++++ .../tests/parameters/test_parameters_class.py | 1 - 3 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index 508165362..7e1b012e3 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -184,6 +184,18 @@ def _read_and_validate_dataframe( return pd.DataFrame.from_records(data) +@overload +def _evaled_expression(val: float) -> float: ... + + +@overload +def _evaled_expression(val: str) -> float: ... + + +@overload +def _evaled_expression(val: T) -> T: ... + + def _evaled_expression(val: float | str | Any) -> float | Any: """ Evaluates an expression and attempts to convert it to a float. @@ -196,6 +208,20 @@ def _evaled_expression(val: float | str | Any) -> float | Any: Raises: ValueError: on parsing errors + + Example: + >>> _evaled_expression("1 + 1") + 2.0 + >>> _evaled_expression("5 / 2") + 2.5 + >>> _evaled_expression(99.5) + 99.5 + >>> _evaled_expression(None) + + >>> _evaled_expression("a * b") + Traceback (most recent call last): + ... + ValueError: can't convert expression to float """ if isinstance(val, float): diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py new file mode 100644 index 000000000..dca886d4b --- /dev/null +++ b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py @@ -0,0 +1,36 @@ +"""Tests for the `gempyor._pydantic_ext._evaled_expression` function.""" + +import pytest + +from gempyor._pydantic_ext import _evaled_expression + + +@pytest.mark.parametrize( + ("val", "expected"), + [ + ("1 + 1", 2.0), + ("5 / 2", 2.5), + ("sqrt(9)", 3.0), + (99.5, 99.5), + (-10.0, -10.0), + (None, None), + (True, True), + ([1, 2, 3], [1, 2, 3]), + ({"a": 1}, {"a": 1}), + ], +) +def test_evaled_expression_returns_expected_results(val, expected): + assert _evaled_expression(val) == expected + + +@pytest.mark.parametrize( + "val", + [ + "a * b", + "not_a_valid_expression", + "1 / 0", + ], +) +def test_evaled_expression_raises_value_error_for_invalid_input(val): + with pytest.raises(ValueError): + _evaled_expression(val) diff --git a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py index 0ca75a98e..7e84f953f 100644 --- a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py +++ b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py @@ -469,7 +469,6 @@ def test_parameters_quick_draw( dist_obj = params.pdata[param_name]["dist"] drawn_values = p_draw[i, :, :] - @pytest.mark.parametrize( "factory,param_df,n_days,nsubpops", [ From 29240a775527b558fb1dc90206218c947407651c Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Thu, 24 Jul 2025 15:52:31 -0400 Subject: [PATCH 44/73] Introduce a `_evaled_expression_int()` function Implement this new function, add unit tests, rename other function to `_evaled_expression_float()` --- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 75 +++++++++++++++---- .../gempyor_pkg/src/gempyor/distributions.py | 4 +- ...on.py => test__evaled_expression_float.py} | 12 +-- .../test__evaled_expression_int.py | 37 +++++++++ 4 files changed, 107 insertions(+), 21 deletions(-) rename flepimop/gempyor_pkg/tests/_pydantic_ext/{test__evaled_expression.py => test__evaled_expression_float.py} (55%) create mode 100644 flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index 7e1b012e3..6eaf5635e 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -185,40 +185,40 @@ def _read_and_validate_dataframe( @overload -def _evaled_expression(val: float) -> float: ... +def _evaled_expression_float(val: float) -> float: ... @overload -def _evaled_expression(val: str) -> float: ... +def _evaled_expression_float(val: str) -> float: ... @overload -def _evaled_expression(val: T) -> T: ... +def _evaled_expression_float(val: T) -> T: ... -def _evaled_expression(val: float | str | Any) -> float | Any: +def _evaled_expression_float(val: float | str | Any) -> float | Any: """ Evaluates an expression and attempts to convert it to a float. Args: - val: Expression to be evaluated: + val: Expression to be evaluated. Returns: Either the expression coerced into a float, or the expression. Raises: - ValueError: on parsing errors + ValueError: on parsing errors. Example: - >>> _evaled_expression("1 + 1") + >>> _evaled_expression_float("1 + 1") 2.0 - >>> _evaled_expression("5 / 2") + >>> _evaled_expression_float("5 / 2") 2.5 - >>> _evaled_expression(99.5) + >>> _evaled_expression_float(99.5) 99.5 - >>> _evaled_expression(None) + >>> _evaled_expression_float(None) - >>> _evaled_expression("a * b") + >>> _evaled_expression_float("a * b") Traceback (most recent call last): ... ValueError: can't convert expression to float @@ -234,5 +234,54 @@ def _evaled_expression(val: float | str | Any) -> float | Any: return val -# TODO: add EvaluatedInt = Annotated[int, BeforeValidator(_evaled_expression_int)] -EvaledFloat = Annotated[float, BeforeValidator(_evaled_expression)] +@overload +def _evaled_expression_int(val: int) -> int: ... + + +@overload +def _evaled_expression_int(val: str) -> int: ... + + +@overload +def _evaled_expression_int(val: T) -> T: ... + + +def _evaled_expression_int(val: int | str | Any) -> int | Any: + """ + Evaluates an expression and attempts to convert it to an int. + + Args: + val: Expression to be evaluated. + + Returns: + Either the expression coerced into an integer, or the original expression. + + Raises: + ValueError: On parsing errors. + + Examples: + >>> _evaled_expression_int("1 + 1") + 2 + >>> _evaled_expression_int("10 / 2") # Note: Result is truncated + 5 + >>> _evaled_expression_int(99) + 99 + >>> _evaled_expression_int(None) + + >>> _evaled_expression_int("a * b") + Traceback (most recent call last): + ... + ValueError: can't convert expression to int + """ + if isinstance(val, int): + return val + if isinstance(val, str): + try: + return int(sympy.parsing.sympy_parser.parse_expr(val)) + except TypeError as e: + raise ValueError(e) from e + return val + + +EvaledInt = Annotated[int, BeforeValidator(_evaled_expression_int)] +EvaledFloat = Annotated[float, BeforeValidator(_evaled_expression_float)] diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index bb7b81e62..b787434f6 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -26,7 +26,7 @@ from pydantic import BaseModel, PrivateAttr, Field, TypeAdapter, model_validator from scipy.stats import truncnorm -from ._pydantic_ext import EvaledFloat +from ._pydantic_ext import EvaledFloat, EvaledInt class DistributionABC(ABC, BaseModel): @@ -350,7 +350,7 @@ class BinomialDistribution(DistributionABC): """ distribution: Literal["binomial"] = "binomial" - n: int = Field(..., ge=0) + n: EvaledInt = Field(..., ge=0) p: EvaledFloat = Field(..., ge=0.0, le=1.0) def _sample_from_generator( diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_float.py similarity index 55% rename from flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py rename to flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_float.py index dca886d4b..16b5ffd8c 100644 --- a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py +++ b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_float.py @@ -1,8 +1,8 @@ -"""Tests for the `gempyor._pydantic_ext._evaled_expression` function.""" +"""Tests for the `gempyor._pydantic_ext._evaled_expression_float` function.""" import pytest -from gempyor._pydantic_ext import _evaled_expression +from gempyor._pydantic_ext import _evaled_expression_float @pytest.mark.parametrize( @@ -19,8 +19,8 @@ ({"a": 1}, {"a": 1}), ], ) -def test_evaled_expression_returns_expected_results(val, expected): - assert _evaled_expression(val) == expected +def test_evaled_expression_float_returns_expected_results(val, expected): + assert _evaled_expression_float(val) == expected @pytest.mark.parametrize( @@ -31,6 +31,6 @@ def test_evaled_expression_returns_expected_results(val, expected): "1 / 0", ], ) -def test_evaled_expression_raises_value_error_for_invalid_input(val): +def test_evaled_expression_float_raises_value_error_for_invalid_input(val): with pytest.raises(ValueError): - _evaled_expression(val) + _evaled_expression_float(val) diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py new file mode 100644 index 000000000..4755f15a2 --- /dev/null +++ b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py @@ -0,0 +1,37 @@ +"""Tests for the `gempyor._pydantic_ext._evaled_expression_int` function.""" + +import pytest + +from gempyor._pydantic_ext import _evaled_expression_int + + +@pytest.mark.parametrize( + ("val", "expected"), + [ + ("2 * 5", 10), + ("sqrt(81)", 9), + ("11 / 2", 5), + ("3.99", 3), + (99, 99), + (-50, -50), + (None, None), + (True, True), + (10.5, 10.5), + ([1, 2], [1, 2]), + ], +) +def test_evaled_expression_int_returns_expected_results(val, expected): + assert _evaled_expression_int(val) == expected + + +@pytest.mark.parametrize( + "val", + [ + "a * b", + "not_a_valid_expression", + "1 / 0", + ], +) +def test_evaled_expression_int_raises_value_error_for_invalid_input(val): + with pytest.raises(ValueError): + _evaled_expression_int(val) From 16a90aaa989032747d8b4a9b8b063935ab5e6408 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Thu, 24 Jul 2025 16:52:32 -0400 Subject: [PATCH 45/73] Outsource distribution construction to `distributions.py` introduce a new function (`build_distribution_from_confuse_config()`) that will assemble distributions accordingly --- .../gempyor_pkg/src/gempyor/distributions.py | 23 +++++++++++++++++++ .../gempyor_pkg/src/gempyor/parameters.py | 20 ++++++++-------- 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index b787434f6..5b846df50 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -16,6 +16,7 @@ ) +import confuse from abc import ABC, abstractmethod from math import isclose from typing import Annotated, Literal @@ -480,3 +481,25 @@ def _sample_from_generator( ] DISTRIBUTION_ADAPTER = TypeAdapter(Distribution) + + +def build_distribution_from_confuse_config( + param_config: confuse.ConfigView, +) -> Distribution: + """ + Creates a Distribution object from a confuse.ConfigView (single value). + + Handles the case where the value is a simple number or string, + interpreting it as a 'fixed' distribution. + + Args: + param_config: A confuse.ConfigView for a single parameter. + + Returns: + A Distribution object. + """ + conf = param_config["value"].get() + if isinstance(conf, float | int | str): + conf = {"distribution": "fixed", "value": conf} + + return DISTRIBUTION_ADAPTER.validate_python(conf) diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index b8003e844..37b335777 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -22,7 +22,11 @@ import pydantic from . import NPI, utils -from .distributions import Distribution, DISTRIBUTION_ADAPTER +from .distributions import ( + Distribution, + DISTRIBUTION_ADAPTER, + build_distribution_from_confuse_config, +) logger = logging.getLogger(__name__) @@ -101,10 +105,9 @@ def __init__( # Parameter characterized by its distribution if self.pconfig[pn]["value"].exists(): - dist_dict = self.pconfig[pn]["value"].get() - if isinstance(dist_dict, float | int | str): - dist_dict = {"distribution": "fixed", "value": dist_dict} - self.pdata[pn]["dist"] = DISTRIBUTION_ADAPTER.validate_python(dist_dict) + self.pdata[pn]["dist"] = build_distribution_from_confuse_config( + self.pconfig[pn] + ) # Parameter given as a file elif self.pconfig[pn]["timeseries"].exists(): @@ -187,10 +190,9 @@ def reinitialize_distributions(self) -> None: """ for pn in self.pnames: if "dist" in self.pdata[pn]: - dist_dict = self.pconfig[pn]["value"].get() - if isinstance(dist_dict, float | int | str): - dist_dict = {"distribution": "fixed", "value": dist_dict} - self.pdata[pn]["dist"] = DISTRIBUTION_ADAPTER.validate_python(dist_dict) + self.pdata[pn]["dist"] = build_distribution_from_confuse_config( + self.pconfig[pn] + ) def get_pnames2pindex(self) -> dict: """ From 9470e3270f8e366a40f1ffb0498fafb89b946fa3 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 25 Jul 2025 14:51:25 -0400 Subject: [PATCH 46/73] Combine `_pydantic_ext` expression evaluation functions, de-duplicate logic in `utils.py` --- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 124 +- .../gempyor_pkg/src/gempyor/utils copy.py | 1454 +++++++++++++++++ flepimop/gempyor_pkg/src/gempyor/utils.py | 57 +- .../_pydantic_ext/test__evaled_expression.py | 48 + .../test__evaled_expression_float.py | 36 - .../test__evaled_expression_int.py | 37 - .../gempyor_pkg/tests/seir/test_parameters.py | 3 +- .../gempyor_pkg/tests/utils/test_utils2.py | 38 - 8 files changed, 1577 insertions(+), 220 deletions(-) create mode 100644 flepimop/gempyor_pkg/src/gempyor/utils copy.py create mode 100644 flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py delete mode 100644 flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_float.py delete mode 100644 flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index 6eaf5635e..8968b3697 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -10,7 +10,8 @@ import numbers from pathlib import Path -from typing import Any, TypeVar, overload, Annotated +from functools import partial +from typing import Any, TypeVar, overload, Annotated, Type import pyarrow.parquet as pq import pandas as pd @@ -20,6 +21,7 @@ T = TypeVar("T") U = TypeVar("U") +EE = TypeVar("EE", int, float) def _ensure_list(value: list[T] | tuple[T] | T | None) -> list[T] | None: @@ -185,103 +187,89 @@ def _read_and_validate_dataframe( @overload -def _evaled_expression_float(val: float) -> float: ... - - +def _evaled_expression(val: str, *, target_type: Type[int]) -> int: ... @overload -def _evaled_expression_float(val: str) -> float: ... - +def _evaled_expression(val: Any, *, target_type: Type[int]) -> int | Any: ... @overload -def _evaled_expression_float(val: T) -> T: ... - +def _evaled_expression(val: str, *, target_type: Type[float]) -> float: ... +@overload +def _evaled_expression(val: Any, *, target_type: Type[float]) -> float | Any: ... -def _evaled_expression_float(val: float | str | Any) -> float | Any: +def _evaled_expression(val: Any, *, target_type: Type[EE]) -> EE | Any: """ - Evaluates an expression and attempts to convert it to a float. + Evaluates a string expression to a target numeric type (int or float). Args: - val: Expression to be evaluated. + val: The input value to process. + target_type: The type (int or float) to convert the expression to. Returns: - Either the expression coerced into a float, or the expression. + The value coerced into the target numeric type, or the original value. Raises: - ValueError: on parsing errors. + ValueError: On parsing errors. - Example: - >>> _evaled_expression_float("1 + 1") - 2.0 - >>> _evaled_expression_float("5 / 2") + Examples: + >>> _evaled_expression("1 + 1", target_type=int) + 2 + >>> _evaled_expression("10 / 4", target_type=float) 2.5 - >>> _evaled_expression_float(99.5) + >>> _evaled_expression("10 / 4", target_type=int) + # note that result is trucnated; probably undesirable if misused + 2 + >>> _evaled_expression(99.5, target_type=float) 99.5 - >>> _evaled_expression_float(None) + >>> _evaled_expression(None, target_type=int) - >>> _evaled_expression_float("a * b") + >>> _evaled_expression("a * b", target_type=float) Traceback (most recent call last): ... - ValueError: can't convert expression to float + ValueError: Can't convert expression to float. """ - - if isinstance(val, float): + if isinstance(val, target_type): return val + if isinstance(val, str): try: - return float(sympy.parsing.sympy_parser.parse_expr(val)) + return target_type(sympy.parsing.sympy_parser.parse_expr(val)) except TypeError as e: - raise ValueError(e) from e + raise ValueError(f"Can't convert expression to {target_type.__name__}.") from e + return val -@overload -def _evaled_expression_int(val: int) -> int: ... +EvaledInt = Annotated[int, BeforeValidator(partial(_evaled_expression, target_type=int))] +EvaledFloat = Annotated[float, BeforeValidator(partial(_evaled_expression, target_type=float))] -@overload -def _evaled_expression_int(val: str) -> int: ... +# TEMP -------- +# --- TEMPORARY DEBUGGING LOGIC --- +# This entire block should be deleted after the refactor is complete. +import functools +import confuse +from typing import Any -@overload -def _evaled_expression_int(val: T) -> T: ... - +def add_method(cls: Any): + """A function which adds a function to a class.""" + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + setattr(cls, func.__name__, wrapper) + return func + return decorator -def _evaled_expression_int(val: int | str | Any) -> int | Any: +@add_method(confuse.ConfigView) +def _as_evaled_expression(self): """ - Evaluates an expression and attempts to convert it to an int. - - Args: - val: Expression to be evaluated. - - Returns: - Either the expression coerced into an integer, or the original expression. - - Raises: - ValueError: On parsing errors. - - Examples: - >>> _evaled_expression_int("1 + 1") - 2 - >>> _evaled_expression_int("10 / 2") # Note: Result is truncated - 5 - >>> _evaled_expression_int(99) - 99 - >>> _evaled_expression_int(None) - - >>> _evaled_expression_int("a * b") - Traceback (most recent call last): - ... - ValueError: can't convert expression to int + This is a temporary function to find all usages of the old method. + It will be removed after the refactoring is complete. """ - if isinstance(val, int): - return val - if isinstance(val, str): - try: - return int(sympy.parsing.sympy_parser.parse_expr(val)) - except TypeError as e: - raise ValueError(e) from e - return val - + raise NotImplementedError( + "DEPRECATED: _as_evaled_expression has been replaced. " + "Use the new pattern: _evaled_expression(view.get(), target_type=...)" + ) -EvaledInt = Annotated[int, BeforeValidator(_evaled_expression_int)] -EvaledFloat = Annotated[float, BeforeValidator(_evaled_expression_float)] +# --- END OF TEMPORARY BLOCK --- \ No newline at end of file diff --git a/flepimop/gempyor_pkg/src/gempyor/utils copy.py b/flepimop/gempyor_pkg/src/gempyor/utils copy.py new file mode 100644 index 000000000..9a46d50d3 --- /dev/null +++ b/flepimop/gempyor_pkg/src/gempyor/utils copy.py @@ -0,0 +1,1454 @@ +""" +Helper functions for interacting with model I/O. +""" + +from collections import Counter +from collections.abc import Iterable +import datetime +import functools +import logging +import numbers +import os +from pathlib import Path +import random +from shlex import quote as shlex_quote +import shutil +import subprocess +import time +from typing import Any, Callable, Literal, overload + +import confuse +import numpy as np +import numpy.typing as npt +import pandas as pd +import pyarrow as pa +import scipy.ndimage +import scipy.stats +import sympy.parsing.sympy_parser +import yaml + +from . import file_paths +from ._pydantic_ext import _evaled_expression + + +logger = logging.getLogger(__name__) + +config = confuse.Configuration("flepiMoP", read=False) + + +def write_df( + fname: str | bytes | os.PathLike, + df: pd.DataFrame, + extension: Literal[None, "", "csv", "parquet"] = "", +) -> None: + """Writes a pandas DataFrame without its index to a file. + + Writes a pandas DataFrame to either a CSV or Parquet file without its index and can + infer which format to use based on the extension given in `fname` or based on + explicit `extension`. + + Args: + fname: The name of the file to write to. + df: A pandas DataFrame whose contents to write, but without its index. + extension: A user specified extension to use for the file if not contained in + `fname` already. + + Returns: + None + + Raises: + NotImplementedError: The given output extension is not supported yet. + """ + # Decipher the path given + fname = fname.decode() if isinstance(fname, bytes) else fname + path = Path(f"{fname}.{extension}") if extension else Path(fname) + # Write df to either a csv or parquet or raise if an invalid extension + if path.suffix == ".csv": + return df.to_csv(path, index=False) + elif path.suffix == ".parquet": + return df.to_parquet(path, index=False, engine="pyarrow") + raise NotImplementedError( + f"Invalid extension provided: '.{path.suffix[1:]}'. Supported extensions are `.csv` or `.parquet`." + ) + + +def read_df( + fname: str | bytes | os.PathLike, + extension: Literal[None, "", "csv", "parquet"] = "", +) -> pd.DataFrame: + """Reads a pandas DataFrame from either a CSV or Parquet file. + + Reads a pandas DataFrame to either a CSV or Parquet file and can infer which format + to use based on the extension given in `fname` or based on explicit `extension`. If + the file being read is a csv with a column called 'subpop' then that column will be + cast as a string. + + Args: + fname: The name of the file to read from. + extension: A user specified extension to use for the file if not contained in + `fname` already. + + Returns: + A pandas DataFrame parsed from the file given. + + Raises: + NotImplementedError: The given output extension is not supported yet. + """ + # Decipher the path given + fname = fname.decode() if isinstance(fname, bytes) else fname + path = Path(f"{fname}.{extension}") if extension else Path(fname) + # Read df from either a csv or parquet or raise if an invalid extension + if path.suffix == ".csv": + return pd.read_csv( + path, converters={"subpop": lambda x: str(x)}, skipinitialspace=True + ) + elif path.suffix == ".parquet": + return pd.read_parquet(path, engine="pyarrow") + raise NotImplementedError( + f"Invalid extension provided: '.{path.suffix[1:]}'. Supported extensions are `.csv` or `.parquet`." + ) + + +def command_safe_run( + command: str, command_name: str = "mycommand", fail_on_fail: bool = True +) -> tuple[int, str, str]: + """ + Runs a shell command and prints diagnostics if command fails. + + Args: + command: The CLI command to be given. + command_name: The reference name for you command. Default value is "mycommand". + fail_on_fail: If True, an exception will be thrown if the command fails (default is True) + + Returns: + As a tuple; the return code, the standard output, and standard error from running the command. + + Raises: + RuntimeError: If fail_on_fail=True and the command fails, an error will be thrown. + """ + import subprocess + import shlex # using shlex to split the command because it's not obvious https://docs.python.org/3/library/subprocess.html#subprocess.Popen + + sr = subprocess.Popen( + shlex.split(command), stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + (stdout, stderr) = sr.communicate() + if sr.returncode != 0: + print(f"'{command_name}' failed failed with returncode '{sr.returncode}'") + print(f"'{command_name}': '{command}'") + print("{command_name} command failed with stdout and stderr:") + + print("{command_name} stdout >>>>>>") + print(stdout.decode()) + print("{command_name} stdout <<<<<<") + + print("{command_name} stderr >>>>>>") + print(stderr.decode()) + print("{command_name} stderr <<<<<<") + if fail_on_fail: + raise RuntimeError(f"The '{command_name}' command failed.") + + return sr.returncode, stdout, stderr + + +def add_method(cls: Any): + """ + A function which adds a function to a class. + + Args: + cls: The class you want to add a method to. + + Returns: + decorator: The decorator. + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + setattr(cls, func.__name__, wrapper) + return func + + return decorator + + +def search_and_import_plugins_class( + plugin_file_path: str, path_prefix: str, class_name: str, **kwargs: dict[str, Any] +) -> Any: + """ + Function serving to create a class that finds and imports the necessary modules. + + Args: + plugin_file_path: Pathway to the module. + path_prefix: Pathway prefix to the module. + class_name: Name of the class. + **kwargs: Further arguments passed to initilization of the class. + + Returns: + The instance of the class that was instantiated with provided **kwargs. + + Examples: + Suppose there is a module called `my_plugin.py with a class `MyClass` located at `/path/to/plugin/`. + + Dynamically import and instantiate the class: + + >>> instance = search_and_import_plugins_class('/path/to/plugin', path_prefix, 'MyClass', **params) + + View the instance: + + >>> print(instance) + <__main__.MyClass object at 0x7f8b2c6b4d60> + + """ + # Look for all possible plugins and import them + # https://stackoverflow.com/questions/67631/how-can-i-import-a-module-dynamically-given-the-full-path + # unfortunatelly very complicated, this is cpython only ?? + import sys, os + + full_path = os.path.join(path_prefix, plugin_file_path) + sys.path.append(os.path.dirname(full_path)) + # the following works, but these above lines seems necessary to pickle // runs + from pydoc import importfile + + module = importfile(full_path) + klass = getattr(module, class_name) + return klass(**kwargs) + + +### Profile configuration +import cProfile +import pstats +from functools import wraps + + +def profile( + output_file: str = None, + sort_by: str = "cumulative", + lines_to_print: int = None, + strip_dirs: bool = False, +): + """ + A time profiler decorator. + Inspired by and modified the profile decorator of Giampaolo Rodola: + http://code.activestate.com/recipes/577817-profile-decorator/ + + Args: + output_file: + Path of the output file. If only name of the file is given, it's + saved in the current directory. + If it's None, the name of the decorated function is used. + sort_by: + Sorting criteria for the Stats object. + For a list of valid string and SortKey refer to: + https://docs.python.org/3/library/profile.html#pstats.Stats.sort_stats + lines_to_print: + Number of lines to print. Default (None) is for all the lines. + This is useful in reducing the size of the printout, especially + that sorting by 'cumulative', the time consuming operations + are printed toward the top of the file. + strip_dirs: + Whether to remove the leading path info from file names. + + Returns: + Profile of the decorated function. + + Examples: + >>> @profile(output_file="my_function.prof") + >>> def my_function(): + # Function body content + pass + >>> my_function() + After running ``my_function``, a file named ``my_function.prof`` will be created in the current WD. + This file contains the profiling data. + """ + + def inner(func): + @wraps(func) + def wrapper(*args, **kwargs): + _output_file = output_file or func.__name__ + ".prof" + pr = cProfile.Profile() + pr.enable() + retval = func(*args, **kwargs) + pr.disable() + pr.dump_stats(_output_file) + return retval + + return wrapper + + return inner + + +def as_list(thing: Any) -> list[Any]: + """ + Returns argument passed as a list. + + Args: + thing: The object that you would like to be converted to a list. + + Returns: + The object converted to a list. + """ + if type(thing) == list: + return thing + return [thing] + + +### A little timer class +class Timer(object): + """ + A timer class that starts, ends, and records time in between. + + Attributes: + name: Name of event. + tstart: Time start. + """ + + def __init__(self, name): + self.name = name + + def __enter__(self): + logging.debug(f"[{self.name}] started") + self.tstart = time.time() + + def __exit__(self, type, value, traceback): + logging.debug(f"[{self.name}] completed in {time.time() - self.tstart:,.2f} s") + + +class ISO8601Date(confuse.Template): + """ + Reads in config dates into datetimes.dates. + """ + + def convert(self, value: any, view: confuse.ConfigView): + """ + Converts the given value to a datetime.date object. + + Args: + value: The value to be converted. Can be datetime.date object or ISO8601 string. + view: A view object from confuse, to be used for error reporting. + + Raises: + confuse.TemplateError: If `value` is neither a datetime.date nor an ISO8601Date string. + """ + if isinstance(value, datetime.date): + return value + elif isinstance(value, str): + return datetime.datetime.strptime(value, "%Y-%m-%d").date() + else: + self.fail("must be a date object or ISO8601 date", True) + + +@add_method(confuse.ConfigView) +def as_date(self) -> datetime.date: + """ + Evaluates an datetime.date or ISO8601 date string. + + Returns: + A datetime.date data type of the date associated with the object. + """ + + return self.get(ISO8601Date()) + + +def get_truncated_normal( + mean: float | int = 0, + sd: float | int = 1, + a: float | int = 0, + b: float | int = 10, +) -> scipy.stats._distn_infrastructure.rv_frozen: + """ + Returns a truncated normal distribution. + + This function constructs a truncated normal distribution with the specified + mean, standard deviation, and bounds. The truncated normal distribution is + a normal distribution bounded within the interval [a, b]. + + Args: + mean: The mean of the truncated normal distribution. Defaults to 0. + sd: The standard deviation of the truncated normal distribution. Defaults to 1. + a: The lower bound of the truncated normal distribution. Defaults to 0. + b: The upper bound of the truncated normal distribution. Defaults to 10. + + Returns: + rv_frozen: A frozen instance of the truncated normal distribution with the specified parameters. + + Examples: + Create a truncated normal distribution with specified parameters (truncated between 1 and 10): + >>> truncated_normal_dist = get_truncated_normal(mean=5, sd=2, a=1, b=10) + >>> print(truncated_normal_dist) + rv_frozen() + """ + lower = (a - mean) / sd + upper = (b - mean) / sd + return scipy.stats.truncnorm(lower, upper, loc=mean, scale=sd) + + +def get_log_normal( + meanlog: float | int, + sdlog: float | int, +) -> scipy.stats._distn_infrastructure.rv_frozen: + """ + Returns a log normal distribution. + + This function constructs a log normal distribution with the specified + log mean and log standard deviation. + + Args: + meanlog: The log of the mean of the log normal distribution. + sdlog: The log of the standard deviation of the log normal distribution. + + Returns: + rv_frozen: A frozen instance of the log normal distribution with the + specified parameters. + + Examples: + Create a log-normal distribution with specified parameters: + >>> log_normal_dist = get_log_normal(meanlog=1, sdlog=0.5) + >>> print(log_normal_dist) + + """ + return scipy.stats.lognorm(s=sdlog, scale=np.exp(meanlog), loc=0) + + +def random_distribution_sampler( + distribution: Literal[ + "fixed", "uniform", "poisson", "binomial", "truncnorm", "lognorm" + ], + **kwargs: dict[str, Any], +) -> Callable[[], float | int]: + """ + Create function to sample from a random distribution. + + Args: + distribution: The type of distribution to generate a sampling function for. + **kwargs: Further parameters that are passed to the underlying function for the + given distribution. + + Notes: + The further args expected by each distribution type are: + - fixed: value, + - uniform: low, high, + - poisson: lam, + - binomial: n, p, + - truncnorm: mean, sd, a, b, + - lognorm: meanlog, sdlog. + + Returns: + A function that can be called to sample from that distribution. + + Raises: + ValueError: If `distribution` is 'binomial' the given `p` must be in (0,1). + NotImplementedError: If `distribution` is not one of the type hinted options. + + Examples: + >>> import numpy as np + >>> np.random.seed(123) + >>> uniform_sampler = random_distribution_sampler("uniform", low=0.0, high=3.0) + >>> uniform_sampler() + 2.089407556793585 + >>> uniform_sampler() + 0.8584180048511384 + """ + if distribution == "fixed": + # Fixed value is the same as uniform on [a, a) + return functools.partial( + np.random.uniform, + kwargs.get("value"), + kwargs.get("value"), + ) + elif distribution == "uniform": + # Uniform on [low, high) + return functools.partial( + np.random.uniform, + kwargs.get("low"), + kwargs.get("high"), + ) + elif distribution == "poisson": + # Poisson with mean lambda + return functools.partial(np.random.poisson, kwargs.get("lam")) + elif distribution == "binomial": + p = kwargs.get("p") + if not (0 < p < 1): + raise ValueError(f"Invalid `p-value`: '{p}' is out of range [0,1].") + return functools.partial(np.random.binomial, kwargs.get("n"), p) + elif distribution == "truncnorm": + # Truncated normal with mean, sd on interval [a, b] + return get_truncated_normal( + mean=kwargs.get("mean"), + sd=kwargs.get("sd"), + a=kwargs.get("a"), + b=kwargs.get("b"), + ).rvs + elif distribution == "lognorm": + # Lognormal distribution with meanlog, sdlog + return get_log_normal(kwargs.get("meanlog"), kwargs.get("sdlog")).rvs + raise NotImplementedError(f"Unknown distribution [received: '{distribution}'].") + + +@add_method(confuse.ConfigView) +def as_random_distribution(self): + """ + Constructs a random distribution object from a distribution config key. + + Args: + self: Class instance (in this case, a config key) to construct the random distribution from. + + Returns: + A partial object containing the random distribution. + + Raises: + ValueError: When values are out of range. + NotImplementedError: If an unknown distribution is found. + + Examples: + Say that ``config`` is a ``confuse.ConfigView`` instance. + + To create a uniform distribution between 1 and 10: + >>> dist_function = config.as_random_distribution() + >>> sample = dist_function() + 5.436789235794546 + + To use a truncated normal distribution: + >>> config_truncnorm = confuse.ConfigView({ + "distribution": "truncnorm", + "mean": 0 + "sd": 1, + "a": -1, + "b": 1 + }) + >>> truncnorm_dist_function = config_truncnorm.as_random_distribution() + >>> truncnorm_sample = truncnorm_dist_function() + 0.312745 + ``` + """ + + if isinstance(self.get(), dict): + dist = self["distribution"].get() + if dist == "fixed": + return functools.partial( + np.random.uniform, + _evaled_expression(self["value"].get(), target_type=float), + _evaled_expression(self["value"].get(), target_type=float), + ) + elif dist == "uniform": + return functools.partial( + np.random.uniform, + _evaled_expression(self["low"].get(), target_type=float), + _evaled_expression(self["high"].get(), target_type=float), + ) + elif dist == "poisson": + return functools.partial(np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float)) + elif dist == "binomial": + p = _evaled_expression(self["p"].get(), target_type=float) + if (p < 0) or (p > 1): + raise ValueError(f"Invalid `p-value`: '{p}' is out of range [0,1].") + # if (self["p"] < 0) or (self["p"] > 1): + # raise ValueError(f"""p value { self["p"] } is out of range [0,1]""") + return functools.partial( + np.random.binomial, + _evaled_expression(self["n"].get(), target_type=float), + # _evaled_expression(self["p"].get(), target_type=float), + p, + ) + elif dist == "truncnorm": + return get_truncated_normal( + mean=_evaled_expression(self["mean"].get(), target_type=float), + sd=_evaled_expression(self["sd"].get(), target_type=float), + a=_evaled_expression(self["a"].get(), target_type=float), + b=_evaled_expression(self["b"].get(), target_type=float), + ).rvs + elif dist == "lognorm": + return get_log_normal( + meanlog=_evaled_expression(self["mean_log"].get(), target_type=float), + sdlog=_evaled_expression(self["sd_log"].get(), target_type=float), + ).rvs + else: + raise NotImplementedError(f"Unknown distribution [received: '{dist}'].") + else: + # we allow a fixed value specified directly: + return functools.partial( + np.random.uniform, + _evaled_expression(self, target_type=float), + _evaled_expression(self, target_type=float), + ) + + +def list_filenames( + folder: str | bytes | os.PathLike = ".", + filters: str | list[str] = [], +) -> list[str]: + """ + Return the list of all filenames and paths in the provided folder. + + This function lists all files in the specified folder and its subdirectories. + If filters are provided, only the files containing each of the substrings + in the filters will be returned. + + Args: + folder: + The directory to search for files. Defaults to the current directory. + filters: + A string or a list of strings to filter filenames. Only files + containing all the provided substrings will be returned. Defaults to an + empty list. + + Returns: + A list of strings representing the paths to the files that match the filters. + + Examples: + To get all files containing "hosp": + >>> gempyor.utils.list_filenames( + folder="model_output/", + filters=["hosp"], + ) + + To get only "hosp" files with a ".parquet" extension: + >>> gempyor.utils.list_filenames( + folder="model_output/", + filters=["hosp", ".parquet"], + ) + """ + filters = [filters] if not isinstance(filters, list) else filters + filters = filters if len(filters) else [""] + folder = Path(folder.decode() if isinstance(folder, bytes) else folder) + files = [ + str(file) + for file in folder.rglob("*") + if file.is_file() and all(f in str(file) for f in filters) + ] + return files + + +def extract_slot(file: Path) -> pd.DataFrame: + """ + Extract the slot number from the filename and add to DataFrame. + + Args: + file: The filename to extract the slot number from. + + Returns: + A pandas DataFrame containing the contents of `file` along with a column called + 'slot' that contains the slot number extracted from the filename. + + Raises: + ValueError: If the 'slot' column already exists in the DataFrame read in. + + Examples: + >>> from pathlib import Path + >>> import pandas as pd + >>> from gempyor.utils import extract_slot + >>> sample = pd.DataFrame(data={"letters": ["a", "b", "c"]}) + >>> file = Path.cwd() / "00017.foobar.csv" + >>> sample.to_csv(file, index=False) + >>> extract_slot(file) + letters slot + 0 a 17 + 1 b 17 + 2 c 17 + >>> extract_slot(file).info() + + RangeIndex: 3 entries, 0 to 2 + Data columns (total 2 columns): + # Column Non-Null Count Dtype + --- ------ -------------- ----- + 0 letters 3 non-null object + 1 slot 3 non-null int64 + dtypes: int64(1), object(1) + memory usage: 180.0+ bytes + """ + df = pd.read_parquet(file) if file.suffix == ".parquet" else pd.read_csv(file) + if "slot" in df.columns: + raise ValueError( + f"Column 'slot' already exists in the DataFrame from file: {file}." + ) + df["slot"] = int(file.name.split(".")[0]) + return df + + +def read_directory( + directory: str | bytes | os.PathLike, filters: str | list[str] | None = None +) -> pd.DataFrame: + """ + Read all files in a directory into a single DataFrame. + + Args: + directory: The directory to read files from. + filters: A string or a list of strings to filter filenames. Only files + containing all the provided substrings will be read. Defaults to `None` for + no filters. + + Returns: + A pandas DataFrame containing the contents of all the files in the directory. + """ + files = [Path(file) for file in list_filenames(folder=directory, filters=filters or [])] + dfs: list[pd.DataFrame] | pd.DataFrame = [] + for file in sorted(files): + dfs.append(extract_slot(file)) + return pd.concat(dfs).reset_index(drop=True) + + +def rolling_mean_pad( + data: npt.NDArray[np.number], + window: int, +) -> npt.NDArray[np.number]: + """ + Calculates the column-wise rolling mean with centered window. + + Args: + data: A two dimensional numpy array, typically the row dimension is time and the + column dimension is subpop. + window: The window size for the rolling mean. + + Returns: + A two dimensional numpy array that is the same shape as `data`. + + Examples: + Below is a brief set of examples showcasing how to smooth a metric, like + hospitalizations, using this function. + ``` + >>> import numpy as np + >>> from gempyor.utils import rolling_mean_pad + >>> hospitalizations = np.arange(1., 29.).reshape((7, 4)) + >>> hospitalizations + array([[ 1., 2., 3., 4.], + [ 5., 6., 7., 8.], + [ 9., 10., 11., 12.], + [13., 14., 15., 16.], + [17., 18., 19., 20.], + [21., 22., 23., 24.], + [25., 26., 27., 28.]]) + >>> rolling_mean_pad(hospitalizations, 5) + array([[ 3.4, 4.4, 5.4, 6.4], + [ 5.8, 6.8, 7.8, 8.8], + [ 9. , 10. , 11. , 12. ], + [13. , 14. , 15. , 16. ], + [17. , 18. , 19. , 20. ], + [20.2, 21.2, 22.2, 23.2], + [22.6, 23.6, 24.6, 25.6]]) + ``` + """ + weights = (1.0 / window) * np.ones(window) + output = scipy.ndimage.convolve1d(data, weights, axis=0, mode="nearest") + if window % 2 == 0: + rows, cols = data.shape + i = rows - 1 + output[i, :] = 0.0 + window -= 1 + weight = 1.0 / window + for l in range(-((window - 1) // 2), 1 + (window // 2)): + i_star = min(max(i + l, 0), i) + for j in range(cols): + output[i, j] += weight * data[i_star, j] + return output + + +def print_disk_diagnosis(): + """ + Reads and prints AWS disk diagnostic information. + """ + import os + from os import path + from shutil import disk_usage + + def bash(command: str) -> str: + """ + Executes a shell command and returns its output. + + Args: + command: The shell command to be executed. + + Returns: + The output of the shell command. + """ + output = os.popen(command).read() + return output + + print("START AWS DIAGNOSIS ================================") + total_bytes, used_bytes, free_bytes = disk_usage(path.realpath("/")) + print( + f"shutil.disk_usage: {total_bytes/ 1000000} Mb total, {used_bytes / 1000000} Mb used, {free_bytes / 1000000} Mb free..." + ) + print("------------") + print(f"df -hT: '{bash('df -hT')}'") + print("------------") + print(f"df -i: '{bash('df -i')}'") + print("------------") + print(f"free -h: '{bash('free -h')}'") + print("------------") + print(f"lsblk: '{bash('lsblk')}'") + print("END AWS DIAGNOSIS ================================") + + +def create_resume_out_filename( + flepi_run_index: str, + flepi_prefix: str, + flepi_slot_index: str, + flepi_block_index: str, + filetype: str, + liketype: str, +) -> str: + """ + Compiles run output information. + + Args: + flepi_run_index: Index of the run. + flepi_prefix: File prefix. + flepi_slot_index: Index of the slot. + flepi_block_index: Index of the block. + filetype: File type. + liketype: Chimeric or global. + + Returns: + The path to a corresponding output file. + + Examples: + Generate an output file with specified parameters: + >>> filename = create_resume_out_filename( + flepi_run_index="test_run", + flepi_prefix="model_output/run_id/", + flepi_slot_index="1", + flepi_block_index="2", + filetype="seed", + liketype="chimeric" + ) + >>> print(filename) + "experiment/001/normal/intermediate/000000123.000000000.1.parquet" + """ + prefix = f"{flepi_prefix}/{flepi_run_index}" + inference_filepath_suffix = f"{liketype}/intermediate" + inference_filename_prefix = "{:09d}.".format(int(flepi_slot_index)) + index = "{:09d}.{:09d}".format(1, int(flepi_block_index) - 1) + extension = "parquet" + if filetype == "seed": + extension = "csv" + return file_paths.create_file_name( + run_id=flepi_run_index, + prefix=prefix, + inference_filename_prefix=inference_filename_prefix, + inference_filepath_suffix=inference_filepath_suffix, + index=index, + ftype=filetype, + extension=extension, + ) + + +def create_resume_input_filename( + resume_run_index: str, + flepi_prefix: str, + flepi_slot_index: str, + filetype: str, + liketype: str, +) -> str: + """ + Compiles run input information. + + Args: + resume_run_index: Index of the run. + flepi_prefix: File prefix. + flepi_slot_index: Index of the slot. + filetype: File type. + liketype: Chimeric or global. + + Returns: + The path to the a corresponding input file. + + Examples: + Generate an input file with specified parameters: + >>> filename = create_resume_input_filename( + resume_run_index="2", + flepi_prefix="model_output/run_id/", + flepi_slot_index="1", + filetype="seed", + liketype="chimeric" + ) + >>> print(filename) + "experiment/002/normal/final/789.csv" + """ + prefix = f"{flepi_prefix}/{resume_run_index}" + inference_filepath_suffix = f"{liketype}/final" + index = flepi_slot_index + extension = "parquet" + if filetype == "seed": + extension = "csv" + return file_paths.create_file_name( + run_id=resume_run_index, + prefix=prefix, + inference_filepath_suffix=inference_filepath_suffix, + index=index, + ftype=filetype, + extension=extension, + ) + + +def get_filetype_for_resume( + resume_discard_seeding: str, flepi_block_index: str +) -> list[str]: + """ + Retrieves a list of parquet file types that are relevant for resuming a process based on + specific environment variable settings. + This function dynamically determines the list + based on the current operational context given by the environment. + + Args: + resume_discard_seeding: Determines whether seeding-related file types should be included (str). + flepi_block_index: Determines a specific operational mode or block of the process (str). + + Returns: + List of file types. + + Examples: + Determine file types for block index 1 with seeding data NOT discarded: + >>> filetypes = get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="1") + >>> print(filetypes) + ["seed", "spar", "snpi", "hpar", "hnpi", "init"] + + Determine file types for block index 2 with seeding data discarded: + >>> filtypes = get_filetype_for_resume(resume_discard_seeding="true", flepi_block_index="2") + >>> print(filetypes) + ["seed", "spar", "snpi", "hpar", "hnpi", "host", "llik", "init"] + """ + if flepi_block_index == "1": + if resume_discard_seeding == "true": + return ["spar", "snpi", "hpar", "hnpi", "init"] + else: + return ["seed", "spar", "snpi", "hpar", "hnpi", "init"] + else: + return ["seed", "spar", "snpi", "hpar", "hnpi", "host", "llik", "init"] + + +def create_resume_file_names_map( + resume_discard_seeding: str, + flepi_block_index: str, + resume_run_index: str, + flepi_prefix: str, + flepi_slot_index: str, + flepi_run_index: str, + last_job_output: str, +) -> dict[str, str]: + """ + Generates a mapping of input file names to output file names for a resume process based on + parquet file types and environmental conditions. The function adjusts the file name mappings + based on the operational block index and the location of the last job output. + + Args: + resume_discard_seeding: Determines whether seeding-related file types should be included. + flepi_block_index: Determines a specific operational mode or block of the process. + resume_run_index: Resume run index. + flepi_prefix: File prefix. + flepi_slot_index: Index of the slot. + flepi_run_index: flepiMoP run index. + last_job_output: Adjusts the keys in the mapping to be prefixed with this path. + + Returns: + A dictionary where keys are input file paths and values are corresponding + output file paths. + + The mappings depend on: + - Parquet file types appropriate for resuming a process, as determined by the environment. + - Whether the files are for 'global' or 'chimeric' types, as these liketypes influence the + file naming convention. + - The operational block index ('FLEPI_BLOCK_INDEX'), which can alter the input file names for + block index '1'. + - The presence and value of 'LAST_JOB_OUTPUT' environment variable, which if set to an S3 path, + adjusts the keys in the mapping to be prefixed with this path. + + Raises: + No explicit exceptions are raised within the function, but it relies heavily on external + functions and environment variables which if improperly configured could lead to unexpected + behavior. + + Examples: + Generate a mapping of file names for a given resume process: + >>> file_names_map = create_resume_file_names_map( + resume_discard_seeding="false", + flepi_block_index="1", + resume_run_index="1", + flepi_prefix="model_output/run_id/", + flepi_slot_index="1", + flepi_run_index="test_run", + last_job_output="s3://bucket/path/") + >>> print(file_names_map) + { + 's3://bucket/path/model_output/run_id/1_type1_global_1.in': 'model_output/run_id/test_run_type1_global_1_1.out', + 's3://bucket/path/model_output/run_id/1_type1_chimeric_1.in': 'model_output/run_id/test_run_type1_chimeric_1_1.out', + 's3://bucket/path/model_output/run_id/1_type2_global_1.in': 'model_output/run_id/test_run_type2_global_1_1.out', + 's3://bucket/path/model_output/run_id/1_type2_chimeric_1.in': 'model_output/run_id/test_run_type2_chimeric_1_1.out' + } + # Note: this output is toy output implemented with toy file names. + + Notes: + - The paths may be modified by the 'LAST_JOB_OUTPUT' if it is set and points to an S3 location. + """ + file_types = get_filetype_for_resume( + resume_discard_seeding=resume_discard_seeding, flepi_block_index=flepi_block_index + ) + resume_file_name_mapping = dict() + liketypes = ["global", "chimeric"] + for filetype in file_types: + for liketype in liketypes: + output_file_name = create_resume_out_filename( + flepi_run_index=flepi_run_index, + flepi_prefix=flepi_prefix, + flepi_slot_index=flepi_slot_index, + flepi_block_index=flepi_block_index, + filetype=filetype, + liketype=liketype, + ) + input_file_name = output_file_name + if flepi_block_index == "1": + input_file_name = create_resume_input_filename( + resume_run_index=resume_run_index, + flepi_prefix=flepi_prefix, + flepi_slot_index=flepi_slot_index, + filetype=filetype, + liketype=liketype, + ) + resume_file_name_mapping[input_file_name] = output_file_name + if last_job_output.find("s3://") >= 0: + old_keys = list(resume_file_name_mapping.keys()) + for k in old_keys: + new_key = os.path.join(last_job_output, k) + resume_file_name_mapping[new_key] = resume_file_name_mapping[k] + del resume_file_name_mapping[k] + return resume_file_name_mapping + + +def download_file_from_s3(name_map: dict[str, str]) -> None: + """ + Downloads files from AWS S3 based on a mapping of S3 URIs to local file paths. The function + checks if the directory for the first output file exists and creates it if necessary. It + then iterates over each S3 URI in the provided mapping, downloads the file to the corresponding + local path, and handles errors if the S3 URI format is incorrect or if the download fails. + + Args: + name_map: + A dictionary where keys are S3 URIs (strings) and values + are the local file paths (strings) where the files should + be saved. + + Returns: + This function does not return a value; its primary effect is the side effect of + downloading files and potentially creating directories. + + Raises: + ValueError: If an S3 URI does not start with 's3://', indicating an invalid format. + ClientError: If an error occurs during the download from S3, such as a permissions issue, + a missing file, or network-related errors. These are caught and logged but not + re-raised, to allow the function to attempt subsequent downloads. + + Examples: + >>> name_map = { + "s3://mybucket/data/file1.txt": "/local/path/to/file1.txt", + "s3://mybucket/data/file2.txt": "/local/path/to/file2.txt" + } + >>> download_file_from_s3(name_map) + # This would download 'file1.txt' and 'file2.txt' from 'mybucket' on S3 to the specified local paths. + + # If an S3 URI is malformed: + >>> name_map = { + "http://wrongurl.com/data/file1.txt": "/local/path/to/file1.txt" + } + >>> download_file_from_s3(name_map) + # This will raise a ValueError indicating the invalid S3 URI format. + """ + try: + import boto3 + from botocore.exceptions import ClientError + except ModuleNotFoundError: + raise ModuleNotFoundError( + ( + "No module named `boto3` found, which is required for " + "`gempyor.utils.download_file_from_s3`. Please install the aws target." + ) + ) + s3 = boto3.client("s3") + first_output_filename = next(iter(name_map.values())) + output_dir = os.path.dirname(first_output_filename) + if not os.path.exists(output_dir): + os.makedirs(output_dir) + for s3_uri in name_map: + try: + if s3_uri.startswith("s3://"): + bucket = s3_uri.split("/")[2] + object = s3_uri[len(bucket) + 6 :] + s3.download_file(bucket, object, name_map[s3_uri]) + else: + raise ValueError(f"Invalid S3 URI format [received: '{s3_uri}'].") + except ClientError as e: + raise Exception(f"'{e}': could not download filefrom S3.") + + +def move_file_at_local(name_map: dict[str, str]) -> None: + """ + Moves files locally according to a given mapping. + This function takes a dictionary where the keys are source file paths and + the values are destination file paths. It ensures that the destination + directories exist and then copies the files from the source paths to the + destination paths. + + Args: + name_map: A dictionary mapping source file paths to destination file paths. + """ + for src, dst in name_map.items(): + os.path.makedirs(os.path.dirname(dst), exist_ok=True) + shutil.copy(src, dst) + + +def _dump_formatted_yaml(cfg: confuse.Configuration) -> str: + """ + Dump confuse configuration to a formatted YAML string. + + Args: + cfg: The confuse configuration object. + + Returns: + A formatted YAML string representation of the configuration. + + Examples: + >>> from gempyor.utils import _dump_formatted_yaml + >>> import confuse + >>> conf = confuse.Configuration("foobar") + >>> data = { + ... "name": "Test Config", + ... "compartments": { + ... "infection_stage": ["S", "E", "I", "R"] + ... }, + ... "seir": { + ... "parameters": { + ... "beta": {"value": 3.4}, + ... "gamma": {"value": 5.6}, + ... }, + ... "transitions": { + ... "source": ["S"], + ... "destination": ["E"], + ... "rate": ["beta * gamma"], + ... "proportional_to": [["S"], ["I"]], + ... "proportion_exponent": [1, 1], + ... }, + ... }, + ... } + >>> conf.set(data) + >>> print(_dump_formatted_yaml(conf)) + name: "Test Config" + compartments: + infection_stage: [S, E, I, R] + seir: + parameters: + beta: + value: 3.4 + gamma: + value: 5.6 + transitions: + source: [S] + destination: [E] + rate: ["beta * gamma"] + proportional_to: [[S], [I]] + proportion_exponent: [1, 1] + """ + + class CustomDumper(yaml.Dumper): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.add_representer(list, self._represent_list) + self.add_representer(str, self._represent_str) + + def _represent_list(self, dumper, data): + return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True) + + def _represent_str(self, dumper, data): + if " " in data: + return dumper.represent_scalar("tag:yaml.org,2002:str", data, style='"') + return dumper.represent_scalar("tag:yaml.org,2002:str", data) + + return yaml.dump( + yaml.safe_load(cfg.dump()), Dumper=CustomDumper, indent=4, sort_keys=False + ) + + +def _shutil_which( + cmd: str, + mode: int = os.F_OK | os.X_OK, + path: str | bytes | os.PathLike | None = None, + check: bool = True, +) -> str | None: + """ + A thin wrapper around `shutil.which` with extra validation. + + Args: + cmd: The name of the command to search for. + mode: The permission mask required of possible files. + path: A path describing the locations to search, or `None` to use + the PATH environment variable. + check: If `True` an `OSError` will be raised if a `cmd` is not found. + Similar in spirit to the `check` arg of `subprocess.run`. + + Returns: + Either the full path to the `cmd` found, or `None` if `cmd` is not + found and `check` is `False`. + + Raises: + OSError: If `cmd` is not found and `check` is `True`. + + Examples: + >>> import os + >>> from shutil import which + >>> _shutil_which("python") == which("python") + True + >>> _shutil_which("does_not_exist", check=False) + >>> try: + ... _shutil_which("does_not_exist", check=True) + ... except Exception as e: + ... print(type(e)) + ... print(str(e).replace(os.environ.get("PATH"), "...")) + ... + + Did not find 'does_not_exist' on path '...'. + + See Also: + [`shutil.which`](https://docs.python.org/3/library/shutil.html#shutil.which) + """ + result = shutil.which(cmd, mode=mode, path=path) + if check and result is None: + path = os.environ.get("PATH") if path is None else path + raise OSError(f"Did not find '{cmd}' on path '{path}'.") + return result + + +def _git_head(repository: Path) -> str: + """ + Get the sha commit hash for the head of a git repository. + + Args: + repository: A directory under version control with git to get the sha commit of. + + Returns: + The sha commit of head for `repository`. + + Examples: + >>> import os + >>> from pathlib import Path + >>> _git_head(Path(os.environ["FLEPI_PATH"])) + 'efe896b1a5e4f8e33667c170cd5319d6ef1e3db5' + """ + git_cmd = _shutil_which("git") + proc = subprocess.run( + [git_cmd, "rev-parse", "HEAD"], + cwd=repository.expanduser().absolute(), + capture_output=True, + check=True, + ) + return proc.stdout.decode().strip() + + +def _git_checkout(repository: Path, branch: str) -> subprocess.CompletedProcess: + """ + Checkout a new branch from a given git repository. + + Args: + repository: A directory under version control with git to + checkout a new branch in. + branch: The name of the new branch to checkout. + + Examples: + >>> import os + >>> from pathlib import Path + >>> _git_checkout(Path(os.environ["FLEPI_PATH"]), "my-new-branch") + """ + git_cmd = _shutil_which("git") + return subprocess.run( + [git_cmd, "checkout", "-b", shlex_quote(branch)], + cwd=repository.expanduser().absolute(), + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + +def _format_cli_options( + options: dict[str, str | Iterable[str]] | None, + always_single: bool = False, + iterable_formatting: Literal["split", "comma"] = "split", +) -> list[str]: + """ + Convert a dictionary of CLI options into a formatted list. + + Args: + options: A dictionary where the keys correspond to the option name and the + values correspond to the option value. Values are escaped for shell. + always_single: If `True` all options will be formatted with a single dash prefix + always, useful for commands that don't support long options like + `pdftotext`. If `False` the option will be formatted with a double dash + prefix if the key is longer than one character. + iterable_formatting: The formatting to use for iterable values. If 'split' each + value will be formatted as a separate option. If 'comma' the values will be + joined with a comma and formatted as a single option. + + Returns: + A list of options that can be passed to + [`subprocess.run`](https://docs.python.org/3/library/subprocess.html#subprocess.run) + or similar functions. + + Examples: + >>> from gempyor.utils import _format_cli_options + >>> _format_cli_options({"name": "foo bar fizz buzz"}) + ["--name='foo bar fizz buzz'"] + >>> _format_cli_options({"o": "/path/to/output.log"}) + ['-o=/path/to/output.log'] + >>> _format_cli_options({"opt1": "```", "opt2": "$( echo 'Hello!')"}) + ["--opt1='```'", '--opt2=\'$( echo \'"\'"\'Hello!\'"\'"\')\''] + >>> _format_cli_options({"output": "/path/to/output.log"}, always_single=True) + ['-output=/path/to/output.log'] + >>> _format_cli_options({"person": ["Alice", "Bob", "Charlie"]}) + ['--person=Alice', '--person=Bob', '--person=Charlie'] + >>> _format_cli_options( + ... {"person": ["Alice", "Bob", "Charlie"]}, + ... iterable_formatting="comma", + ... ) + ['--person=Alice,Bob,Charlie'] + """ + opts = [] + for k, v in (options or {}).items(): + new_opts = [] + opt_name = f"{'-' if (always_single or len(k) == 1) else '--'}{k}" + if isinstance(v, str): + new_opts.append(f"{opt_name}={shlex_quote(v)}") + elif iterable_formatting == "comma": + new_opts.append(f"{opt_name}={','.join(shlex_quote(w) for w in v)}") + else: + new_opts.extend(f"{opt_name}={shlex_quote(w)}" for w in v) + opts.extend(new_opts) + return opts + + +def _random_seeds_list(a: int, b: int, n: int) -> list[int]: + """ + Generate a list of unique random integers within a specified range. + + Args: + a: The lower bound of the range, inclusive. + b: The upper bound of the range, inclusive. + n: The number of unique random integers to generate. + + Returns: + A list of `n` unique integers between `a` and `b`. + + Raises: + ValueError: If the lower bound is not less than the upper bound. + ValueError: If the range is too small to accommodate `n` unique integers. + + Examples: + >>> import random + >>> from gempyor.utils import _random_seeds_list + >>> random.seed(42) + >>> _random_seeds_list(1, 100, 5) + [82, 15, 4, 95, 36] + >>> _random_seeds_list(1, 100, 5) + [32, 29, 18, 95, 14] + >>> _random_seeds_list(1, 5, 5) + [5, 1, 4, 2, 3] + >>> _random_seeds_list(1, 20, 0) + [] + >>> try: + ... _random_seeds_list(1, 5, 6) + ... except Exception as e: + ... print(e) + ... + Range [1, 5] is too small for 6 unique random integers. + """ + random_seeds = [] + if a >= b: + raise ValueError(f"The lower bound, {a}, must be less than the upper bound, {b}.") + if n < 0: + raise ValueError("The number of random integers to generate must be non-negative.") + if (b - a) + 1 < n: + raise ValueError(f"Range [{a}, {b}] is too small for {n} unique random integers.") + while len(random_seeds) < n: + seed = random.randint(a, b) + if seed not in random_seeds: + random_seeds.append(seed) + return random_seeds + + +def _nslots_random_seeds(nslots: int) -> list[int]: + """ + Generate a list of unique random integers for the number of slots. + + This function generates a list of `nslots` unique random integers between + [`nslots` + 1, 2^32 + 1]. Intended for use in generating random seeds for individual + slots when simulations are done in parallel. + + Args: + nslots: The number of slots to generate random integers for. + + Examples: + >>> import random + >>> from gempyor.utils import _nslots_random_seeds + >>> random.seed(42) + >>> _nslots_random_seeds(5) + [2746317219, 478163333, 107420375, 3184935169, 1181241949] + >>> _nslots_random_seeds(5) + [1051802518, 958682852, 599310831, 3163119791, 440213421] + >>> _nslots_random_seeds(1) + [2906402159] + >>> _nslots_random_seeds(1) + [3181143733] + """ + return _random_seeds_list(nslots + 1, 2**32 - 1, nslots) + + +def _duplicate_strings(it: Iterable[str]) -> list[str]: + """ + Efficiently find duplicate strings in an iterable. + + Args: + it: An iterable of strings. + + Returns: + A list of strings that appear more than once in `it`. + + Examples: + >>> from gempyor.utils import _duplicate_strings + >>> _duplicate_strings(["a", "b", "c"]) + [] + >>> _duplicate_strings(["a", "b", "c", "b"]) + ['b'] + >>> _duplicate_strings(["a", "b", "c", "b", "a"]) + ['a', 'b'] + >>> _duplicate_strings("foobar") + ['o'] + """ + return sorted([k for k, v in Counter(it).items() if v > 1]) + + +@overload +def _trim_s3_path(path: str) -> str: ... + + +@overload +def _trim_s3_path(path: Path) -> Path: ... + + +def _trim_s3_path(path: str | Path) -> str | Path: + """ + Strips the 's3:' prefix from a given path. + + Args: + path: The path to be trimmed. Can be a string or a Path object. + + Returns: + The trimmed path as the same type as the input. + + Examples: + >>> from pathlib import Path + >>> from gempyor.utils import _trim_s3_path + >>> _trim_s3_path("/abc/def/ghi.txt") + '/abc/def/ghi.txt' + >>> _trim_s3_path(Path("/abc/def/ghi.txt")) + PosixPath('/abc/def/ghi.txt') + >>> _trim_s3_path("s3://foo/bar.txt") + '//foo/bar.txt' + >>> _trim_s3_path(Path("s3://foo/bar.txt")) + PosixPath('s3:/foo/bar.txt') + """ + return path.lstrip("s3:") if isinstance(path, str) else path diff --git a/flepimop/gempyor_pkg/src/gempyor/utils.py b/flepimop/gempyor_pkg/src/gempyor/utils.py index e59ac16f8..9a46d50d3 100644 --- a/flepimop/gempyor_pkg/src/gempyor/utils.py +++ b/flepimop/gempyor_pkg/src/gempyor/utils.py @@ -28,6 +28,7 @@ import yaml from . import file_paths +from ._pydantic_ext import _evaled_expression logger = logging.getLogger(__name__) @@ -350,30 +351,6 @@ def as_date(self) -> datetime.date: return self.get(ISO8601Date()) -@add_method(confuse.ConfigView) -def as_evaled_expression(self): - """ - Evaluates an expression string, returning a float. - - Returns: - A float data type of the value associated with the object. - - Raises: - ValueError: On parsing errors. - """ - - value = self.get() - if isinstance(value, numbers.Number): - return value - elif isinstance(value, str): - try: - return float(sympy.parsing.sympy_parser.parse_expr(value)) - except TypeError as e: - raise ValueError(e) from e - else: - raise ValueError(f"Expected numeric or string expression [received: '{value}'].") - - def get_truncated_normal( mean: float | int = 0, sd: float | int = 1, @@ -551,40 +528,40 @@ def as_random_distribution(self): if dist == "fixed": return functools.partial( np.random.uniform, - self["value"].as_evaled_expression(), - self["value"].as_evaled_expression(), + _evaled_expression(self["value"].get(), target_type=float), + _evaled_expression(self["value"].get(), target_type=float), ) elif dist == "uniform": return functools.partial( np.random.uniform, - self["low"].as_evaled_expression(), - self["high"].as_evaled_expression(), + _evaled_expression(self["low"].get(), target_type=float), + _evaled_expression(self["high"].get(), target_type=float), ) elif dist == "poisson": - return functools.partial(np.random.poisson, self["lam"].as_evaled_expression()) + return functools.partial(np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float)) elif dist == "binomial": - p = self["p"].as_evaled_expression() + p = _evaled_expression(self["p"].get(), target_type=float) if (p < 0) or (p > 1): raise ValueError(f"Invalid `p-value`: '{p}' is out of range [0,1].") # if (self["p"] < 0) or (self["p"] > 1): # raise ValueError(f"""p value { self["p"] } is out of range [0,1]""") return functools.partial( np.random.binomial, - self["n"].as_evaled_expression(), - # self["p"].as_evaled_expression(), + _evaled_expression(self["n"].get(), target_type=float), + # _evaled_expression(self["p"].get(), target_type=float), p, ) elif dist == "truncnorm": return get_truncated_normal( - mean=self["mean"].as_evaled_expression(), - sd=self["sd"].as_evaled_expression(), - a=self["a"].as_evaled_expression(), - b=self["b"].as_evaled_expression(), + mean=_evaled_expression(self["mean"].get(), target_type=float), + sd=_evaled_expression(self["sd"].get(), target_type=float), + a=_evaled_expression(self["a"].get(), target_type=float), + b=_evaled_expression(self["b"].get(), target_type=float), ).rvs elif dist == "lognorm": return get_log_normal( - meanlog=self["meanlog"].as_evaled_expression(), - sdlog=self["sdlog"].as_evaled_expression(), + meanlog=_evaled_expression(self["mean_log"].get(), target_type=float), + sdlog=_evaled_expression(self["sd_log"].get(), target_type=float), ).rvs else: raise NotImplementedError(f"Unknown distribution [received: '{dist}'].") @@ -592,8 +569,8 @@ def as_random_distribution(self): # we allow a fixed value specified directly: return functools.partial( np.random.uniform, - self.as_evaled_expression(), - self.as_evaled_expression(), + _evaled_expression(self, target_type=float), + _evaled_expression(self, target_type=float), ) diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py new file mode 100644 index 000000000..57ca86d3f --- /dev/null +++ b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py @@ -0,0 +1,48 @@ +"""Tests for the gempyor._pydantic_ext._evaled_expression function.""" + +import pytest + +from gempyor._pydantic_ext import _evaled_expression + + +@pytest.mark.parametrize( + ("val", "target_type", "expected"), + [ + ("1 + 1", float, 2.0), + ("5 / 2", float, 2.5), + ("sqrt(9)", float, 3.0), + (99.5, float, 99.5), + (-10.0, float, -10.0), + ("2 * 5", int, 10), + ("sqrt(81)", int, 9), + ("11 / 2", int, 5), + ("3.99", int, 3), + (99, int, 99), + (-50, int, -50), + (None, float, None), + (None, int, None), + (True, float, True), + (True, int, True), + (10.5, int, 10.5), + ([1, 2, 3], float, [1, 2, 3]), + ({"a": 1}, float, {"a": 1}), + ], +) +def test_evaled_expression_returns_expected_results(val, target_type, expected): + assert _evaled_expression(val, target_type=target_type) == expected + + +@pytest.mark.parametrize( + ("val", "target_type"), + [ + ("a * b", int), + ("a * b", float), + ("not_a_valid_expression", int), + ("not_a_valid_expression", float), + ("1 / 0", int), + ("1 / 0", float), + ], +) +def test_evaled_expression_raises_value_error_for_invalid_input(val, target_type): + with pytest.raises(ValueError): + _evaled_expression(val, target_type=target_type) \ No newline at end of file diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_float.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_float.py deleted file mode 100644 index 16b5ffd8c..000000000 --- a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_float.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Tests for the `gempyor._pydantic_ext._evaled_expression_float` function.""" - -import pytest - -from gempyor._pydantic_ext import _evaled_expression_float - - -@pytest.mark.parametrize( - ("val", "expected"), - [ - ("1 + 1", 2.0), - ("5 / 2", 2.5), - ("sqrt(9)", 3.0), - (99.5, 99.5), - (-10.0, -10.0), - (None, None), - (True, True), - ([1, 2, 3], [1, 2, 3]), - ({"a": 1}, {"a": 1}), - ], -) -def test_evaled_expression_float_returns_expected_results(val, expected): - assert _evaled_expression_float(val) == expected - - -@pytest.mark.parametrize( - "val", - [ - "a * b", - "not_a_valid_expression", - "1 / 0", - ], -) -def test_evaled_expression_float_raises_value_error_for_invalid_input(val): - with pytest.raises(ValueError): - _evaled_expression_float(val) diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py deleted file mode 100644 index 4755f15a2..000000000 --- a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression_int.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Tests for the `gempyor._pydantic_ext._evaled_expression_int` function.""" - -import pytest - -from gempyor._pydantic_ext import _evaled_expression_int - - -@pytest.mark.parametrize( - ("val", "expected"), - [ - ("2 * 5", 10), - ("sqrt(81)", 9), - ("11 / 2", 5), - ("3.99", 3), - (99, 99), - (-50, -50), - (None, None), - (True, True), - (10.5, 10.5), - ([1, 2], [1, 2]), - ], -) -def test_evaled_expression_int_returns_expected_results(val, expected): - assert _evaled_expression_int(val) == expected - - -@pytest.mark.parametrize( - "val", - [ - "a * b", - "not_a_valid_expression", - "1 / 0", - ], -) -def test_evaled_expression_int_raises_value_error_for_invalid_input(val): - with pytest.raises(ValueError): - _evaled_expression_int(val) diff --git a/flepimop/gempyor_pkg/tests/seir/test_parameters.py b/flepimop/gempyor_pkg/tests/seir/test_parameters.py index bc0937720..7c674f34b 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_parameters.py +++ b/flepimop/gempyor_pkg/tests/seir/test_parameters.py @@ -13,6 +13,7 @@ from gempyor import model_info, seir, NPI, file_paths, parameters, subpopulation_structure from gempyor.utils import config, write_df, read_df +from gempyor._pydantic_ext import _evaled_expression DATA_DIR = os.path.dirname(__file__) + "/data" os.chdir(os.path.dirname(__file__)) @@ -134,7 +135,7 @@ def test_parameters_quick_draw_old(): assert sigma.shape == (modinf.n_days, modinf.nsubpops) assert ( sigma - == config["seir"]["parameters"]["sigma"]["value"]["value"].as_evaled_expression() + == _evaled_expression(config["seir"]["parameters"]["sigma"]["value"]["value"].get(), target_type=float) ).all() assert gamma.shape == (modinf.n_days, modinf.nsubpops) diff --git a/flepimop/gempyor_pkg/tests/utils/test_utils2.py b/flepimop/gempyor_pkg/tests/utils/test_utils2.py index 0822604ed..cfaeb2943 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_utils2.py +++ b/flepimop/gempyor_pkg/tests/utils/test_utils2.py @@ -108,50 +108,12 @@ def test_as_date_with_valid_date_string(): assert result == datetime.date(2022, 1, 15) -def test_as_evaled_expression_with_valid_expression(): - # ConfigViewオブジェクトをモック化 - mock_config_view = MagicMock(spec=confuse.ConfigView) - mock_config_view.as_evaled_expression.return_value = 7.5 - - # as_evaled_expressionメソッドを呼び出し、正しい結果を確認 - result = mock_config_view.as_evaled_expression() - - assert result == 7.5 - - @pytest.fixture def config(): config = confuse.Configuration("myapp", __name__) return config -def test_as_evaled_expression_number(config): - config.add({"myvalue": 123}) - assert config["myvalue"].as_evaled_expression() == 123 - - -def test_as_evaled_expression_number(config): - config.add({"myvalue": 1.10}) - assert config["myvalue"].as_evaled_expression() == 1.1 - - -def test_as_evaled_expression_string(config): - config.add({"myvalue": "2 + 3"}) - assert config["myvalue"].as_evaled_expression() == 5.0 - - -def test_as_evaled_expression_other(config): - config.add({"myvalue": [1, 2, 3]}) - with pytest.raises(ValueError): - config["myvalue"].as_evaled_expression() - - -def test_as_evaled_expression_Invalid_string(config): - config.add({"myvalue": "invalid"}) - with pytest.raises(ValueError): - config["myvalue"].as_evaled_expression() - - def test_as_date(config): config.add({"myvalue": "2022-01-15"}) assert config["myvalue"].as_date() == datetime.date(2022, 1, 15) From 7f35376b73ca8072f3610de785a2b634f2097882 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 25 Jul 2025 17:16:20 -0400 Subject: [PATCH 47/73] Fix tests --- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 42 ++++--------------- .../gempyor_pkg/src/gempyor/utils copy.py | 4 +- flepimop/gempyor_pkg/src/gempyor/utils.py | 12 +++--- .../_pydantic_ext/test__evaled_expression.py | 8 ++-- .../gempyor_pkg/tests/seir/test_parameters.py | 4 +- 5 files changed, 24 insertions(+), 46 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index 8968b3697..c76f2388c 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -191,11 +191,13 @@ def _evaled_expression(val: str, *, target_type: Type[int]) -> int: ... @overload def _evaled_expression(val: Any, *, target_type: Type[int]) -> int | Any: ... + @overload def _evaled_expression(val: str, *, target_type: Type[float]) -> float: ... @overload def _evaled_expression(val: Any, *, target_type: Type[float]) -> float | Any: ... + def _evaled_expression(val: Any, *, target_type: Type[EE]) -> EE | Any: """ Evaluates a string expression to a target numeric type (int or float). @@ -215,8 +217,8 @@ def _evaled_expression(val: Any, *, target_type: Type[EE]) -> EE | Any: 2 >>> _evaled_expression("10 / 4", target_type=float) 2.5 - >>> _evaled_expression("10 / 4", target_type=int) - # note that result is trucnated; probably undesirable if misused + >>> _evaled_expression("10 / 4", target_type=int) + # note that result is trucnated; probably undesirable if misused 2 >>> _evaled_expression(99.5, target_type=float) 99.5 @@ -240,36 +242,6 @@ def _evaled_expression(val: Any, *, target_type: Type[EE]) -> EE | Any: EvaledInt = Annotated[int, BeforeValidator(partial(_evaled_expression, target_type=int))] -EvaledFloat = Annotated[float, BeforeValidator(partial(_evaled_expression, target_type=float))] - - -# TEMP -------- -# --- TEMPORARY DEBUGGING LOGIC --- -# This entire block should be deleted after the refactor is complete. - -import functools -import confuse -from typing import Any - -def add_method(cls: Any): - """A function which adds a function to a class.""" - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - return func(*args, **kwargs) - setattr(cls, func.__name__, wrapper) - return func - return decorator - -@add_method(confuse.ConfigView) -def _as_evaled_expression(self): - """ - This is a temporary function to find all usages of the old method. - It will be removed after the refactoring is complete. - """ - raise NotImplementedError( - "DEPRECATED: _as_evaled_expression has been replaced. " - "Use the new pattern: _evaled_expression(view.get(), target_type=...)" - ) - -# --- END OF TEMPORARY BLOCK --- \ No newline at end of file +EvaledFloat = Annotated[ + float, BeforeValidator(partial(_evaled_expression, target_type=float)) +] diff --git a/flepimop/gempyor_pkg/src/gempyor/utils copy.py b/flepimop/gempyor_pkg/src/gempyor/utils copy.py index 9a46d50d3..876acc712 100644 --- a/flepimop/gempyor_pkg/src/gempyor/utils copy.py +++ b/flepimop/gempyor_pkg/src/gempyor/utils copy.py @@ -538,7 +538,9 @@ def as_random_distribution(self): _evaled_expression(self["high"].get(), target_type=float), ) elif dist == "poisson": - return functools.partial(np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float)) + return functools.partial( + np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float) + ) elif dist == "binomial": p = _evaled_expression(self["p"].get(), target_type=float) if (p < 0) or (p > 1): diff --git a/flepimop/gempyor_pkg/src/gempyor/utils.py b/flepimop/gempyor_pkg/src/gempyor/utils.py index 9a46d50d3..423cd647c 100644 --- a/flepimop/gempyor_pkg/src/gempyor/utils.py +++ b/flepimop/gempyor_pkg/src/gempyor/utils.py @@ -538,7 +538,9 @@ def as_random_distribution(self): _evaled_expression(self["high"].get(), target_type=float), ) elif dist == "poisson": - return functools.partial(np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float)) + return functools.partial( + np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float) + ) elif dist == "binomial": p = _evaled_expression(self["p"].get(), target_type=float) if (p < 0) or (p > 1): @@ -560,8 +562,8 @@ def as_random_distribution(self): ).rvs elif dist == "lognorm": return get_log_normal( - meanlog=_evaled_expression(self["mean_log"].get(), target_type=float), - sdlog=_evaled_expression(self["sd_log"].get(), target_type=float), + meanlog=_evaled_expression(self["meanlog"].get(), target_type=float), + sdlog=_evaled_expression(self["sdlog"].get(), target_type=float), ).rvs else: raise NotImplementedError(f"Unknown distribution [received: '{dist}'].") @@ -569,8 +571,8 @@ def as_random_distribution(self): # we allow a fixed value specified directly: return functools.partial( np.random.uniform, - _evaled_expression(self, target_type=float), - _evaled_expression(self, target_type=float), + _evaled_expression(self.get(), target_type=float), + _evaled_expression(self.get(), target_type=float), ) diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py index 57ca86d3f..0a20de6d8 100644 --- a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py +++ b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py @@ -15,15 +15,15 @@ (-10.0, float, -10.0), ("2 * 5", int, 10), ("sqrt(81)", int, 9), - ("11 / 2", int, 5), - ("3.99", int, 3), + ("11 / 2", int, 5), + ("3.99", int, 3), (99, int, 99), (-50, int, -50), (None, float, None), (None, int, None), (True, float, True), (True, int, True), - (10.5, int, 10.5), + (10.5, int, 10.5), ([1, 2, 3], float, [1, 2, 3]), ({"a": 1}, float, {"a": 1}), ], @@ -45,4 +45,4 @@ def test_evaled_expression_returns_expected_results(val, target_type, expected): ) def test_evaled_expression_raises_value_error_for_invalid_input(val, target_type): with pytest.raises(ValueError): - _evaled_expression(val, target_type=target_type) \ No newline at end of file + _evaled_expression(val, target_type=target_type) diff --git a/flepimop/gempyor_pkg/tests/seir/test_parameters.py b/flepimop/gempyor_pkg/tests/seir/test_parameters.py index 7c674f34b..932e8b4c2 100644 --- a/flepimop/gempyor_pkg/tests/seir/test_parameters.py +++ b/flepimop/gempyor_pkg/tests/seir/test_parameters.py @@ -135,7 +135,9 @@ def test_parameters_quick_draw_old(): assert sigma.shape == (modinf.n_days, modinf.nsubpops) assert ( sigma - == _evaled_expression(config["seir"]["parameters"]["sigma"]["value"]["value"].get(), target_type=float) + == _evaled_expression( + config["seir"]["parameters"]["sigma"]["value"]["value"].get(), target_type=float + ) ).all() assert gamma.shape == (modinf.n_days, modinf.nsubpops) From 6e4e39c5716268355c4d233b17a86f58a104cc9c Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Fri, 25 Jul 2025 17:29:21 -0400 Subject: [PATCH 48/73] Remove accidental file duplicate --- .../gempyor_pkg/src/gempyor/utils copy.py | 1456 ----------------- 1 file changed, 1456 deletions(-) delete mode 100644 flepimop/gempyor_pkg/src/gempyor/utils copy.py diff --git a/flepimop/gempyor_pkg/src/gempyor/utils copy.py b/flepimop/gempyor_pkg/src/gempyor/utils copy.py deleted file mode 100644 index 876acc712..000000000 --- a/flepimop/gempyor_pkg/src/gempyor/utils copy.py +++ /dev/null @@ -1,1456 +0,0 @@ -""" -Helper functions for interacting with model I/O. -""" - -from collections import Counter -from collections.abc import Iterable -import datetime -import functools -import logging -import numbers -import os -from pathlib import Path -import random -from shlex import quote as shlex_quote -import shutil -import subprocess -import time -from typing import Any, Callable, Literal, overload - -import confuse -import numpy as np -import numpy.typing as npt -import pandas as pd -import pyarrow as pa -import scipy.ndimage -import scipy.stats -import sympy.parsing.sympy_parser -import yaml - -from . import file_paths -from ._pydantic_ext import _evaled_expression - - -logger = logging.getLogger(__name__) - -config = confuse.Configuration("flepiMoP", read=False) - - -def write_df( - fname: str | bytes | os.PathLike, - df: pd.DataFrame, - extension: Literal[None, "", "csv", "parquet"] = "", -) -> None: - """Writes a pandas DataFrame without its index to a file. - - Writes a pandas DataFrame to either a CSV or Parquet file without its index and can - infer which format to use based on the extension given in `fname` or based on - explicit `extension`. - - Args: - fname: The name of the file to write to. - df: A pandas DataFrame whose contents to write, but without its index. - extension: A user specified extension to use for the file if not contained in - `fname` already. - - Returns: - None - - Raises: - NotImplementedError: The given output extension is not supported yet. - """ - # Decipher the path given - fname = fname.decode() if isinstance(fname, bytes) else fname - path = Path(f"{fname}.{extension}") if extension else Path(fname) - # Write df to either a csv or parquet or raise if an invalid extension - if path.suffix == ".csv": - return df.to_csv(path, index=False) - elif path.suffix == ".parquet": - return df.to_parquet(path, index=False, engine="pyarrow") - raise NotImplementedError( - f"Invalid extension provided: '.{path.suffix[1:]}'. Supported extensions are `.csv` or `.parquet`." - ) - - -def read_df( - fname: str | bytes | os.PathLike, - extension: Literal[None, "", "csv", "parquet"] = "", -) -> pd.DataFrame: - """Reads a pandas DataFrame from either a CSV or Parquet file. - - Reads a pandas DataFrame to either a CSV or Parquet file and can infer which format - to use based on the extension given in `fname` or based on explicit `extension`. If - the file being read is a csv with a column called 'subpop' then that column will be - cast as a string. - - Args: - fname: The name of the file to read from. - extension: A user specified extension to use for the file if not contained in - `fname` already. - - Returns: - A pandas DataFrame parsed from the file given. - - Raises: - NotImplementedError: The given output extension is not supported yet. - """ - # Decipher the path given - fname = fname.decode() if isinstance(fname, bytes) else fname - path = Path(f"{fname}.{extension}") if extension else Path(fname) - # Read df from either a csv or parquet or raise if an invalid extension - if path.suffix == ".csv": - return pd.read_csv( - path, converters={"subpop": lambda x: str(x)}, skipinitialspace=True - ) - elif path.suffix == ".parquet": - return pd.read_parquet(path, engine="pyarrow") - raise NotImplementedError( - f"Invalid extension provided: '.{path.suffix[1:]}'. Supported extensions are `.csv` or `.parquet`." - ) - - -def command_safe_run( - command: str, command_name: str = "mycommand", fail_on_fail: bool = True -) -> tuple[int, str, str]: - """ - Runs a shell command and prints diagnostics if command fails. - - Args: - command: The CLI command to be given. - command_name: The reference name for you command. Default value is "mycommand". - fail_on_fail: If True, an exception will be thrown if the command fails (default is True) - - Returns: - As a tuple; the return code, the standard output, and standard error from running the command. - - Raises: - RuntimeError: If fail_on_fail=True and the command fails, an error will be thrown. - """ - import subprocess - import shlex # using shlex to split the command because it's not obvious https://docs.python.org/3/library/subprocess.html#subprocess.Popen - - sr = subprocess.Popen( - shlex.split(command), stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) - (stdout, stderr) = sr.communicate() - if sr.returncode != 0: - print(f"'{command_name}' failed failed with returncode '{sr.returncode}'") - print(f"'{command_name}': '{command}'") - print("{command_name} command failed with stdout and stderr:") - - print("{command_name} stdout >>>>>>") - print(stdout.decode()) - print("{command_name} stdout <<<<<<") - - print("{command_name} stderr >>>>>>") - print(stderr.decode()) - print("{command_name} stderr <<<<<<") - if fail_on_fail: - raise RuntimeError(f"The '{command_name}' command failed.") - - return sr.returncode, stdout, stderr - - -def add_method(cls: Any): - """ - A function which adds a function to a class. - - Args: - cls: The class you want to add a method to. - - Returns: - decorator: The decorator. - """ - - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - return func(*args, **kwargs) - - setattr(cls, func.__name__, wrapper) - return func - - return decorator - - -def search_and_import_plugins_class( - plugin_file_path: str, path_prefix: str, class_name: str, **kwargs: dict[str, Any] -) -> Any: - """ - Function serving to create a class that finds and imports the necessary modules. - - Args: - plugin_file_path: Pathway to the module. - path_prefix: Pathway prefix to the module. - class_name: Name of the class. - **kwargs: Further arguments passed to initilization of the class. - - Returns: - The instance of the class that was instantiated with provided **kwargs. - - Examples: - Suppose there is a module called `my_plugin.py with a class `MyClass` located at `/path/to/plugin/`. - - Dynamically import and instantiate the class: - - >>> instance = search_and_import_plugins_class('/path/to/plugin', path_prefix, 'MyClass', **params) - - View the instance: - - >>> print(instance) - <__main__.MyClass object at 0x7f8b2c6b4d60> - - """ - # Look for all possible plugins and import them - # https://stackoverflow.com/questions/67631/how-can-i-import-a-module-dynamically-given-the-full-path - # unfortunatelly very complicated, this is cpython only ?? - import sys, os - - full_path = os.path.join(path_prefix, plugin_file_path) - sys.path.append(os.path.dirname(full_path)) - # the following works, but these above lines seems necessary to pickle // runs - from pydoc import importfile - - module = importfile(full_path) - klass = getattr(module, class_name) - return klass(**kwargs) - - -### Profile configuration -import cProfile -import pstats -from functools import wraps - - -def profile( - output_file: str = None, - sort_by: str = "cumulative", - lines_to_print: int = None, - strip_dirs: bool = False, -): - """ - A time profiler decorator. - Inspired by and modified the profile decorator of Giampaolo Rodola: - http://code.activestate.com/recipes/577817-profile-decorator/ - - Args: - output_file: - Path of the output file. If only name of the file is given, it's - saved in the current directory. - If it's None, the name of the decorated function is used. - sort_by: - Sorting criteria for the Stats object. - For a list of valid string and SortKey refer to: - https://docs.python.org/3/library/profile.html#pstats.Stats.sort_stats - lines_to_print: - Number of lines to print. Default (None) is for all the lines. - This is useful in reducing the size of the printout, especially - that sorting by 'cumulative', the time consuming operations - are printed toward the top of the file. - strip_dirs: - Whether to remove the leading path info from file names. - - Returns: - Profile of the decorated function. - - Examples: - >>> @profile(output_file="my_function.prof") - >>> def my_function(): - # Function body content - pass - >>> my_function() - After running ``my_function``, a file named ``my_function.prof`` will be created in the current WD. - This file contains the profiling data. - """ - - def inner(func): - @wraps(func) - def wrapper(*args, **kwargs): - _output_file = output_file or func.__name__ + ".prof" - pr = cProfile.Profile() - pr.enable() - retval = func(*args, **kwargs) - pr.disable() - pr.dump_stats(_output_file) - return retval - - return wrapper - - return inner - - -def as_list(thing: Any) -> list[Any]: - """ - Returns argument passed as a list. - - Args: - thing: The object that you would like to be converted to a list. - - Returns: - The object converted to a list. - """ - if type(thing) == list: - return thing - return [thing] - - -### A little timer class -class Timer(object): - """ - A timer class that starts, ends, and records time in between. - - Attributes: - name: Name of event. - tstart: Time start. - """ - - def __init__(self, name): - self.name = name - - def __enter__(self): - logging.debug(f"[{self.name}] started") - self.tstart = time.time() - - def __exit__(self, type, value, traceback): - logging.debug(f"[{self.name}] completed in {time.time() - self.tstart:,.2f} s") - - -class ISO8601Date(confuse.Template): - """ - Reads in config dates into datetimes.dates. - """ - - def convert(self, value: any, view: confuse.ConfigView): - """ - Converts the given value to a datetime.date object. - - Args: - value: The value to be converted. Can be datetime.date object or ISO8601 string. - view: A view object from confuse, to be used for error reporting. - - Raises: - confuse.TemplateError: If `value` is neither a datetime.date nor an ISO8601Date string. - """ - if isinstance(value, datetime.date): - return value - elif isinstance(value, str): - return datetime.datetime.strptime(value, "%Y-%m-%d").date() - else: - self.fail("must be a date object or ISO8601 date", True) - - -@add_method(confuse.ConfigView) -def as_date(self) -> datetime.date: - """ - Evaluates an datetime.date or ISO8601 date string. - - Returns: - A datetime.date data type of the date associated with the object. - """ - - return self.get(ISO8601Date()) - - -def get_truncated_normal( - mean: float | int = 0, - sd: float | int = 1, - a: float | int = 0, - b: float | int = 10, -) -> scipy.stats._distn_infrastructure.rv_frozen: - """ - Returns a truncated normal distribution. - - This function constructs a truncated normal distribution with the specified - mean, standard deviation, and bounds. The truncated normal distribution is - a normal distribution bounded within the interval [a, b]. - - Args: - mean: The mean of the truncated normal distribution. Defaults to 0. - sd: The standard deviation of the truncated normal distribution. Defaults to 1. - a: The lower bound of the truncated normal distribution. Defaults to 0. - b: The upper bound of the truncated normal distribution. Defaults to 10. - - Returns: - rv_frozen: A frozen instance of the truncated normal distribution with the specified parameters. - - Examples: - Create a truncated normal distribution with specified parameters (truncated between 1 and 10): - >>> truncated_normal_dist = get_truncated_normal(mean=5, sd=2, a=1, b=10) - >>> print(truncated_normal_dist) - rv_frozen() - """ - lower = (a - mean) / sd - upper = (b - mean) / sd - return scipy.stats.truncnorm(lower, upper, loc=mean, scale=sd) - - -def get_log_normal( - meanlog: float | int, - sdlog: float | int, -) -> scipy.stats._distn_infrastructure.rv_frozen: - """ - Returns a log normal distribution. - - This function constructs a log normal distribution with the specified - log mean and log standard deviation. - - Args: - meanlog: The log of the mean of the log normal distribution. - sdlog: The log of the standard deviation of the log normal distribution. - - Returns: - rv_frozen: A frozen instance of the log normal distribution with the - specified parameters. - - Examples: - Create a log-normal distribution with specified parameters: - >>> log_normal_dist = get_log_normal(meanlog=1, sdlog=0.5) - >>> print(log_normal_dist) - - """ - return scipy.stats.lognorm(s=sdlog, scale=np.exp(meanlog), loc=0) - - -def random_distribution_sampler( - distribution: Literal[ - "fixed", "uniform", "poisson", "binomial", "truncnorm", "lognorm" - ], - **kwargs: dict[str, Any], -) -> Callable[[], float | int]: - """ - Create function to sample from a random distribution. - - Args: - distribution: The type of distribution to generate a sampling function for. - **kwargs: Further parameters that are passed to the underlying function for the - given distribution. - - Notes: - The further args expected by each distribution type are: - - fixed: value, - - uniform: low, high, - - poisson: lam, - - binomial: n, p, - - truncnorm: mean, sd, a, b, - - lognorm: meanlog, sdlog. - - Returns: - A function that can be called to sample from that distribution. - - Raises: - ValueError: If `distribution` is 'binomial' the given `p` must be in (0,1). - NotImplementedError: If `distribution` is not one of the type hinted options. - - Examples: - >>> import numpy as np - >>> np.random.seed(123) - >>> uniform_sampler = random_distribution_sampler("uniform", low=0.0, high=3.0) - >>> uniform_sampler() - 2.089407556793585 - >>> uniform_sampler() - 0.8584180048511384 - """ - if distribution == "fixed": - # Fixed value is the same as uniform on [a, a) - return functools.partial( - np.random.uniform, - kwargs.get("value"), - kwargs.get("value"), - ) - elif distribution == "uniform": - # Uniform on [low, high) - return functools.partial( - np.random.uniform, - kwargs.get("low"), - kwargs.get("high"), - ) - elif distribution == "poisson": - # Poisson with mean lambda - return functools.partial(np.random.poisson, kwargs.get("lam")) - elif distribution == "binomial": - p = kwargs.get("p") - if not (0 < p < 1): - raise ValueError(f"Invalid `p-value`: '{p}' is out of range [0,1].") - return functools.partial(np.random.binomial, kwargs.get("n"), p) - elif distribution == "truncnorm": - # Truncated normal with mean, sd on interval [a, b] - return get_truncated_normal( - mean=kwargs.get("mean"), - sd=kwargs.get("sd"), - a=kwargs.get("a"), - b=kwargs.get("b"), - ).rvs - elif distribution == "lognorm": - # Lognormal distribution with meanlog, sdlog - return get_log_normal(kwargs.get("meanlog"), kwargs.get("sdlog")).rvs - raise NotImplementedError(f"Unknown distribution [received: '{distribution}'].") - - -@add_method(confuse.ConfigView) -def as_random_distribution(self): - """ - Constructs a random distribution object from a distribution config key. - - Args: - self: Class instance (in this case, a config key) to construct the random distribution from. - - Returns: - A partial object containing the random distribution. - - Raises: - ValueError: When values are out of range. - NotImplementedError: If an unknown distribution is found. - - Examples: - Say that ``config`` is a ``confuse.ConfigView`` instance. - - To create a uniform distribution between 1 and 10: - >>> dist_function = config.as_random_distribution() - >>> sample = dist_function() - 5.436789235794546 - - To use a truncated normal distribution: - >>> config_truncnorm = confuse.ConfigView({ - "distribution": "truncnorm", - "mean": 0 - "sd": 1, - "a": -1, - "b": 1 - }) - >>> truncnorm_dist_function = config_truncnorm.as_random_distribution() - >>> truncnorm_sample = truncnorm_dist_function() - 0.312745 - ``` - """ - - if isinstance(self.get(), dict): - dist = self["distribution"].get() - if dist == "fixed": - return functools.partial( - np.random.uniform, - _evaled_expression(self["value"].get(), target_type=float), - _evaled_expression(self["value"].get(), target_type=float), - ) - elif dist == "uniform": - return functools.partial( - np.random.uniform, - _evaled_expression(self["low"].get(), target_type=float), - _evaled_expression(self["high"].get(), target_type=float), - ) - elif dist == "poisson": - return functools.partial( - np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float) - ) - elif dist == "binomial": - p = _evaled_expression(self["p"].get(), target_type=float) - if (p < 0) or (p > 1): - raise ValueError(f"Invalid `p-value`: '{p}' is out of range [0,1].") - # if (self["p"] < 0) or (self["p"] > 1): - # raise ValueError(f"""p value { self["p"] } is out of range [0,1]""") - return functools.partial( - np.random.binomial, - _evaled_expression(self["n"].get(), target_type=float), - # _evaled_expression(self["p"].get(), target_type=float), - p, - ) - elif dist == "truncnorm": - return get_truncated_normal( - mean=_evaled_expression(self["mean"].get(), target_type=float), - sd=_evaled_expression(self["sd"].get(), target_type=float), - a=_evaled_expression(self["a"].get(), target_type=float), - b=_evaled_expression(self["b"].get(), target_type=float), - ).rvs - elif dist == "lognorm": - return get_log_normal( - meanlog=_evaled_expression(self["mean_log"].get(), target_type=float), - sdlog=_evaled_expression(self["sd_log"].get(), target_type=float), - ).rvs - else: - raise NotImplementedError(f"Unknown distribution [received: '{dist}'].") - else: - # we allow a fixed value specified directly: - return functools.partial( - np.random.uniform, - _evaled_expression(self, target_type=float), - _evaled_expression(self, target_type=float), - ) - - -def list_filenames( - folder: str | bytes | os.PathLike = ".", - filters: str | list[str] = [], -) -> list[str]: - """ - Return the list of all filenames and paths in the provided folder. - - This function lists all files in the specified folder and its subdirectories. - If filters are provided, only the files containing each of the substrings - in the filters will be returned. - - Args: - folder: - The directory to search for files. Defaults to the current directory. - filters: - A string or a list of strings to filter filenames. Only files - containing all the provided substrings will be returned. Defaults to an - empty list. - - Returns: - A list of strings representing the paths to the files that match the filters. - - Examples: - To get all files containing "hosp": - >>> gempyor.utils.list_filenames( - folder="model_output/", - filters=["hosp"], - ) - - To get only "hosp" files with a ".parquet" extension: - >>> gempyor.utils.list_filenames( - folder="model_output/", - filters=["hosp", ".parquet"], - ) - """ - filters = [filters] if not isinstance(filters, list) else filters - filters = filters if len(filters) else [""] - folder = Path(folder.decode() if isinstance(folder, bytes) else folder) - files = [ - str(file) - for file in folder.rglob("*") - if file.is_file() and all(f in str(file) for f in filters) - ] - return files - - -def extract_slot(file: Path) -> pd.DataFrame: - """ - Extract the slot number from the filename and add to DataFrame. - - Args: - file: The filename to extract the slot number from. - - Returns: - A pandas DataFrame containing the contents of `file` along with a column called - 'slot' that contains the slot number extracted from the filename. - - Raises: - ValueError: If the 'slot' column already exists in the DataFrame read in. - - Examples: - >>> from pathlib import Path - >>> import pandas as pd - >>> from gempyor.utils import extract_slot - >>> sample = pd.DataFrame(data={"letters": ["a", "b", "c"]}) - >>> file = Path.cwd() / "00017.foobar.csv" - >>> sample.to_csv(file, index=False) - >>> extract_slot(file) - letters slot - 0 a 17 - 1 b 17 - 2 c 17 - >>> extract_slot(file).info() - - RangeIndex: 3 entries, 0 to 2 - Data columns (total 2 columns): - # Column Non-Null Count Dtype - --- ------ -------------- ----- - 0 letters 3 non-null object - 1 slot 3 non-null int64 - dtypes: int64(1), object(1) - memory usage: 180.0+ bytes - """ - df = pd.read_parquet(file) if file.suffix == ".parquet" else pd.read_csv(file) - if "slot" in df.columns: - raise ValueError( - f"Column 'slot' already exists in the DataFrame from file: {file}." - ) - df["slot"] = int(file.name.split(".")[0]) - return df - - -def read_directory( - directory: str | bytes | os.PathLike, filters: str | list[str] | None = None -) -> pd.DataFrame: - """ - Read all files in a directory into a single DataFrame. - - Args: - directory: The directory to read files from. - filters: A string or a list of strings to filter filenames. Only files - containing all the provided substrings will be read. Defaults to `None` for - no filters. - - Returns: - A pandas DataFrame containing the contents of all the files in the directory. - """ - files = [Path(file) for file in list_filenames(folder=directory, filters=filters or [])] - dfs: list[pd.DataFrame] | pd.DataFrame = [] - for file in sorted(files): - dfs.append(extract_slot(file)) - return pd.concat(dfs).reset_index(drop=True) - - -def rolling_mean_pad( - data: npt.NDArray[np.number], - window: int, -) -> npt.NDArray[np.number]: - """ - Calculates the column-wise rolling mean with centered window. - - Args: - data: A two dimensional numpy array, typically the row dimension is time and the - column dimension is subpop. - window: The window size for the rolling mean. - - Returns: - A two dimensional numpy array that is the same shape as `data`. - - Examples: - Below is a brief set of examples showcasing how to smooth a metric, like - hospitalizations, using this function. - ``` - >>> import numpy as np - >>> from gempyor.utils import rolling_mean_pad - >>> hospitalizations = np.arange(1., 29.).reshape((7, 4)) - >>> hospitalizations - array([[ 1., 2., 3., 4.], - [ 5., 6., 7., 8.], - [ 9., 10., 11., 12.], - [13., 14., 15., 16.], - [17., 18., 19., 20.], - [21., 22., 23., 24.], - [25., 26., 27., 28.]]) - >>> rolling_mean_pad(hospitalizations, 5) - array([[ 3.4, 4.4, 5.4, 6.4], - [ 5.8, 6.8, 7.8, 8.8], - [ 9. , 10. , 11. , 12. ], - [13. , 14. , 15. , 16. ], - [17. , 18. , 19. , 20. ], - [20.2, 21.2, 22.2, 23.2], - [22.6, 23.6, 24.6, 25.6]]) - ``` - """ - weights = (1.0 / window) * np.ones(window) - output = scipy.ndimage.convolve1d(data, weights, axis=0, mode="nearest") - if window % 2 == 0: - rows, cols = data.shape - i = rows - 1 - output[i, :] = 0.0 - window -= 1 - weight = 1.0 / window - for l in range(-((window - 1) // 2), 1 + (window // 2)): - i_star = min(max(i + l, 0), i) - for j in range(cols): - output[i, j] += weight * data[i_star, j] - return output - - -def print_disk_diagnosis(): - """ - Reads and prints AWS disk diagnostic information. - """ - import os - from os import path - from shutil import disk_usage - - def bash(command: str) -> str: - """ - Executes a shell command and returns its output. - - Args: - command: The shell command to be executed. - - Returns: - The output of the shell command. - """ - output = os.popen(command).read() - return output - - print("START AWS DIAGNOSIS ================================") - total_bytes, used_bytes, free_bytes = disk_usage(path.realpath("/")) - print( - f"shutil.disk_usage: {total_bytes/ 1000000} Mb total, {used_bytes / 1000000} Mb used, {free_bytes / 1000000} Mb free..." - ) - print("------------") - print(f"df -hT: '{bash('df -hT')}'") - print("------------") - print(f"df -i: '{bash('df -i')}'") - print("------------") - print(f"free -h: '{bash('free -h')}'") - print("------------") - print(f"lsblk: '{bash('lsblk')}'") - print("END AWS DIAGNOSIS ================================") - - -def create_resume_out_filename( - flepi_run_index: str, - flepi_prefix: str, - flepi_slot_index: str, - flepi_block_index: str, - filetype: str, - liketype: str, -) -> str: - """ - Compiles run output information. - - Args: - flepi_run_index: Index of the run. - flepi_prefix: File prefix. - flepi_slot_index: Index of the slot. - flepi_block_index: Index of the block. - filetype: File type. - liketype: Chimeric or global. - - Returns: - The path to a corresponding output file. - - Examples: - Generate an output file with specified parameters: - >>> filename = create_resume_out_filename( - flepi_run_index="test_run", - flepi_prefix="model_output/run_id/", - flepi_slot_index="1", - flepi_block_index="2", - filetype="seed", - liketype="chimeric" - ) - >>> print(filename) - "experiment/001/normal/intermediate/000000123.000000000.1.parquet" - """ - prefix = f"{flepi_prefix}/{flepi_run_index}" - inference_filepath_suffix = f"{liketype}/intermediate" - inference_filename_prefix = "{:09d}.".format(int(flepi_slot_index)) - index = "{:09d}.{:09d}".format(1, int(flepi_block_index) - 1) - extension = "parquet" - if filetype == "seed": - extension = "csv" - return file_paths.create_file_name( - run_id=flepi_run_index, - prefix=prefix, - inference_filename_prefix=inference_filename_prefix, - inference_filepath_suffix=inference_filepath_suffix, - index=index, - ftype=filetype, - extension=extension, - ) - - -def create_resume_input_filename( - resume_run_index: str, - flepi_prefix: str, - flepi_slot_index: str, - filetype: str, - liketype: str, -) -> str: - """ - Compiles run input information. - - Args: - resume_run_index: Index of the run. - flepi_prefix: File prefix. - flepi_slot_index: Index of the slot. - filetype: File type. - liketype: Chimeric or global. - - Returns: - The path to the a corresponding input file. - - Examples: - Generate an input file with specified parameters: - >>> filename = create_resume_input_filename( - resume_run_index="2", - flepi_prefix="model_output/run_id/", - flepi_slot_index="1", - filetype="seed", - liketype="chimeric" - ) - >>> print(filename) - "experiment/002/normal/final/789.csv" - """ - prefix = f"{flepi_prefix}/{resume_run_index}" - inference_filepath_suffix = f"{liketype}/final" - index = flepi_slot_index - extension = "parquet" - if filetype == "seed": - extension = "csv" - return file_paths.create_file_name( - run_id=resume_run_index, - prefix=prefix, - inference_filepath_suffix=inference_filepath_suffix, - index=index, - ftype=filetype, - extension=extension, - ) - - -def get_filetype_for_resume( - resume_discard_seeding: str, flepi_block_index: str -) -> list[str]: - """ - Retrieves a list of parquet file types that are relevant for resuming a process based on - specific environment variable settings. - This function dynamically determines the list - based on the current operational context given by the environment. - - Args: - resume_discard_seeding: Determines whether seeding-related file types should be included (str). - flepi_block_index: Determines a specific operational mode or block of the process (str). - - Returns: - List of file types. - - Examples: - Determine file types for block index 1 with seeding data NOT discarded: - >>> filetypes = get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="1") - >>> print(filetypes) - ["seed", "spar", "snpi", "hpar", "hnpi", "init"] - - Determine file types for block index 2 with seeding data discarded: - >>> filtypes = get_filetype_for_resume(resume_discard_seeding="true", flepi_block_index="2") - >>> print(filetypes) - ["seed", "spar", "snpi", "hpar", "hnpi", "host", "llik", "init"] - """ - if flepi_block_index == "1": - if resume_discard_seeding == "true": - return ["spar", "snpi", "hpar", "hnpi", "init"] - else: - return ["seed", "spar", "snpi", "hpar", "hnpi", "init"] - else: - return ["seed", "spar", "snpi", "hpar", "hnpi", "host", "llik", "init"] - - -def create_resume_file_names_map( - resume_discard_seeding: str, - flepi_block_index: str, - resume_run_index: str, - flepi_prefix: str, - flepi_slot_index: str, - flepi_run_index: str, - last_job_output: str, -) -> dict[str, str]: - """ - Generates a mapping of input file names to output file names for a resume process based on - parquet file types and environmental conditions. The function adjusts the file name mappings - based on the operational block index and the location of the last job output. - - Args: - resume_discard_seeding: Determines whether seeding-related file types should be included. - flepi_block_index: Determines a specific operational mode or block of the process. - resume_run_index: Resume run index. - flepi_prefix: File prefix. - flepi_slot_index: Index of the slot. - flepi_run_index: flepiMoP run index. - last_job_output: Adjusts the keys in the mapping to be prefixed with this path. - - Returns: - A dictionary where keys are input file paths and values are corresponding - output file paths. - - The mappings depend on: - - Parquet file types appropriate for resuming a process, as determined by the environment. - - Whether the files are for 'global' or 'chimeric' types, as these liketypes influence the - file naming convention. - - The operational block index ('FLEPI_BLOCK_INDEX'), which can alter the input file names for - block index '1'. - - The presence and value of 'LAST_JOB_OUTPUT' environment variable, which if set to an S3 path, - adjusts the keys in the mapping to be prefixed with this path. - - Raises: - No explicit exceptions are raised within the function, but it relies heavily on external - functions and environment variables which if improperly configured could lead to unexpected - behavior. - - Examples: - Generate a mapping of file names for a given resume process: - >>> file_names_map = create_resume_file_names_map( - resume_discard_seeding="false", - flepi_block_index="1", - resume_run_index="1", - flepi_prefix="model_output/run_id/", - flepi_slot_index="1", - flepi_run_index="test_run", - last_job_output="s3://bucket/path/") - >>> print(file_names_map) - { - 's3://bucket/path/model_output/run_id/1_type1_global_1.in': 'model_output/run_id/test_run_type1_global_1_1.out', - 's3://bucket/path/model_output/run_id/1_type1_chimeric_1.in': 'model_output/run_id/test_run_type1_chimeric_1_1.out', - 's3://bucket/path/model_output/run_id/1_type2_global_1.in': 'model_output/run_id/test_run_type2_global_1_1.out', - 's3://bucket/path/model_output/run_id/1_type2_chimeric_1.in': 'model_output/run_id/test_run_type2_chimeric_1_1.out' - } - # Note: this output is toy output implemented with toy file names. - - Notes: - - The paths may be modified by the 'LAST_JOB_OUTPUT' if it is set and points to an S3 location. - """ - file_types = get_filetype_for_resume( - resume_discard_seeding=resume_discard_seeding, flepi_block_index=flepi_block_index - ) - resume_file_name_mapping = dict() - liketypes = ["global", "chimeric"] - for filetype in file_types: - for liketype in liketypes: - output_file_name = create_resume_out_filename( - flepi_run_index=flepi_run_index, - flepi_prefix=flepi_prefix, - flepi_slot_index=flepi_slot_index, - flepi_block_index=flepi_block_index, - filetype=filetype, - liketype=liketype, - ) - input_file_name = output_file_name - if flepi_block_index == "1": - input_file_name = create_resume_input_filename( - resume_run_index=resume_run_index, - flepi_prefix=flepi_prefix, - flepi_slot_index=flepi_slot_index, - filetype=filetype, - liketype=liketype, - ) - resume_file_name_mapping[input_file_name] = output_file_name - if last_job_output.find("s3://") >= 0: - old_keys = list(resume_file_name_mapping.keys()) - for k in old_keys: - new_key = os.path.join(last_job_output, k) - resume_file_name_mapping[new_key] = resume_file_name_mapping[k] - del resume_file_name_mapping[k] - return resume_file_name_mapping - - -def download_file_from_s3(name_map: dict[str, str]) -> None: - """ - Downloads files from AWS S3 based on a mapping of S3 URIs to local file paths. The function - checks if the directory for the first output file exists and creates it if necessary. It - then iterates over each S3 URI in the provided mapping, downloads the file to the corresponding - local path, and handles errors if the S3 URI format is incorrect or if the download fails. - - Args: - name_map: - A dictionary where keys are S3 URIs (strings) and values - are the local file paths (strings) where the files should - be saved. - - Returns: - This function does not return a value; its primary effect is the side effect of - downloading files and potentially creating directories. - - Raises: - ValueError: If an S3 URI does not start with 's3://', indicating an invalid format. - ClientError: If an error occurs during the download from S3, such as a permissions issue, - a missing file, or network-related errors. These are caught and logged but not - re-raised, to allow the function to attempt subsequent downloads. - - Examples: - >>> name_map = { - "s3://mybucket/data/file1.txt": "/local/path/to/file1.txt", - "s3://mybucket/data/file2.txt": "/local/path/to/file2.txt" - } - >>> download_file_from_s3(name_map) - # This would download 'file1.txt' and 'file2.txt' from 'mybucket' on S3 to the specified local paths. - - # If an S3 URI is malformed: - >>> name_map = { - "http://wrongurl.com/data/file1.txt": "/local/path/to/file1.txt" - } - >>> download_file_from_s3(name_map) - # This will raise a ValueError indicating the invalid S3 URI format. - """ - try: - import boto3 - from botocore.exceptions import ClientError - except ModuleNotFoundError: - raise ModuleNotFoundError( - ( - "No module named `boto3` found, which is required for " - "`gempyor.utils.download_file_from_s3`. Please install the aws target." - ) - ) - s3 = boto3.client("s3") - first_output_filename = next(iter(name_map.values())) - output_dir = os.path.dirname(first_output_filename) - if not os.path.exists(output_dir): - os.makedirs(output_dir) - for s3_uri in name_map: - try: - if s3_uri.startswith("s3://"): - bucket = s3_uri.split("/")[2] - object = s3_uri[len(bucket) + 6 :] - s3.download_file(bucket, object, name_map[s3_uri]) - else: - raise ValueError(f"Invalid S3 URI format [received: '{s3_uri}'].") - except ClientError as e: - raise Exception(f"'{e}': could not download filefrom S3.") - - -def move_file_at_local(name_map: dict[str, str]) -> None: - """ - Moves files locally according to a given mapping. - This function takes a dictionary where the keys are source file paths and - the values are destination file paths. It ensures that the destination - directories exist and then copies the files from the source paths to the - destination paths. - - Args: - name_map: A dictionary mapping source file paths to destination file paths. - """ - for src, dst in name_map.items(): - os.path.makedirs(os.path.dirname(dst), exist_ok=True) - shutil.copy(src, dst) - - -def _dump_formatted_yaml(cfg: confuse.Configuration) -> str: - """ - Dump confuse configuration to a formatted YAML string. - - Args: - cfg: The confuse configuration object. - - Returns: - A formatted YAML string representation of the configuration. - - Examples: - >>> from gempyor.utils import _dump_formatted_yaml - >>> import confuse - >>> conf = confuse.Configuration("foobar") - >>> data = { - ... "name": "Test Config", - ... "compartments": { - ... "infection_stage": ["S", "E", "I", "R"] - ... }, - ... "seir": { - ... "parameters": { - ... "beta": {"value": 3.4}, - ... "gamma": {"value": 5.6}, - ... }, - ... "transitions": { - ... "source": ["S"], - ... "destination": ["E"], - ... "rate": ["beta * gamma"], - ... "proportional_to": [["S"], ["I"]], - ... "proportion_exponent": [1, 1], - ... }, - ... }, - ... } - >>> conf.set(data) - >>> print(_dump_formatted_yaml(conf)) - name: "Test Config" - compartments: - infection_stage: [S, E, I, R] - seir: - parameters: - beta: - value: 3.4 - gamma: - value: 5.6 - transitions: - source: [S] - destination: [E] - rate: ["beta * gamma"] - proportional_to: [[S], [I]] - proportion_exponent: [1, 1] - """ - - class CustomDumper(yaml.Dumper): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.add_representer(list, self._represent_list) - self.add_representer(str, self._represent_str) - - def _represent_list(self, dumper, data): - return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True) - - def _represent_str(self, dumper, data): - if " " in data: - return dumper.represent_scalar("tag:yaml.org,2002:str", data, style='"') - return dumper.represent_scalar("tag:yaml.org,2002:str", data) - - return yaml.dump( - yaml.safe_load(cfg.dump()), Dumper=CustomDumper, indent=4, sort_keys=False - ) - - -def _shutil_which( - cmd: str, - mode: int = os.F_OK | os.X_OK, - path: str | bytes | os.PathLike | None = None, - check: bool = True, -) -> str | None: - """ - A thin wrapper around `shutil.which` with extra validation. - - Args: - cmd: The name of the command to search for. - mode: The permission mask required of possible files. - path: A path describing the locations to search, or `None` to use - the PATH environment variable. - check: If `True` an `OSError` will be raised if a `cmd` is not found. - Similar in spirit to the `check` arg of `subprocess.run`. - - Returns: - Either the full path to the `cmd` found, or `None` if `cmd` is not - found and `check` is `False`. - - Raises: - OSError: If `cmd` is not found and `check` is `True`. - - Examples: - >>> import os - >>> from shutil import which - >>> _shutil_which("python") == which("python") - True - >>> _shutil_which("does_not_exist", check=False) - >>> try: - ... _shutil_which("does_not_exist", check=True) - ... except Exception as e: - ... print(type(e)) - ... print(str(e).replace(os.environ.get("PATH"), "...")) - ... - - Did not find 'does_not_exist' on path '...'. - - See Also: - [`shutil.which`](https://docs.python.org/3/library/shutil.html#shutil.which) - """ - result = shutil.which(cmd, mode=mode, path=path) - if check and result is None: - path = os.environ.get("PATH") if path is None else path - raise OSError(f"Did not find '{cmd}' on path '{path}'.") - return result - - -def _git_head(repository: Path) -> str: - """ - Get the sha commit hash for the head of a git repository. - - Args: - repository: A directory under version control with git to get the sha commit of. - - Returns: - The sha commit of head for `repository`. - - Examples: - >>> import os - >>> from pathlib import Path - >>> _git_head(Path(os.environ["FLEPI_PATH"])) - 'efe896b1a5e4f8e33667c170cd5319d6ef1e3db5' - """ - git_cmd = _shutil_which("git") - proc = subprocess.run( - [git_cmd, "rev-parse", "HEAD"], - cwd=repository.expanduser().absolute(), - capture_output=True, - check=True, - ) - return proc.stdout.decode().strip() - - -def _git_checkout(repository: Path, branch: str) -> subprocess.CompletedProcess: - """ - Checkout a new branch from a given git repository. - - Args: - repository: A directory under version control with git to - checkout a new branch in. - branch: The name of the new branch to checkout. - - Examples: - >>> import os - >>> from pathlib import Path - >>> _git_checkout(Path(os.environ["FLEPI_PATH"]), "my-new-branch") - """ - git_cmd = _shutil_which("git") - return subprocess.run( - [git_cmd, "checkout", "-b", shlex_quote(branch)], - cwd=repository.expanduser().absolute(), - check=True, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - - -def _format_cli_options( - options: dict[str, str | Iterable[str]] | None, - always_single: bool = False, - iterable_formatting: Literal["split", "comma"] = "split", -) -> list[str]: - """ - Convert a dictionary of CLI options into a formatted list. - - Args: - options: A dictionary where the keys correspond to the option name and the - values correspond to the option value. Values are escaped for shell. - always_single: If `True` all options will be formatted with a single dash prefix - always, useful for commands that don't support long options like - `pdftotext`. If `False` the option will be formatted with a double dash - prefix if the key is longer than one character. - iterable_formatting: The formatting to use for iterable values. If 'split' each - value will be formatted as a separate option. If 'comma' the values will be - joined with a comma and formatted as a single option. - - Returns: - A list of options that can be passed to - [`subprocess.run`](https://docs.python.org/3/library/subprocess.html#subprocess.run) - or similar functions. - - Examples: - >>> from gempyor.utils import _format_cli_options - >>> _format_cli_options({"name": "foo bar fizz buzz"}) - ["--name='foo bar fizz buzz'"] - >>> _format_cli_options({"o": "/path/to/output.log"}) - ['-o=/path/to/output.log'] - >>> _format_cli_options({"opt1": "```", "opt2": "$( echo 'Hello!')"}) - ["--opt1='```'", '--opt2=\'$( echo \'"\'"\'Hello!\'"\'"\')\''] - >>> _format_cli_options({"output": "/path/to/output.log"}, always_single=True) - ['-output=/path/to/output.log'] - >>> _format_cli_options({"person": ["Alice", "Bob", "Charlie"]}) - ['--person=Alice', '--person=Bob', '--person=Charlie'] - >>> _format_cli_options( - ... {"person": ["Alice", "Bob", "Charlie"]}, - ... iterable_formatting="comma", - ... ) - ['--person=Alice,Bob,Charlie'] - """ - opts = [] - for k, v in (options or {}).items(): - new_opts = [] - opt_name = f"{'-' if (always_single or len(k) == 1) else '--'}{k}" - if isinstance(v, str): - new_opts.append(f"{opt_name}={shlex_quote(v)}") - elif iterable_formatting == "comma": - new_opts.append(f"{opt_name}={','.join(shlex_quote(w) for w in v)}") - else: - new_opts.extend(f"{opt_name}={shlex_quote(w)}" for w in v) - opts.extend(new_opts) - return opts - - -def _random_seeds_list(a: int, b: int, n: int) -> list[int]: - """ - Generate a list of unique random integers within a specified range. - - Args: - a: The lower bound of the range, inclusive. - b: The upper bound of the range, inclusive. - n: The number of unique random integers to generate. - - Returns: - A list of `n` unique integers between `a` and `b`. - - Raises: - ValueError: If the lower bound is not less than the upper bound. - ValueError: If the range is too small to accommodate `n` unique integers. - - Examples: - >>> import random - >>> from gempyor.utils import _random_seeds_list - >>> random.seed(42) - >>> _random_seeds_list(1, 100, 5) - [82, 15, 4, 95, 36] - >>> _random_seeds_list(1, 100, 5) - [32, 29, 18, 95, 14] - >>> _random_seeds_list(1, 5, 5) - [5, 1, 4, 2, 3] - >>> _random_seeds_list(1, 20, 0) - [] - >>> try: - ... _random_seeds_list(1, 5, 6) - ... except Exception as e: - ... print(e) - ... - Range [1, 5] is too small for 6 unique random integers. - """ - random_seeds = [] - if a >= b: - raise ValueError(f"The lower bound, {a}, must be less than the upper bound, {b}.") - if n < 0: - raise ValueError("The number of random integers to generate must be non-negative.") - if (b - a) + 1 < n: - raise ValueError(f"Range [{a}, {b}] is too small for {n} unique random integers.") - while len(random_seeds) < n: - seed = random.randint(a, b) - if seed not in random_seeds: - random_seeds.append(seed) - return random_seeds - - -def _nslots_random_seeds(nslots: int) -> list[int]: - """ - Generate a list of unique random integers for the number of slots. - - This function generates a list of `nslots` unique random integers between - [`nslots` + 1, 2^32 + 1]. Intended for use in generating random seeds for individual - slots when simulations are done in parallel. - - Args: - nslots: The number of slots to generate random integers for. - - Examples: - >>> import random - >>> from gempyor.utils import _nslots_random_seeds - >>> random.seed(42) - >>> _nslots_random_seeds(5) - [2746317219, 478163333, 107420375, 3184935169, 1181241949] - >>> _nslots_random_seeds(5) - [1051802518, 958682852, 599310831, 3163119791, 440213421] - >>> _nslots_random_seeds(1) - [2906402159] - >>> _nslots_random_seeds(1) - [3181143733] - """ - return _random_seeds_list(nslots + 1, 2**32 - 1, nslots) - - -def _duplicate_strings(it: Iterable[str]) -> list[str]: - """ - Efficiently find duplicate strings in an iterable. - - Args: - it: An iterable of strings. - - Returns: - A list of strings that appear more than once in `it`. - - Examples: - >>> from gempyor.utils import _duplicate_strings - >>> _duplicate_strings(["a", "b", "c"]) - [] - >>> _duplicate_strings(["a", "b", "c", "b"]) - ['b'] - >>> _duplicate_strings(["a", "b", "c", "b", "a"]) - ['a', 'b'] - >>> _duplicate_strings("foobar") - ['o'] - """ - return sorted([k for k, v in Counter(it).items() if v > 1]) - - -@overload -def _trim_s3_path(path: str) -> str: ... - - -@overload -def _trim_s3_path(path: Path) -> Path: ... - - -def _trim_s3_path(path: str | Path) -> str | Path: - """ - Strips the 's3:' prefix from a given path. - - Args: - path: The path to be trimmed. Can be a string or a Path object. - - Returns: - The trimmed path as the same type as the input. - - Examples: - >>> from pathlib import Path - >>> from gempyor.utils import _trim_s3_path - >>> _trim_s3_path("/abc/def/ghi.txt") - '/abc/def/ghi.txt' - >>> _trim_s3_path(Path("/abc/def/ghi.txt")) - PosixPath('/abc/def/ghi.txt') - >>> _trim_s3_path("s3://foo/bar.txt") - '//foo/bar.txt' - >>> _trim_s3_path(Path("s3://foo/bar.txt")) - PosixPath('s3:/foo/bar.txt') - """ - return path.lstrip("s3:") if isinstance(path, str) else path From f33aa8b988b983ad930b8fff993c51895f2de4ba Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 28 Jul 2025 16:26:20 -0400 Subject: [PATCH 49/73] Refactoring all uses of `.as_random_distribution` to `build_distribution_from_confuse_config()` This code is presently broken. --- .../src/gempyor/NPI/MultiPeriodModifier.py | 15 ++++++++------- .../src/gempyor/NPI/SinglePeriodModifier.py | 9 +++++---- .../gempyor_pkg/src/gempyor/distributions.py | 7 ++++++- .../src/gempyor/inference_parameter.py | 17 +++++++++++++---- flepimop/gempyor_pkg/src/gempyor/outcomes.py | 11 ++++++++--- flepimop/gempyor_pkg/src/gempyor/parameters.py | 6 +----- 6 files changed, 41 insertions(+), 24 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py index 82c227336..0f4013070 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py @@ -5,6 +5,7 @@ from . import helpers from .base import NPIBase +from ..distributions import build_distribution_from_confuse_config class MultiPeriodModifier(NPIBase): @@ -149,7 +150,7 @@ def __createFromConfig(self, npi_config): self.affected_subpops = self.__get_affected_subpops(npi_config) self.parameters = self.parameters[self.parameters.index.isin(self.affected_subpops)] - dist = npi_config["value"].as_random_distribution() + dist = build_distribution_from_confuse_config(npi_config["value"]) self.parameters["modifier_name"] = self.name self.parameters["parameter"] = self.param_name @@ -177,9 +178,9 @@ def __createFromConfig(self, npi_config): for subpop in this_spatial_group["ungrouped"]: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates - self.parameters.at[subpop, "value"] = dist(size=1) + self.parameters.at[subpop, "value"] = dist() for group in this_spatial_group["grouped"]: - drawn_value = dist(size=1) + drawn_value = dist() for subpop in group: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates @@ -235,8 +236,8 @@ def __createFromDf(self, loaded_df, npi_config): if not subpop in loaded_df.index: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates - dist = npi_config["value"].as_random_distribution() - self.parameters.at[subpop, "value"] = dist(size=1) + dist = build_distribution_from_confuse_config(npi_config["value"]) + self.parameters.at[subpop, "value"] = dist() else: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates @@ -250,8 +251,8 @@ def __createFromDf(self, loaded_df, npi_config): ",".join(group), "value" ] else: - dist = npi_config["value"].as_random_distribution() - drawn_value = dist(size=1) + dist = build_distribution_from_confuse_config(npi_config["value"]) + drawn_value = dist() for subpop in group: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py index 7aece6306..81680422a 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py @@ -5,6 +5,7 @@ from . import helpers from .base import NPIBase +from ..distributions import build_distribution_from_confuse_config class SinglePeriodModifier(NPIBase): @@ -134,7 +135,7 @@ def __createFromConfig(self, npi_config): self.parameters = self.parameters[self.parameters.index.isin(self.affected_subpops)] # Create reduction - self.dist = npi_config["value"].as_random_distribution() + self.dist = build_distribution_from_confuse_config(npi_config["value"]) self.parameters["modifier_name"] = self.name self.parameters["start_date"] = ( @@ -152,12 +153,12 @@ def __createFromConfig(self, npi_config): npi_config, list(self.affected_subpops) ) if self.spatial_groups["ungrouped"]: - self.parameters.loc[self.spatial_groups["ungrouped"], "value"] = self.dist( - size=len(self.spatial_groups["ungrouped"]) + self.parameters.loc[self.spatial_groups["ungrouped"], "value"] = ( + self.dist.sample(size=len(self.spatial_groups["ungrouped"])) ) if self.spatial_groups["grouped"]: for group in self.spatial_groups["grouped"]: - drawn_value = self.dist(size=1) * np.ones(len(group)) + drawn_value = self.dist() * np.ones(len(group)) self.parameters.loc[group, "value"] = drawn_value def __createFromDf(self, loaded_df, npi_config): diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 5b846df50..4ba9d2d16 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -54,6 +54,10 @@ def sample( """ rng = rng if rng is not None else self._rng return self._sample_from_generator(size=size, rng=rng) + + def __call__(self) -> float | int: + """A shortcut for `self.sample(size=1)`.""" + return self.sample(size=1).item() @abstractmethod def _sample_from_generator( @@ -498,7 +502,8 @@ def build_distribution_from_confuse_config( Returns: A Distribution object. """ - conf = param_config["value"].get() + # previously: conf = param_config["value"].get() + conf = param_config.get() if isinstance(conf, float | int | str): conf = {"distribution": "fixed", "value": conf} diff --git a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py index b4ee24274..b5707103a 100644 --- a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py +++ b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py @@ -7,6 +7,7 @@ import numpy as np import confuse from . import NPI +from .distributions import build_distribution_from_confuse_config # TODO cast uper and lower bound as arrays @@ -58,7 +59,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=sp, - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) @@ -70,7 +73,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=",".join(group), - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) @@ -100,7 +105,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=sp, - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) @@ -112,7 +119,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=",".join(group), - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) diff --git a/flepimop/gempyor_pkg/src/gempyor/outcomes.py b/flepimop/gempyor_pkg/src/gempyor/outcomes.py index 4ff727d34..357488b23 100644 --- a/flepimop/gempyor_pkg/src/gempyor/outcomes.py +++ b/flepimop/gempyor_pkg/src/gempyor/outcomes.py @@ -17,6 +17,7 @@ from .utils import Timer, _nslots_random_seeds, config, read_df from . import NPI, model_info +from .distributions import build_distribution_from_confuse_config logger = logging.getLogger(__name__) @@ -466,13 +467,15 @@ def compute_all_multioutcomes( else: # One draw for all subpops probabilities = np.repeat( - parameters[new_comp]["probability"].as_random_distribution()(), + build_distribution_from_confuse_config( + parameters[new_comp]["probability"] + )(), len(modinf.subpop_struct.subpop_names), ) if "rel_probability" in parameters[new_comp]: probabilities = probabilities * parameters[new_comp]["rel_probability"] delays = np.repeat( - parameters[new_comp]["delay"].as_random_distribution()(), + build_distribution_from_confuse_config(parameters[new_comp]["delay"])(), len(modinf.subpop_struct.subpop_names), ) probabilities[probabilities > 1] = 1 @@ -553,7 +556,9 @@ def compute_all_multioutcomes( ]["value"].to_numpy() else: durations = np.repeat( - parameters[new_comp]["duration"].as_random_distribution()(), + build_distribution_from_confuse_config( + parameters[new_comp]["duration"] + )(), len(modinf.subpop_struct.subpop_names), ) durations = np.repeat( diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index 37b335777..32f2fa7bf 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -22,11 +22,7 @@ import pydantic from . import NPI, utils -from .distributions import ( - Distribution, - DISTRIBUTION_ADAPTER, - build_distribution_from_confuse_config, -) +from .distributions import build_distribution_from_confuse_config logger = logging.getLogger(__name__) From 45f857f4120f0196d9a4697d33e84118462b559d Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Mon, 28 Jul 2025 16:26:20 -0400 Subject: [PATCH 50/73] Refactoring all uses of `.as_random_distribution` to `build_distribution_from_confuse_config()` This code is presently broken. --- .../src/gempyor/NPI/MultiPeriodModifier.py | 15 ++++++++------- .../src/gempyor/NPI/SinglePeriodModifier.py | 9 +++++---- .../gempyor_pkg/src/gempyor/distributions.py | 15 ++++++++------- .../src/gempyor/inference_parameter.py | 17 +++++++++++++---- flepimop/gempyor_pkg/src/gempyor/outcomes.py | 11 ++++++++--- flepimop/gempyor_pkg/src/gempyor/parameters.py | 10 +++------- 6 files changed, 45 insertions(+), 32 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py index 82c227336..0f4013070 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py @@ -5,6 +5,7 @@ from . import helpers from .base import NPIBase +from ..distributions import build_distribution_from_confuse_config class MultiPeriodModifier(NPIBase): @@ -149,7 +150,7 @@ def __createFromConfig(self, npi_config): self.affected_subpops = self.__get_affected_subpops(npi_config) self.parameters = self.parameters[self.parameters.index.isin(self.affected_subpops)] - dist = npi_config["value"].as_random_distribution() + dist = build_distribution_from_confuse_config(npi_config["value"]) self.parameters["modifier_name"] = self.name self.parameters["parameter"] = self.param_name @@ -177,9 +178,9 @@ def __createFromConfig(self, npi_config): for subpop in this_spatial_group["ungrouped"]: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates - self.parameters.at[subpop, "value"] = dist(size=1) + self.parameters.at[subpop, "value"] = dist() for group in this_spatial_group["grouped"]: - drawn_value = dist(size=1) + drawn_value = dist() for subpop in group: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates @@ -235,8 +236,8 @@ def __createFromDf(self, loaded_df, npi_config): if not subpop in loaded_df.index: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates - dist = npi_config["value"].as_random_distribution() - self.parameters.at[subpop, "value"] = dist(size=1) + dist = build_distribution_from_confuse_config(npi_config["value"]) + self.parameters.at[subpop, "value"] = dist() else: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates @@ -250,8 +251,8 @@ def __createFromDf(self, loaded_df, npi_config): ",".join(group), "value" ] else: - dist = npi_config["value"].as_random_distribution() - drawn_value = dist(size=1) + dist = build_distribution_from_confuse_config(npi_config["value"]) + drawn_value = dist() for subpop in group: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py index 7aece6306..81680422a 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py @@ -5,6 +5,7 @@ from . import helpers from .base import NPIBase +from ..distributions import build_distribution_from_confuse_config class SinglePeriodModifier(NPIBase): @@ -134,7 +135,7 @@ def __createFromConfig(self, npi_config): self.parameters = self.parameters[self.parameters.index.isin(self.affected_subpops)] # Create reduction - self.dist = npi_config["value"].as_random_distribution() + self.dist = build_distribution_from_confuse_config(npi_config["value"]) self.parameters["modifier_name"] = self.name self.parameters["start_date"] = ( @@ -152,12 +153,12 @@ def __createFromConfig(self, npi_config): npi_config, list(self.affected_subpops) ) if self.spatial_groups["ungrouped"]: - self.parameters.loc[self.spatial_groups["ungrouped"], "value"] = self.dist( - size=len(self.spatial_groups["ungrouped"]) + self.parameters.loc[self.spatial_groups["ungrouped"], "value"] = ( + self.dist.sample(size=len(self.spatial_groups["ungrouped"])) ) if self.spatial_groups["grouped"]: for group in self.spatial_groups["grouped"]: - drawn_value = self.dist(size=1) * np.ones(len(group)) + drawn_value = self.dist() * np.ones(len(group)) self.parameters.loc[group, "value"] = drawn_value def __createFromDf(self, loaded_df, npi_config): diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 5b846df50..6c6200c7b 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -55,6 +55,10 @@ def sample( rng = rng if rng is not None else self._rng return self._sample_from_generator(size=size, rng=rng) + def __call__(self) -> float | int: + """A shortcut for `self.sample(size=1)`.""" + return self.sample(size=1).item() + @abstractmethod def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator @@ -483,23 +487,20 @@ def _sample_from_generator( DISTRIBUTION_ADAPTER = TypeAdapter(Distribution) -def build_distribution_from_confuse_config( - param_config: confuse.ConfigView, -) -> Distribution: +def build_distribution_from_confuse_config(config: confuse.ConfigView) -> Distribution: """ - Creates a Distribution object from a confuse.ConfigView (single value). + Creates a Distribution object from a `confuse.ConfigView`. Handles the case where the value is a simple number or string, interpreting it as a 'fixed' distribution. Args: - param_config: A confuse.ConfigView for a single parameter. + config: A `confuse.ConfigView` for a single parameter. Returns: A Distribution object. """ - conf = param_config["value"].get() + conf = config.get() if isinstance(conf, float | int | str): conf = {"distribution": "fixed", "value": conf} - return DISTRIBUTION_ADAPTER.validate_python(conf) diff --git a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py index b4ee24274..b5707103a 100644 --- a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py +++ b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py @@ -7,6 +7,7 @@ import numpy as np import confuse from . import NPI +from .distributions import build_distribution_from_confuse_config # TODO cast uper and lower bound as arrays @@ -58,7 +59,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=sp, - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) @@ -70,7 +73,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=",".join(group), - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) @@ -100,7 +105,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=sp, - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) @@ -112,7 +119,9 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=",".join(group), - pdist=parameter_config["value"].as_random_distribution(), + pdist=build_distribution_from_confuse_config( + parameter_config["value"] + ), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) diff --git a/flepimop/gempyor_pkg/src/gempyor/outcomes.py b/flepimop/gempyor_pkg/src/gempyor/outcomes.py index 4ff727d34..357488b23 100644 --- a/flepimop/gempyor_pkg/src/gempyor/outcomes.py +++ b/flepimop/gempyor_pkg/src/gempyor/outcomes.py @@ -17,6 +17,7 @@ from .utils import Timer, _nslots_random_seeds, config, read_df from . import NPI, model_info +from .distributions import build_distribution_from_confuse_config logger = logging.getLogger(__name__) @@ -466,13 +467,15 @@ def compute_all_multioutcomes( else: # One draw for all subpops probabilities = np.repeat( - parameters[new_comp]["probability"].as_random_distribution()(), + build_distribution_from_confuse_config( + parameters[new_comp]["probability"] + )(), len(modinf.subpop_struct.subpop_names), ) if "rel_probability" in parameters[new_comp]: probabilities = probabilities * parameters[new_comp]["rel_probability"] delays = np.repeat( - parameters[new_comp]["delay"].as_random_distribution()(), + build_distribution_from_confuse_config(parameters[new_comp]["delay"])(), len(modinf.subpop_struct.subpop_names), ) probabilities[probabilities > 1] = 1 @@ -553,7 +556,9 @@ def compute_all_multioutcomes( ]["value"].to_numpy() else: durations = np.repeat( - parameters[new_comp]["duration"].as_random_distribution()(), + build_distribution_from_confuse_config( + parameters[new_comp]["duration"] + )(), len(modinf.subpop_struct.subpop_names), ) durations = np.repeat( diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index 37b335777..dc27e2bc3 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -22,11 +22,7 @@ import pydantic from . import NPI, utils -from .distributions import ( - Distribution, - DISTRIBUTION_ADAPTER, - build_distribution_from_confuse_config, -) +from .distributions import build_distribution_from_confuse_config logger = logging.getLogger(__name__) @@ -106,7 +102,7 @@ def __init__( # Parameter characterized by its distribution if self.pconfig[pn]["value"].exists(): self.pdata[pn]["dist"] = build_distribution_from_confuse_config( - self.pconfig[pn] + self.pconfig[pn]["value"] ) # Parameter given as a file @@ -191,7 +187,7 @@ def reinitialize_distributions(self) -> None: for pn in self.pnames: if "dist" in self.pdata[pn]: self.pdata[pn]["dist"] = build_distribution_from_confuse_config( - self.pconfig[pn] + self.pconfig[pn]["value"] ) def get_pnames2pindex(self) -> dict: From 3dccedb56a148909280ab67a4abebcdca585e3d8 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 29 Jul 2025 11:14:30 -0400 Subject: [PATCH 51/73] De-duplicate code by removing three four functions from `gempyor.utils` Add new `__call__` syntactic sugar to `Distributions`. --- .../gempyor_pkg/src/gempyor/distributions.py | 2 +- flepimop/gempyor_pkg/src/gempyor/utils.py | 225 ------------------ .../tests/parameters/test_parameters_class.py | 1 - .../tests/utils/test_get_log_normal.py | 49 ---- .../tests/utils/test_get_truncated_normal.py | 49 ---- .../utils/test_random_distribution_sampler.py | 79 ------ .../gempyor_pkg/tests/utils/test_utils.py | 8 - .../gempyor_pkg/tests/utils/test_utils2.py | 71 ------ 8 files changed, 1 insertion(+), 483 deletions(-) delete mode 100644 flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py delete mode 100644 flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py delete mode 100644 flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 98e9dc6ca..4a641c3c8 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -54,7 +54,7 @@ def sample( """ rng = rng if rng is not None else self._rng return self._sample_from_generator(size=size, rng=rng) - + def __call__(self) -> float | int: """A shortcut for `self.sample(size=1)`.""" return self.sample(size=1).item() diff --git a/flepimop/gempyor_pkg/src/gempyor/utils.py b/flepimop/gempyor_pkg/src/gempyor/utils.py index 423cd647c..f0c7916eb 100644 --- a/flepimop/gempyor_pkg/src/gempyor/utils.py +++ b/flepimop/gempyor_pkg/src/gempyor/utils.py @@ -351,231 +351,6 @@ def as_date(self) -> datetime.date: return self.get(ISO8601Date()) -def get_truncated_normal( - mean: float | int = 0, - sd: float | int = 1, - a: float | int = 0, - b: float | int = 10, -) -> scipy.stats._distn_infrastructure.rv_frozen: - """ - Returns a truncated normal distribution. - - This function constructs a truncated normal distribution with the specified - mean, standard deviation, and bounds. The truncated normal distribution is - a normal distribution bounded within the interval [a, b]. - - Args: - mean: The mean of the truncated normal distribution. Defaults to 0. - sd: The standard deviation of the truncated normal distribution. Defaults to 1. - a: The lower bound of the truncated normal distribution. Defaults to 0. - b: The upper bound of the truncated normal distribution. Defaults to 10. - - Returns: - rv_frozen: A frozen instance of the truncated normal distribution with the specified parameters. - - Examples: - Create a truncated normal distribution with specified parameters (truncated between 1 and 10): - >>> truncated_normal_dist = get_truncated_normal(mean=5, sd=2, a=1, b=10) - >>> print(truncated_normal_dist) - rv_frozen() - """ - lower = (a - mean) / sd - upper = (b - mean) / sd - return scipy.stats.truncnorm(lower, upper, loc=mean, scale=sd) - - -def get_log_normal( - meanlog: float | int, - sdlog: float | int, -) -> scipy.stats._distn_infrastructure.rv_frozen: - """ - Returns a log normal distribution. - - This function constructs a log normal distribution with the specified - log mean and log standard deviation. - - Args: - meanlog: The log of the mean of the log normal distribution. - sdlog: The log of the standard deviation of the log normal distribution. - - Returns: - rv_frozen: A frozen instance of the log normal distribution with the - specified parameters. - - Examples: - Create a log-normal distribution with specified parameters: - >>> log_normal_dist = get_log_normal(meanlog=1, sdlog=0.5) - >>> print(log_normal_dist) - - """ - return scipy.stats.lognorm(s=sdlog, scale=np.exp(meanlog), loc=0) - - -def random_distribution_sampler( - distribution: Literal[ - "fixed", "uniform", "poisson", "binomial", "truncnorm", "lognorm" - ], - **kwargs: dict[str, Any], -) -> Callable[[], float | int]: - """ - Create function to sample from a random distribution. - - Args: - distribution: The type of distribution to generate a sampling function for. - **kwargs: Further parameters that are passed to the underlying function for the - given distribution. - - Notes: - The further args expected by each distribution type are: - - fixed: value, - - uniform: low, high, - - poisson: lam, - - binomial: n, p, - - truncnorm: mean, sd, a, b, - - lognorm: meanlog, sdlog. - - Returns: - A function that can be called to sample from that distribution. - - Raises: - ValueError: If `distribution` is 'binomial' the given `p` must be in (0,1). - NotImplementedError: If `distribution` is not one of the type hinted options. - - Examples: - >>> import numpy as np - >>> np.random.seed(123) - >>> uniform_sampler = random_distribution_sampler("uniform", low=0.0, high=3.0) - >>> uniform_sampler() - 2.089407556793585 - >>> uniform_sampler() - 0.8584180048511384 - """ - if distribution == "fixed": - # Fixed value is the same as uniform on [a, a) - return functools.partial( - np.random.uniform, - kwargs.get("value"), - kwargs.get("value"), - ) - elif distribution == "uniform": - # Uniform on [low, high) - return functools.partial( - np.random.uniform, - kwargs.get("low"), - kwargs.get("high"), - ) - elif distribution == "poisson": - # Poisson with mean lambda - return functools.partial(np.random.poisson, kwargs.get("lam")) - elif distribution == "binomial": - p = kwargs.get("p") - if not (0 < p < 1): - raise ValueError(f"Invalid `p-value`: '{p}' is out of range [0,1].") - return functools.partial(np.random.binomial, kwargs.get("n"), p) - elif distribution == "truncnorm": - # Truncated normal with mean, sd on interval [a, b] - return get_truncated_normal( - mean=kwargs.get("mean"), - sd=kwargs.get("sd"), - a=kwargs.get("a"), - b=kwargs.get("b"), - ).rvs - elif distribution == "lognorm": - # Lognormal distribution with meanlog, sdlog - return get_log_normal(kwargs.get("meanlog"), kwargs.get("sdlog")).rvs - raise NotImplementedError(f"Unknown distribution [received: '{distribution}'].") - - -@add_method(confuse.ConfigView) -def as_random_distribution(self): - """ - Constructs a random distribution object from a distribution config key. - - Args: - self: Class instance (in this case, a config key) to construct the random distribution from. - - Returns: - A partial object containing the random distribution. - - Raises: - ValueError: When values are out of range. - NotImplementedError: If an unknown distribution is found. - - Examples: - Say that ``config`` is a ``confuse.ConfigView`` instance. - - To create a uniform distribution between 1 and 10: - >>> dist_function = config.as_random_distribution() - >>> sample = dist_function() - 5.436789235794546 - - To use a truncated normal distribution: - >>> config_truncnorm = confuse.ConfigView({ - "distribution": "truncnorm", - "mean": 0 - "sd": 1, - "a": -1, - "b": 1 - }) - >>> truncnorm_dist_function = config_truncnorm.as_random_distribution() - >>> truncnorm_sample = truncnorm_dist_function() - 0.312745 - ``` - """ - - if isinstance(self.get(), dict): - dist = self["distribution"].get() - if dist == "fixed": - return functools.partial( - np.random.uniform, - _evaled_expression(self["value"].get(), target_type=float), - _evaled_expression(self["value"].get(), target_type=float), - ) - elif dist == "uniform": - return functools.partial( - np.random.uniform, - _evaled_expression(self["low"].get(), target_type=float), - _evaled_expression(self["high"].get(), target_type=float), - ) - elif dist == "poisson": - return functools.partial( - np.random.poisson, _evaled_expression(self["lam"].get(), target_type=float) - ) - elif dist == "binomial": - p = _evaled_expression(self["p"].get(), target_type=float) - if (p < 0) or (p > 1): - raise ValueError(f"Invalid `p-value`: '{p}' is out of range [0,1].") - # if (self["p"] < 0) or (self["p"] > 1): - # raise ValueError(f"""p value { self["p"] } is out of range [0,1]""") - return functools.partial( - np.random.binomial, - _evaled_expression(self["n"].get(), target_type=float), - # _evaled_expression(self["p"].get(), target_type=float), - p, - ) - elif dist == "truncnorm": - return get_truncated_normal( - mean=_evaled_expression(self["mean"].get(), target_type=float), - sd=_evaled_expression(self["sd"].get(), target_type=float), - a=_evaled_expression(self["a"].get(), target_type=float), - b=_evaled_expression(self["b"].get(), target_type=float), - ).rvs - elif dist == "lognorm": - return get_log_normal( - meanlog=_evaled_expression(self["meanlog"].get(), target_type=float), - sdlog=_evaled_expression(self["sdlog"].get(), target_type=float), - ).rvs - else: - raise NotImplementedError(f"Unknown distribution [received: '{dist}'].") - else: - # we allow a fixed value specified directly: - return functools.partial( - np.random.uniform, - _evaled_expression(self.get(), target_type=float), - _evaled_expression(self.get(), target_type=float), - ) - - def list_filenames( folder: str | bytes | os.PathLike = ".", filters: str | list[str] = [], diff --git a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py index 7e84f953f..b70d3c75b 100644 --- a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py +++ b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py @@ -31,7 +31,6 @@ partials_are_similar, sample_fits_distribution, ) -from gempyor.utils import random_distribution_sampler class MockParametersInput: diff --git a/flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py b/flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py deleted file mode 100644 index b12989ec4..000000000 --- a/flepimop/gempyor_pkg/tests/utils/test_get_log_normal.py +++ /dev/null @@ -1,49 +0,0 @@ -import numpy as np -import pytest -import scipy.stats - -from gempyor.utils import get_log_normal - - -class TestGetLogNormal: - """Unit tests for the `gempyor.utils.get_log_normal` function.""" - - @pytest.mark.parametrize( - "meanlog,sdlog", - [ - (1.0, 1.0), - (0.0, 2.0), - (10.0, 30.0), - (0.33, 4.56), - (9.87, 4.21), - (1, 1), - (0, 2), - (10, 30), - ], - ) - def test_construct_distribution( - self, - meanlog: float | int, - sdlog: float | int, - ) -> None: - """Test the construction of a log normal distribution. - - This test checks whether the `get_log_normal` function correctly constructs - a log normal distribution with the specified parameters. It verifies that - the returned object is an instance of `rv_frozen`, and that its support and - parameters (log mean and log standard deviation) are correctly set. - - Args: - mean: The mean of the truncated normal distribution. - sd: The standard deviation of the truncated normal distribution. - a: The lower bound of the truncated normal distribution. - b: The upper bound of the truncated normal distribution. - """ - dist = get_log_normal(meanlog=meanlog, sdlog=sdlog) - assert isinstance(dist, scipy.stats._distn_infrastructure.rv_frozen) - lower, upper = dist.support() - assert np.isclose(lower, 0.0) - assert np.isclose(upper, np.inf) - assert np.isclose(dist.kwds.get("s"), sdlog) - assert np.isclose(dist.kwds.get("scale"), np.exp(meanlog)) - assert np.isclose(dist.kwds.get("loc"), 0.0) diff --git a/flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py b/flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py deleted file mode 100644 index c3ccbca79..000000000 --- a/flepimop/gempyor_pkg/tests/utils/test_get_truncated_normal.py +++ /dev/null @@ -1,49 +0,0 @@ -import numpy as np -import pytest -import scipy.stats - -from gempyor.utils import get_truncated_normal - - -class TestGetTruncatedNormal: - """Unit tests for the `gempyor.utils.get_truncated_normal` function.""" - - @pytest.mark.parametrize( - "mean,sd,a,b", - [ - (0.0, 1.0, 0.0, 10.0), - (0.0, 2.0, -4.0, 4.0), - (-5.0, 3.0, -5.0, 10.0), - (-3.25, 1.4, -8.74, 4.89), - (0, 1, 0, 10), - (0, 2, -4, 4), - (-5, 3, -5, 10), - ], - ) - def test_construct_distribution( - self, - mean: float | int, - sd: float | int, - a: float | int, - b: float | int, - ) -> None: - """Test the construction of a truncated normal distribution. - - This test checks whether the `get_truncated_normal` function correctly - constructs a truncated normal distribution with the specified parameters. - It verifies that the returned object is an instance of `rv_frozen`, and that - its support and parameters (mean and standard deviation) are correctly set. - - Args: - mean: The mean of the truncated normal distribution. - sd: The standard deviation of the truncated normal distribution. - a: The lower bound of the truncated normal distribution. - b: The upper bound of the truncated normal distribution. - """ - dist = get_truncated_normal(mean=mean, sd=sd, a=a, b=b) - assert isinstance(dist, scipy.stats._distn_infrastructure.rv_frozen) - lower, upper = dist.support() - assert np.isclose(lower, a) - assert np.isclose(upper, b) - assert np.isclose(dist.kwds.get("loc"), mean) - assert np.isclose(dist.kwds.get("scale"), sd) diff --git a/flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py b/flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py deleted file mode 100644 index 0dcd110fb..000000000 --- a/flepimop/gempyor_pkg/tests/utils/test_random_distribution_sampler.py +++ /dev/null @@ -1,79 +0,0 @@ -from functools import partial -import inspect -from typing import Any - -import numpy as np -import pytest -import re - -from gempyor.testing import partials_are_similar -from gempyor.utils import random_distribution_sampler - - -class TestRandomDistributionSampler: - @pytest.mark.parametrize("distribution", [("abc"), ("def"), ("ghi")]) - def test_not_implemented_error_exception(self, distribution: str) -> None: - with pytest.raises( - NotImplementedError, - match=re.escape(rf"Unknown distribution [received: '{distribution}']."), - ): - random_distribution_sampler(distribution) - - @pytest.mark.parametrize("p", [(-0.5), (1.2), (0.0), (1.0)]) - def test_binomial_p_value_error(self, p: float) -> None: - with pytest.raises( - ValueError, - match=re.escape(rf"Invalid `p-value`: '{p}' is out of range [0,1]."), - ): - random_distribution_sampler("binomial", n=100, p=p) - - @pytest.mark.parametrize( - "distribution,kwargs", - [ - ("fixed", {"value": 0.12}), - ("fixed", {"value": -3.45}), - ("fixed", {"value": 0.0}), - ("uniform", {"low": 0.0, "high": 1.0}), - ("uniform", {"low": 50.0, "high": 200.0}), - ("uniform", {"low": -1.0, "high": 1.0}), - ("uniform", {"low": 1.0, "high": -1.0}), - ("poisson", {"lam": 0.1}), - ("poisson", {"lam": 1.23}), - ("poisson", {"lam": -0.1}), - ("binomial", {"n": 10, "p": 0.1}), - ("binomial", {"n": -10, "p": 0.1}), - ("binomial", {"n": 50, "p": 0.67}), - ("truncnorm", {"mean": 0.0, "sd": 1.0, "a": -2.0, "b": 2.0}), - ("truncnorm", {"mean": 1.4, "sd": 0.34, "a": -0.3, "b": 22.8}), - ("lognorm", {"sdlog": 1.0, "meanlog": 1.0}), - ("lognorm", {"sdlog": 3.4, "meanlog": -0.56}), - ], - ) - def test_output_validation(self, distribution: str, kwargs: dict[str, Any]) -> None: - actual = random_distribution_sampler(distribution, **kwargs) - if distribution == "fixed": - expected = partial(np.random.uniform, kwargs.get("value"), kwargs.get("value")) - assert partials_are_similar(actual, expected) - elif distribution == "uniform": - expected = partial(np.random.uniform, kwargs.get("low"), kwargs.get("high")) - assert partials_are_similar(actual, expected) - elif distribution == "poisson": - expected = partial(np.random.poisson, kwargs.get("lam")) - assert partials_are_similar(actual, expected) - elif distribution == "binomial": - expected = partial(np.random.binomial, kwargs.get("n"), kwargs.get("p")) - assert partials_are_similar(actual, expected) - elif distribution == "truncnorm": - assert inspect.ismethod(actual) - assert actual.__self__.kwds.get("loc") == kwargs.get("mean") - assert actual.__self__.kwds.get("scale") == kwargs.get("sd") - assert actual.__self__.a == (kwargs.get("a") - kwargs.get("mean")) / kwargs.get( - "sd" - ) - assert actual.__self__.b == (kwargs.get("b") - kwargs.get("mean")) / kwargs.get( - "sd" - ) - elif distribution == "lognorm": - assert inspect.ismethod(actual) - assert actual.__self__.kwds.get("s") == kwargs.get("sdlog") - assert actual.__self__.kwds.get("scale") == np.exp(kwargs.get("meanlog")) diff --git a/flepimop/gempyor_pkg/tests/utils/test_utils.py b/flepimop/gempyor_pkg/tests/utils/test_utils.py index 4753ac142..ec1f608f4 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_utils.py +++ b/flepimop/gempyor_pkg/tests/utils/test_utils.py @@ -88,14 +88,6 @@ def test_ISO8601Date_success(): # assert t == datetime.datetime("2020-02-01").strftime("%Y-%m-%d") -def test_get_truncated_normal_success(): - utils.get_truncated_normal(mean=0, sd=1, a=-2, b=2) - - -def test_get_log_normal_success(): - utils.get_log_normal(meanlog=0, sdlog=1) - - def test_create_resume_out_filename(): result = utils.create_resume_out_filename( flepi_run_index="123", diff --git a/flepimop/gempyor_pkg/tests/utils/test_utils2.py b/flepimop/gempyor_pkg/tests/utils/test_utils2.py index cfaeb2943..36ca2dc1e 100644 --- a/flepimop/gempyor_pkg/tests/utils/test_utils2.py +++ b/flepimop/gempyor_pkg/tests/utils/test_utils2.py @@ -84,12 +84,7 @@ def test_ISO8601Date_success(): #assert t == datetime.datetime("2020-02-01").strftime("%Y-%m-%d") -def test_get_truncated_normal_success(): - utils.get_truncated_normal(mean=0, sd=1, a=-2, b=2) - -def test_get_log_normal_success(): - utils.get_log_normal(meanlog=0, sdlog=1) """ @@ -117,69 +112,3 @@ def config(): def test_as_date(config): config.add({"myvalue": "2022-01-15"}) assert config["myvalue"].as_date() == datetime.date(2022, 1, 15) - - -def test_as_random_distribution_fixed(config): - config.add({"value": {"distribution": "fixed", "value": 1}}) - dist = config["value"].as_random_distribution() - assert dist() == 1 - - -def test_as_random_distribution_uniform(config): - config.add({"value": {"distribution": "uniform", "low": 1, "high": 2.6}}) - dist = config["value"].as_random_distribution() - assert 1 <= dist() <= 2.6 - - -def test_as_random_distribution_poisson(config): - config.add({"value": {"distribution": "poisson", "lam": 1}}) - dist = config["value"].as_random_distribution() - assert isinstance(dist(), int) - - -def test_as_random_distribution_binomial(config): - config.add({"value": {"distribution": "binomial", "n": 10, "p": 0.5}}) - dist = config["value"].as_random_distribution() - assert 0 <= dist() <= 10 - - -def test_as_random_distribution_binomial_w_fraction(config): - config.add({"value": {"distribution": "binomial", "n": 10, "p": "1/2"}}) - dist = config["value"].as_random_distribution() - assert 0 <= dist() <= 10 - - -def test_as_random_distribution_binomial_error(config): - config.add({"value": {"distribution": "binomial", "n": 10, "p": 1.1}}) - with pytest.raises(ValueError, match=r".*p.*value.*"): - dist = config["value"].as_random_distribution() - - -def test_as_random_distribution_binomial_w_fraction_error(config): - config.add({"value": {"distribution": "binomial", "n": 10, "p": "5/4"}}) - with pytest.raises(ValueError, match=r".*p.*value.*"): - dist = config["value"].as_random_distribution() - - -def test_as_random_distribution_truncnorm(config): - config.add( - {"value": {"distribution": "truncnorm", "mean": 0, "sd": 1, "a": -1, "b": 1}} - ) - dist = config["value"].as_random_distribution() - rvs = dist(size=1000) - assert len(rvs) == 1000 - assert all(-1 <= x <= 1 for x in rvs) - - -def test_as_random_distribution_lognorm(config): - config.add({"value": {"distribution": "lognorm", "meanlog": 0, "sdlog": 1}}) - dist = config["value"].as_random_distribution() - rvs = dist(size=1000) - assert len(rvs) == 1000 - assert all(x > 0 for x in rvs) - - -def test_as_random_distribution_unknown(config): - config.add({"value": {"distribution": "unknown", "mean": 0, "sd": 1}}) - with pytest.raises(NotImplementedError): - config["value"].as_random_distribution() From 51ec305e2334e0eac0360331c92640d49f5630cb Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Tue, 29 Jul 2025 12:15:02 -0400 Subject: [PATCH 52/73] Update flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index c76f2388c..e67294211 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -187,18 +187,12 @@ def _read_and_validate_dataframe( @overload -def _evaled_expression(val: str, *, target_type: Type[int]) -> int: ... -@overload -def _evaled_expression(val: Any, *, target_type: Type[int]) -> int | Any: ... - +def _evaled_expression(val: EE | str, *, target_type: Type[EE]) -> EE: ... @overload -def _evaled_expression(val: str, *, target_type: Type[float]) -> float: ... -@overload -def _evaled_expression(val: Any, *, target_type: Type[float]) -> float | Any: ... - +def _evaled_expression(val: T, *, target_type: Type[EE]) -> T: ... -def _evaled_expression(val: Any, *, target_type: Type[EE]) -> EE | Any: +def _evaled_expression(val: EE | str | Any, *, target_type: Type[EE]) -> EE | Any: """ Evaluates a string expression to a target numeric type (int or float). From e5f3220f60d04cef27a2b1d93610171be3ba6a4e Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 29 Jul 2025 12:17:07 -0400 Subject: [PATCH 53/73] Remove `*` keyword arg for `target_type` in `_evaled_expression()` --- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index e67294211..95c32226d 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -187,12 +187,12 @@ def _read_and_validate_dataframe( @overload -def _evaled_expression(val: EE | str, *, target_type: Type[EE]) -> EE: ... +def _evaled_expression(val: EE | str, target_type: Type[EE]) -> EE: ... @overload -def _evaled_expression(val: T, *, target_type: Type[EE]) -> T: ... +def _evaled_expression(val: T, target_type: Type[EE]) -> T: ... -def _evaled_expression(val: EE | str | Any, *, target_type: Type[EE]) -> EE | Any: +def _evaled_expression(val: EE | str | Any, target_type: Type[EE]) -> EE | Any: """ Evaluates a string expression to a target numeric type (int or float). @@ -207,18 +207,18 @@ def _evaled_expression(val: EE | str | Any, *, target_type: Type[EE]) -> EE | An ValueError: On parsing errors. Examples: - >>> _evaled_expression("1 + 1", target_type=int) + >>> _evaled_expression("1 + 1", int) 2 - >>> _evaled_expression("10 / 4", target_type=float) + >>> _evaled_expression("10 / 4", float) 2.5 - >>> _evaled_expression("10 / 4", target_type=int) + >>> _evaled_expression("10 / 4", int) # note that result is trucnated; probably undesirable if misused 2 - >>> _evaled_expression(99.5, target_type=float) + >>> _evaled_expression(99.5, float) 99.5 - >>> _evaled_expression(None, target_type=int) + >>> _evaled_expression(None, int) - >>> _evaled_expression("a * b", target_type=float) + >>> _evaled_expression("a * b", float) Traceback (most recent call last): ... ValueError: Can't convert expression to float. From 88640bb066e672d89d6a8cb33f0e07c3ea39a914 Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Tue, 29 Jul 2025 12:17:33 -0400 Subject: [PATCH 54/73] Update flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index e67294211..a5146576a 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -16,7 +16,7 @@ import pyarrow.parquet as pq import pandas as pd from pydantic import BaseModel, RootModel, BeforeValidator -import sympy +from sympy.parsing.sympy_parser import parse_expr T = TypeVar("T") From 89f960b163f6ce1803d383efb05a59377bb4ce7a Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Tue, 29 Jul 2025 12:22:24 -0400 Subject: [PATCH 55/73] Update flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py Co-authored-by: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> --- flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index a5146576a..d52e05a92 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -227,10 +227,10 @@ def _evaled_expression(val: EE | str | Any, *, target_type: Type[EE]) -> EE | An return val if isinstance(val, str): - try: - return target_type(sympy.parsing.sympy_parser.parse_expr(val)) - except TypeError as e: - raise ValueError(f"Can't convert expression to {target_type.__name__}.") from e + expr = parse_expr(val) + if not expr.is_Number: + raise ValueError(f"Cannot convert expression '{expr}' to {target_type}.") + return target_type(expr) return val From 1172363a9bc8cef4ffbbc33b14db5eeff0b6cca1 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 29 Jul 2025 12:34:31 -0400 Subject: [PATCH 56/73] Naming suggestions, deleting duplication --- .../gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py | 8 ++++---- .../src/gempyor/NPI/SinglePeriodModifier.py | 4 ++-- flepimop/gempyor_pkg/src/gempyor/distributions.py | 5 +---- .../gempyor_pkg/src/gempyor/inference_parameter.py | 10 +++++----- flepimop/gempyor_pkg/src/gempyor/outcomes.py | 8 ++++---- 5 files changed, 16 insertions(+), 19 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py index 0f4013070..f5b6aa0cb 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py @@ -5,7 +5,7 @@ from . import helpers from .base import NPIBase -from ..distributions import build_distribution_from_confuse_config +from ..distributions import distribution_from_confuse_config class MultiPeriodModifier(NPIBase): @@ -150,7 +150,7 @@ def __createFromConfig(self, npi_config): self.affected_subpops = self.__get_affected_subpops(npi_config) self.parameters = self.parameters[self.parameters.index.isin(self.affected_subpops)] - dist = build_distribution_from_confuse_config(npi_config["value"]) + dist = distribution_from_confuse_config(npi_config["value"]) self.parameters["modifier_name"] = self.name self.parameters["parameter"] = self.param_name @@ -236,7 +236,7 @@ def __createFromDf(self, loaded_df, npi_config): if not subpop in loaded_df.index: self.parameters.at[subpop, "start_date"] = start_dates self.parameters.at[subpop, "end_date"] = end_dates - dist = build_distribution_from_confuse_config(npi_config["value"]) + dist = distribution_from_confuse_config(npi_config["value"]) self.parameters.at[subpop, "value"] = dist() else: self.parameters.at[subpop, "start_date"] = start_dates @@ -251,7 +251,7 @@ def __createFromDf(self, loaded_df, npi_config): ",".join(group), "value" ] else: - dist = build_distribution_from_confuse_config(npi_config["value"]) + dist = distribution_from_confuse_config(npi_config["value"]) drawn_value = dist() for subpop in group: self.parameters.at[subpop, "start_date"] = start_dates diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py index 81680422a..36035ef16 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py @@ -5,7 +5,7 @@ from . import helpers from .base import NPIBase -from ..distributions import build_distribution_from_confuse_config +from ..distributions import distribution_from_confuse_config class SinglePeriodModifier(NPIBase): @@ -135,7 +135,7 @@ def __createFromConfig(self, npi_config): self.parameters = self.parameters[self.parameters.index.isin(self.affected_subpops)] # Create reduction - self.dist = build_distribution_from_confuse_config(npi_config["value"]) + self.dist = distribution_from_confuse_config(npi_config["value"]) self.parameters["modifier_name"] = self.name self.parameters["start_date"] = ( diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 4a641c3c8..d8bac6ea3 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -59,9 +59,6 @@ def __call__(self) -> float | int: """A shortcut for `self.sample(size=1)`.""" return self.sample(size=1).item() - def __call__(self) -> float | int: - """A shortcut for `self.sample(size=1)`.""" - return self.sample(size=1).item() @abstractmethod def _sample_from_generator( @@ -491,7 +488,7 @@ def _sample_from_generator( DISTRIBUTION_ADAPTER = TypeAdapter(Distribution) -def build_distribution_from_confuse_config(config: confuse.ConfigView) -> Distribution: +def distribution_from_confuse_config(config: confuse.ConfigView) -> Distribution: """ Creates a Distribution object from a `confuse.ConfigView`. diff --git a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py index b5707103a..494a9800f 100644 --- a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py +++ b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py @@ -7,7 +7,7 @@ import numpy as np import confuse from . import NPI -from .distributions import build_distribution_from_confuse_config +from .distributions import distribution_from_confuse_config # TODO cast uper and lower bound as arrays @@ -59,7 +59,7 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=sp, - pdist=build_distribution_from_confuse_config( + pdist=distribution_from_confuse_config( parameter_config["value"] ), lb=parameter_config["value"]["a"].get(float), @@ -73,7 +73,7 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=",".join(group), - pdist=build_distribution_from_confuse_config( + pdist=distribution_from_confuse_config( parameter_config["value"] ), lb=parameter_config["value"]["a"].get(float), @@ -105,7 +105,7 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=sp, - pdist=build_distribution_from_confuse_config( + pdist=istribution_from_confuse_config( parameter_config["value"] ), lb=parameter_config["value"]["a"].get(float), @@ -119,7 +119,7 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=",".join(group), - pdist=build_distribution_from_confuse_config( + pdist=istribution_from_confuse_config( parameter_config["value"] ), lb=parameter_config["value"]["a"].get(float), diff --git a/flepimop/gempyor_pkg/src/gempyor/outcomes.py b/flepimop/gempyor_pkg/src/gempyor/outcomes.py index 357488b23..6d593b140 100644 --- a/flepimop/gempyor_pkg/src/gempyor/outcomes.py +++ b/flepimop/gempyor_pkg/src/gempyor/outcomes.py @@ -17,7 +17,7 @@ from .utils import Timer, _nslots_random_seeds, config, read_df from . import NPI, model_info -from .distributions import build_distribution_from_confuse_config +from .distributions import distribution_from_confuse_config logger = logging.getLogger(__name__) @@ -467,7 +467,7 @@ def compute_all_multioutcomes( else: # One draw for all subpops probabilities = np.repeat( - build_distribution_from_confuse_config( + distribution_from_confuse_config( parameters[new_comp]["probability"] )(), len(modinf.subpop_struct.subpop_names), @@ -475,7 +475,7 @@ def compute_all_multioutcomes( if "rel_probability" in parameters[new_comp]: probabilities = probabilities * parameters[new_comp]["rel_probability"] delays = np.repeat( - build_distribution_from_confuse_config(parameters[new_comp]["delay"])(), + distribution_from_confuse_config(parameters[new_comp]["delay"])(), len(modinf.subpop_struct.subpop_names), ) probabilities[probabilities > 1] = 1 @@ -556,7 +556,7 @@ def compute_all_multioutcomes( ]["value"].to_numpy() else: durations = np.repeat( - build_distribution_from_confuse_config( + distribution_from_confuse_config( parameters[new_comp]["duration"] )(), len(modinf.subpop_struct.subpop_names), From 606cb4cf76d9a527e7820282c63dd9af6dea834f Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 29 Jul 2025 13:19:28 -0400 Subject: [PATCH 57/73] Implement review suggestions --- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 4 ++- .../gempyor_pkg/src/gempyor/distributions.py | 1 - .../src/gempyor/inference_parameter.py | 8 ++--- flepimop/gempyor_pkg/src/gempyor/outcomes.py | 4 +-- .../gempyor_pkg/src/gempyor/parameters.py | 6 ++-- .../_pydantic_ext/test__evaled_expression.py | 2 +- .../test_distributions_common.py | 12 +++++++ .../tests/parameters/test_parameters_class.py | 36 +------------------ 8 files changed, 23 insertions(+), 50 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index 650321797..a8e59b187 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -189,9 +189,11 @@ def _read_and_validate_dataframe( @overload def _evaled_expression(val: EE | str, target_type: Type[EE]) -> EE: ... + @overload def _evaled_expression(val: T, target_type: Type[EE]) -> T: ... + def _evaled_expression(val: EE | str | Any, target_type: Type[EE]) -> EE | Any: """ Evaluates a string expression to a target numeric type (int or float). @@ -229,7 +231,7 @@ def _evaled_expression(val: EE | str | Any, target_type: Type[EE]) -> EE | Any: if isinstance(val, str): expr = parse_expr(val) if not expr.is_Number: - raise ValueError(f"Cannot convert expression '{expr}' to {target_type}.") + raise ValueError(f"Cannot convert expression '{expr}' to {target_type}.") return target_type(expr) return val diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index d8bac6ea3..917a03282 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -59,7 +59,6 @@ def __call__(self) -> float | int: """A shortcut for `self.sample(size=1)`.""" return self.sample(size=1).item() - @abstractmethod def _sample_from_generator( self, size: int | tuple[int, ...], rng: Generator diff --git a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py index 494a9800f..9bed34792 100644 --- a/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py +++ b/flepimop/gempyor_pkg/src/gempyor/inference_parameter.py @@ -59,9 +59,7 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=sp, - pdist=distribution_from_confuse_config( - parameter_config["value"] - ), + pdist=distribution_from_confuse_config(parameter_config["value"]), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) @@ -73,9 +71,7 @@ def add_modifier(self, pname, ptype, parameter_config, subpops): ptype=ptype, pname=pname, subpop=",".join(group), - pdist=distribution_from_confuse_config( - parameter_config["value"] - ), + pdist=distribution_from_confuse_config(parameter_config["value"]), lb=parameter_config["value"]["a"].get(float), ub=parameter_config["value"]["b"].get(float), ) diff --git a/flepimop/gempyor_pkg/src/gempyor/outcomes.py b/flepimop/gempyor_pkg/src/gempyor/outcomes.py index 6d593b140..6ba061bf5 100644 --- a/flepimop/gempyor_pkg/src/gempyor/outcomes.py +++ b/flepimop/gempyor_pkg/src/gempyor/outcomes.py @@ -467,9 +467,7 @@ def compute_all_multioutcomes( else: # One draw for all subpops probabilities = np.repeat( - distribution_from_confuse_config( - parameters[new_comp]["probability"] - )(), + distribution_from_confuse_config(parameters[new_comp]["probability"])(), len(modinf.subpop_struct.subpop_names), ) if "rel_probability" in parameters[new_comp]: diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index dc27e2bc3..1e2d5bac3 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -22,7 +22,7 @@ import pydantic from . import NPI, utils -from .distributions import build_distribution_from_confuse_config +from .distributions import distribution_from_confuse_config logger = logging.getLogger(__name__) @@ -101,7 +101,7 @@ def __init__( # Parameter characterized by its distribution if self.pconfig[pn]["value"].exists(): - self.pdata[pn]["dist"] = build_distribution_from_confuse_config( + self.pdata[pn]["dist"] = distribution_from_confuse_config( self.pconfig[pn]["value"] ) @@ -186,7 +186,7 @@ def reinitialize_distributions(self) -> None: """ for pn in self.pnames: if "dist" in self.pdata[pn]: - self.pdata[pn]["dist"] = build_distribution_from_confuse_config( + self.pdata[pn]["dist"] = distribution_from_confuse_config( self.pconfig[pn]["value"] ) diff --git a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py index 0a20de6d8..682adbba7 100644 --- a/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py +++ b/flepimop/gempyor_pkg/tests/_pydantic_ext/test__evaled_expression.py @@ -44,5 +44,5 @@ def test_evaled_expression_returns_expected_results(val, target_type, expected): ], ) def test_evaled_expression_raises_value_error_for_invalid_input(val, target_type): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match=r"Cannot convert expression '.*' to .*"): _evaled_expression(val, target_type=target_type) diff --git a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py index e95a0bcf8..62a4ee20b 100644 --- a/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py +++ b/flepimop/gempyor_pkg/tests/distributions/test_distributions_common.py @@ -8,6 +8,10 @@ class DummyDistribution(DistributionABC): distribution: str = "dummy" + def __call__(self) -> float | int: + """A shortcut for `self.sample(size=1)`.""" + return self.sample(size=1).item() + def _sample_from_generator( self, size: int | tuple[int, ...], rng: np.random.Generator ) -> np.ndarray: @@ -47,3 +51,11 @@ def test_distribution_abc_sample_properties(size, expected_shape) -> None: assert isinstance(sample, np.ndarray) assert sample.shape == expected_shape assert np.issubdtype(sample.dtype, np.number) + + +def test_distribution_abc_callable() -> None: + dist = DummyDistribution() + sample_callable_result = dist() + assert isinstance( + sample_callable_result, float + ) # b/c rng.random always returns a float diff --git a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py index b70d3c75b..d2a3ac849 100644 --- a/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py +++ b/flepimop/gempyor_pkg/tests/parameters/test_parameters_class.py @@ -600,41 +600,7 @@ def test_parameters_load( else: dist_obj = params.pdata[param_name]["dist"] - drawn_values = p_draw[i, :, :] - - # For a fixed distribution, all values must be equal to the specified value - if isinstance(dist_obj, FixedDistribution): - assert np.allclose(drawn_values, dist_obj.value) - - # For bounded distributions, check that all values are within the domain - elif isinstance(dist_obj, UniformDistribution): - assert np.all(drawn_values >= dist_obj.low) - assert np.all(drawn_values <= dist_obj.high) - elif isinstance(dist_obj, TruncatedNormalDistribution): - assert np.all(drawn_values >= dist_obj.a) - assert np.all(drawn_values <= dist_obj.b) - elif isinstance(dist_obj, BetaDistribution): - assert np.all(drawn_values >= 0) - assert np.all(drawn_values <= 1) - - # For discrete distributions, check for integer type and domain - elif isinstance(dist_obj, BinomialDistribution): - assert np.all(drawn_values >= 0) - assert np.all(drawn_values <= dist_obj.n) - elif isinstance(dist_obj, PoissonDistribution): - assert np.all(drawn_values >= 0) - - # For distributions with a non-negative domain - elif isinstance( - dist_obj, - (LognormalDistribution, GammaDistribution, WeibullDistribution), - ): - assert np.all(drawn_values >= 0) - - # For distributions without a simple domain to check (Normal) - else: - assert isinstance(dist_obj, NormalDistribution) - assert not np.isnan(drawn_values).any() + assert isinstance(dist_obj, DistributionABC) @pytest.mark.parametrize( "factory,n_days,nsubpops", From 549b6c556d854cd57d3c728374c15a61a145b7c0 Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Tue, 29 Jul 2025 14:25:10 -0400 Subject: [PATCH 58/73] Update flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py Co-authored-by: Carl A. B. Pearson --- flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py index 36035ef16..ede2ce848 100644 --- a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py +++ b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py @@ -158,7 +158,7 @@ def __createFromConfig(self, npi_config): ) if self.spatial_groups["grouped"]: for group in self.spatial_groups["grouped"]: - drawn_value = self.dist() * np.ones(len(group)) + drawn_value = np.repeat(self.dist(), len(group)) self.parameters.loc[group, "value"] = drawn_value def __createFromDf(self, loaded_df, npi_config): From 260dcb254597e4cc0c817bb7e8ca53bf9b67fe97 Mon Sep 17 00:00:00 2001 From: Emily Przykucki <100221052+emprzy@users.noreply.github.com> Date: Tue, 29 Jul 2025 14:27:10 -0400 Subject: [PATCH 59/73] Update flepimop/gempyor_pkg/src/gempyor/parameters.py Co-authored-by: Carl A. B. Pearson --- flepimop/gempyor_pkg/src/gempyor/parameters.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index 1e2d5bac3..31742ab7d 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -231,7 +231,7 @@ class via `subpop_names`. for idx, pn in enumerate(self.pnames): if "dist" in self.pdata[pn]: param_arr[idx] = np.full( - (n_days, nsubpops), self.pdata[pn]["dist"].sample().item() + (n_days, nsubpops), self.pdata[pn]["dist"]() ) else: param_arr[idx] = self.pdata[pn]["ts"].values From 9a2daea06a8ccd0eaabb0d625d11a28e2152fad3 Mon Sep 17 00:00:00 2001 From: Emily Przykucki Date: Tue, 29 Jul 2025 14:31:14 -0400 Subject: [PATCH 60/73] Update parameters.py --- flepimop/gempyor_pkg/src/gempyor/parameters.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/parameters.py b/flepimop/gempyor_pkg/src/gempyor/parameters.py index 31742ab7d..a7f164d75 100644 --- a/flepimop/gempyor_pkg/src/gempyor/parameters.py +++ b/flepimop/gempyor_pkg/src/gempyor/parameters.py @@ -230,9 +230,7 @@ class via `subpop_names`. for idx, pn in enumerate(self.pnames): if "dist" in self.pdata[pn]: - param_arr[idx] = np.full( - (n_days, nsubpops), self.pdata[pn]["dist"]() - ) + param_arr[idx] = np.full((n_days, nsubpops), self.pdata[pn]["dist"]()) else: param_arr[idx] = self.pdata[pn]["ts"].values @@ -280,9 +278,7 @@ def parameters_load( f"PARAM: parameter {pn} NOT found in loadID file. " "Drawing from config distribution" ) - param_arr[idx] = np.full( - (n_days, nsubpops), self.pdata[pn]["dist"].sample().item() - ) + param_arr[idx] = np.full((n_days, nsubpops), self.pdata[pn]["dist"]()) return param_arr From b24e9a3802fe2d0f61c55d0d944de8507bbae622 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Wed, 30 Jul 2025 16:49:39 -0400 Subject: [PATCH 61/73] add vectorization expeirments to dev --- .../src/gempyor/vectorization_experiments.py | 266 ++++++++++++++++++ .../tests/steps_rk4/test_vector_build.py | 229 +++++++++++++++ 2 files changed, 495 insertions(+) create mode 100644 flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py create mode 100644 flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py diff --git a/flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py b/flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py new file mode 100644 index 000000000..2e63a0c13 --- /dev/null +++ b/flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py @@ -0,0 +1,266 @@ +import numpy as np +from numba import njit, prange +from scipy.interpolate import interp1d + +# === Constants === +FLOAT_TOLERANCE = 1e-9 +_PARALLEL_THRESHOLD = 1e7 + +# === Helper for symbolic expression resolution === + +def resolve_parameter_expression(expr, param_lookup, today): + if "*" in expr: + terms = [s.strip() for s in expr.split("*")] + result = param_lookup[terms[0]].copy() + if result.ndim > 1: + result = result[:, today] + for term in terms[1:]: + val = param_lookup[term] + result *= val if val.ndim == 1 else val[:, today] + return result + else: + val = param_lookup[expr] + return val if val.ndim == 1 else val[:, today] + +# === 1. Core Proportion Logic === + +@njit(fastmath=False) +def prod_along_axis0(arr_2d): + n_cols = arr_2d.shape[1] + result = np.ones(n_cols, dtype=arr_2d.dtype) + for i in range(arr_2d.shape[0]): + for j in range(n_cols): + result[j] *= arr_2d[i, j] + return result + +@njit(fastmath=False) +def compute_proportion_sums_exponents(states_current, transitions, proportion_info, + transition_sum_compartments, parameters, today): + n_transitions = transitions.shape[1] + n_nodes = states_current.shape[1] + n_props = proportion_info.shape[1] + + total_rates = np.ones((n_transitions, n_nodes)) + source_numbers = np.zeros((n_transitions, n_nodes)) + proportion_contribs = np.zeros((n_props, n_nodes)) + + for t_idx in range(n_transitions): + p_start = transitions[3, t_idx] + p_stop = transitions[4, t_idx] + n_p = p_stop - p_start + first = True + + for i, p_idx in enumerate(range(p_start, p_stop)): + sum_start = proportion_info[0, p_idx] + sum_stop = proportion_info[1, p_idx] + expnt = parameters[proportion_info[2, p_idx], today] + summed = states_current[transition_sum_compartments[sum_start:sum_stop], :].sum(axis=0) + summed_exp = summed ** expnt + + if first: + source_numbers[t_idx] = summed + safe_src = np.where(summed > 0, summed, 1.0) + proportion_contribs[i, :] = summed_exp / safe_src + first = False + else: + proportion_contribs[i, :] = summed_exp + + if n_p > 0: + total_rates[t_idx, :] *= prod_along_axis0(proportion_contribs[p_start:p_stop, :]) + + return total_rates, source_numbers + +# === 2. Transition Rate Computation === + +def compute_transition_rates(total_rates_base, source_numbers, transitions, parameters, today, + percent_day_away, proportion_who_move, + mobility_data, mobility_data_indices, + mobility_row_indices, population, + param_expr_lookup=None): + n_transitions, n_nodes = total_rates_base.shape + total_rates = total_rates_base.copy() + + for t_idx in range(n_transitions): + p_start = transitions[3, t_idx] + p_stop = transitions[4, t_idx] + param_idx = transitions[2, t_idx] + + if (p_stop - p_start) == 1: + if param_expr_lookup is None: + param_val = parameters[param_idx, today] + else: + param_val = resolve_parameter_expression(param_expr_lookup[param_idx], parameters, today) + total_rates[t_idx] *= param_val + else: + for node in range(n_nodes): + pop_n = population[node] + pop_n_safe = pop_n if pop_n > 0 else 1.0 + + if param_expr_lookup is None: + param_node = parameters[param_idx, today][node] + else: + param_node = resolve_parameter_expression(param_expr_lookup[param_idx], parameters, today)[node] + + source_node = source_numbers[t_idx, node] + prop_keep = 1.0 - percent_day_away * proportion_who_move[node] + visitors_idx = slice(mobility_data_indices[node], mobility_data_indices[node + 1]) + visitors = mobility_row_indices[visitors_idx] + mobility_vals = mobility_data[visitors_idx] + prop_change = percent_day_away * mobility_vals / pop_n_safe + + rate_keep = prop_keep * source_node * param_node / pop_n if pop_n > 0 and param_node >= 0 else 0.0 + rate_change = 0.0 + for i in range(visitors.size): + v = visitors[i] + pop_v = population[v] + pop_v_safe = pop_v if pop_v > 0 else 1.0 + src_v = source_numbers[t_idx, v] + param_v = resolve_parameter_expression(param_expr_lookup[param_idx], parameters, today)[v] if param_expr_lookup else parameters[param_idx, today][v] + if pop_v > 0 and param_v >= 0: + rate_change += prop_change[i] * src_v * param_v / pop_v_safe + + total = rate_keep + rate_change + total_rates[t_idx, node] *= total if np.isfinite(total) else 0.0 + + return total_rates + +# === 3. Binomial Stochastic (NumPy) === + +def compute_transition_amounts_numpy_binomial(source_numbers, total_rates, dt): + probs = 1.0 - np.exp(-dt * total_rates) + draws = np.random.binomial(source_numbers.astype(np.int32), probs) + return draws.astype(np.float32) + +# === 4. Transition Amount Dispatch === + +@njit +def compute_transition_amounts_meta(source_numbers, total_rates, method, dt): + workload = source_numbers.shape[0] * source_numbers.shape[1] + if method == "stochastic": + raise RuntimeError("Stochastic method must be handled outside JIT") + elif workload >= _PARALLEL_THRESHOLD: + return compute_transition_amounts_parallel(source_numbers, total_rates, dt) + else: + return compute_transition_amounts_serial(source_numbers, total_rates, method, dt) + +@njit(fastmath=False) +def compute_transition_amounts_serial(source_numbers, total_rates, method, dt): + n_transitions, n_nodes = total_rates.shape + amounts = np.zeros((n_transitions, n_nodes)) + for t_idx in range(n_transitions): + for node in range(n_nodes): + if method == "rk4": + amounts[t_idx, node] = source_numbers[t_idx, node] * total_rates[t_idx, node] + elif method == "euler": + rate = 1 - np.exp(-dt * total_rates[t_idx, node]) + amounts[t_idx, node] = source_numbers[t_idx, node] * rate + return amounts + +@njit(parallel=True, fastmath=False) +def compute_transition_amounts_parallel(source_numbers, total_rates, dt): + n_transitions, n_nodes = total_rates.shape + amounts = np.zeros((n_transitions, n_nodes)) + for t_idx in prange(n_transitions): + for node in range(n_nodes): + rate = 1 - np.exp(-dt * total_rates[t_idx, node]) + amounts[t_idx, node] = source_numbers[t_idx, node] * rate + return amounts + +# === 5. Assemble Flux Vector === + +@njit(fastmath=False) +def assemble_flux(amounts, transitions, ncompartments, nspatial_nodes): + dy_dt = np.zeros((ncompartments, nspatial_nodes)) + for t_idx in range(amounts.shape[0]): + src = transitions[0, t_idx] + dst = transitions[1, t_idx] + for node in range(nspatial_nodes): + dy_dt[src, node] -= amounts[t_idx, node] + dy_dt[dst, node] += amounts[t_idx, node] + return dy_dt.ravel() + +# === 6. RHS Builder === + +def build_rhs(ncompartments, nspatial_nodes, transitions, proportion_info, + transition_sum_compartments, parameters, method, dt, + percent_day_away, proportion_who_move, + mobility_data, mobility_data_indices, + mobility_row_indices, population, + param_expr_lookup=None): + def rhs(t, y): + today = int(t) + states_current = y.reshape((ncompartments, nspatial_nodes)) + + total_base, source_numbers = compute_proportion_sums_exponents( + states_current, transitions, proportion_info, + transition_sum_compartments, parameters, today + ) + + total_rates = compute_transition_rates( + total_base, source_numbers, transitions, parameters, today, + percent_day_away, proportion_who_move, + mobility_data, mobility_data_indices, + mobility_row_indices, population, + param_expr_lookup=param_expr_lookup + ) + + if method == "stochastic": + amounts = compute_transition_amounts_numpy_binomial(source_numbers, total_rates, dt) + else: + amounts = compute_transition_amounts_meta(source_numbers, total_rates, method, dt) + + return assemble_flux(amounts, transitions, ncompartments, nspatial_nodes) + + return rhs + +# === 7. Solver Step Functions === + +@njit(inline='always', fastmath=False) +def update(y, delta_t, dy): + return y + delta_t * dy + +def rk4_step(rhs_fn, y, t, dt): + k1 = rhs_fn(t, y) + k2 = rhs_fn(t + dt / 2, update(y, dt / 2, k1)) + k3 = rhs_fn(t + dt / 2, update(y, dt / 2, k2)) + k4 = rhs_fn(t + dt, update(y, dt, k3)) + return update(y, dt / 6, k1 + 2 * k2 + 2 * k3 + k4) + +def euler_or_stochastic_step(rhs_fn, y, t, dt): + return update(y, dt, rhs_fn(t, y)) + +# === 8. Solver Interface === + +def run_solver(rhs_fn, y0, t_grid, method="rk4", + record_daily=False, ncompartments=None, nspatial_nodes=None): + state_shape = y0.shape + y = y0.copy().ravel() + n_steps = len(t_grid) + states = np.zeros((n_steps, *state_shape)) + incid = np.zeros((n_steps, *state_shape)) if record_daily else None + + for i, t in enumerate(t_grid): + today = int(np.floor(t)) + states[i] = y.reshape(state_shape) + + if record_daily and method != "rk4": + dy = rhs_fn(t, y).reshape(state_shape) + incid[i] = np.maximum(dy, 0.0) + + if i < n_steps - 1: + dt = t_grid[i + 1] - t + if method == "rk4": + y = rk4_step(rhs_fn, y, t, dt) + elif method in {"euler", "stochastic"}: + y = euler_or_stochastic_step(rhs_fn, y, t, dt) + else: + raise ValueError(f"Unknown method: {method}") + + if t_grid[1] - t_grid[0] == 2.0: + interp = interp1d(t_grid[::2], states[::2], axis=0, kind="linear", fill_value="extrapolate") + states = interp(t_grid) + if record_daily: + incid /= 2 + incid[1::2] = incid[:-1:2] + + return states, incid diff --git a/flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py b/flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py new file mode 100644 index 000000000..95eecd1ce --- /dev/null +++ b/flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py @@ -0,0 +1,229 @@ +import pytest +import numpy as np +import shutil +from pathlib import Path +import confuse +from scipy.sparse import csr_matrix + +from gempyor.model_info import ModelInfo +from gempyor.vectorization_experiments import ( + compute_proportion_sums_exponents, + compute_transition_rates, + compute_transition_amounts_meta, + assemble_flux, + run_solver, + build_rhs, +) + +@pytest.fixture(scope="module") +def modelinfo_from_config(tmp_path_factory): + tmp_path = tmp_path_factory.mktemp("model_input") + + original_dir = Path(__file__).resolve() + while original_dir.name != "flepimop": + original_dir = original_dir.parent + tutorial_dir = original_dir.parent / "examples/tutorials" + + input_dir = tmp_path / "model_input" + input_dir.mkdir() + + input_files = [ + "geodata_sample_2pop.csv", + "ic_2pop.csv", + "mobility_sample_2pop.csv", + ] + for fname in input_files: + shutil.copyfile(tutorial_dir / "model_input" / fname, input_dir / fname) + + config_file = "config_sample_2pop_modifiers.yml" + config_path = tmp_path / config_file + shutil.copyfile(tutorial_dir / config_file, config_path) + + config = confuse.Configuration("TestModel", __name__) + config.set_file(str(config_path)) + + return ModelInfo( + config=config, + config_filepath=str(config_path), + path_prefix=str(tmp_path), + setup_name="sample_2pop", + seir_modifiers_scenario="Ro_all", + ), config + +@pytest.fixture +def model_and_inputs(modelinfo_from_config): + model, config = modelinfo_from_config + + initial_array = model.initial_conditions.get_from_config(sim_id=0, modinf=model) + + (unique_strings, + transitions, + transition_sum_compartments, + proportion_info + ) = model.compartments.get_transition_array() + + + # This is the correct call — no .parameter_names used + parsed_params = model.compartments.parse_parameters( + model.parameters.parameters_quick_draw(model.n_days, model.nsubpops), + config["seir"]["parameters"].get(), + unique_strings + ) + + # === Mobility setup === + mobility_csr: csr_matrix = model.mobility + population = model.subpop_pop + + mobility_data = mobility_csr.data + mobility_data_indices = mobility_csr.indptr + mobility_row_indices = mobility_csr.indices + + proportion_who_move = np.zeros(model.nsubpops) + for i in range(model.nsubpops): + total_flux = mobility_data[mobility_data_indices[i]:mobility_data_indices[i + 1]].sum() + proportion_who_move[i] = min(total_flux / population[i], 1.0) + + # === Time grid === + offset = model.ti.toordinal() + t0 = 0 # start at day 0 + t1 = model.tf.toordinal() - offset # number of days between tf and ti + dt = 0.1 + n_steps = int((t1 - t0) / dt) + 1 + time_grid = np.linspace(t0, t1, n_steps).astype(np.float32) + + out= { + "model": model, + "initial_array": initial_array, + "params": parsed_params, + "transitions": transitions, + "proportion_info": proportion_info, + "transition_sum_compartments": transition_sum_compartments, + "mobility_data": mobility_data, + "mobility_data_indices": mobility_data_indices, + "mobility_row_indices": mobility_row_indices, + "proportion_who_move": proportion_who_move, + "population": population, + "time_grid": time_grid, + "dt": dt, + "percent_day_away": 0.5, + "today_idx": 0, + } + + return out + + + +import numpy as np +import pytest + +from gempyor.vectorization_experiments import ( + compute_proportion_sums_exponents, + compute_transition_rates, + compute_transition_amounts_meta, + assemble_flux, + run_solver, + build_rhs, +) + + +def test_proportion_sums_and_sources(model_and_inputs): + out = model_and_inputs + prop_sums, source_nums = compute_proportion_sums_exponents( + out["initial_array"], out["transitions"], out["proportion_info"], + out["transition_sum_compartments"], out["params"], out["today_idx"] + ) + assert prop_sums.shape == source_nums.shape + assert np.all(np.isfinite(prop_sums)) + assert np.all(source_nums >= 0) + + +def test_transition_rate_shape_and_finiteness(model_and_inputs): + out = model_and_inputs + + + prop_sums, source_nums = compute_proportion_sums_exponents( + out["initial_array"], out["transitions"], out["proportion_info"], + out["transition_sum_compartments"], out["params"], out["today_idx"] + ) + rates = compute_transition_rates( + prop_sums, source_nums, out["transitions"], out["params"], out["today_idx"], + out["percent_day_away"], out["proportion_who_move"], + out["mobility_data"], out["mobility_data_indices"], + out["mobility_row_indices"], out["population"] + ) + assert rates.shape == source_nums.shape + assert np.all(np.isfinite(rates)) + + +def test_transition_amounts_meta_rk4(model_and_inputs): + out = model_and_inputs + _, source_nums = compute_proportion_sums_exponents( + out["initial_array"], out["transitions"], out["proportion_info"], + out["transition_sum_compartments"], out["params"], out["today_idx"] + ) + rates = np.ones_like(source_nums) + amounts = compute_transition_amounts_meta(source_nums, rates, method="rk4", dt=out["dt"]) + assert amounts.shape == rates.shape + assert np.all(amounts >= 0) + + +def test_flux_assembly_validity(model_and_inputs): + out = model_and_inputs + ncomp, nloc = out["initial_array"].shape + ntrans = len(out["transitions"]) + mock_amounts = np.random.rand(ntrans, nloc).astype(np.float32) + flux = assemble_flux(mock_amounts, out["transitions"], ncomp, nloc) + assert flux.shape == (ncomp * nloc,) + assert np.all(np.isfinite(flux)) + + +def test_run_solver_vectorized_backend(model_and_inputs): + out = model_and_inputs + ncomp, nloc = out["initial_array"].shape + rhs_fn = build_rhs( + ncomp, nloc, out["transitions"], out["proportion_info"], + out["transition_sum_compartments"], out["params"], "rk4", out["dt"], + out["percent_day_away"], out["proportion_who_move"], + out["mobility_data"], out["mobility_data_indices"], + out["mobility_row_indices"], out["population"] + ) + states, fluxes = run_solver(rhs_fn, out["initial_array"], out["time_grid"], + method="rk4", record_daily=True, + ncompartments=ncomp, nspatial_nodes=nloc) + + + + assert states.shape == (len(out["time_grid"]), ncomp, nloc) + assert fluxes.shape == states.shape + assert np.all(np.isfinite(states)) + assert np.all(np.isfinite(fluxes)) + + +def test_rhs_output_matches_input_shape(model_and_inputs): + out = model_and_inputs + ncomp, nloc = out["initial_array"].shape + rhs_fn = build_rhs( + ncomp, nloc, out["transitions"], out["proportion_info"], + out["transition_sum_compartments"], out["params"], "rk4", out["dt"], + out["percent_day_away"], out["proportion_who_move"], + out["mobility_data"], out["mobility_data_indices"], + out["mobility_row_indices"], out["population"] + ) + flat_y = out["initial_array"].ravel() + dy = rhs_fn(out["today_idx"], flat_y) + assert dy.shape == flat_y.shape + assert np.all(np.isfinite(dy)) + + +def test_mass_conservation_for_dummy_flux(model_and_inputs): + out = model_and_inputs + ncomp, nloc = out["initial_array"].shape + _, source_nums = compute_proportion_sums_exponents( + out["initial_array"], out["transitions"], out["proportion_info"], + out["transition_sum_compartments"], out["params"], out["today_idx"] + ) + dummy_rates = np.ones_like(source_nums) + dummy_amounts = compute_transition_amounts_meta(source_nums, dummy_rates, method="rk4", dt=out["dt"]) + flux = assemble_flux(dummy_amounts, out["transitions"], ncomp, nloc) + total_flux_per_node = flux.reshape((ncomp, nloc)).sum(axis=0) + assert np.allclose(total_flux_per_node, 0.0, atol=1e-5) From 9fc2b6a960deede31f2d43088c111eec1efa1568 Mon Sep 17 00:00:00 2001 From: "J.C. Macdonald" Date: Wed, 30 Jul 2025 17:05:08 -0400 Subject: [PATCH 62/73] Revert "add vectorization expeirments to dev" This reverts commit b24e9a3802fe2d0f61c55d0d944de8507bbae622. --- .../src/gempyor/vectorization_experiments.py | 266 ------------------ .../tests/steps_rk4/test_vector_build.py | 229 --------------- 2 files changed, 495 deletions(-) delete mode 100644 flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py delete mode 100644 flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py diff --git a/flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py b/flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py deleted file mode 100644 index 2e63a0c13..000000000 --- a/flepimop/gempyor_pkg/src/gempyor/vectorization_experiments.py +++ /dev/null @@ -1,266 +0,0 @@ -import numpy as np -from numba import njit, prange -from scipy.interpolate import interp1d - -# === Constants === -FLOAT_TOLERANCE = 1e-9 -_PARALLEL_THRESHOLD = 1e7 - -# === Helper for symbolic expression resolution === - -def resolve_parameter_expression(expr, param_lookup, today): - if "*" in expr: - terms = [s.strip() for s in expr.split("*")] - result = param_lookup[terms[0]].copy() - if result.ndim > 1: - result = result[:, today] - for term in terms[1:]: - val = param_lookup[term] - result *= val if val.ndim == 1 else val[:, today] - return result - else: - val = param_lookup[expr] - return val if val.ndim == 1 else val[:, today] - -# === 1. Core Proportion Logic === - -@njit(fastmath=False) -def prod_along_axis0(arr_2d): - n_cols = arr_2d.shape[1] - result = np.ones(n_cols, dtype=arr_2d.dtype) - for i in range(arr_2d.shape[0]): - for j in range(n_cols): - result[j] *= arr_2d[i, j] - return result - -@njit(fastmath=False) -def compute_proportion_sums_exponents(states_current, transitions, proportion_info, - transition_sum_compartments, parameters, today): - n_transitions = transitions.shape[1] - n_nodes = states_current.shape[1] - n_props = proportion_info.shape[1] - - total_rates = np.ones((n_transitions, n_nodes)) - source_numbers = np.zeros((n_transitions, n_nodes)) - proportion_contribs = np.zeros((n_props, n_nodes)) - - for t_idx in range(n_transitions): - p_start = transitions[3, t_idx] - p_stop = transitions[4, t_idx] - n_p = p_stop - p_start - first = True - - for i, p_idx in enumerate(range(p_start, p_stop)): - sum_start = proportion_info[0, p_idx] - sum_stop = proportion_info[1, p_idx] - expnt = parameters[proportion_info[2, p_idx], today] - summed = states_current[transition_sum_compartments[sum_start:sum_stop], :].sum(axis=0) - summed_exp = summed ** expnt - - if first: - source_numbers[t_idx] = summed - safe_src = np.where(summed > 0, summed, 1.0) - proportion_contribs[i, :] = summed_exp / safe_src - first = False - else: - proportion_contribs[i, :] = summed_exp - - if n_p > 0: - total_rates[t_idx, :] *= prod_along_axis0(proportion_contribs[p_start:p_stop, :]) - - return total_rates, source_numbers - -# === 2. Transition Rate Computation === - -def compute_transition_rates(total_rates_base, source_numbers, transitions, parameters, today, - percent_day_away, proportion_who_move, - mobility_data, mobility_data_indices, - mobility_row_indices, population, - param_expr_lookup=None): - n_transitions, n_nodes = total_rates_base.shape - total_rates = total_rates_base.copy() - - for t_idx in range(n_transitions): - p_start = transitions[3, t_idx] - p_stop = transitions[4, t_idx] - param_idx = transitions[2, t_idx] - - if (p_stop - p_start) == 1: - if param_expr_lookup is None: - param_val = parameters[param_idx, today] - else: - param_val = resolve_parameter_expression(param_expr_lookup[param_idx], parameters, today) - total_rates[t_idx] *= param_val - else: - for node in range(n_nodes): - pop_n = population[node] - pop_n_safe = pop_n if pop_n > 0 else 1.0 - - if param_expr_lookup is None: - param_node = parameters[param_idx, today][node] - else: - param_node = resolve_parameter_expression(param_expr_lookup[param_idx], parameters, today)[node] - - source_node = source_numbers[t_idx, node] - prop_keep = 1.0 - percent_day_away * proportion_who_move[node] - visitors_idx = slice(mobility_data_indices[node], mobility_data_indices[node + 1]) - visitors = mobility_row_indices[visitors_idx] - mobility_vals = mobility_data[visitors_idx] - prop_change = percent_day_away * mobility_vals / pop_n_safe - - rate_keep = prop_keep * source_node * param_node / pop_n if pop_n > 0 and param_node >= 0 else 0.0 - rate_change = 0.0 - for i in range(visitors.size): - v = visitors[i] - pop_v = population[v] - pop_v_safe = pop_v if pop_v > 0 else 1.0 - src_v = source_numbers[t_idx, v] - param_v = resolve_parameter_expression(param_expr_lookup[param_idx], parameters, today)[v] if param_expr_lookup else parameters[param_idx, today][v] - if pop_v > 0 and param_v >= 0: - rate_change += prop_change[i] * src_v * param_v / pop_v_safe - - total = rate_keep + rate_change - total_rates[t_idx, node] *= total if np.isfinite(total) else 0.0 - - return total_rates - -# === 3. Binomial Stochastic (NumPy) === - -def compute_transition_amounts_numpy_binomial(source_numbers, total_rates, dt): - probs = 1.0 - np.exp(-dt * total_rates) - draws = np.random.binomial(source_numbers.astype(np.int32), probs) - return draws.astype(np.float32) - -# === 4. Transition Amount Dispatch === - -@njit -def compute_transition_amounts_meta(source_numbers, total_rates, method, dt): - workload = source_numbers.shape[0] * source_numbers.shape[1] - if method == "stochastic": - raise RuntimeError("Stochastic method must be handled outside JIT") - elif workload >= _PARALLEL_THRESHOLD: - return compute_transition_amounts_parallel(source_numbers, total_rates, dt) - else: - return compute_transition_amounts_serial(source_numbers, total_rates, method, dt) - -@njit(fastmath=False) -def compute_transition_amounts_serial(source_numbers, total_rates, method, dt): - n_transitions, n_nodes = total_rates.shape - amounts = np.zeros((n_transitions, n_nodes)) - for t_idx in range(n_transitions): - for node in range(n_nodes): - if method == "rk4": - amounts[t_idx, node] = source_numbers[t_idx, node] * total_rates[t_idx, node] - elif method == "euler": - rate = 1 - np.exp(-dt * total_rates[t_idx, node]) - amounts[t_idx, node] = source_numbers[t_idx, node] * rate - return amounts - -@njit(parallel=True, fastmath=False) -def compute_transition_amounts_parallel(source_numbers, total_rates, dt): - n_transitions, n_nodes = total_rates.shape - amounts = np.zeros((n_transitions, n_nodes)) - for t_idx in prange(n_transitions): - for node in range(n_nodes): - rate = 1 - np.exp(-dt * total_rates[t_idx, node]) - amounts[t_idx, node] = source_numbers[t_idx, node] * rate - return amounts - -# === 5. Assemble Flux Vector === - -@njit(fastmath=False) -def assemble_flux(amounts, transitions, ncompartments, nspatial_nodes): - dy_dt = np.zeros((ncompartments, nspatial_nodes)) - for t_idx in range(amounts.shape[0]): - src = transitions[0, t_idx] - dst = transitions[1, t_idx] - for node in range(nspatial_nodes): - dy_dt[src, node] -= amounts[t_idx, node] - dy_dt[dst, node] += amounts[t_idx, node] - return dy_dt.ravel() - -# === 6. RHS Builder === - -def build_rhs(ncompartments, nspatial_nodes, transitions, proportion_info, - transition_sum_compartments, parameters, method, dt, - percent_day_away, proportion_who_move, - mobility_data, mobility_data_indices, - mobility_row_indices, population, - param_expr_lookup=None): - def rhs(t, y): - today = int(t) - states_current = y.reshape((ncompartments, nspatial_nodes)) - - total_base, source_numbers = compute_proportion_sums_exponents( - states_current, transitions, proportion_info, - transition_sum_compartments, parameters, today - ) - - total_rates = compute_transition_rates( - total_base, source_numbers, transitions, parameters, today, - percent_day_away, proportion_who_move, - mobility_data, mobility_data_indices, - mobility_row_indices, population, - param_expr_lookup=param_expr_lookup - ) - - if method == "stochastic": - amounts = compute_transition_amounts_numpy_binomial(source_numbers, total_rates, dt) - else: - amounts = compute_transition_amounts_meta(source_numbers, total_rates, method, dt) - - return assemble_flux(amounts, transitions, ncompartments, nspatial_nodes) - - return rhs - -# === 7. Solver Step Functions === - -@njit(inline='always', fastmath=False) -def update(y, delta_t, dy): - return y + delta_t * dy - -def rk4_step(rhs_fn, y, t, dt): - k1 = rhs_fn(t, y) - k2 = rhs_fn(t + dt / 2, update(y, dt / 2, k1)) - k3 = rhs_fn(t + dt / 2, update(y, dt / 2, k2)) - k4 = rhs_fn(t + dt, update(y, dt, k3)) - return update(y, dt / 6, k1 + 2 * k2 + 2 * k3 + k4) - -def euler_or_stochastic_step(rhs_fn, y, t, dt): - return update(y, dt, rhs_fn(t, y)) - -# === 8. Solver Interface === - -def run_solver(rhs_fn, y0, t_grid, method="rk4", - record_daily=False, ncompartments=None, nspatial_nodes=None): - state_shape = y0.shape - y = y0.copy().ravel() - n_steps = len(t_grid) - states = np.zeros((n_steps, *state_shape)) - incid = np.zeros((n_steps, *state_shape)) if record_daily else None - - for i, t in enumerate(t_grid): - today = int(np.floor(t)) - states[i] = y.reshape(state_shape) - - if record_daily and method != "rk4": - dy = rhs_fn(t, y).reshape(state_shape) - incid[i] = np.maximum(dy, 0.0) - - if i < n_steps - 1: - dt = t_grid[i + 1] - t - if method == "rk4": - y = rk4_step(rhs_fn, y, t, dt) - elif method in {"euler", "stochastic"}: - y = euler_or_stochastic_step(rhs_fn, y, t, dt) - else: - raise ValueError(f"Unknown method: {method}") - - if t_grid[1] - t_grid[0] == 2.0: - interp = interp1d(t_grid[::2], states[::2], axis=0, kind="linear", fill_value="extrapolate") - states = interp(t_grid) - if record_daily: - incid /= 2 - incid[1::2] = incid[:-1:2] - - return states, incid diff --git a/flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py b/flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py deleted file mode 100644 index 95eecd1ce..000000000 --- a/flepimop/gempyor_pkg/tests/steps_rk4/test_vector_build.py +++ /dev/null @@ -1,229 +0,0 @@ -import pytest -import numpy as np -import shutil -from pathlib import Path -import confuse -from scipy.sparse import csr_matrix - -from gempyor.model_info import ModelInfo -from gempyor.vectorization_experiments import ( - compute_proportion_sums_exponents, - compute_transition_rates, - compute_transition_amounts_meta, - assemble_flux, - run_solver, - build_rhs, -) - -@pytest.fixture(scope="module") -def modelinfo_from_config(tmp_path_factory): - tmp_path = tmp_path_factory.mktemp("model_input") - - original_dir = Path(__file__).resolve() - while original_dir.name != "flepimop": - original_dir = original_dir.parent - tutorial_dir = original_dir.parent / "examples/tutorials" - - input_dir = tmp_path / "model_input" - input_dir.mkdir() - - input_files = [ - "geodata_sample_2pop.csv", - "ic_2pop.csv", - "mobility_sample_2pop.csv", - ] - for fname in input_files: - shutil.copyfile(tutorial_dir / "model_input" / fname, input_dir / fname) - - config_file = "config_sample_2pop_modifiers.yml" - config_path = tmp_path / config_file - shutil.copyfile(tutorial_dir / config_file, config_path) - - config = confuse.Configuration("TestModel", __name__) - config.set_file(str(config_path)) - - return ModelInfo( - config=config, - config_filepath=str(config_path), - path_prefix=str(tmp_path), - setup_name="sample_2pop", - seir_modifiers_scenario="Ro_all", - ), config - -@pytest.fixture -def model_and_inputs(modelinfo_from_config): - model, config = modelinfo_from_config - - initial_array = model.initial_conditions.get_from_config(sim_id=0, modinf=model) - - (unique_strings, - transitions, - transition_sum_compartments, - proportion_info - ) = model.compartments.get_transition_array() - - - # This is the correct call — no .parameter_names used - parsed_params = model.compartments.parse_parameters( - model.parameters.parameters_quick_draw(model.n_days, model.nsubpops), - config["seir"]["parameters"].get(), - unique_strings - ) - - # === Mobility setup === - mobility_csr: csr_matrix = model.mobility - population = model.subpop_pop - - mobility_data = mobility_csr.data - mobility_data_indices = mobility_csr.indptr - mobility_row_indices = mobility_csr.indices - - proportion_who_move = np.zeros(model.nsubpops) - for i in range(model.nsubpops): - total_flux = mobility_data[mobility_data_indices[i]:mobility_data_indices[i + 1]].sum() - proportion_who_move[i] = min(total_flux / population[i], 1.0) - - # === Time grid === - offset = model.ti.toordinal() - t0 = 0 # start at day 0 - t1 = model.tf.toordinal() - offset # number of days between tf and ti - dt = 0.1 - n_steps = int((t1 - t0) / dt) + 1 - time_grid = np.linspace(t0, t1, n_steps).astype(np.float32) - - out= { - "model": model, - "initial_array": initial_array, - "params": parsed_params, - "transitions": transitions, - "proportion_info": proportion_info, - "transition_sum_compartments": transition_sum_compartments, - "mobility_data": mobility_data, - "mobility_data_indices": mobility_data_indices, - "mobility_row_indices": mobility_row_indices, - "proportion_who_move": proportion_who_move, - "population": population, - "time_grid": time_grid, - "dt": dt, - "percent_day_away": 0.5, - "today_idx": 0, - } - - return out - - - -import numpy as np -import pytest - -from gempyor.vectorization_experiments import ( - compute_proportion_sums_exponents, - compute_transition_rates, - compute_transition_amounts_meta, - assemble_flux, - run_solver, - build_rhs, -) - - -def test_proportion_sums_and_sources(model_and_inputs): - out = model_and_inputs - prop_sums, source_nums = compute_proportion_sums_exponents( - out["initial_array"], out["transitions"], out["proportion_info"], - out["transition_sum_compartments"], out["params"], out["today_idx"] - ) - assert prop_sums.shape == source_nums.shape - assert np.all(np.isfinite(prop_sums)) - assert np.all(source_nums >= 0) - - -def test_transition_rate_shape_and_finiteness(model_and_inputs): - out = model_and_inputs - - - prop_sums, source_nums = compute_proportion_sums_exponents( - out["initial_array"], out["transitions"], out["proportion_info"], - out["transition_sum_compartments"], out["params"], out["today_idx"] - ) - rates = compute_transition_rates( - prop_sums, source_nums, out["transitions"], out["params"], out["today_idx"], - out["percent_day_away"], out["proportion_who_move"], - out["mobility_data"], out["mobility_data_indices"], - out["mobility_row_indices"], out["population"] - ) - assert rates.shape == source_nums.shape - assert np.all(np.isfinite(rates)) - - -def test_transition_amounts_meta_rk4(model_and_inputs): - out = model_and_inputs - _, source_nums = compute_proportion_sums_exponents( - out["initial_array"], out["transitions"], out["proportion_info"], - out["transition_sum_compartments"], out["params"], out["today_idx"] - ) - rates = np.ones_like(source_nums) - amounts = compute_transition_amounts_meta(source_nums, rates, method="rk4", dt=out["dt"]) - assert amounts.shape == rates.shape - assert np.all(amounts >= 0) - - -def test_flux_assembly_validity(model_and_inputs): - out = model_and_inputs - ncomp, nloc = out["initial_array"].shape - ntrans = len(out["transitions"]) - mock_amounts = np.random.rand(ntrans, nloc).astype(np.float32) - flux = assemble_flux(mock_amounts, out["transitions"], ncomp, nloc) - assert flux.shape == (ncomp * nloc,) - assert np.all(np.isfinite(flux)) - - -def test_run_solver_vectorized_backend(model_and_inputs): - out = model_and_inputs - ncomp, nloc = out["initial_array"].shape - rhs_fn = build_rhs( - ncomp, nloc, out["transitions"], out["proportion_info"], - out["transition_sum_compartments"], out["params"], "rk4", out["dt"], - out["percent_day_away"], out["proportion_who_move"], - out["mobility_data"], out["mobility_data_indices"], - out["mobility_row_indices"], out["population"] - ) - states, fluxes = run_solver(rhs_fn, out["initial_array"], out["time_grid"], - method="rk4", record_daily=True, - ncompartments=ncomp, nspatial_nodes=nloc) - - - - assert states.shape == (len(out["time_grid"]), ncomp, nloc) - assert fluxes.shape == states.shape - assert np.all(np.isfinite(states)) - assert np.all(np.isfinite(fluxes)) - - -def test_rhs_output_matches_input_shape(model_and_inputs): - out = model_and_inputs - ncomp, nloc = out["initial_array"].shape - rhs_fn = build_rhs( - ncomp, nloc, out["transitions"], out["proportion_info"], - out["transition_sum_compartments"], out["params"], "rk4", out["dt"], - out["percent_day_away"], out["proportion_who_move"], - out["mobility_data"], out["mobility_data_indices"], - out["mobility_row_indices"], out["population"] - ) - flat_y = out["initial_array"].ravel() - dy = rhs_fn(out["today_idx"], flat_y) - assert dy.shape == flat_y.shape - assert np.all(np.isfinite(dy)) - - -def test_mass_conservation_for_dummy_flux(model_and_inputs): - out = model_and_inputs - ncomp, nloc = out["initial_array"].shape - _, source_nums = compute_proportion_sums_exponents( - out["initial_array"], out["transitions"], out["proportion_info"], - out["transition_sum_compartments"], out["params"], out["today_idx"] - ) - dummy_rates = np.ones_like(source_nums) - dummy_amounts = compute_transition_amounts_meta(source_nums, dummy_rates, method="rk4", dt=out["dt"]) - flux = assemble_flux(dummy_amounts, out["transitions"], ncomp, nloc) - total_flux_per_node = flux.reshape((ncomp, nloc)).sum(axis=0) - assert np.allclose(total_flux_per_node, 0.0, atol=1e-5) From 1c0370271f5b476119eed1466b4165f73295bd20 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Thu, 17 Jul 2025 11:34:55 -0400 Subject: [PATCH 63/73] Pull flepiMoP install info out of install Pulled the script that generates info on the flepiMoP install out of `bin/flepimop-install` into its own `bin/flepimop-install-info` for standalone use. Also updated installation script to also install the tests for the R packages. --- bin/flepimop-install | 24 ++------------------ bin/flepimop-install-info | 41 ++++++++++++++++++++++++++++++++++ build/create_environment_yml.R | 1 + 3 files changed, 44 insertions(+), 22 deletions(-) create mode 100755 bin/flepimop-install-info diff --git a/bin/flepimop-install b/bin/flepimop-install index faecc7d01..5b874d9ca 100755 --- a/bin/flepimop-install +++ b/bin/flepimop-install @@ -179,7 +179,7 @@ cd $FLEPI_PATH/flepimop/R_packages/ R_PACKAGES=("flepicommon" "flepiconfig" "inference") for pkg in "${R_PACKAGES[@]}"; do if [ -d "$pkg" ]; then - R_LIBS=$CONDA_R_LIBS R CMD INSTALL $pkg + R_LIBS=$CONDA_R_LIBS R CMD INSTALL --install-tests $pkg else echo "WARNING: The R package '$pkg' does not exist in '$FLEPI_PATH/flepimop/R_packages/'. Skipping installation." >&2 fi @@ -188,27 +188,7 @@ cd $RETURNTO R -e "library(inference); inference::install_cli()" # Output a summary of the installation -WHICH_R=$( which R ) -WHICH_PYTHON=$( which python ) -CONDA_VERSION=$( conda --version | cut -d' ' -f2 ) -R_VERSION=$( R --version | head -n 1 | cut -d' ' -f3 ) -PYTHON_VERSION=$( python --version | cut -d' ' -f2 ) -GEMPYOR_VERSION=$( pip show gempyor | grep Version | cut -d' ' -f2 ) - -cat << EOF -flepiMoP installation summary: -> flepiMoP version: $( git -C $FLEPI_PATH rev-parse HEAD ) -> flepiMoP path: $FLEPI_PATH -> flepiMoP conda env: $FLEPI_CONDA -> conda: $CONDA_VERSION -> R $R_VERSION: $WHICH_R -> Python $PYTHON_VERSION: $WHICH_PYTHON -> gempyor version: $GEMPYOR_VERSION -EOF -for pkg in "${R_PACKAGES[@]}"; do - PKG_VERSION=$( R -s -e "print(packageVersion('$pkg'))" | grep -oE "[0-9]+\.[0-9]+\.[0-9]+" ) - echo "> R $pkg version: $PKG_VERSION" -done +FLEPI_PATH=$FLEPI_PATH FLEPI_CONDA=$FLEPI_CONDA $FLEPI_PATH/bin/flepimop-install-info cat << EOF To activate the flepimop conda environment, run: conda activate $FLEPI_CONDA diff --git a/bin/flepimop-install-info b/bin/flepimop-install-info new file mode 100755 index 000000000..0ce001d17 --- /dev/null +++ b/bin/flepimop-install-info @@ -0,0 +1,41 @@ +#!/usr/bin/env bash + +WHICH_R=$( which R ) +WHICH_PYTHON=$( which python ) +CONDA_VERSION=$( conda --version | cut -d' ' -f2 ) +R_VERSION=$( R --version | head -n 1 | cut -d' ' -f3 ) +PYTHON_VERSION=$( python --version | cut -d' ' -f2 ) +GEMPYOR_VERSION=$( pip show gempyor | grep Version | cut -d' ' -f2 ) +R_PACKAGES=("flepicommon" "flepiconfig" "inference") + +echo "flepiMoP installation summary:" + +if command -v git &> /dev/null && [ -d "$FLEPI_PATH/.git" ]; then + FLEPI_VERSION=$( git -C "$FLEPI_PATH" rev-parse HEAD ) + echo "> flepiMoP version: $FLEPI_VERSION" +else + echo "> flepiMoP version: *could not determine version*" +fi + +if test -n "$FLEPI_PATH"; then + echo "> flepiMoP path: $FLEPI_PATH" +else + echo "> flepiMoP path: *not set*" +fi + +if test -n "$FLEPI_CONDA"; then + echo "> flepiMoP conda env: $FLEPI_CONDA" +else + echo "> flepiMoP conda env: *not set*" +fi + +cat << EOF +> conda: $CONDA_VERSION +> R $R_VERSION: $WHICH_R +> Python $PYTHON_VERSION: $WHICH_PYTHON +> gempyor version: $GEMPYOR_VERSION +EOF +for pkg in "${R_PACKAGES[@]}"; do + PKG_VERSION=$( R -s -e "print(packageVersion('$pkg'))" | grep -oE "[0-9]+\.[0-9]+\.[0-9]+" ) + echo "> R $pkg version: $PKG_VERSION" +done diff --git a/build/create_environment_yml.R b/build/create_environment_yml.R index 961c233b0..8d812d17e 100644 --- a/build/create_environment_yml.R +++ b/build/create_environment_yml.R @@ -43,6 +43,7 @@ new_environment_yml <- c( "- pyarrow=17.0.0", "- r-arrow=17.0.0", "- r-sf", + "- r-testthat", paste0("- r-", dependencies) ) if (file.exists(environment_yml)) { From 82f82f74686dd1ba145092e14fcf29d51a1c2269 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Thu, 17 Jul 2025 10:24:18 -0400 Subject: [PATCH 64/73] Use conda and install scripts in CI workflows Created a custom `flepimop-install` action and incorporated it into the gempyor, flepicommon, and inference CI workflows. This custom action takes care of: * Installing apt packages required by some of the R dependencies, * Installing miniconda, * Either pulling a conda venv from cache or installing flepiMoP with the standard install scripts with a user specified python/R version, and * Displaying install information for debugging purposes. --- .github/actions/flepimop-install/action.yml | 53 ++++++++++++ .github/workflows/flepicommon-ci.yml | 66 +++------------ .github/workflows/gempyor-ci.yml | 59 ++++++------- .github/workflows/inference-ci.yml | 93 +++------------------ .github/workflows/lint.yml | 27 +++--- 5 files changed, 124 insertions(+), 174 deletions(-) create mode 100644 .github/actions/flepimop-install/action.yml diff --git a/.github/actions/flepimop-install/action.yml b/.github/actions/flepimop-install/action.yml new file mode 100644 index 000000000..ee59210a1 --- /dev/null +++ b/.github/actions/flepimop-install/action.yml @@ -0,0 +1,53 @@ +name: flepiMoP Install +description: Install flepiMoP and its dependencies +inputs: + python-version: + description: Python version to use + required: true + default: "3.11" + R-version: + description: R version to use + required: true + default: "4.3" +runs: + using: "composite" + steps: + - name: Install system dependencies + run: | + sudo apt update + sudo apt install \ + libcurl4-openssl-dev \ + libharfbuzz-dev \ + libfribidi-dev \ + libtiff5-dev \ + libudunits2-dev \ + libgdal-dev \ + libgeos-dev \ + libproj-dev \ + libfontconfig1-dev \ + pandoc + shell: bash + - name: Setup miniconda + uses: conda-incubator/setup-miniconda@v3 + with: + auto-update-conda: true + - name: Conda venv cache + id: conda-venv-cache + uses: actions/cache@v4 + with: + key: flepimop-conda-${{ runner.os }}-${{ inputs.R-version }}-${{ inputs.python-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE', 'flepimop/R_packages/flepiconfig/DESCRIPTION', 'flepimop/R_packages/flepiconfig/NAMESPACE', 'flepimop/R_packages/inference/DESCRIPTION', 'flepimop/R_packages/inference/NAMESPACE', 'flepimop/gempyor_pkg/pyproject.toml', 'environment.yml') }} + path: './venv' + - name: Install flepiMoP + if: steps.conda-venv-cache.outputs.cache-hit != 'true' + run: | + rm -rf venv/ + sed -i 's/^- python.*$/- python=${{ inputs.python-version}}/g' environment.yml + sed -i 's/^- r-base.*$/- r-base=${{ inputs.R-version}}/g' environment.yml + ./bin/flepimop-install-dev + shell: bash -el {0} + - name: flepiMoP install info + run: | + conda activate venv/ + export FLEPI_PATH=$(pwd) + ./bin/flepimop-install-info + shell: bash -el {0} diff --git a/.github/workflows/flepicommon-ci.yml b/.github/workflows/flepicommon-ci.yml index 0a0049089..7594f3e95 100644 --- a/.github/workflows/flepicommon-ci.yml +++ b/.github/workflows/flepicommon-ci.yml @@ -26,9 +26,10 @@ jobs: if: github.event_name != 'pull_request' || github.event.pull_request.draft == false strategy: matrix: - R-version: ["4.3.3"] + python-version: ["3.11"] + R-version: ["4.3"] steps: - - name: Checkout + - name: Checkout flepiMoP uses: actions/checkout@v4 with: lfs: true @@ -36,56 +37,15 @@ jobs: * !documentation/ sparse-checkout-cone-mode: false - - name: Setup R - uses: r-lib/actions/setup-r@v2 + persist-credentials: false + - name: Install flepiMoP + uses: ./.github/actions/flepimop-install with: - r-version: ${{ matrix.R-version }} - update-rtools: true - - name: Install System Dependencies - run: sudo apt install libcurl4-openssl-dev libharfbuzz-dev libfribidi-dev libtiff5-dev libudunits2-dev libgdal-dev libgeos-dev libproj-dev libfontconfig1-dev - - name: Determine R Library Location + python-version: ${{ matrix.python-version }} + R-version: ${{ matrix.R-version }} + - name: Run tests run: | - R_LIBPATH=$( R -s -e "cat(.libPaths()[1L])" | xargs ) - echo "R_LIBPATH=$R_LIBPATH" >> $GITHUB_ENV - R_LIBPATH_CKSUM=$( echo "$R_LIBPATH" | cksum | cut -d ' ' -f 1 ) - echo "R_LIBPATH_CKSUM=$R_LIBPATH_CKSUM" >> $GITHUB_ENV - CACHE_DATE=$( date -d "last Sunday" +%Y%m%d ) - echo "CACHE_DATE=$CACHE_DATE" >> $GITHUB_ENV - - name: R Library Cache - uses: actions/cache@v4 - with: - key: flepicommon-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE') }}-${{ env.R_LIBPATH_CKSUM }}-${{ env.CACHE_DATE }} - path: ${{ env.R_LIBPATH }} - restore-keys: | - flepicommon-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE') }}-${{ env.R_LIBPATH_CKSUM }}- - flepicommon-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE') }}- - flepicommon-rlibs-${{ runner.os }}-${{ matrix.R-version }}- - - name: Install R Dependencies - if: steps.r-library-cache.outputs.cache-hit != 'true' - run: | - install.packages( - "devtools", - repos = "https://cloud.r-project.org", - ) - library(devtools) - devtools::install_deps( - pkg = "flepimop/R_packages/flepicommon", - dependencies = TRUE - ) - install.packages("epidatr", repos = "https://cloud.r-project.org") - shell: Rscript {0} - - name: Install The flepicommon Package - run: | - devtools::install( - pkg = "flepimop/R_packages/flepicommon", - args = c(getOption("devtools.install.args"), "--install-tests"), - quick = TRUE, - dependencies = TRUE, - force = TRUE - ) - shell: Rscript {0} - - name: Run Tests - run: | - library(testthat) - test_package("flepicommon") - shell: Rscript {0} + conda activate venv/ + export FLEPI_PATH=$(pwd) + Rscript -e "library(testthat); test_package('flepicommon')" + shell: bash -el {0} diff --git a/.github/workflows/gempyor-ci.yml b/.github/workflows/gempyor-ci.yml index 1af373e62..6ae15f13b 100644 --- a/.github/workflows/gempyor-ci.yml +++ b/.github/workflows/gempyor-ci.yml @@ -29,8 +29,9 @@ jobs: strategy: matrix: python-version: ["3.10", "3.11"] + R-version: ["4.3"] steps: - - name: Checkout + - name: Checkout flepiMoP uses: actions/checkout@v4 with: lfs: true @@ -38,67 +39,69 @@ jobs: * !documentation/ sparse-checkout-cone-mode: false - - name: Setup Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + persist-credentials: false + - name: Install flepiMoP + uses: ./.github/actions/flepimop-install with: python-version: ${{ matrix.python-version }} - - name: Install gempyor - run: | - python -m pip install --upgrade pip - python -m pip install "flepimop/gempyor_pkg[dev]" - shell: bash - - name: Run not slow tests + R-version: ${{ matrix.R-version }} + - name: Run fast unit tests run: | + conda activate venv/ export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg pytest -m "not slow" --exitfirst - shell: bash - - name: Run slow tests + shell: bash -el {0} + - name: Run slow unit tests run: | + conda activate venv/ export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg pytest -m "slow" --exitfirst - shell: bash + shell: bash -el {0} - name: Run gempyor-cli integration tests from examples run: | + conda activate venv/ cd examples pytest --exitfirst - shell: bash + shell: bash -el {0} doc-updates: runs-on: ubuntu-latest if: github.event_name != 'pull_request' || github.event.pull_request.draft == false + strategy: + matrix: + python-version: ["3.11"] + R-version: ["4.3"] steps: - - name: Checkout + - name: Checkout flepiMoP uses: actions/checkout@v4 with: - ref: ${{ github.head_ref || github.ref }} lfs: true sparse-checkout: | * - !tests/ + !documentation/ sparse-checkout-cone-mode: false - - name: Setup Python 3.11 - uses: actions/setup-python@v5 + persist-credentials: false + - name: Install flepiMoP + uses: ./.github/actions/flepimop-install with: - python-version: '3.11' - - name: Install gempyor - run: | - python -m pip install --upgrade pip - python -m pip install "flepimop/gempyor_pkg[dev]" - shell: bash + python-version: ${{ matrix.python-version }} + R-version: ${{ matrix.R-version }} - name: Build gempyor documentation with Sphinx run: | + conda activate venv/ + export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg/docs make fullbuild - shell: bash - - name: Install pandoc - run: | - sudo apt-get update && sudo apt-get install -y pandoc + shell: bash -el {0} - name: Build click documentation with click-man run: | + conda activate venv/ + export FLEPI_PATH=$(pwd) cd flepimop/gempyor_pkg click-man flepimop --target docs/cli_man --man-version 1 python docs/click-man_parsing.py + shell: bash -el {0} - name: Remove temporal metadata from man page headers run: | sed -i 's/^\(\.TH [^"]* \)"[^"]*" "\(.*\)"/\1"" "\2"/' flepimop/gempyor_pkg/docs/cli_man/*.1 diff --git a/.github/workflows/inference-ci.yml b/.github/workflows/inference-ci.yml index b5bd53467..ac27beb48 100644 --- a/.github/workflows/inference-ci.yml +++ b/.github/workflows/inference-ci.yml @@ -32,10 +32,10 @@ jobs: if: github.event_name != 'pull_request' || github.event.pull_request.draft == false strategy: matrix: - R-version: ["4.3.3"] python-version: ["3.10", "3.11"] + R-version: ["4.3"] steps: - - name: Checkout + - name: Checkout flepiMoP uses: actions/checkout@v4 with: lfs: true @@ -43,86 +43,15 @@ jobs: * !documentation/ sparse-checkout-cone-mode: false - - name: Setup Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + persist-credentials: false + - name: Install flepiMoP + uses: ./.github/actions/flepimop-install with: python-version: ${{ matrix.python-version }} - - name: Install gempyor + R-version: ${{ matrix.R-version }} + - name: Run tests run: | - python -m pip install --upgrade pip - python -m pip install "flepimop/gempyor_pkg[test]" - shell: bash - - name: Setup R - uses: r-lib/actions/setup-r@v2 - with: - r-version: ${{ matrix.R-version }} - update-rtools: true - - name: Install System Dependencies - run: sudo apt install libcurl4-openssl-dev libharfbuzz-dev libfribidi-dev libtiff5-dev libudunits2-dev libgdal-dev libgeos-dev libproj-dev libfontconfig1-dev - - name: Determine R Library Location - run: | - R_LIBPATH=$( R -s -e "cat(.libPaths()[1L])" | xargs ) - echo "R_LIBPATH=$R_LIBPATH" >> $GITHUB_ENV - R_LIBPATH_CKSUM=$( echo "$R_LIBPATH" | cksum | cut -d ' ' -f 1 ) - echo "R_LIBPATH_CKSUM=$R_LIBPATH_CKSUM" >> $GITHUB_ENV - CACHE_DATE=$( date -d "last Sunday" +%Y%m%d ) - echo "CACHE_DATE=$CACHE_DATE" >> $GITHUB_ENV - - name: R Library Cache - uses: actions/cache@v4 - with: - key: inference-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE', 'flepimop/R_packages/inference/DESCRIPTION', 'flepimop/R_packages/inference/NAMESPACE') }}-${{ env.R_LIBPATH_CKSUM }}-${{ env.CACHE_DATE }} - path: ${{ env.R_LIBPATH }} - restore-keys: | - inference-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE', 'flepimop/R_packages/inference/DESCRIPTION', 'flepimop/R_packages/inference/NAMESPACE') }}-${{ env.R_LIBPATH_CKSUM }}- - inference-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE', 'flepimop/R_packages/inference/DESCRIPTION', 'flepimop/R_packages/inference/NAMESPACE') }}- - inference-rlibs-${{ runner.os }}-${{ matrix.R-version }}- - flepicommon-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE') }}-${{ env.R_LIBPATH_CKSUM }}- - flepicommon-rlibs-${{ runner.os }}-${{ matrix.R-version }}-${{ hashFiles('flepimop/R_packages/flepicommon/DESCRIPTION', 'flepimop/R_packages/flepicommon/NAMESPACE') }}- - flepicommon-rlibs-${{ runner.os }}-${{ matrix.R-version }}- - - name: Install R Dependencies For flepicommon - if: steps.r-library-cache.outputs.cache-hit != 'true' - run: | - install.packages( - "devtools", - repos = "https://cloud.r-project.org", - ) - library(devtools) - devtools::install_deps( - pkg = "flepimop/R_packages/flepicommon", - dependencies = TRUE - ) - install.packages("epidatr", repos = "https://cloud.r-project.org") - shell: Rscript {0} - - name: Install The flepicommon Package - run: | - devtools::install( - pkg = "flepimop/R_packages/flepicommon", - quick = TRUE, - dependencies = TRUE, - force = TRUE - ) - shell: Rscript {0} - - name: Install R Dependencies For inference - if: steps.r-library-cache.outputs.cache-hit != 'true' - run: | - devtools::install_deps( - pkg = "flepimop/R_packages/inference", - dependencies = TRUE - ) - install.packages("epidatr", repos = "https://cloud.r-project.org") - shell: Rscript {0} - - name: Install The inference Package - run: | - devtools::install( - pkg = "flepimop/R_packages/inference", - args = c(getOption("devtools.install.args"), "--install-tests"), - quick = TRUE, - dependencies = TRUE, - force = TRUE - ) - shell: Rscript {0} - - name: Run Tests - run: | - library(testthat) - test_package("inference") - shell: Rscript {0} + conda activate venv/ + export FLEPI_PATH=$(pwd) + Rscript -e "library(testthat); test_package('inference')" + shell: bash -el {0} diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 077cd66f1..e36fff5a1 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -53,30 +53,35 @@ jobs: if: github.event_name != 'pull_request' || github.event.pull_request.draft == false env: PYLINT_SRC: 'flepimop/gempyor_pkg/src/gempyor' - PYLINT_FAIL_UNDER: '4' + PYLINT_FAIL_UNDER: '5' PYLINT_RCFILE: 'flepimop/gempyor_pkg/.pylintrc' + strategy: + matrix: + python-version: ["3.11"] + R-version: ["4.3"] steps: - - name: Checkout + - name: Checkout flepiMoP uses: actions/checkout@v4 with: lfs: true sparse-checkout: | - flepimop/gempyor_pkg/ + * + !documentation/ sparse-checkout-cone-mode: false - - name: Setup Python - uses: actions/setup-python@v5 + persist-credentials: false + - name: Install flepiMoP + uses: ./.github/actions/flepimop-install with: - python-version: '3.11' - - name: Install Pylint - run: | - python -m pip install --upgrade pip - python -m pip install pylint - - name: Linting with Pylint + python-version: ${{ matrix.python-version }} + R-version: ${{ matrix.R-version }} + - name: Run pylint run: | + conda activate venv/ pylint ${{ env.PYLINT_SRC }} \ --fail-under ${{ env.PYLINT_FAIL_UNDER }} \ --rcfile ${{ env.PYLINT_RCFILE }} \ --verbose + shell: bash -el {0} check-info-json-schema: runs-on: ubuntu-latest if: github.event_name != 'pull_request' || github.event.pull_request.draft == false From 1c0eff26c88f0cb1fa6d4db2e4ca955454f2d0a1 Mon Sep 17 00:00:00 2001 From: TimothyWillard Date: Thu, 17 Jul 2025 19:26:17 +0000 Subject: [PATCH 65/73] Update `environment.yml` via GitHub action --- environment.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/environment.yml b/environment.yml index 19441ca8d..bef6b0394 100644 --- a/environment.yml +++ b/environment.yml @@ -1,4 +1,4 @@ -# Mon Jun 16 21:15:32 2025 UTC +# Thu Jul 17 19:26:17 2025 UTC channels: - conda-forge - defaults @@ -11,6 +11,7 @@ dependencies: - pyarrow=17.0.0 - r-arrow=17.0.0 - r-sf +- r-testthat - r-data.table - r-doParallel - r-dplyr From b0825df2ea8dccd11e1f8992286d036a38550fb7 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Thu, 17 Jul 2025 16:01:12 -0400 Subject: [PATCH 66/73] Add doctests to gempyor CI Added a step to the gempyor CI workflow to run doctests as well as fixed the current state of doctests throughout the gempyor package. Mostly quick fixes, more detailed and higher quality examples will be required in follow ups. --- .github/workflows/gempyor-ci.yml | 7 + .../python-guidelines-for-developers.md | 24 +- flepimop/gempyor_pkg/conftest.py | 18 ++ flepimop/gempyor_pkg/src/gempyor/_jinja.py | 2 +- .../gempyor_pkg/src/gempyor/_pydantic_ext.py | 19 +- .../gempyor_pkg/src/gempyor/batch/_helpers.py | 8 +- .../gempyor_pkg/src/gempyor/batch/manifest.py | 16 -- .../gempyor_pkg/src/gempyor/batch/systems.py | 27 +- .../gempyor_pkg/src/gempyor/batch/types.py | 77 +++--- .../src/gempyor/config_validator.py | 1 + .../gempyor_pkg/src/gempyor/distributions.py | 45 ++-- .../gempyor_pkg/src/gempyor/file_paths.py | 4 +- flepimop/gempyor_pkg/src/gempyor/info.py | 14 +- flepimop/gempyor_pkg/src/gempyor/logging.py | 2 +- .../gempyor_pkg/src/gempyor/shared_cli.py | 24 +- .../gempyor_pkg/src/gempyor/sync/_sync.py | 4 +- flepimop/gempyor_pkg/src/gempyor/testing.py | 2 +- flepimop/gempyor_pkg/src/gempyor/utils.py | 243 ++++++++---------- 18 files changed, 242 insertions(+), 295 deletions(-) create mode 100644 flepimop/gempyor_pkg/conftest.py diff --git a/.github/workflows/gempyor-ci.yml b/.github/workflows/gempyor-ci.yml index 6ae15f13b..7a5a3a054 100644 --- a/.github/workflows/gempyor-ci.yml +++ b/.github/workflows/gempyor-ci.yml @@ -45,6 +45,13 @@ jobs: with: python-version: ${{ matrix.python-version }} R-version: ${{ matrix.R-version }} + - name: Run doctests + run: | + conda activate venv/ + export FLEPI_PATH=$(pwd) + cd flepimop/gempyor_pkg + pytest --doctest-modules --exitfirst src/gempyor/ + shell: bash -el {0} - name: Run fast unit tests run: | conda activate venv/ diff --git a/documentation/gitbook/development/python-guidelines-for-developers.md b/documentation/gitbook/development/python-guidelines-for-developers.md index fe1a62f7b..87d476ceb 100644 --- a/documentation/gitbook/development/python-guidelines-for-developers.md +++ b/documentation/gitbook/development/python-guidelines-for-developers.md @@ -34,31 +34,31 @@ One of the current focus is to switch internal data types from dataframes and nu To run the tests suite locally, you'll need to install the gempyor package with build dependencies: ```bash -pip install "flepimop/gempyor_pkg[test]" +pip install "flepimop/gempyor_pkg[dev]" ``` -which installs the `pytest` and `mock` packages in addition to all other gempyor dependencies so that one can run tests. - -If you are running from a conda environment and installing with \`--no-deps\`, then you should make sure that these two packages are installed. - -Now you can try to run the gempyor test suite by running, from the `flepimop/gempyor_pkg` folder: +which installs the `pytest` package in addition to all other gempyor dependencies so that one can run tests. Now you can try to run the gempyor test suite by running, from the `flepimop/gempyor_pkg` folder: ```bash pytest ``` -If that works, then you are ready to develop gempyor. Feel free to open your first pull request. +If that works, then you are ready to develop gempyor. Feel free to open your first pull request. If you want more output on tests, e.g capturing standard output (print), you can use: + +```bash +pytest -vvv +``` -If you want more output on tests, e.g capturing standard output (print), you can use: +And to run just some subset of the tests (e.g here just the outcome tests), use: ```bash -pytest -vvvv +pytest -vvv -k outcomes ``` -and to run just some subset of the tests (e.g here just the outcome tests), use: +Furthermore, to ensure that the examples provided in the documentation are high quality we run the doctests in CI with: ```bash -pytest -vvvv -k outcomes +pytest --doctest-modules src/gempyor/ ``` For more details on how to use `pytest` please refer to their [usage guide](https://docs.pytest.org/en/latest/how-to/usage.html). @@ -86,4 +86,4 @@ For those using a Mac or Linux system for development, these commands are also a ```bash cp -f bin/pre-commit .git/hooks/ -``` \ No newline at end of file +``` diff --git a/flepimop/gempyor_pkg/conftest.py b/flepimop/gempyor_pkg/conftest.py new file mode 100644 index 000000000..2fb80203e --- /dev/null +++ b/flepimop/gempyor_pkg/conftest.py @@ -0,0 +1,18 @@ +"""Pytest configuration file.""" + +from collections.abc import Generator + +import pytest + + +@pytest.fixture(autouse=True) +def _docdir(request: pytest.FixtureRequest) -> Generator[None, None, None]: + # Run doctests in the temporary directory. + # https://stackoverflow.com/q/46962007 + doctest_plugin = request.config.pluginmanager.getplugin("doctest") + if isinstance(request.node, doctest_plugin.DoctestItem): + tmpdir = request.getfixturevalue("tmpdir") + with tmpdir.as_cwd(): + yield + else: + yield diff --git a/flepimop/gempyor_pkg/src/gempyor/_jinja.py b/flepimop/gempyor_pkg/src/gempyor/_jinja.py index 847e68382..fc397963e 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_jinja.py +++ b/flepimop/gempyor_pkg/src/gempyor/_jinja.py @@ -77,7 +77,7 @@ def _render_template_to_temp_file( ... "test_template.j2", {"name": "John"}, suffix=".txt", prefix="foo_" ... ) >>> file - PosixPath('/var/folders/2z/h3pc0p7s3ng1tvxrgsw5kr680000gp/T/foo_ocaomg4k.txt') + PosixPath(...) >>> file.read_text() 'Hello John!' """ diff --git a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py index a8e59b187..bf373b71c 100644 --- a/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py +++ b/flepimop/gempyor_pkg/src/gempyor/_pydantic_ext.py @@ -134,16 +134,20 @@ def _read_and_validate_dataframe( ... name: str ... age: int >>> _read_and_validate_dataframe(file, model=Person) - name age + name age 0 Jack 23 1 Jill 25 >>> pd.DataFrame(data={"name": [32], "age": ["Jane"]}).to_csv(file, index=False) >>> _read_and_validate_dataframe(file, model=Person) Traceback (most recent call last): ... - pydantic_core._pydantic_core.ValidationError: 1 validation error for RootModel[list[Person]] + pydantic_core._pydantic_core.ValidationError: 2 validation errors for RootModel[list[Person]] + 0.name + Input should be a valid string [type=string_type, input_value=32, input_type=int] + For further information visit https://errors.pydantic.dev/2.11/v/string_type 0.age - Input should be a valid integer, unable to parse string as an integer ... + Input should be a valid integer, unable to parse string as an integer [type=int_parsing, input_value='Jane', input_type=str] + For further information visit https://errors.pydantic.dev/2.11/v/int_parsing >>> class PartitionSlice(BaseModel): ... name: str ... amount: Annotated[float, Field(gt=0.0, lt=1.0)] @@ -159,7 +163,7 @@ def _read_and_validate_dataframe( ... data={"name": ["A", "B"], "amount": [0.5, 0.5]}, ... ).to_csv(file, index=False) >>> _read_and_validate_dataframe(file, model=Partition) - name amount + name amount 0 A 0.5 1 B 0.5 >>> pd.DataFrame( @@ -169,7 +173,8 @@ def _read_and_validate_dataframe( Traceback (most recent call last): ... pydantic_core._pydantic_core.ValidationError: 1 validation error for Partition - Value error, The sum of the amounts must be equal to 1.0 ... + Value error, The sum of the amounts must be equal to 1.0 [type=value_error, input_value=[{'name': 'A', 'amount': ...e': 'B', 'amount': 0.1}], input_type=list] + For further information visit https://errors.pydantic.dev/2.11/v/value_error """ if file.suffix == ".csv": @@ -213,8 +218,8 @@ def _evaled_expression(val: EE | str | Any, target_type: Type[EE]) -> EE | Any: 2 >>> _evaled_expression("10 / 4", float) 2.5 + >>> # Note that result is truncated, probably undesirable if misused >>> _evaled_expression("10 / 4", int) - # note that result is trucnated; probably undesirable if misused 2 >>> _evaled_expression(99.5, float) 99.5 @@ -223,7 +228,7 @@ def _evaled_expression(val: EE | str | Any, target_type: Type[EE]) -> EE | Any: >>> _evaled_expression("a * b", float) Traceback (most recent call last): ... - ValueError: Can't convert expression to float. + ValueError: Cannot convert expression 'a*b' to . """ if isinstance(val, target_type): return val diff --git a/flepimop/gempyor_pkg/src/gempyor/batch/_helpers.py b/flepimop/gempyor_pkg/src/gempyor/batch/_helpers.py index 2a62cde47..cb48c5a93 100644 --- a/flepimop/gempyor_pkg/src/gempyor/batch/_helpers.py +++ b/flepimop/gempyor_pkg/src/gempyor/batch/_helpers.py @@ -39,12 +39,8 @@ def _job_name(name: str | None, timestamp: datetime | None) -> str: ValueError: If `name` does not start with a letter and contains characters other than the alphabet, numbers, underscores or dashes. Examples: - >>> from gempyor.batch import _job_name - >>> _job_name(None, None) - '20241105T153818' - >>> _job_name("foobar", None) - 'foobar-20241105T153831' >>> from datetime import datetime, timezone + >>> from gempyor.batch._helpers import _job_name >>> _job_name(None, datetime(2024, 1, 1, tzinfo=timezone.utc)) '20240101T000000' """ @@ -66,7 +62,7 @@ def _parse_extra_options(extra: Iterable[str] | None) -> dict[str, str]: A dictionary of the parsed extra options. Examples: - >>> from gempyor.batch import _parse_extra_options + >>> from gempyor.batch._helpers import _parse_extra_options >>> _parse_extra_options(["abc=def", "ghi=jkl"]) {'abc': 'def', 'ghi': 'jkl'} >>> _parse_extra_options([ diff --git a/flepimop/gempyor_pkg/src/gempyor/batch/manifest.py b/flepimop/gempyor_pkg/src/gempyor/batch/manifest.py index 392122021..aa6faa167 100644 --- a/flepimop/gempyor_pkg/src/gempyor/batch/manifest.py +++ b/flepimop/gempyor_pkg/src/gempyor/batch/manifest.py @@ -42,22 +42,6 @@ def write_manifest( Returns: The path to the written json file. - - Examples: - >>> import os - >>> from pathlib import Path - >>> flepi_path = Path(os.environ["FLEPI_PATH"]) - >>> project_path = flepi_path / "examples" / "tutorials" - >>> manifest = write_manifest("Foobar", flepi_path, project_path) - >>> manifest.name - 'manifest.json' - >>> print(manifest.read_text()) - { - "cmd": "", - "job_name": "Foobar", - "data_sha": "59fe36d13fe34b6c1fb5c92bf8c53b83bd3ba593", - "flepimop_sha": "2bdfbc74e69bdd0243ef8340dda238f5504f1ad9" - } """ flepimop_sha = _git_head(flepi_path) data_sha = _git_head(project_path) diff --git a/flepimop/gempyor_pkg/src/gempyor/batch/systems.py b/flepimop/gempyor_pkg/src/gempyor/batch/systems.py index 0d677cea8..f183c61d7 100644 --- a/flepimop/gempyor_pkg/src/gempyor/batch/systems.py +++ b/flepimop/gempyor_pkg/src/gempyor/batch/systems.py @@ -67,7 +67,7 @@ class BatchSystem(ABC): >>> batch_system.format_time_limit(time_limit) '1595' >>> batch_system.size_from_jobs_simulations_blocks(2, 10, None, 5) - JobSize(blocks=2, chains=10, samples=None, simulations=5) + JobSize(blocks=2, chains=10, samples=None, simulations=5, samples_per_chain=None, simulations_per_chain=10, total_samples=None, total_simulations=100) """ needs_cluster: bool = False @@ -300,7 +300,7 @@ class LocalBatchSystem(BatchSystem): >>> from gempyor.batch import LocalBatchSystem >>> batch_system = LocalBatchSystem() >>> batch_system.size_from_jobs_simulations_blocks(1, 1, 25, 50) - JobSize(blocks=1, chains=1, samples=25, simulations=50) + JobSize(blocks=1, chains=1, samples=25, simulations=50, samples_per_chain=25, simulations_per_chain=50, total_samples=25, total_simulations=50) >>> with warnings.catch_warnings(record=True) as warns: ... size = batch_system.size_from_jobs_simulations_blocks(2, 4, 50, 100) ... for warn in warns: @@ -308,10 +308,10 @@ class LocalBatchSystem(BatchSystem): ... Local batch system only supports 1 chain but was given 4, overriding. Local batch system only supports 1 block but was given 2, overriding. - Local batch system only supports 50 total simulations but was given 800, overriding. Local batch system only supports 25 total samples but was given 400, overriding. + Local batch system only supports 50 total simulations but was given 800, overriding. >>> size - JobSize(blocks=1, chains=1, samples=25, simulations=50) + JobSize(blocks=1, chains=1, samples=25, simulations=50, samples_per_chain=25, simulations_per_chain=50, total_samples=25, total_simulations=50) """ name = "local" @@ -625,19 +625,6 @@ def register_batch_system(batch_system: BatchSystem) -> None: Raises: ValueError: If the batch system is already registered, checks the `name` attribute to determine this. - - Examples: - >>> from gempyor.batch import BatchSystem, register_batch_system - >>> class CustomBatchSystem(BatchSystem): - ... name = "custom" - >>> register_batch_system(CustomBatchSystem()) is None - True - >>> try: - ... register_batch_system(CustomBatchSystem()) is None - ... except Exception as e: - ... print(e) - ... - Batch system 'custom' already registered. """ if any(batch_system.name == bs.name for bs in _batch_systems): raise ValueError(f"Batch system '{batch_system.name}' already registered.") @@ -665,9 +652,9 @@ def get_batch_system(name: str, raise_on_missing: bool = True) -> BatchSystem | Examples: >>> from gempyor.batch import get_batch_system >>> get_batch_system("local") - + >>> get_batch_system("slurm") - + >>> try: ... get_batch_system("does not exist") ... except Exception as e: @@ -715,7 +702,7 @@ def _resolve_batch_system_name(name: str | None, local: bool, slurm: bool) -> st ValueError: If the given name conflicts with the boolean flags. Examples: - >>> from gempyor.batch import _resolve_batch_system_name + >>> from gempyor.batch.systems import _resolve_batch_system_name >>> _resolve_batch_system_name("abc", False, False) 'abc' >>> _resolve_batch_system_name("SLURM", False, False) diff --git a/flepimop/gempyor_pkg/src/gempyor/batch/types.py b/flepimop/gempyor_pkg/src/gempyor/batch/types.py index fd92bc81f..026ca17a3 100644 --- a/flepimop/gempyor_pkg/src/gempyor/batch/types.py +++ b/flepimop/gempyor_pkg/src/gempyor/batch/types.py @@ -71,22 +71,20 @@ class JobResources(BaseModel): 51200 >>> resources.total_resources() (5, 50, 51200) - >>> try: - ... JobResources(nodes=0, cpus=1, memory=1024) - ... except Exception as e: - ... print(e) - 1 validation error for JobResources + >>> JobResources(nodes=0, cpus=1, memory=1024) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for JobResources nodes - Input should be greater than 0 [type=greater_than, input_value=0, input_type=int] - For further information visit https://errors.pydantic.dev/2.10/v/greater_than - >>> try: - ... JobResources(nodes=2, cpus=4.5, memory=1024) - ... except Exception as e: - ... print(e) - 1 validation error for JobResources + Input should be greater than 0 [type=greater_than, input_value=0, input_type=int] + For further information visit https://errors.pydantic.dev/2.11/v/greater_than + >>> JobResources(nodes=2, cpus=4.5, memory=1024) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for JobResources cpus - Input should be a valid integer, got a number with a fractional part [type=int_from_float, input_value=4.5, input_type=float] - For further information visit https://errors.pydantic.dev/2.10/v/int_from_float + Input should be a valid integer, got a number with a fractional part [type=int_from_float, input_value=4.5, input_type=float] + For further information visit https://errors.pydantic.dev/2.11/v/int_from_float """ nodes: PositiveInt @@ -190,44 +188,37 @@ class JobSize(BaseModel): >>> import warnings >>> from gempyor.batch import JobSize >>> JobSize(blocks=5, chains=10, samples=100, simulations=200) - JobSize(blocks=5, chains=10, samples=100, simulations=200) + JobSize(blocks=5, chains=10, samples=100, simulations=200, samples_per_chain=500, simulations_per_chain=1000, total_samples=5000, total_simulations=10000) >>> JobSize(chains=32, simulations=500) - JobSize(blocks=None, chains=32, samples=None, simulations=500) - >>> try: - ... JobSize(blocks=0, chains=12, simulations=100) - ... except Exception as e: - ... print(e) - ... - 1 validation error for JobSize + JobSize(blocks=None, chains=32, samples=None, simulations=500, samples_per_chain=None, simulations_per_chain=500, total_samples=None, total_simulations=16000) + >>> JobSize(blocks=0, chains=12, simulations=100) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for JobSize blocks - Input should be greater than 0 [type=greater_than, input_value=0, input_type=int] - For further information visit https://errors.pydantic.dev/2.10/v/greater_than - >>> try: - ... JobSize(chains=10, samples=50.5, simulations=200) - ... except Exception as e: - ... print(e) - ... - 1 validation error for JobSize + Input should be greater than 0 [type=greater_than, input_value=0, input_type=int] + For further information visit https://errors.pydantic.dev/2.11/v/greater_than + >>> JobSize(chains=10, samples=50.5, simulations=200) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for JobSize samples - Input should be a valid integer, got a number with a fractional part [type=int_from_float, input_value=50.5, input_type=float] - For further information visit https://errors.pydantic.dev/2.10/v/int_from_float - >>> try: - ... JobSize(samples=100, simulations=50) - ... except Exception as e: - ... print(e) - ... - 1 validation error for JobSize - Value error, The number of samples, 100, must be less than or equal to the number of simulations, 50, per a block. [type=value_error, input_value={'samples': 100, 'simulations': 50}, input_type=dict] - For further information visit https://errors.pydantic.dev/2.10/v/value_error + Input should be a valid integer, got a number with a fractional part [type=int_from_float, input_value=50.5, input_type=float] + For further information visit https://errors.pydantic.dev/2.11/v/int_from_float + >>> JobSize(samples=100, simulations=50) + Traceback (most recent call last): + ... + pydantic_core._pydantic_core.ValidationError: 1 validation error for JobSize + Value error, The number of samples, 100, must be less than or equal to the number of simulations, 50, per a block. [type=value_error, input_value={'samples': 100, 'simulations': 50}, input_type=dict] + For further information visit https://errors.pydantic.dev/2.11/v/value_error >>> with warnings.catch_warnings(record=True) as warns: ... JobSize(samples=75, simulations=100) ... for warn in warns: ... print(warn.message) - ... - JobSize(blocks=None, chains=None, samples=75, simulations=100) + JobSize(blocks=None, chains=None, samples=75, simulations=100, samples_per_chain=75, simulations_per_chain=100, total_samples=75, total_simulations=100) The samples to simulations ratio is 75%, which is higher than the recommended limit of 60%. >>> JobSize() - JobSize(blocks=None, chains=None, samples=None, simulations=None) + JobSize(blocks=None, chains=None, samples=None, simulations=None, samples_per_chain=None, simulations_per_chain=None, total_samples=None, total_simulations=None) >>> size = JobSize(blocks=4, chains=3, samples=10, simulations=25) >>> size.samples_per_chain 40 diff --git a/flepimop/gempyor_pkg/src/gempyor/config_validator.py b/flepimop/gempyor_pkg/src/gempyor/config_validator.py index 610c60b46..68ff5343d 100644 --- a/flepimop/gempyor_pkg/src/gempyor/config_validator.py +++ b/flepimop/gempyor_pkg/src/gempyor/config_validator.py @@ -6,6 +6,7 @@ Field, AfterValidator, validator, + field_validator, ) from datetime import date from typing import Dict, List, Union, Literal, Optional, Annotated, Any diff --git a/flepimop/gempyor_pkg/src/gempyor/distributions.py b/flepimop/gempyor_pkg/src/gempyor/distributions.py index 917a03282..fa983dcb6 100644 --- a/flepimop/gempyor_pkg/src/gempyor/distributions.py +++ b/flepimop/gempyor_pkg/src/gempyor/distributions.py @@ -108,7 +108,7 @@ class NormalDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = NormalDistribution(mu=2.3, sigma=4.5) >>> dist - NormalDistribution(distribution='norm', mu=2.3, sigma=4.5) + NormalDistribution(distribution='norm', allow_edge_cases=False, mu=2.3, sigma=4.5) >>> dist.sample(rng=rng) array([3.67122686]) >>> dist.sample(size=(3, 5), rng=rng) @@ -138,7 +138,7 @@ class UniformDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = UniformDistribution(low=-0.5, high=1.5) >>> dist - UniformDistribution(distribution='uniform', low=-0.5, high=1.5) + UniformDistribution(distribution='uniform', allow_edge_cases=False, low=-0.5, high=1.5) >>> dist.sample(rng=rng) array([1.0479121]) >>> dist.sample(size=(3, 5), rng=rng) @@ -190,7 +190,7 @@ class LognormalDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = LognormalDistribution(meanlog=0.0, sdlog=1.0) >>> dist - LognormalDistribution(distribution='lognorm', meanlog=0.0, sdlog=1.0) + LognormalDistribution(distribution='lognorm', allow_edge_cases=False, meanlog=0.0, sdlog=1.0) >>> dist.sample(rng=rng) array([1.35624124]) >>> dist.sample(size=(3, 5), rng=rng) @@ -221,7 +221,7 @@ class TruncatedNormalDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = TruncatedNormalDistribution(mean=1.0, sd=1.0, a=0.0, b=10.0) >>> dist - TruncatedNormalDistribution(distribution='truncnorm', mean=1.0, sd=1.0, a=0.0, b=10.0) + TruncatedNormalDistribution(distribution='truncnorm', allow_edge_cases=False, mean=1.0, sd=1.0, a=0.0, b=10.0) >>> dist.sample(rng=rng) array([1.87722989]) >>> dist.sample(size=(3, 5), rng=rng) @@ -288,7 +288,7 @@ class PoissonDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = PoissonDistribution(lam=3.0) >>> dist - PoissonDistribution(distribution='poisson', lam=3.0) + PoissonDistribution(distribution='poisson', allow_edge_cases=False, lam=3.0) >>> dist.sample(rng=rng) array([4]) >>> dist.sample(size=(3, 5), rng=rng) @@ -335,7 +335,7 @@ class BinomialDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = BinomialDistribution(n=10, p=0.5) >>> dist - BinomialDistribution(distribution='binomial', n=10, p=0.5) + BinomialDistribution(distribution='binomial', allow_edge_cases=False, n=10, p=0.5) >>> dist.sample(rng=rng) array([6]) >>> dist.sample(size=(3, 5), rng=rng) @@ -351,7 +351,8 @@ class BinomialDistribution(DistributionABC): Traceback (most recent call last): ... pydantic_core._pydantic_core.ValidationError: 1 validation error for BinomialDistribution - Value error, `p` cannot be 0 or 1 when `allow_edge_cases` is `False`. [type=value_error, ... + Value error, Input for `p` cannot be 0 or 1 when `allow_edge_cases` is `False`. [type=value_error, input_value={'n': 10, 'p': 0.0}, input_type=dict] + For further information visit https://errors.pydantic.dev/2.11/v/value_error """ distribution: Literal["binomial"] = "binomial" @@ -389,13 +390,13 @@ class GammaDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = GammaDistribution(shape=2.0, scale=1.5) >>> dist - GammaDistribution(distribution='gamma', shape=2.0, scale=1.5) + GammaDistribution(distribution='gamma', allow_edge_cases=False, shape=2.0, scale=1.5) >>> dist.sample(rng=rng) - array([2.78453141]) + array([3.13772591]) >>> dist.sample(size=(3, 5), rng=rng) - array([[1.0116844 , 5.1632733 , 6.51682397, 0.49053229, 0.82522731], - [2.45524838, 1.4870423 , 1.95460596, 0.94191942, 5.92679803], - [5.08819234, 2.21976694, 7.82570081, 3.42761858, 1.1070266 ]]) + array([[4.25301838, 2.75582337, 2.46760563, 4.61888299, 2.63006031], + [3.51900762, 3.28422915, 4.61612039, 2.15883076, 5.69337578], + [1.75889742, 3.67897959, 3.38745296, 1.6334585 , 3.89261212]]) """ distribution: Literal["gamma"] = "gamma" @@ -419,13 +420,13 @@ class WeibullDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = WeibullDistribution(shape=2.5, scale=5.0) >>> dist - WeibullDistribution(distribution='weibull', shape=2.5, scale=5.0) + WeibullDistribution(distribution='weibull', allow_edge_cases=False, shape=2.5, scale=5.0) >>> dist.sample(rng=rng) - array([5.18664161]) + array([7.10164344]) >>> dist.sample(size=(3, 5), rng=rng) - array([[4.09267151, 6.25764353, 6.57560416, 2.72348545, 3.73147743], - [5.08637841, 4.63044996, 4.96639599, 3.99658252, 6.42564251], - [6.29525049, 5.03450379, 6.88371131, 5.65602432, 3.99307222]]) + array([[7.02058405, 7.0786094 , 3.00403884, 1.87780582, 5.8054469 ], + [5.73657673, 7.886256 , 1.81412232, 5.0918523 , 1.73016943], + [5.17350562, 6.22761392, 3.4198509 , 5.43445321, 2.36440749]]) """ distribution: Literal["weibull"] = "weibull" @@ -450,13 +451,13 @@ class BetaDistribution(DistributionABC): >>> rng = np.random.default_rng(42) >>> dist = BetaDistribution(alpha=2.0, beta=5.0) >>> dist - BetaDistribution(distribution='beta', alpha=2.0, beta=5.0) + BetaDistribution(distribution='beta', allow_edge_cases=False, alpha=2.0, beta=5.0) >>> dist.sample(rng=rng) - array([0.31153835]) + array([0.24395464]) >>> dist.sample(size=(3, 5), rng=rng) - array([[0.18343469, 0.43545995, 0.48512838, 0.08625927, 0.22223849], - [0.34020943, 0.28014294, 0.3069152 , 0.201726 , 0.4729195 ], - [0.44331908, 0.32357912, 0.5282496 , 0.38318282, 0.20141687]]) + array([[0.28406092, 0.39027204, 0.29864681, 0.41835336, 0.49963165], + [0.30396328, 0.15089427, 0.32937986, 0.52373987, 0.16127411], + [0.32746504, 0.48761242, 0.2162056 , 0.29178583, 0.22819733]]) """ distribution: Literal["beta"] = "beta" diff --git a/flepimop/gempyor_pkg/src/gempyor/file_paths.py b/flepimop/gempyor_pkg/src/gempyor/file_paths.py index eba3acaf4..7916ebf8f 100644 --- a/flepimop/gempyor_pkg/src/gempyor/file_paths.py +++ b/flepimop/gempyor_pkg/src/gempyor/file_paths.py @@ -55,7 +55,7 @@ def create_file_name( ... "jkl", ... create_directory=False, ... ) - PosixPath('model_output/abc/hosp/global/jkl000000001.20240101_000000.hosp.parquet') + 'model_output/abc/hosp/global/jkl000000001.20240101_000000.hosp.parquet' """ if create_directory: os.makedirs( @@ -157,8 +157,6 @@ def run_id(timestamp: None | datetime = None) -> str: Examples: >>> from datetime import datetime, timezone >>> from gempyor.file_paths import run_id - >>> run_id() - '20240711_160059' >>> run_id(timestamp=datetime(2024, 1, 1)) '20240101_000000' >>> run_id(timestamp=datetime(2024, 1, 1, tzinfo=timezone.utc)) diff --git a/flepimop/gempyor_pkg/src/gempyor/info.py b/flepimop/gempyor_pkg/src/gempyor/info.py index 45961caeb..f3f6d6aaa 100644 --- a/flepimop/gempyor_pkg/src/gempyor/info.py +++ b/flepimop/gempyor_pkg/src/gempyor/info.py @@ -30,6 +30,10 @@ function being used via the `search_paths` argument. Examples: + >>> import os + >>> import pytest + >>> if "FLEPI_PATH" not in os.environ: + ... pytest.skip("FLEPI_PATH environment variable is not set.") >>> from pprint import pprint >>> from gempyor.info import get_cluster_info >>> cluster_info = get_cluster_info("longleaf") @@ -37,9 +41,9 @@ 'longleaf' >>> pprint(cluster_info.modules) [Module(name='gcc', version='9.1.0'), - Module(name='anaconda', version='2023.03'), - Module(name='git', version=None), - Module(name='aws', version=None)] + Module(name='anaconda', version='2023.03'), + Module(name='git', version=None), + Module(name='aws', version=None)] """ __all__ = ["Cluster", "Module", "PathExport", "get_cluster_info"] @@ -188,6 +192,10 @@ def get_cluster_info( An object containing the information about the `name` cluster. Examples: + >>> import os + >>> import pytest + >>> if "FLEPI_PATH" not in os.environ: + ... pytest.skip("FLEPI_PATH environment variable is not set.") >>> from gempyor.info import get_cluster_info >>> cluster_info = get_cluster_info("longleaf") >>> cluster_info.name diff --git a/flepimop/gempyor_pkg/src/gempyor/logging.py b/flepimop/gempyor_pkg/src/gempyor/logging.py index 12fc0a636..5f4e28c6f 100644 --- a/flepimop/gempyor_pkg/src/gempyor/logging.py +++ b/flepimop/gempyor_pkg/src/gempyor/logging.py @@ -106,7 +106,7 @@ def get_script_logger( Examples: >>> from gempyor.logging import get_script_logger >>> logger = get_script_logger(__name__, 3) - >>> logger.info("This is a log info message") + >>> logger.info("This is a log info message") # doctest: +SKIP 2024-10-29 16:07:20,272:INFO:__main__> This is a log info message. """ logger = logging.getLogger(name) diff --git a/flepimop/gempyor_pkg/src/gempyor/shared_cli.py b/flepimop/gempyor_pkg/src/gempyor/shared_cli.py index c2535a179..611938e5b 100644 --- a/flepimop/gempyor_pkg/src/gempyor/shared_cli.py +++ b/flepimop/gempyor_pkg/src/gempyor/shared_cli.py @@ -169,13 +169,11 @@ def click_helpstring( >>> @click_helpstring(opt) ... def test_cli(num: int) -> None: ... print(f"Your favorite number is {num}!") - ... - >>> print(test_cli.__doc__) - - + >>> print(test_cli.__doc__) # doctest: +SKIP Command Line Interface arguments: + + n: Int - n: Int """ if not isinstance(params, list): params = [params] @@ -334,24 +332,12 @@ def log_cli_inputs(kwargs: dict[str, Any], verbosity: int | None = None) -> None Examples: >>> from gempyor.shared_cli import log_cli_inputs - >>> log_cli_inputs({"abc": 123, "def": True}, 3) + >>> log_cli_inputs({"abc": 123, "def": True}, 3) # doctest: +SKIP 2024-11-05 09:27:58,884:DEBUG:gempyor.shared_cli> CLI was given 2 arguments: 2024-11-05 09:27:58,885:DEBUG:gempyor.shared_cli> abc = 123. 2024-11-05 09:27:58,885:DEBUG:gempyor.shared_cli> def = True. >>> log_cli_inputs({"abc": 123, "def": True}, 2) - >>> from pathlib import Path - >>> kwargs = { - ... "input_file": Path("config.in"), - ... "method": "stochastic", - ... "cluster": "longleaf", - ... "verbosity": 3, - ... } - >>> log_cli_inputs(kwargs) - 2024-11-05 09:29:21,666:DEBUG:gempyor.shared_cli> CLI was given 4 arguments: - 2024-11-05 09:29:21,667:DEBUG:gempyor.shared_cli> input_file = /Users/twillard/Desktop/GitHub/HopkinsIDD/flepiMoP/flepimop/gempyor_pkg/config.in. - 2024-11-05 09:29:21,667:DEBUG:gempyor.shared_cli> method = stochastic. - 2024-11-05 09:29:21,668:DEBUG:gempyor.shared_cli> cluster = longleaf. - 2024-11-05 09:29:21,668:DEBUG:gempyor.shared_cli> verbosity = 3. + """ verbosity = kwargs.get("verbosity") if verbosity is None else verbosity if verbosity is None: diff --git a/flepimop/gempyor_pkg/src/gempyor/sync/_sync.py b/flepimop/gempyor_pkg/src/gempyor/sync/_sync.py index 98b62f6a5..c86c72139 100644 --- a/flepimop/gempyor_pkg/src/gempyor/sync/_sync.py +++ b/flepimop/gempyor_pkg/src/gempyor/sync/_sync.py @@ -181,10 +181,8 @@ def _rsync_ensure_path(target: str, verbosity: int, dry_run: bool) -> CompletedP Examples: >>> from gempyor.sync._sync import _rsync_ensure_path >>> _rsync_ensure_path("/foo/bar", 0, True) - (DRY RUN): mkdir -p /foo/bar CompletedProcess(args=['echo', '(DRY RUN): mkdir -p /foo/bar'], returncode=0) >>> _rsync_ensure_path("user@host:/fizz/buzz", 0, True) - (DRY RUN): ssh user@host mkdir -p /fizz/buzz CompletedProcess(args=['echo', '(DRY RUN): ssh user@host mkdir -p /fizz/buzz'], returncode=0) """ # pylint: enable=line-too-long @@ -349,7 +347,7 @@ class S3SyncModel(SyncABC, WithFilters): Traceback (most recent call last): ... pydantic_core._pydantic_core.ValidationError: 1 validation error for S3SyncModel - Value error, At least one of `source` or `target` must be an s3 bucket, as indicated by a `s3://` prefix [type=value_error, input_value={'type': 's3sync', 'targe...et', 'source': 'source'}, input_type=dict] + Value error, At least one of `source` or `target` must be an s3 bucket, as indicated by a `s3://` prefix [type=value_error, input_value={'type': 's3sync', 'targe...et', 'source': 'source'}, input_type=dict] For further information visit https://errors.pydantic.dev/2.11/v/value_error """ # pylint: enable=line-too-long diff --git a/flepimop/gempyor_pkg/src/gempyor/testing.py b/flepimop/gempyor_pkg/src/gempyor/testing.py index 69018034c..0dccc8d84 100644 --- a/flepimop/gempyor_pkg/src/gempyor/testing.py +++ b/flepimop/gempyor_pkg/src/gempyor/testing.py @@ -304,7 +304,7 @@ def run_test_in_separate_process( Examples: Typically this function is used in a test like so: - >>> assert run_test_in_separate_process( + >>> assert run_test_in_separate_process( # doctest: +SKIP ... Path(__file__).parent / "external_script.py", ... tmp_path / "test.py", ... args=["arg1", "arg2"], diff --git a/flepimop/gempyor_pkg/src/gempyor/utils.py b/flepimop/gempyor_pkg/src/gempyor/utils.py index f0c7916eb..411224ee6 100644 --- a/flepimop/gempyor_pkg/src/gempyor/utils.py +++ b/flepimop/gempyor_pkg/src/gempyor/utils.py @@ -188,18 +188,6 @@ def search_and_import_plugins_class( Returns: The instance of the class that was instantiated with provided **kwargs. - Examples: - Suppose there is a module called `my_plugin.py with a class `MyClass` located at `/path/to/plugin/`. - - Dynamically import and instantiate the class: - - >>> instance = search_and_import_plugins_class('/path/to/plugin', path_prefix, 'MyClass', **params) - - View the instance: - - >>> print(instance) - <__main__.MyClass object at 0x7f8b2c6b4d60> - """ # Look for all possible plugins and import them # https://stackoverflow.com/questions/67631/how-can-i-import-a-module-dynamically-given-the-full-path @@ -253,14 +241,6 @@ def profile( Returns: Profile of the decorated function. - Examples: - >>> @profile(output_file="my_function.prof") - >>> def my_function(): - # Function body content - pass - >>> my_function() - After running ``my_function``, a file named ``my_function.prof`` will be created in the current WD. - This file contains the profiling data. """ def inner(func): @@ -373,18 +353,6 @@ def list_filenames( Returns: A list of strings representing the paths to the files that match the filters. - Examples: - To get all files containing "hosp": - >>> gempyor.utils.list_filenames( - folder="model_output/", - filters=["hosp"], - ) - - To get only "hosp" files with a ".parquet" extension: - >>> gempyor.utils.list_filenames( - folder="model_output/", - filters=["hosp", ".parquet"], - ) """ filters = [filters] if not isinstance(filters, list) else filters filters = filters if len(filters) else [""] @@ -419,20 +387,11 @@ def extract_slot(file: Path) -> pd.DataFrame: >>> file = Path.cwd() / "00017.foobar.csv" >>> sample.to_csv(file, index=False) >>> extract_slot(file) - letters slot + letters slot 0 a 17 1 b 17 2 c 17 - >>> extract_slot(file).info() - - RangeIndex: 3 entries, 0 to 2 - Data columns (total 2 columns): - # Column Non-Null Count Dtype - --- ------ -------------- ----- - 0 letters 3 non-null object - 1 slot 3 non-null int64 - dtypes: int64(1), object(1) - memory usage: 180.0+ bytes + """ df = pd.read_parquet(file) if file.suffix == ".parquet" else pd.read_csv(file) if "slot" in df.columns: @@ -483,27 +442,27 @@ def rolling_mean_pad( Examples: Below is a brief set of examples showcasing how to smooth a metric, like hospitalizations, using this function. - ``` + >>> import numpy as np >>> from gempyor.utils import rolling_mean_pad >>> hospitalizations = np.arange(1., 29.).reshape((7, 4)) >>> hospitalizations array([[ 1., 2., 3., 4.], - [ 5., 6., 7., 8.], - [ 9., 10., 11., 12.], - [13., 14., 15., 16.], - [17., 18., 19., 20.], - [21., 22., 23., 24.], - [25., 26., 27., 28.]]) + [ 5., 6., 7., 8.], + [ 9., 10., 11., 12.], + [13., 14., 15., 16.], + [17., 18., 19., 20.], + [21., 22., 23., 24.], + [25., 26., 27., 28.]]) >>> rolling_mean_pad(hospitalizations, 5) array([[ 3.4, 4.4, 5.4, 6.4], - [ 5.8, 6.8, 7.8, 8.8], - [ 9. , 10. , 11. , 12. ], - [13. , 14. , 15. , 16. ], - [17. , 18. , 19. , 20. ], - [20.2, 21.2, 22.2, 23.2], - [22.6, 23.6, 24.6, 25.6]]) - ``` + [ 5.8, 6.8, 7.8, 8.8], + [ 9. , 10. , 11. , 12. ], + [13. , 14. , 15. , 16. ], + [17. , 18. , 19. , 20. ], + [20.2, 21.2, 22.2, 23.2], + [22.6, 23.6, 24.6, 25.6]]) + """ weights = (1.0 / window) * np.ones(window) output = scipy.ndimage.convolve1d(data, weights, axis=0, mode="nearest") @@ -580,17 +539,16 @@ def create_resume_out_filename( The path to a corresponding output file. Examples: - Generate an output file with specified parameters: - >>> filename = create_resume_out_filename( - flepi_run_index="test_run", - flepi_prefix="model_output/run_id/", - flepi_slot_index="1", - flepi_block_index="2", - filetype="seed", - liketype="chimeric" - ) - >>> print(filename) - "experiment/001/normal/intermediate/000000123.000000000.1.parquet" + >>> from gempyor.utils import create_resume_out_filename + >>> create_resume_out_filename( + ... flepi_run_index="test_run", + ... flepi_prefix="model_output/run_id/", + ... flepi_slot_index="1", + ... flepi_block_index="2", + ... filetype="seed", + ... liketype="chimeric", + ... ) + 'model_output/model_output/run_id/test_run/seed/chimeric/intermediate/000000001.000000001.000000001.test_run.seed.csv' """ prefix = f"{flepi_prefix}/{flepi_run_index}" inference_filepath_suffix = f"{liketype}/intermediate" @@ -631,16 +589,15 @@ def create_resume_input_filename( The path to the a corresponding input file. Examples: - Generate an input file with specified parameters: - >>> filename = create_resume_input_filename( - resume_run_index="2", - flepi_prefix="model_output/run_id/", - flepi_slot_index="1", - filetype="seed", - liketype="chimeric" - ) - >>> print(filename) - "experiment/002/normal/final/789.csv" + >>> from gempyor.utils import create_resume_input_filename + >>> create_resume_input_filename( + ... resume_run_index="2", + ... flepi_prefix="model_output/run_id/", + ... flepi_slot_index="1", + ... filetype="seed", + ... liketype="chimeric", + ... ) + 'model_output/model_output/run_id/2/seed/chimeric/final/000000001.2.seed.csv' """ prefix = f"{flepi_prefix}/{resume_run_index}" inference_filepath_suffix = f"{liketype}/final" @@ -676,14 +633,20 @@ def get_filetype_for_resume( Examples: Determine file types for block index 1 with seeding data NOT discarded: - >>> filetypes = get_filetype_for_resume(resume_discard_seeding="false", flepi_block_index="1") - >>> print(filetypes) - ["seed", "spar", "snpi", "hpar", "hnpi", "init"] + + >>> from gempyor.utils import get_filetype_for_resume + >>> get_filetype_for_resume( + ... resume_discard_seeding="false", flepi_block_index="1" + ... ) + ['seed', 'spar', 'snpi', 'hpar', 'hnpi', 'init'] Determine file types for block index 2 with seeding data discarded: - >>> filtypes = get_filetype_for_resume(resume_discard_seeding="true", flepi_block_index="2") - >>> print(filetypes) - ["seed", "spar", "snpi", "hpar", "hnpi", "host", "llik", "init"] + + >>> get_filetype_for_resume( + ... resume_discard_seeding="true", flepi_block_index="2" + ... ) + ['seed', 'spar', 'snpi', 'hpar', 'hnpi', 'host', 'llik', 'init'] + """ if flepi_block_index == "1": if resume_discard_seeding == "true": @@ -704,13 +667,18 @@ def create_resume_file_names_map( last_job_output: str, ) -> dict[str, str]: """ - Generates a mapping of input file names to output file names for a resume process based on - parquet file types and environmental conditions. The function adjusts the file name mappings - based on the operational block index and the location of the last job output. + Generate a mapping of input file names to output file names for a resume process. + + Generates a mapping of input file names to output file names for a resume process + based on parquet file types and environmental conditions. The function adjusts the + file name mappings based on the operational block index and the location of the + last job output. Args: - resume_discard_seeding: Determines whether seeding-related file types should be included. - flepi_block_index: Determines a specific operational mode or block of the process. + resume_discard_seeding: Determines whether seeding-related file types should + be included. + flepi_block_index: Determines a specific operational mode or block of the + process. resume_run_index: Resume run index. flepi_prefix: File prefix. flepi_slot_index: Index of the slot. @@ -722,40 +690,49 @@ def create_resume_file_names_map( output file paths. The mappings depend on: - - Parquet file types appropriate for resuming a process, as determined by the environment. - - Whether the files are for 'global' or 'chimeric' types, as these liketypes influence the - file naming convention. - - The operational block index ('FLEPI_BLOCK_INDEX'), which can alter the input file names for - block index '1'. - - The presence and value of 'LAST_JOB_OUTPUT' environment variable, which if set to an S3 path, - adjusts the keys in the mapping to be prefixed with this path. + - Parquet file types appropriate for resuming a process, as determined by the + environment. + - Whether the files are for 'global' or 'chimeric' types, as these like types + influence the file naming convention. + - The operational block index ('FLEPI_BLOCK_INDEX'), which can alter the input + file names for block index '1'. + - The presence and value of 'LAST_JOB_OUTPUT' environment variable, which if set + to an S3 path, adjusts the keys in the mapping to be prefixed with this path. Raises: - No explicit exceptions are raised within the function, but it relies heavily on external - functions and environment variables which if improperly configured could lead to unexpected - behavior. + No explicit exceptions are raised within the function, but it relies heavily + on external functions and environment variables which if improperly configured + could lead to unexpected behavior. Examples: - Generate a mapping of file names for a given resume process: + >>> from pprint import pprint + >>> from gempyor.utils import create_resume_file_names_map >>> file_names_map = create_resume_file_names_map( - resume_discard_seeding="false", - flepi_block_index="1", - resume_run_index="1", - flepi_prefix="model_output/run_id/", - flepi_slot_index="1", - flepi_run_index="test_run", - last_job_output="s3://bucket/path/") - >>> print(file_names_map) - { - 's3://bucket/path/model_output/run_id/1_type1_global_1.in': 'model_output/run_id/test_run_type1_global_1_1.out', - 's3://bucket/path/model_output/run_id/1_type1_chimeric_1.in': 'model_output/run_id/test_run_type1_chimeric_1_1.out', - 's3://bucket/path/model_output/run_id/1_type2_global_1.in': 'model_output/run_id/test_run_type2_global_1_1.out', - 's3://bucket/path/model_output/run_id/1_type2_chimeric_1.in': 'model_output/run_id/test_run_type2_chimeric_1_1.out' - } - # Note: this output is toy output implemented with toy file names. + ... resume_discard_seeding="false", + ... flepi_block_index="1", + ... resume_run_index="1", + ... flepi_prefix="model_output/run_id/", + ... flepi_slot_index="1", + ... flepi_run_index="test_run", + ... last_job_output="s3://bucket/path/", + ... ) + >>> pprint(file_names_map) + {'s3://bucket/path/model_output/model_output/run_id/1/hnpi/chimeric/final/000000001.1.hnpi.parquet': 'model_output/model_output/run_id/test_run/hnpi/chimeric/intermediate/000000001.000000001.000000000.test_run.hnpi.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/hnpi/global/final/000000001.1.hnpi.parquet': 'model_output/model_output/run_id/test_run/hnpi/global/intermediate/000000001.000000001.000000000.test_run.hnpi.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/hpar/chimeric/final/000000001.1.hpar.parquet': 'model_output/model_output/run_id/test_run/hpar/chimeric/intermediate/000000001.000000001.000000000.test_run.hpar.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/hpar/global/final/000000001.1.hpar.parquet': 'model_output/model_output/run_id/test_run/hpar/global/intermediate/000000001.000000001.000000000.test_run.hpar.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/init/chimeric/final/000000001.1.init.parquet': 'model_output/model_output/run_id/test_run/init/chimeric/intermediate/000000001.000000001.000000000.test_run.init.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/init/global/final/000000001.1.init.parquet': 'model_output/model_output/run_id/test_run/init/global/intermediate/000000001.000000001.000000000.test_run.init.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/seed/chimeric/final/000000001.1.seed.csv': 'model_output/model_output/run_id/test_run/seed/chimeric/intermediate/000000001.000000001.000000000.test_run.seed.csv', + 's3://bucket/path/model_output/model_output/run_id/1/seed/global/final/000000001.1.seed.csv': 'model_output/model_output/run_id/test_run/seed/global/intermediate/000000001.000000001.000000000.test_run.seed.csv', + 's3://bucket/path/model_output/model_output/run_id/1/snpi/chimeric/final/000000001.1.snpi.parquet': 'model_output/model_output/run_id/test_run/snpi/chimeric/intermediate/000000001.000000001.000000000.test_run.snpi.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/snpi/global/final/000000001.1.snpi.parquet': 'model_output/model_output/run_id/test_run/snpi/global/intermediate/000000001.000000001.000000000.test_run.snpi.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/spar/chimeric/final/000000001.1.spar.parquet': 'model_output/model_output/run_id/test_run/spar/chimeric/intermediate/000000001.000000001.000000000.test_run.spar.parquet', + 's3://bucket/path/model_output/model_output/run_id/1/spar/global/final/000000001.1.spar.parquet': 'model_output/model_output/run_id/test_run/spar/global/intermediate/000000001.000000001.000000000.test_run.spar.parquet'} Notes: - - The paths may be modified by the 'LAST_JOB_OUTPUT' if it is set and points to an S3 location. + - The paths may be modified by the 'LAST_JOB_OUTPUT' if it is set and points to + an S3 location. """ file_types = get_filetype_for_resume( resume_discard_seeding=resume_discard_seeding, flepi_block_index=flepi_block_index @@ -813,21 +790,6 @@ def download_file_from_s3(name_map: dict[str, str]) -> None: ClientError: If an error occurs during the download from S3, such as a permissions issue, a missing file, or network-related errors. These are caught and logged but not re-raised, to allow the function to attempt subsequent downloads. - - Examples: - >>> name_map = { - "s3://mybucket/data/file1.txt": "/local/path/to/file1.txt", - "s3://mybucket/data/file2.txt": "/local/path/to/file2.txt" - } - >>> download_file_from_s3(name_map) - # This would download 'file1.txt' and 'file2.txt' from 'mybucket' on S3 to the specified local paths. - - # If an S3 URI is malformed: - >>> name_map = { - "http://wrongurl.com/data/file1.txt": "/local/path/to/file1.txt" - } - >>> download_file_from_s3(name_map) - # This will raise a ValueError indicating the invalid S3 URI format. """ try: import boto3 @@ -922,6 +884,7 @@ def _dump_formatted_yaml(cfg: confuse.Configuration) -> str: rate: ["beta * gamma"] proportional_to: [[S], [I]] proportion_exponent: [1, 1] + """ class CustomDumper(yaml.Dumper): @@ -1004,9 +967,13 @@ def _git_head(repository: Path) -> str: Examples: >>> import os + >>> import pytest + >>> if os.environ.get("FLEPI_PATH") is None: + ... pytest.skip("FLEPI_PATH environment variable is not set.") >>> from pathlib import Path - >>> _git_head(Path(os.environ["FLEPI_PATH"])) - 'efe896b1a5e4f8e33667c170cd5319d6ef1e3db5' + >>> _git_head(Path(os.environ["FLEPI_PATH"])) # doctest: +ELLIPSIS + '...' + """ git_cmd = _shutil_which("git") proc = subprocess.run( @@ -1027,10 +994,10 @@ def _git_checkout(repository: Path, branch: str) -> subprocess.CompletedProcess: checkout a new branch in. branch: The name of the new branch to checkout. - Examples: - >>> import os - >>> from pathlib import Path - >>> _git_checkout(Path(os.environ["FLEPI_PATH"]), "my-new-branch") + Returns: + A `subprocess.CompletedProcess` object containing the result of the checkout + command. + """ git_cmd = _shutil_which("git") return subprocess.run( @@ -1073,7 +1040,7 @@ def _format_cli_options( >>> _format_cli_options({"o": "/path/to/output.log"}) ['-o=/path/to/output.log'] >>> _format_cli_options({"opt1": "```", "opt2": "$( echo 'Hello!')"}) - ["--opt1='```'", '--opt2=\'$( echo \'"\'"\'Hello!\'"\'"\')\''] + ["--opt1='```'", '--opt2=\\'$( echo \\'"\\'"\\'Hello!\\'"\\'"\\')\\''] >>> _format_cli_options({"output": "/path/to/output.log"}, always_single=True) ['-output=/path/to/output.log'] >>> _format_cli_options({"person": ["Alice", "Bob", "Charlie"]}) From 641b6a1c82a9c251f2a0596b434a8e52e8798053 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Tue, 22 Jul 2025 16:16:53 -0400 Subject: [PATCH 67/73] Improve `environment.yml` file formatting Add spacing to list item entries so IDEs are a bit less confused when viewing the file. Addresses #485. --- build/create_environment_yml.R | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/build/create_environment_yml.R b/build/create_environment_yml.R index 8d812d17e..356be9263 100644 --- a/build/create_environment_yml.R +++ b/build/create_environment_yml.R @@ -32,19 +32,19 @@ dependencies <- dependencies[!grepl("^R(\\(.*\\))?$", dependencies)] environment_yml <- file.path(flepi_path, "environment.yml") new_environment_yml <- c( "channels:", - "- conda-forge", - "- defaults", - "- r", - "- dnachun", + " - conda-forge", + " - defaults", + " - r", + " - dnachun", "dependencies:", - "- python=3.11", - "- pip", - "- r-base>=4.3", - "- pyarrow=17.0.0", - "- r-arrow=17.0.0", - "- r-sf", - "- r-testthat", - paste0("- r-", dependencies) + " - python=3.11", + " - pip", + " - r-base>=4.3", + " - pyarrow=17.0.0", + " - r-arrow=17.0.0", + " - r-sf", + " - r-testthat", + paste0(" - r-", dependencies) ) if (file.exists(environment_yml)) { old_environment_yml <- readLines(environment_yml) From 6dc69f4c6dabd19c41a0d095cdfca0f6ae5233bd Mon Sep 17 00:00:00 2001 From: TimothyWillard Date: Tue, 22 Jul 2025 20:21:27 +0000 Subject: [PATCH 68/73] Update `environment.yml` via GitHub action --- environment.yml | 76 ++++++++++++++++++++++++------------------------- 1 file changed, 38 insertions(+), 38 deletions(-) diff --git a/environment.yml b/environment.yml index bef6b0394..498a4762f 100644 --- a/environment.yml +++ b/environment.yml @@ -1,40 +1,40 @@ -# Thu Jul 17 19:26:17 2025 UTC +# Tue Jul 22 20:21:27 2025 UTC channels: -- conda-forge -- defaults -- r -- dnachun + - conda-forge + - defaults + - r + - dnachun dependencies: -- python=3.11 -- pip -- r-base>=4.3 -- pyarrow=17.0.0 -- r-arrow=17.0.0 -- r-sf -- r-testthat -- r-data.table -- r-doParallel -- r-dplyr -- r-foreach -- r-ggplot2 -- r-ggraph -- r-httr -- r-jsonlite -- r-lubridate -- r-magrittr -- r-MMWRweek -- r-optparse -- r-purrr -- r-readr -- r-reticulate -- r-rlang -- r-stringr -- r-tibble -- r-tidygraph -- r-tidyr -- r-tidyselect -- r-tidyverse -- r-truncnorm -- r-vroom -- r-xts -- r-yaml + - python=3.11 + - pip + - r-base>=4.3 + - pyarrow=17.0.0 + - r-arrow=17.0.0 + - r-sf + - r-testthat + - r-data.table + - r-doParallel + - r-dplyr + - r-foreach + - r-ggplot2 + - r-ggraph + - r-httr + - r-jsonlite + - r-lubridate + - r-magrittr + - r-MMWRweek + - r-optparse + - r-purrr + - r-readr + - r-reticulate + - r-rlang + - r-stringr + - r-tibble + - r-tidygraph + - r-tidyr + - r-tidyselect + - r-tidyverse + - r-truncnorm + - r-vroom + - r-xts + - r-yaml From ce8942755d2e95531c9857a714f8d81c6a035bd1 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Thu, 31 Jul 2025 08:47:13 -0400 Subject: [PATCH 69/73] Remove push to `dev` trigger, split out slow tests * Removed the trigger whenever there is a push to the `dev` branch for the GitHub actions, in practice results in actions being ran twice on PR from `dev` into `main`. * Split the slow gempyor tests out into their own job that depends on the fast tests job and only runs with python 3.11. Running the slow tests with all python versions is not needed. --- .github/workflows/conda-env.yml | 6 ---- .github/workflows/flepicommon-ci.yml | 5 ---- .github/workflows/gempyor-ci.yml | 45 +++++++++++++++++++--------- .github/workflows/inference-ci.yml | 8 ----- 4 files changed, 31 insertions(+), 33 deletions(-) diff --git a/.github/workflows/conda-env.yml b/.github/workflows/conda-env.yml index 711cd0968..4005c4105 100644 --- a/.github/workflows/conda-env.yml +++ b/.github/workflows/conda-env.yml @@ -2,12 +2,6 @@ name: Generate Conda Environment on: workflow_dispatch: - push: - paths: - - build/create_environment_yml.R - - flepimop/R_packages/*/DESCRIPTION - branches: - - dev pull_request: types: - edited diff --git a/.github/workflows/flepicommon-ci.yml b/.github/workflows/flepicommon-ci.yml index 7594f3e95..f6ca5373f 100644 --- a/.github/workflows/flepicommon-ci.yml +++ b/.github/workflows/flepicommon-ci.yml @@ -2,11 +2,6 @@ name: flepicommon CI on: workflow_dispatch: - push: - paths: - - flepimop/R_packages/flepicommon/**/* - branches: - - dev pull_request: types: - edited diff --git a/.github/workflows/gempyor-ci.yml b/.github/workflows/gempyor-ci.yml index 7a5a3a054..330f597aa 100644 --- a/.github/workflows/gempyor-ci.yml +++ b/.github/workflows/gempyor-ci.yml @@ -2,12 +2,6 @@ name: gempyor CI on: workflow_dispatch: - push: - paths: - - examples/**/* - - flepimop/gempyor_pkg/**/* - branches: - - dev pull_request: types: - edited @@ -23,7 +17,7 @@ on: - main jobs: - tests: + fast-tests: runs-on: ubuntu-latest if: github.event_name != 'pull_request' || github.event.pull_request.draft == false strategy: @@ -59,19 +53,42 @@ jobs: cd flepimop/gempyor_pkg pytest -m "not slow" --exitfirst shell: bash -el {0} - - name: Run slow unit tests - run: | - conda activate venv/ - export FLEPI_PATH=$(pwd) - cd flepimop/gempyor_pkg - pytest -m "slow" --exitfirst - shell: bash -el {0} - name: Run gempyor-cli integration tests from examples run: | conda activate venv/ cd examples pytest --exitfirst shell: bash -el {0} + slow-tests: + needs: fast-tests + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.draft == false + strategy: + matrix: + python-version: ["3.11"] + R-version: ["4.3"] + steps: + - name: Checkout flepiMoP + uses: actions/checkout@v4 + with: + lfs: true + sparse-checkout: | + * + !documentation/ + sparse-checkout-cone-mode: false + persist-credentials: false + - name: Install flepiMoP + uses: ./.github/actions/flepimop-install + with: + python-version: ${{ matrix.python-version }} + R-version: ${{ matrix.R-version }} + - name: Run slow unit tests + run: | + conda activate venv/ + export FLEPI_PATH=$(pwd) + cd flepimop/gempyor_pkg + pytest -m "slow" + shell: bash -el {0} doc-updates: runs-on: ubuntu-latest if: github.event_name != 'pull_request' || github.event.pull_request.draft == false diff --git a/.github/workflows/inference-ci.yml b/.github/workflows/inference-ci.yml index ac27beb48..a9f403067 100644 --- a/.github/workflows/inference-ci.yml +++ b/.github/workflows/inference-ci.yml @@ -2,14 +2,6 @@ name: inference CI on: workflow_dispatch: - push: - paths: - - flepimop/R_packages/inference/**/* - - flepimop/R_packages/flepicommon/**/* - - examples/**/* - - flepimop/gempyor_pkg/**/* - branches: - - dev pull_request: types: - edited From fda910dfe243fe490c94ee1548eff6b3ad328dec Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Thu, 31 Jul 2025 10:24:36 -0400 Subject: [PATCH 70/73] Extract out R pkg installer Extracted out the custom R pkg installer into its own `bin/flepimop-install-r-pkgs` script. This is helpful, because: 1. It helps to modularize `bin/flepimop-install` a little bit more for maintainability, and 2. Assists with development of R packages by making them easier to reinstall for testing. R does not have an `--editable` install like python does so edits to `flepimop-inference-*` require reinstall. --- bin/flepimop-install | 15 +-------------- bin/flepimop-install-r-pkgs | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 14 deletions(-) create mode 100755 bin/flepimop-install-r-pkgs diff --git a/bin/flepimop-install b/bin/flepimop-install index 5b874d9ca..d725d4377 100755 --- a/bin/flepimop-install +++ b/bin/flepimop-install @@ -172,20 +172,7 @@ else fi # Install the local R packages -CONDA_R_LIBS=$( realpath "$CONDA_PREFIX/lib/R/library" ) -R -e "install.packages('epidatr', '$CONDA_R_LIBS', repos='https://cloud.r-project.org')" -RETURNTO=$( pwd ) -cd $FLEPI_PATH/flepimop/R_packages/ -R_PACKAGES=("flepicommon" "flepiconfig" "inference") -for pkg in "${R_PACKAGES[@]}"; do - if [ -d "$pkg" ]; then - R_LIBS=$CONDA_R_LIBS R CMD INSTALL --install-tests $pkg - else - echo "WARNING: The R package '$pkg' does not exist in '$FLEPI_PATH/flepimop/R_packages/'. Skipping installation." >&2 - fi -done -cd $RETURNTO -R -e "library(inference); inference::install_cli()" +FLEPI_PATH=$FLEPI_PATH CONDA_PREFIX=$CONDA_PREFIX $FLEPI_PATH/bin/flepimop-install-r-pkgs # Output a summary of the installation FLEPI_PATH=$FLEPI_PATH FLEPI_CONDA=$FLEPI_CONDA $FLEPI_PATH/bin/flepimop-install-info diff --git a/bin/flepimop-install-r-pkgs b/bin/flepimop-install-r-pkgs new file mode 100755 index 000000000..e3db98a2f --- /dev/null +++ b/bin/flepimop-install-r-pkgs @@ -0,0 +1,19 @@ +#!/usr/bin/env bash + +if [[ -z "$FLEPI_PATH" ]]; then + FLEPI_PATH=$( pwd ) +fi +CONDA_R_LIBS=$( realpath "$CONDA_PREFIX/lib/R/library" ) +R -e "install.packages('epidatr', '$CONDA_R_LIBS', repos='https://cloud.r-project.org')" +RETURNTO=$( pwd ) +cd $FLEPI_PATH/flepimop/R_packages/ +R_PACKAGES=("flepicommon" "flepiconfig" "inference") +for pkg in "${R_PACKAGES[@]}"; do + if [ -d "$pkg" ]; then + R_LIBS=$CONDA_R_LIBS R CMD INSTALL --install-tests $pkg + else + echo "WARNING: The R package '$pkg' does not exist in '$FLEPI_PATH/flepimop/R_packages/'. Skipping installation." >&2 + fi +done +cd $RETURNTO +R -e "library(inference); inference::install_cli()" From 17565d5a17dc96cc3517cb694eb79ea7cb50a87d Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Fri, 25 Jul 2025 14:17:33 -0400 Subject: [PATCH 71/73] Adjustment for `gempyor.seir.write_seir` args Follow up to #585 to account for the arguments of `gempyor.seir.write_seir` taking the out dataframe instead of the states, leaving it to the user to call `gempyor.seir.states2DF`. --- .../emcee_vs_r_inference/model_output.ipynb | 258 ++++++++++++++++++ flepimop/gempyor_pkg/src/gempyor/inference.py | 11 +- 2 files changed, 265 insertions(+), 4 deletions(-) create mode 100644 examples/emcee_vs_r_inference/model_output.ipynb diff --git a/examples/emcee_vs_r_inference/model_output.ipynb b/examples/emcee_vs_r_inference/model_output.ipynb new file mode 100644 index 000000000..532596973 --- /dev/null +++ b/examples/emcee_vs_r_inference/model_output.ipynb @@ -0,0 +1,258 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "c4db2e93", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/twillard/Desktop/GitHub/HopkinsIDD/flepiMoP/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], + "source": [ + "from pathlib import Path\n", + "\n", + "import pandas as pd\n", + "import seaborn as sns\n", + "\n", + "from gempyor.utils import read_directory" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "8a71d019", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.microsoft.datawrangler.viewer.v0+json": { + "columns": [ + { + "name": "index", + "rawType": "int64", + "type": "integer" + }, + { + "name": "date", + "rawType": "datetime64[ns]", + "type": "datetime" + }, + { + "name": "subpop", + "rawType": "object", + "type": "string" + }, + { + "name": "hospitalizations", + "rawType": "float64", + "type": "float" + }, + { + "name": "hospitalizations_curr", + "rawType": "float64", + "type": "float" + }, + { + "name": "slot", + "rawType": "int64", + "type": "integer" + } + ], + "ref": "2bc2bdd3-a4f3-4ad9-aadc-3742fd50daee", + "rows": [ + [ + "0", + "2024-01-01 00:00:00", + "NC", + "0.0", + "0.0", + "1" + ], + [ + "1", + "2024-01-02 00:00:00", + "NC", + "0.0", + "0.0", + "1" + ], + [ + "2", + "2024-01-03 00:00:00", + "NC", + "0.0", + "0.0", + "1" + ], + [ + "3", + "2024-01-04 00:00:00", + "NC", + "0.0", + "0.0", + "1" + ], + [ + "4", + "2024-01-05 00:00:00", + "NC", + "0.0", + "0.0", + "1" + ] + ], + "shape": { + "columns": 5, + "rows": 5 + } + }, + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
datesubpophospitalizationshospitalizations_currslot
02024-01-01NC0.00.01
12024-01-02NC0.00.01
22024-01-03NC0.00.01
32024-01-04NC0.00.01
42024-01-05NC0.00.01
\n", + "
" + ], + "text/plain": [ + " date subpop hospitalizations hospitalizations_curr slot\n", + "0 2024-01-01 NC 0.0 0.0 1\n", + "1 2024-01-02 NC 0.0 0.0 1\n", + "2 2024-01-03 NC 0.0 0.0 1\n", + "3 2024-01-04 NC 0.0 0.0 1\n", + "4 2024-01-05 NC 0.0 0.0 1" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df = read_directory(Path.cwd() / \"model_output\" / \"three_state\" / \"20250804_081545\", filters=[\"hosp\"])\n", + "df.head()" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "95d36500", + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'sns' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[31m---------------------------------------------------------------------------\u001b[39m", + "\u001b[31mNameError\u001b[39m Traceback (most recent call last)", + "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43msns\u001b[49m.lineplot(df, x=\u001b[33m\"\u001b[39m\u001b[33mdate\u001b[39m\u001b[33m\"\u001b[39m, y=\u001b[33m\"\u001b[39m\u001b[33mhospitalizations_curr\u001b[39m\u001b[33m\"\u001b[39m, hue=\u001b[33m\"\u001b[39m\u001b[33mslot\u001b[39m\u001b[33m\"\u001b[39m)\n", + "\u001b[31mNameError\u001b[39m: name 'sns' is not defined" + ] + } + ], + "source": [ + "sns.lineplot(df, x=\"date\", y=\"hospitalizations_curr\", hue=\"slot\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fa1a8075", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/flepimop/gempyor_pkg/src/gempyor/inference.py b/flepimop/gempyor_pkg/src/gempyor/inference.py index 0eda9ec94..a2014bd19 100644 --- a/flepimop/gempyor_pkg/src/gempyor/inference.py +++ b/flepimop/gempyor_pkg/src/gempyor/inference.py @@ -125,8 +125,9 @@ def simulation_atomic( seeding_amounts, ) if save: - seir.write_spar_snpi(sim_id=random_id, modinf=modinf, p_draw=p_draw, npi=npi_seir) - seir.write_seir(sim_id=random_id, modinf=modinf, states=states) + out_df = seir.states2Df(modinf, states) + seir.write_spar_snpi(random_id, modinf, p_draw, npi_seir) + seir.write_seir(random_id, modinf, out_df) # Compute outcomes outcomes_df, hpar_df = outcomes.compute_all_multioutcomes( @@ -616,7 +617,8 @@ def build_structure(self): def write_last_seir(self, sim_id2write=None): if sim_id2write is None: sim_id2write = self.lastsim_sim_id2write - out_df = seir.write_seir(sim_id2write, self.modinf, self.lastsim_states) + out_df = seir.states2Df(self.modinf, self.lastsim_states) + seir.write_seir(sim_id2write, self.modinf, out_df) return out_df # @profile() @@ -783,7 +785,8 @@ def one_simulation( if self.modinf.write_csv or self.modinf.write_parquet: seir.write_spar_snpi(sim_id2write, self.modinf, p_draw, npi_seir) if self.autowrite_seir: - out_df = seir.write_seir(sim_id2write, self.modinf, states) + out_df = seir.states2Df(self.modinf, states) + seir.write_seir(sim_id2write, self.modinf, out_df) self.lastsim_out_df = out_df loaded_values = None From 43a6d4b32c05adeae89bf3562efd8048cddd0647 Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Wed, 6 Aug 2025 16:04:08 -0400 Subject: [PATCH 72/73] Remove spuriously committed jupyter notebook --- .../emcee_vs_r_inference/model_output.ipynb | 258 ------------------ 1 file changed, 258 deletions(-) delete mode 100644 examples/emcee_vs_r_inference/model_output.ipynb diff --git a/examples/emcee_vs_r_inference/model_output.ipynb b/examples/emcee_vs_r_inference/model_output.ipynb deleted file mode 100644 index 532596973..000000000 --- a/examples/emcee_vs_r_inference/model_output.ipynb +++ /dev/null @@ -1,258 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "c4db2e93", - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/twillard/Desktop/GitHub/HopkinsIDD/flepiMoP/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] - } - ], - "source": [ - "from pathlib import Path\n", - "\n", - "import pandas as pd\n", - "import seaborn as sns\n", - "\n", - "from gempyor.utils import read_directory" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "8a71d019", - "metadata": {}, - "outputs": [ - { - "data": { - "application/vnd.microsoft.datawrangler.viewer.v0+json": { - "columns": [ - { - "name": "index", - "rawType": "int64", - "type": "integer" - }, - { - "name": "date", - "rawType": "datetime64[ns]", - "type": "datetime" - }, - { - "name": "subpop", - "rawType": "object", - "type": "string" - }, - { - "name": "hospitalizations", - "rawType": "float64", - "type": "float" - }, - { - "name": "hospitalizations_curr", - "rawType": "float64", - "type": "float" - }, - { - "name": "slot", - "rawType": "int64", - "type": "integer" - } - ], - "ref": "2bc2bdd3-a4f3-4ad9-aadc-3742fd50daee", - "rows": [ - [ - "0", - "2024-01-01 00:00:00", - "NC", - "0.0", - "0.0", - "1" - ], - [ - "1", - "2024-01-02 00:00:00", - "NC", - "0.0", - "0.0", - "1" - ], - [ - "2", - "2024-01-03 00:00:00", - "NC", - "0.0", - "0.0", - "1" - ], - [ - "3", - "2024-01-04 00:00:00", - "NC", - "0.0", - "0.0", - "1" - ], - [ - "4", - "2024-01-05 00:00:00", - "NC", - "0.0", - "0.0", - "1" - ] - ], - "shape": { - "columns": 5, - "rows": 5 - } - }, - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
datesubpophospitalizationshospitalizations_currslot
02024-01-01NC0.00.01
12024-01-02NC0.00.01
22024-01-03NC0.00.01
32024-01-04NC0.00.01
42024-01-05NC0.00.01
\n", - "
" - ], - "text/plain": [ - " date subpop hospitalizations hospitalizations_curr slot\n", - "0 2024-01-01 NC 0.0 0.0 1\n", - "1 2024-01-02 NC 0.0 0.0 1\n", - "2 2024-01-03 NC 0.0 0.0 1\n", - "3 2024-01-04 NC 0.0 0.0 1\n", - "4 2024-01-05 NC 0.0 0.0 1" - ] - }, - "execution_count": 2, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df = read_directory(Path.cwd() / \"model_output\" / \"three_state\" / \"20250804_081545\", filters=[\"hosp\"])\n", - "df.head()" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "95d36500", - "metadata": {}, - "outputs": [ - { - "ename": "NameError", - "evalue": "name 'sns' is not defined", - "output_type": "error", - "traceback": [ - "\u001b[31m---------------------------------------------------------------------------\u001b[39m", - "\u001b[31mNameError\u001b[39m Traceback (most recent call last)", - "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43msns\u001b[49m.lineplot(df, x=\u001b[33m\"\u001b[39m\u001b[33mdate\u001b[39m\u001b[33m\"\u001b[39m, y=\u001b[33m\"\u001b[39m\u001b[33mhospitalizations_curr\u001b[39m\u001b[33m\"\u001b[39m, hue=\u001b[33m\"\u001b[39m\u001b[33mslot\u001b[39m\u001b[33m\"\u001b[39m)\n", - "\u001b[31mNameError\u001b[39m: name 'sns' is not defined" - ] - } - ], - "source": [ - "sns.lineplot(df, x=\"date\", y=\"hospitalizations_curr\", hue=\"slot\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fa1a8075", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.13" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} From 4e0e96b8ab3e3d8ec1003b662f24aaf4355b548f Mon Sep 17 00:00:00 2001 From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com> Date: Tue, 19 Aug 2025 09:45:41 -0400 Subject: [PATCH 73/73] Allow explicit 0 mobility Prior to #549 providing an explicit mobility of 0 was allowed, this reallows this previous behaivor. Minor change to the `MobilityFileRow` pydantic model was required and then adjust ments to the unit tests. Should address #605. --- .../src/gempyor/subpopulation_structure.py | 8 +- .../test_subpopulation_structure_class.py | 93 +++++++++++++++++-- 2 files changed, 91 insertions(+), 10 deletions(-) diff --git a/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py b/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py index 5cd8bb7cc..2a75c51af 100644 --- a/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py +++ b/flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py @@ -133,7 +133,7 @@ class GeodataFileTable(RootModel): root: list[GeodataFileRow] @model_validator(mode="after") - def subpop_is_primary_key(self) -> "GeodataFileTable": + def _subpop_is_primary_key(self) -> "GeodataFileTable": """ Validate that the subpopulation names are unique. @@ -163,10 +163,10 @@ class MobilityFileRow(BaseModel): ori: str dest: str - amount: Annotated[int, Field(gt=0)] + amount: Annotated[int, Field(ge=0)] @model_validator(mode="after") - def ori_and_dest_are_different(self) -> "MobilityFileRow": + def _ori_and_dest_are_different(self) -> "MobilityFileRow": """ Validate that the origin and destination subpopulation names are different. @@ -194,7 +194,7 @@ class MobilityFileTable(RootModel): root: list[MobilityFileRow] @model_validator(mode="after") - def ori_and_dest_are_primary_key(self) -> "MobilityFileTable": + def _ori_and_dest_are_primary_key(self) -> "MobilityFileTable": """ Validate that the origin and destination subpopulation names are unique. diff --git a/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py b/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py index 7ae75760a..cfbad73a9 100644 --- a/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py +++ b/flepimop/gempyor_pkg/tests/subpopulation_structure/test_subpopulation_structure_class.py @@ -365,6 +365,51 @@ def valid_2pop_with_csv_mobility_factory( ) +def valid_2pop_with_csv_explicit_no_mobility_factory( + tmp_path: Path, +) -> MockSubpopulationStructureInput: + """ + Factory for geodata file with two valid subpopulations and explicit no mobility. + + Returns: + A `MockSubpopulationStructureInput` instance with a geodata file that contains + two valid subpopulations and a mobility matrix. + """ + return MockSubpopulationStructureInput.create_mock_input( + tmp_path, + { + "geodata": "geodata.csv", + "mobility": "mobility.csv", + }, + pd.DataFrame.from_records( + [ + { + "subpop": "USA", + "population": 100, + }, + { + "subpop": "Canada", + "population": 50, + }, + ] + ), + mobility=pd.DataFrame.from_records( + [ + { + "ori": "USA", + "dest": "Canada", + "amount": 0, + }, + { + "ori": "Canada", + "dest": "USA", + "amount": 2, + }, + ] + ), + ) + + def valid_2pop_with_parquet_mobility_factory( tmp_path: Path, ) -> MockSubpopulationStructureInput: @@ -684,6 +729,29 @@ def mobility_zero_or_less_factory(tmp_path: Path) -> MockSubpopulationStructureI ) +def larger_mobility_zero_or_less_factory(tmp_path: Path) -> MockSubpopulationStructureInput: + return MockSubpopulationStructureInput.create_mock_input( + tmp_path, + { + "geodata": "geodata.parquet", + "mobility": "mobility.parquet", + }, + pd.DataFrame( + data={ + "subpop": ["USA", "Canada", "Mexico"], + "population": [100, 50, 25], + }, + ), + mobility=pd.DataFrame( + data={ + "ori": ["USA", "Canada", "Mexico", "Canada", "USA", "Mexico"], + "dest": ["Canada", "USA", "Canada", "Mexico", "Mexico", "USA"], + "amount": [0, -1, 1, 0, -1, 0], + }, + ), + ) + + def selected_missing_from_geodata(tmp_path: Path) -> MockSubpopulationStructureInput: return MockSubpopulationStructureInput.create_mock_input( tmp_path, @@ -711,6 +779,7 @@ def selected_missing_from_geodata(tmp_path: Path) -> MockSubpopulationStructureI valid_2pop_geodata_parquet_only_factory, valid_2pop_with_txt_mobility_factory, valid_2pop_with_csv_mobility_factory, + valid_2pop_with_csv_explicit_no_mobility_factory, valid_2pop_with_parquet_mobility_factory, valid_2pop_with_npz_mobility_factory, valid_selected_pop_with_no_mobility_factory, @@ -810,13 +879,25 @@ def test_subpopulation_structure_instance_attributes( assert len(caplog.records) == int(mock_input.mobility is None) -def test_mobility_zero_or_less_raises_validation_error(tmp_path: Path) -> None: - """Test that a ValueError is raised when mobility is zero or less.""" - mock_input = mobility_zero_or_less_factory(tmp_path) - assert (zero_or_less := (mock_input.mobility["amount"] <= 0).sum()) > 0 - raises_match = rf"^{zero_or_less} validation errors.*" +@pytest.mark.parametrize( + "factory", [mobility_zero_or_less_factory, larger_mobility_zero_or_less_factory] +) +def test_mobility_less_than_zero_raises_validation_error( + tmp_path: Path, factory: Callable[[Path], MockSubpopulationStructureInput] +) -> None: + """Test that a ValueError is raised when mobility is less than zero.""" + mock_input = factory(tmp_path) + assert ( + less_than_zero_indices := ( + mock_input.mobility.index[mock_input.mobility["amount"] < 0] + ).tolist() + ) + raises_match = rf"^{len(less_than_zero_indices)} validation error.*" raises_match += ".*".join( - [rf"{i}\.amount\s+Input should be greater than 0" for i in range(zero_or_less)] + [ + rf"{i}\.amount\s+Input should be greater than or equal to 0" + for i in less_than_zero_indices + ] ) with pytest.raises( ValidationError,