Merge pull request #29 from pstefanou12/waterboi/migrate-distributions-to-truncated

Migrate distributions/ to truncated/distributions/ and refactor

Author: pstefanou12

delphi/delphi.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import torch as ch
+from pydantic import BaseModel
 from torch.optim import SGD, LBFGS, Adam, AdamW, lr_scheduler
 
 from delphi.utils.constants import (
@@ -15,15 +16,7 @@
     PythonFrameworks,
     SchedulerType,
 )
-from delphi.utils.defaults import (
-    check_and_fill_args,
-    DELPHI_DEFAULTS,
-    SGD_DEFAULTS,
-    LBFGS_DEFAULTS,
-    ADAM_DEFAULTS,
-    ADAMW_DEFAULTS,
-)
-from delphi.utils.helpers import AverageMeter, Parameters
+from delphi.utils.helpers import AverageMeter
 
 
 class delphi(ch.nn.Module):  # pylint: disable=invalid-name,too-many-instance-attributes,abstract-method
@@ -40,22 +33,23 @@ class delphi(ch.nn.Module):  # pylint: disable=invalid-name,too-many-instance-at
 
     _OPTIMIZER_REGISTRY: ClassVar[dict[str, Callable]] = {}
 
-    def __init__(self, args: Parameters):
+    def __init__(self, args: BaseModel):
         """Initialize the delphi model.
 
         Args:
-            args: Hyperparameter object; see DELPHI_DEFAULTS for supported keys.
+            args: Fully constructed Pydantic config. Concrete subclasses are
+                responsible for converting a user-supplied dict to their
+                specific config before calling super().__init__.
 
         Raises:
-            TypeError: If args is not a Parameters instance.
+            TypeError: If args is not a Pydantic BaseModel.
         """
         super().__init__()
-        if not isinstance(args, Parameters):
+        if not isinstance(args, BaseModel):
             raise TypeError(
-                f"args is type {type(args).__name__}; "
-                "expected delphi.utils.helpers.Parameters"
+                f"args is type {type(args).__name__}; expected a pydantic.BaseModel."
             )
-        self.args: Parameters = check_and_fill_args(args, DELPHI_DEFAULTS)
+        self.args = args
 
         self.optimizer: ch.optim.Optimizer | None = None
         self.schedule: lr_scheduler.LRScheduler | None = None
@@ -149,7 +143,6 @@ def _remove_none_config(self, config: dict) -> dict:
 
     def _create_sgd(self, params: list[dict]) -> SGD:
         """Create an SGD optimizer from args."""
-        check_and_fill_args(self.args, SGD_DEFAULTS)
         config = {
             "lr": self.args.lr,
             "momentum": getattr(self.args, "momentum", 0),
@@ -165,7 +158,6 @@ def _create_sgd(self, params: list[dict]) -> SGD:
 
     def _create_lbfgs(self, params: list[dict]) -> LBFGS:
         """Create an L-BFGS optimizer from args."""
-        check_and_fill_args(self.args, LBFGS_DEFAULTS)
         config = {
             "lr": getattr(self.args, "lr", 1.0),
             "max_iter": getattr(self.args, "max_iter", 20),
@@ -179,7 +171,6 @@ def _create_lbfgs(self, params: list[dict]) -> LBFGS:
 
     def _create_adam(self, params: list[dict]) -> Adam:
         """Create an Adam optimizer from args."""
-        check_and_fill_args(self.args, ADAM_DEFAULTS)
         config = {
             "lr": getattr(self.args, "lr", 1e-1),
             "betas": (
@@ -199,7 +190,6 @@ def _create_adam(self, params: list[dict]) -> Adam:
 
     def _create_adamw(self, params: list[dict]) -> AdamW:
         """Create an AdamW optimizer from args."""
-        check_and_fill_args(self.args, ADAMW_DEFAULTS)
         config = {
             "lr": getattr(self.args, "lr", 1e-3),
             "betas": (

delphi/distributions/__init__.py

@@ -1,28 +1,2 @@
 # Author: pstefanou12@
-"""Distribution models subpackage for delphi."""
-
-from delphi.distributions.truncated_normal import TruncatedNormal
-from delphi.distributions.truncated_multivariate_normal import (
-    TruncatedMultivariateNormal,
-)
-from delphi.distributions.unknown_truncated_normal import UnknownTruncationNormal
-from delphi.distributions.unknown_truncated_multivariate_normal import (
-    UnknownTruncationMultivariateNormal,
-    Exp_h,
-)
-from delphi.distributions.truncated_boolean_product import TruncatedBooleanProduct
-from delphi.distributions.truncated_exponential import TruncatedExponential
-from delphi.distributions.truncated_poisson import TruncatedPoisson
-from delphi.distributions.truncated_weibull import TruncatedWeibull
-
-__all__ = [
-    "TruncatedNormal",
-    "TruncatedMultivariateNormal",
-    "UnknownTruncationNormal",
-    "UnknownTruncationMultivariateNormal",
-    "Exp_h",
-    "TruncatedBooleanProduct",
-    "TruncatedExponential",
-    "TruncatedPoisson",
-    "TruncatedWeibull",
-]
+"""Exponential family distributions in natural parameterization for delphi."""

delphi/distributions/boolean_product.py

@@ -0,0 +1,25 @@
+# Author: pstefanou12@
+"""Boolean product distribution in natural parameterization."""
+
+import torch as ch
+from torch.distributions import Bernoulli
+
+
+def calc_bool_prod_suff_stat(x):
+    """Return sufficient statistics for boolean product distribution."""
+    return x
+
+
+class ExponentialFamilyBooleanProduct(Bernoulli):  # pylint: disable=abstract-method
+    """Boolean product distribution parameterized by natural parameters."""
+
+    def __init__(self, theta: ch.Tensor, dims: int):
+        """Initialize with natural parameter theta and dimension."""
+        self.dims = dims
+        p = ch.exp(theta) / (1 + ch.exp(theta))
+        super().__init__(p)
+
+    def log_prob(self, value):
+        """Compute summed log probability over all dimensions."""
+        result = super().log_prob(value)
+        return result.sum(-1)

delphi/distributions/exponential.py

@@ -0,0 +1,25 @@
+# Author: pstefanou12@
+"""Exponential distribution in natural parameterization."""
+
+import torch as ch
+from torch.distributions import Exponential
+
+
+def calc_exp_suff_stat(x):
+    """Return sufficient statistics for exponential distribution."""
+    return x
+
+
+class ExponentialFamilyExponential(Exponential):  # pylint: disable=abstract-method
+    """Exponential distribution parameterized by natural parameters."""
+
+    def __init__(self, theta: ch.Tensor, dims: int):
+        """Initialize with natural parameter theta and dimension."""
+        self.dims = dims
+        lambda_ = -theta
+        super().__init__(lambda_)
+
+    def log_prob(self, value):
+        """Compute summed log probability over all dimensions."""
+        result = super().log_prob(value)
+        return result.sum(-1)

delphi/distributions/multivariate_normal.py

@@ -0,0 +1,72 @@
+# Author: pstefanou12@
+"""Multivariate normal distribution in natural parameterization."""
+
+import torch as ch
+from torch import Tensor
+from torch.distributions import MultivariateNormal
+
+
+class ExponentialFamilyMultivariateNormalKnownCovariance(  # pylint: disable=abstract-method
+    MultivariateNormal
+):
+    """Multivariate normal parameterized by natural parameters with known covariance."""
+
+    def __init__(self, covariance_matrix: Tensor, theta: Tensor, dims: int):
+        """Initialize with covariance matrix, natural parameter theta, and dimension."""
+        self.dims = dims
+        v = theta
+        mu = (covariance_matrix @ v).view(self.dims)
+        super().__init__(mu, covariance_matrix)
+
+    @staticmethod
+    def calc_suff_stat(x: Tensor) -> Tensor:
+        """Return sufficient statistics for multivariate normal with known covariance."""
+        return x
+
+    @staticmethod
+    def to_natural(theta: Tensor, covariance_matrix: Tensor) -> Tensor:
+        """Convert canonical mean to natural parameters."""
+        inv_cov = covariance_matrix.inverse()
+        v = theta @ inv_cov  # pylint: disable=invalid-name
+        return v.flatten()
+
+    @staticmethod
+    def to_canonical(theta: Tensor, covariance_matrix: Tensor) -> Tensor:
+        """Convert natural parameters to canonical mean."""
+        loc = theta @ covariance_matrix
+        return loc.flatten()
+
+
+class ExponentialFamilyMultivariateNormal(MultivariateNormal):  # pylint: disable=abstract-method
+    """Multivariate normal parameterized by natural parameters."""
+
+    def __init__(self, theta: Tensor, dims: int):
+        """Initialize with natural parameter theta and dimension."""
+        self.dims = dims
+        T, v = theta[: self.dims**2], theta[self.dims**2 :]  # pylint: disable=invalid-name
+        covariance_matrix = ch.inverse(-2 * T.view(self.dims, self.dims))
+        mu = (covariance_matrix @ v).view(self.dims)
+        super().__init__(mu, covariance_matrix)
+
+    @staticmethod
+    def calc_suff_stat(x: Tensor) -> Tensor:
+        """Return sufficient statistics for multivariate normal."""
+        return ch.cat([ch.bmm(x.unsqueeze(2), x.unsqueeze(1)).flatten(1), x], 1)
+
+    @staticmethod
+    def to_natural(theta: Tensor, dims: int) -> Tensor:
+        """Convert canonical parameters to natural form."""
+        cov_matrix = theta[: dims**2].view(dims, dims)
+        loc = theta[dims**2 :]
+        mat_t = cov_matrix.inverse()  # pylint: disable=invalid-name
+        v = loc @ mat_t  # pylint: disable=invalid-name
+        return ch.cat([-0.5 * mat_t.flatten(), v.flatten()])
+
+    @staticmethod
+    def to_canonical(theta: Tensor, dims: int) -> Tensor:
+        """Convert natural parameters to canonical form."""
+        mat_t = theta[: dims**2].view(dims, dims)  # pylint: disable=invalid-name
+        v = theta[dims**2 :]
+        covariance_matrix = (-2 * mat_t).inverse()
+        loc = v @ covariance_matrix
+        return ch.cat([covariance_matrix.flatten(), loc.flatten()])

delphi/distributions/poisson.py

@@ -0,0 +1,25 @@
+# Author: pstefanou12@
+"""Poisson distribution in natural parameterization."""
+
+import torch as ch
+from torch.distributions import Poisson
+
+
+def calc_poiss_suff_stat(x):
+    """Return sufficient statistics for Poisson distribution."""
+    return x
+
+
+class ExponentialFamilyPoisson(Poisson):  # pylint: disable=abstract-method
+    """Poisson distribution parameterized by natural parameters."""
+
+    def __init__(self, theta: ch.Tensor, dims: int):
+        """Initialize with natural parameter theta and dimension."""
+        self.dims = dims
+        lambda_ = ch.exp(theta)
+        super().__init__(lambda_)
+
+    def log_prob(self, value):
+        """Compute summed log probability over all dimensions."""
+        result = super().log_prob(value)
+        return result.sum(-1)