drisk 0.1.0__py3-none-any.whl

drisk/__init__.py

@@ -0,0 +1,89 @@
+"""Convenient tools for quick Monte Carlo modelling."""
+
+from . import _style as _style
+from .copulas import Copula, GaussianCopula, StudentTCopula
+from .correlations import CorrelationMatrix
+from .decision import (
+    ChanceBranch,
+    ChanceNode,
+    DecisionBranch,
+    DecisionNode,
+    DTree,
+    DTreeNode,
+    OutcomeNode,
+)
+from .distributions import (
+    PERT,
+    ArrayLike,
+    Bernoulli,
+    Beta,
+    Binomial,
+    DataFrameLike,
+    Distribution,
+    Exponential,
+    Gamma,
+    Geometric,
+    LogitNormal,
+    LogNormal,
+    NegativeBinomial,
+    Normal,
+    Poisson,
+    StretchedBeta,
+    UvBoundedContinuous,
+    UvContinuous,
+    UvCountDiscrete,
+    UvDiscrete,
+    UvDistribution,
+    UvFiniteDiscrete,
+    UvMixture,
+    UvPositiveContinuous,
+    UvRealContinuous,
+    UvUnitBoundedContinuous,
+)
+from .models import MCModel, MCOperation, where
+from .sensitivity import OneAtATimeSensitivity, one_at_a_time
+
+__all__ = [
+    "ArrayLike",
+    "Bernoulli",
+    "Beta",
+    "Binomial",
+    "ChanceBranch",
+    "ChanceNode",
+    "Copula",
+    "CorrelationMatrix",
+    "DataFrameLike",
+    "DecisionBranch",
+    "DecisionNode",
+    "Distribution",
+    "DTree",
+    "DTreeNode",
+    "Exponential",
+    "GaussianCopula",
+    "Gamma",
+    "Geometric",
+    "LogitNormal",
+    "MCModel",
+    "MCOperation",
+    "LogNormal",
+    "UvMixture",
+    "NegativeBinomial",
+    "Normal",
+    "OneAtATimeSensitivity",
+    "OutcomeNode",
+    "one_at_a_time",
+    "PERT",
+    "Poisson",
+    "StudentTCopula",
+    "StretchedBeta",
+    "UvBoundedContinuous",
+    "UvContinuous",
+    "UvCountDiscrete",
+    "UvDiscrete",
+    "UvDistribution",
+    "UvFiniteDiscrete",
+    "UvPositiveContinuous",
+    "UvRealContinuous",
+    "UvUnitBoundedContinuous",
+    "where",
+]

drisk/_style.py

@@ -0,0 +1,5 @@
+"""Package-wide Matplotlib defaults."""
+
+import matplotlib as mpl
+
+mpl.rcParams.update({"axes.grid": True, "grid.alpha": 0.25})

drisk/arithmetic.py

@@ -0,0 +1,111 @@
+"""Arithmetic helpers for composable Monte Carlo expressions."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class ArithmeticMixin:
+    """Mixin that turns arithmetic into lazy Monte Carlo model expressions."""
+
+    def __add__(self, other: Any) -> Any:
+        """Create a lazy model for ``self + other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.ADD, self, other)
+
+    def __radd__(self, other: Any) -> Any:
+        """Create a lazy model for ``other + self``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.ADD, other, self)
+
+    def __sub__(self, other: Any) -> Any:
+        """Create a lazy model for ``self - other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.SUBTRACT, self, other)
+
+    def __rsub__(self, other: Any) -> Any:
+        """Create a lazy model for ``other - self``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.SUBTRACT, other, self)
+
+    def __mul__(self, other: Any) -> Any:
+        """Create a lazy model for ``self * other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.MULTIPLY, self, other)
+
+    def __rmul__(self, other: Any) -> Any:
+        """Create a lazy model for ``other * self``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.MULTIPLY, other, self)
+
+    def __truediv__(self, other: Any) -> Any:
+        """Create a lazy model for ``self / other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.DIVIDE, self, other)
+
+    def __rtruediv__(self, other: Any) -> Any:
+        """Create a lazy model for ``other / self``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.DIVIDE, other, self)
+
+    def __pow__(self, other: Any) -> Any:
+        """Create a lazy model for ``self ** other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.POWER, self, other)
+
+    def __rpow__(self, other: Any) -> Any:
+        """Create a lazy model for ``other ** self``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.POWER, other, self)
+
+    def __lt__(self, other: Any) -> Any:
+        """Create a lazy model for ``self < other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.LESS, self, other)
+
+    def __le__(self, other: Any) -> Any:
+        """Create a lazy model for ``self <= other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.LESS_EQUAL, self, other)
+
+    def __gt__(self, other: Any) -> Any:
+        """Create a lazy model for ``self > other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.GREATER, self, other)
+
+    def __ge__(self, other: Any) -> Any:
+        """Create a lazy model for ``self >= other``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.GREATER_EQUAL, self, other)
+
+    def __neg__(self) -> Any:
+        """Create a lazy model for ``-self``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.NEGATIVE, self)
+
+    def __pos__(self) -> Any:
+        """Create a lazy model for ``+self``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.POSITIVE, self)
+
+    def __abs__(self) -> Any:
+        """Create a lazy model for ``abs(self)``."""
+        from drisk.models import MCModel, MCOperation
+
+        return MCModel.from_operation(MCOperation.ABS, self)

drisk/copulas/__init__.py

@@ -0,0 +1,7 @@
+"""Copula models for jointly sampling marginal distributions."""
+
+from .base import Copula
+from .gaussian import GaussianCopula
+from .student_t import StudentTCopula
+
+__all__ = ["Copula", "GaussianCopula", "StudentTCopula"]

drisk/copulas/base.py

@@ -0,0 +1,95 @@
+"""Base interfaces for copula models."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from inspect import isabstract
+from typing import Any, Self
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, GetCoreSchemaHandler, model_validator
+from pydantic_core import CoreSchema, core_schema
+
+from drisk.correlations import CorrelationMatrix
+from drisk.distributions.univariate import UvDistribution
+from drisk.random import SeedLike
+
+
+class Copula(BaseModel, ABC):
+    """Base class for copulas that jointly sample marginal distributions."""
+
+    distributions: tuple[UvDistribution, ...]
+    corr_matrix: CorrelationMatrix
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """Use ``copula_type`` to validate abstract copula-typed fields."""
+        if not getattr(cls, "__pydantic_complete__", False) or not isabstract(cls):
+            return handler(source_type)
+
+        try:
+            from drisk.copulas.registry import concrete_copula_types_for
+        except ImportError:
+            return handler(source_type)
+
+        try:
+            choices = {
+                copula_cls.model_fields["copula_type"].default: handler.generate_schema(
+                    copula_cls
+                )
+                for copula_cls in concrete_copula_types_for(cls)
+            }
+        except ImportError:
+            return handler(source_type)
+
+        if not choices:
+            return handler(source_type)
+
+        return core_schema.tagged_union_schema(
+            choices=choices,
+            discriminator="copula_type",
+            from_attributes=True,
+        )
+
+    @property
+    def dims(self) -> int:
+        """Number of marginal distributions."""
+        return len(self.distributions)
+
+    @model_validator(mode="after")
+    def validate_dimensions(self) -> Self:
+        """Ensure the correlation matrix dimension matches the marginals."""
+        n = len(self.distributions)
+        matrix_n = len(self.corr_matrix.matrix)
+        if matrix_n != n:
+            raise ValueError(
+                f"Correlation matrix size ({matrix_n}) does not match number of distributions ({n})."
+            )
+        return self
+
+    @classmethod
+    def from_distributions_and_correlation(
+        cls,
+        distributions: Sequence[UvDistribution],
+        correlation: float,
+        **kwargs: object,
+    ) -> Self:
+        """Create a copula from marginals and one shared pairwise correlation."""
+        corr_matrix = CorrelationMatrix.from_n_corr(len(distributions), correlation)
+        return cls(distributions=distributions, corr_matrix=corr_matrix, **kwargs)
+
+    @abstractmethod
+    def sample(
+        self, size: int | tuple[int, ...] = 1, *, seed: SeedLike = None
+    ) -> np.ndarray:
+        """Jointly sample marginals, returning an array shaped ``(dims, *size)``."""
+        pass
+
+    def rvs(
+        self, size: int | tuple[int, ...] = 1, *, seed: SeedLike = None
+    ) -> np.ndarray:
+        """Alias for :meth:`sample` for users familiar with SciPy naming."""
+        return self.sample(size=size, seed=seed)

drisk/copulas/gaussian.py

@@ -0,0 +1,38 @@
+"""Gaussian copula."""
+
+from typing import Literal
+
+import numpy as np
+from scipy import stats
+
+from drisk.random import SeedLike, get_rng
+
+from .base import Copula
+
+
+class GaussianCopula(Copula):
+    """Sample marginal distributions with dependence induced by a Gaussian copula."""
+
+    copula_type: Literal["gaussian"] = "gaussian"
+
+    def sample(
+        self, size: int | tuple[int, ...] = 1, *, seed: SeedLike = None
+    ) -> np.ndarray:
+        """Jointly sample marginals, returning an array shaped ``(dims, *size)``."""
+        if isinstance(size, int):
+            size = (size,)
+
+        rng = get_rng(seed)
+        normal_samples = rng.multivariate_normal(
+            mean=np.zeros(self.dims),
+            cov=self.corr_matrix.to_numpy(),
+            size=size,
+        )
+        normal_samples = np.moveaxis(normal_samples, -1, 0)
+        uniform_samples = stats.norm.cdf(normal_samples)
+
+        samples = np.empty_like(uniform_samples)
+        for i, dist in enumerate(self.distributions):
+            samples[i, ...] = dist.ppf(uniform_samples[i, ...])
+
+        return samples

drisk/copulas/registry.py

@@ -0,0 +1,26 @@
+"""Registry of concrete copula implementations for Pydantic polymorphism."""
+
+from functools import cache
+from typing import cast
+
+from drisk.copulas.base import Copula
+
+
+@cache
+def concrete_copula_types() -> tuple[type[Copula], ...]:
+    """Return all concrete copula classes supported by Drisk."""
+    from drisk.copulas.gaussian import GaussianCopula
+    from drisk.copulas.student_t import StudentTCopula
+
+    return (GaussianCopula, StudentTCopula)
+
+
+def concrete_copula_types_for[CopulaT: Copula](
+    base_cls: type[CopulaT],
+) -> tuple[type[CopulaT], ...]:
+    """Return concrete registered copulas that are subclasses of ``base_cls``."""
+    return tuple(
+        cast(type[CopulaT], copula_cls)
+        for copula_cls in concrete_copula_types()
+        if issubclass(copula_cls, base_cls)
+    )

drisk/copulas/student_t.py

@@ -0,0 +1,51 @@
+"""Student-t copula."""
+
+from typing import Literal
+
+import numpy as np
+from pydantic import field_validator
+from scipy import stats
+
+from drisk.random import SeedLike, get_rng
+
+from .base import Copula
+
+
+class StudentTCopula(Copula):
+    """Sample marginals with dependence induced by a Student-t copula."""
+
+    copula_type: Literal["student_t"] = "student_t"
+    nu: float = 4.0
+
+    @field_validator("nu")
+    @classmethod
+    def validate_nu(cls, nu: float) -> float:
+        """Validate degrees of freedom."""
+        if nu <= 0:
+            raise ValueError(f"nu must be positive, got {nu}")
+        return nu
+
+    def sample(
+        self, size: int | tuple[int, ...] = 1, *, seed: SeedLike = None
+    ) -> np.ndarray:
+        """Jointly sample marginals, returning an array shaped ``(dims, *size)``."""
+        if isinstance(size, int):
+            size = (size,)
+
+        rng = get_rng(seed)
+        normal_samples = rng.multivariate_normal(
+            mean=np.zeros(self.dims),
+            cov=self.corr_matrix.to_numpy(),
+            size=size,
+        )
+        normal_samples = np.moveaxis(normal_samples, -1, 0)
+
+        chi_square_samples = rng.chisquare(df=self.nu, size=size) / self.nu
+        t_samples = normal_samples / np.sqrt(chi_square_samples)
+        uniform_samples = stats.t.cdf(t_samples, df=self.nu)
+
+        samples = np.empty_like(uniform_samples)
+        for i, dist in enumerate(self.distributions):
+            samples[i, ...] = dist.ppf(uniform_samples[i, ...])
+
+        return samples

drisk/correlations/__init__.py

@@ -0,0 +1,5 @@
+"""Correlation structures for Monte Carlo modelling."""
+
+from .matrix import CorrelationMatrix
+
+__all__ = ["CorrelationMatrix"]

drisk/correlations/matrix.py

@@ -0,0 +1,151 @@
+"""Correlation matrix validation and helpers."""
+
+from typing import Any, Self
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, field_validator
+
+
+class CorrelationMatrix(BaseModel):
+    """Represent and validate a numeric correlation matrix."""
+
+    matrix: list[list[float]]
+
+    model_config = ConfigDict(extra="forbid")
+
+    @field_validator("matrix")
+    @classmethod
+    def validate_matrix(cls, matrix: list[list[float]]) -> list[list[float]]:
+        """Validate shape, correlation bounds, symmetry, and PSD-ness."""
+        if not matrix:
+            raise ValueError("Matrix cannot be empty")
+
+        arr = np.asarray(matrix, dtype=float)
+
+        if arr.ndim != 2 or arr.shape[0] != arr.shape[1]:
+            raise ValueError(f"Matrix must be square, got shape {arr.shape}")
+
+        diagonal = np.diag(arr)
+        if not np.allclose(diagonal, 1.0):
+            bad_indices = np.where(~np.isclose(diagonal, 1.0))[0]
+            i = int(bad_indices[0])
+            raise ValueError(
+                f"Diagonal element at ({i}, {i}) must be 1.0, got {diagonal[i]}"
+            )
+
+        if not np.allclose(arr, arr.T):
+            diff = np.abs(arr - arr.T)
+            i, j = np.unravel_index(np.argmax(diff), diff.shape)
+            if i > j:
+                i, j = j, i
+            raise ValueError(
+                f"Matrix is not symmetric: ({i}, {j})={arr[i, j]} != ({j}, {i})={arr[j, i]}"
+            )
+
+        if not np.all((arr >= -1.0) & (arr <= 1.0)):
+            bad_mask = (arr < -1.0) | (arr > 1.0)
+            i, j = np.unravel_index(np.argmax(bad_mask), arr.shape)
+            raise ValueError(
+                f"Correlation value at ({i}, {j}) must be between -1 and 1, got {arr[i, j]}"
+            )
+
+        eigenvalues = np.linalg.eigvalsh(arr)
+        min_eigenvalue = float(np.min(eigenvalues))
+        if min_eigenvalue < -1e-10:
+            raise ValueError(
+                "Correlation matrix must be positive semidefinite; "
+                f"minimum eigenvalue is {min_eigenvalue}"
+            )
+
+        return matrix
+
+    @classmethod
+    def from_n_corr(cls, n: int, corr: float) -> Self:
+        """Create an ``n`` by ``n`` matrix with a shared off-diagonal correlation."""
+        if n <= 0:
+            raise ValueError(f"n must be positive, got {n}")
+        if not (-1.0 <= corr <= 1.0):
+            raise ValueError(f"Correlation value must be between -1 and 1, got {corr}")
+
+        matrix = [[1.0 if i == j else float(corr) for j in range(n)] for i in range(n)]
+        return cls(matrix=matrix)
+
+    @classmethod
+    def from_numpy(cls, arr: np.ndarray) -> Self:
+        """Create a correlation matrix from a NumPy array."""
+        if arr.ndim != 2:
+            raise ValueError(f"Array must be 2-dimensional, got {arr.ndim}")
+        return cls(matrix=arr.tolist())
+
+    def to_numpy(self) -> np.ndarray:
+        """Return the correlation matrix as a NumPy array."""
+        return np.asarray(self.matrix, dtype=float)
+
+    def plot(
+        self,
+        ax: Any = None,
+        *,
+        labels: list[str] | None = None,
+        cmap: str = "Spectral",
+        show: bool = False,
+        colorbar: bool = True,
+        **imshow_kwargs: Any,
+    ) -> Any:
+        """
+        Plot the correlation matrix as an annotated heatmap.
+
+        Returns the Matplotlib ``Axes`` object. Importing Matplotlib is deferred
+        so non-plotting use stays lightweight. Extra keyword arguments are
+        passed to ``imshow``.
+        """
+        if ax is None:
+            import matplotlib.pyplot as plt
+
+            _, ax = plt.subplots()
+
+        arr = self.to_numpy()
+        n = arr.shape[0]
+
+        if labels is not None and len(labels) != n:
+            raise ValueError(f"labels must have length {n}, got {len(labels)}")
+
+        image = ax.imshow(
+            arr,
+            cmap=cmap,
+            vmin=-1,
+            vmax=1,
+            **imshow_kwargs,
+        )
+
+        tick_labels = labels if labels is not None else [str(i) for i in range(n)]
+        ax.set_xticks(np.arange(n), labels=tick_labels, rotation=90)
+        ax.tick_params(
+            axis="x", bottom=True, labelbottom=True, top=False, labeltop=False
+        )
+        ax.set_yticks(np.arange(n), labels=tick_labels)
+        ax.grid(False)
+        ax.set_title("Correlation matrix")
+
+        for i in range(n):
+            for j in range(n):
+                text_color = "white" if abs(arr[i, j]) > 0.5 else "black"
+                ax.text(
+                    j,
+                    i,
+                    f"{arr[i, j]:.2f}",
+                    ha="center",
+                    va="center",
+                    color=text_color,
+                )
+
+        if colorbar:
+            ax.figure.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
+
+        ax.figure.tight_layout()
+
+        if show:
+            import matplotlib.pyplot as plt
+
+            plt.show()
+
+        return ax

drisk/decision/__init__.py

@@ -0,0 +1,23 @@
+"""Decision analysis support."""
+
+from .dtree import (
+    ChanceBranch,
+    ChanceNode,
+    DecisionBranch,
+    DecisionNode,
+    DTree,
+    DTreeNode,
+    OutcomeNode,
+    as_node,
+)
+
+__all__ = [
+    "ChanceBranch",
+    "ChanceNode",
+    "DecisionBranch",
+    "DecisionNode",
+    "DTree",
+    "DTreeNode",
+    "OutcomeNode",
+    "as_node",
+]

drisk/decision/dtree/__init__.py

@@ -0,0 +1,17 @@
+"""Decision tree support."""
+
+from .chance_branch import ChanceBranch
+from .decision_branch import DecisionBranch
+from .nodes import ChanceNode, DecisionNode, DTreeNode, OutcomeNode, as_node
+from .tree import DTree
+
+__all__ = [
+    "ChanceBranch",
+    "ChanceNode",
+    "DecisionBranch",
+    "DecisionNode",
+    "DTree",
+    "DTreeNode",
+    "OutcomeNode",
+    "as_node",
+]

drisk/decision/dtree/_coercion.py

@@ -0,0 +1,48 @@
+"""Coercion helpers for decision tree inputs."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .chance_branch import ChanceBranch
+from .decision_branch import DecisionBranch
+
+
+def coerce_decision_branches(branches: Any) -> list[DecisionBranch]:
+    from .nodes.factory import as_node
+
+    if isinstance(branches, dict):
+        return [
+            DecisionBranch(name=str(name), node=as_node(value))
+            for name, value in branches.items()
+        ]
+    return [
+        branch
+        if isinstance(branch, DecisionBranch)
+        else DecisionBranch.model_validate(branch)
+        for branch in branches
+    ]
+
+
+def coerce_chance_branches(branches: Any) -> list[ChanceBranch]:
+    from .nodes.factory import as_node
+
+    if isinstance(branches, dict):
+        coerced = []
+        for name, spec in branches.items():
+            if isinstance(spec, tuple) and len(spec) == 2:
+                probability, value = spec
+                coerced.append(
+                    ChanceBranch(
+                        name=str(name), probability=probability, node=as_node(value)
+                    )
+                )
+            else:
+                coerced.append(ChanceBranch.model_validate({"name": name, **spec}))
+        return coerced
+    return [
+        branch
+        if isinstance(branch, ChanceBranch)
+        else ChanceBranch.model_validate(branch)
+        for branch in branches
+    ]