Fix truncated exponential family distributions: use configs, correct canonical/natural conversions.

- Replace Parameters + check_and_fill_args with configs.make_config throughout
- Use to_canonical/to_natural staticmethods for parameter conversions in
  TruncatedExponential, TruncatedBooleanProduct, TruncatedPoisson, TruncatedWeibull
- Add _calc_emp_model to TruncatedExponential for empirical initialization
- Fix best_*/ema_*/avg_* properties to return canonical parameters
- Update tests to use dict args instead of Parameters

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: pstefanou12

delphi/truncated/distributions/truncated_boolean_product.py

@@ -4,12 +4,11 @@
 import logging
 from collections.abc import Callable
 
-import torch as ch
 
 from delphi import delphi_logger
 from delphi.exponential_family import boolean_product
 from delphi.truncated.distributions import truncated_exponential_family_distributions
-from delphi.utils import defaults, helpers
+from delphi.utils import configs
 
 
 class TruncatedBooleanProduct(
@@ -20,22 +19,23 @@ class TruncatedBooleanProduct(
     dist = boolean_product.ExponentialFamilyBooleanProduct
 
     def __init__(
-        self, args: helpers.Parameters, phi: Callable, alpha: float, dims: int
+        self,
+        args: dict | configs.TruncatedExponentialFamilyDistributionConfig,
+        phi: Callable,
+        alpha: float,
+        dims: int,
     ):
         """Initialize TruncatedBooleanProduct.
 
         Args:
-            args: Parameter object holding hyperparameters.
+            args: Hyperparameter dict or Pydantic config.
             phi: Truncation set oracle.
             alpha: Survival probability lower bound.
             dims: Number of dimensions.
-
-        Raises:
-            TypeError: If args is not a Parameters instance.
         """
-        if not isinstance(args, helpers.Parameters):
-            raise TypeError(f"args is type {type(args).__name__}; expected Parameters.")
-        args = defaults.check_and_fill_args(args, defaults.TRUNC_BOOL_PROD_DEFAULTS)
+        args = configs.make_config(
+            args, configs.TruncatedExponentialFamilyDistributionConfig
+        )
 
         logger = (
             delphi_logger.delphiLogger()
@@ -50,33 +50,33 @@ def __init__(
             logger,
         )
 
-    def _reparameterize_nat_form(self, theta):
-        """Convert canonical probability parameter to natural log-odds form."""
-        return ch.log(theta / (1 - theta))
-
-    def _reparameterize_canon_form(self, theta):
-        """Convert natural log-odds to canonical probability parameter."""
-        return ch.exp(theta) / (1 + ch.exp(theta))
-
     @property
     def best_p_(self):
         """Return the best probability parameter estimate."""
-        return self.best_params
+        return boolean_product.ExponentialFamilyBooleanProduct.to_canonical(
+            self.best_params
+        )
 
     @property
     def final_p_(self):
         """Return the final probability parameter estimate."""
-        return self.final_params
+        return boolean_product.ExponentialFamilyBooleanProduct.to_canonical(
+            self.final_params
+        )
 
     @property
     def ema_p_(self):
         """Return the EMA probability parameter estimate."""
-        return self.ema_params
+        return boolean_product.ExponentialFamilyBooleanProduct.to_canonical(
+            self.ema_params
+        )
 
     @property
     def avg_p_(self):
         """Return the averaged probability parameter estimate."""
-        return self.avg_params
+        return boolean_product.ExponentialFamilyBooleanProduct.to_canonical(
+            self.avg_params
+        )
 
     def __str__(self):
         """Return a human-readable name for this distribution."""

delphi/truncated/distributions/truncated_exponential.py

@@ -5,11 +5,12 @@
 from collections.abc import Callable
 
 import torch as ch
+from torch import nn
 
 from delphi import delphi_logger
 from delphi.exponential_family import exponential
 from delphi.truncated.distributions import truncated_exponential_family_distributions
-from delphi.utils import defaults, helpers
+from delphi.utils import configs
 
 
 class TruncatedExponential(
@@ -21,25 +22,22 @@ class TruncatedExponential(
 
     def __init__(
         self,
-        args: helpers.Parameters,
+        args: dict | configs.TruncatedExponentialFamilyDistributionConfig,
         phi: Callable,
         alpha: float,
         dims: int,
     ):
         """Initialize TruncatedExponential.
 
         Args:
-            args: Parameter object holding hyperparameters.
+            args: Hyperparameter dict or Pydantic config.
             phi: Truncation set oracle.
             alpha: Survival probability lower bound.
             dims: Number of dimensions.
-
-        Raises:
-            TypeError: If args is not a Parameters instance.
         """
-        if not isinstance(args, helpers.Parameters):
-            raise TypeError(f"args is type {type(args).__name__}; expected Parameters.")
-        args = defaults.check_and_fill_args(args, defaults.TRUNC_EXP_DEFAULTS)
+        args = configs.make_config(
+            args, configs.TruncatedExponentialFamilyDistributionConfig
+        )
 
         logger = (
             delphi_logger.delphiLogger()
@@ -54,37 +52,36 @@ def __init__(
             logger,
         )
 
+    def _calc_emp_model(self):
+        """Initialize theta at the natural parameter corresponding to the empirical rate."""
+        dataset_s = self.train_loader_.dataset.S  # pylint: disable=invalid-name
+        emp_rate = 1.0 / dataset_s.mean(0)
+        self.emp_theta = exponential.ExponentialFamilyExponential.to_natural(emp_rate)
+        self.register_parameter("theta", nn.Parameter(self.emp_theta.clone()))
+
     def _constraints(self, theta):
         """Clamp theta to be strictly negative."""
         return ch.clamp(theta, max=-1e-6)
 
-    def _reparameterize_nat_form(self, theta):
-        """Convert canonical rate parameter to natural form."""
-        return -theta
-
-    def _reparameterize_canon_form(self, theta):
-        """Convert natural parameters to canonical rate parameter."""
-        return -theta
-
     @property
     def best_lambda_(self):
-        """Return the best rate parameter estimate."""
-        return self.best_params
+        """Return the best rate (canonical) parameter estimate."""
+        return exponential.ExponentialFamilyExponential.to_canonical(self.best_params)
 
     @property
     def final_lambda_(self):
-        """Return the final rate parameter estimate."""
-        return self.final_params
+        """Return the final rate (canonical) parameter estimate."""
+        return exponential.ExponentialFamilyExponential.to_canonical(self.final_params)
 
     @property
     def ema_lambda_(self):
-        """Return the EMA rate parameter estimate."""
-        return self.ema_params
+        """Return the EMA rate (canonical) parameter estimate."""
+        return exponential.ExponentialFamilyExponential.to_canonical(self.ema_params)
 
     @property
     def avg_lambda_(self):
-        """Return the averaged rate parameter estimate."""
-        return self.avg_params
+        """Return the averaged rate (canonical) parameter estimate."""
+        return exponential.ExponentialFamilyExponential.to_canonical(self.avg_params)
 
     def __str__(self):
         """Return a human-readable name for this distribution."""

delphi/truncated/distributions/truncated_exponential_family_distributions.py

@@ -2,7 +2,6 @@
 """Parent class for truncated exponential distribution model classes."""
 
 import abc
-import functools
 from collections.abc import Callable
 
 import torch as ch
@@ -63,15 +62,12 @@ def __init__(
         self.alpha = alpha
         self.dims = dims
 
-        dist_cls = (
-            self.dist.func if isinstance(self.dist, functools.partial) else self.dist
-        )
         self.criterion = losses.TruncatedExponentialFamilyDistributionNLL.apply
         self.criterion_params = [
             self.phi,
             self.dims,
             self.dist,
-            dist_cls.calc_suff_stat,
+            self.dist.calc_suff_stat,
             self.args.num_samples,
             self.args.eps,
         ]
@@ -115,20 +111,18 @@ def fit(self, S: ch.Tensor):  # pylint: disable=invalid-name
                 f"or equal to the number of samples ({self.args.num_samples})."
             )
 
-        dist_cls = (
-            self.dist.func if isinstance(self.dist, functools.partial) else self.dist
-        )
         self.train_loader_, self.val_loader_ = datasets.make_train_and_val_distr(
             self.args,
             S,
             datasets.TruncatedExponentialDistributionDataset,
-            {"calc_suff_stat": dist_cls.calc_suff_stat},
+            {"calc_suff_stat": self.dist.calc_suff_stat},
         )
 
         self.prev_best_loss = None
         self.radius_history = []
         self.loss_history = []
         self._calc_emp_model()
+        self.nll_init = self._compute_nll(self.emp_theta)
         self.radius = self.args.min_radius
 
         phase = 0
@@ -224,13 +218,8 @@ def nll_threshold(self) -> float:
     def _calc_emp_model(self):
         """Calculate empirical natural parameters from training data."""
         dataset_s = self.train_loader_.dataset.S  # pylint: disable=invalid-name
-        dist_cls = (
-            self.dist.func if isinstance(self.dist, functools.partial) else self.dist
-        )
-        self.emp_theta = dist_cls.calc_suff_stat(dataset_s).mean(0)
+        self.emp_theta = self.dist.calc_suff_stat(dataset_s).mean(0)
         self.register_parameter("theta", nn.Parameter(self.emp_theta.clone()))
-        with ch.no_grad():
-            self.nll_init = self._compute_nll(self.emp_theta)
 
     def _compute_nll(self, theta: ch.Tensor) -> float:
         """Compute non-truncated NLL of training samples under theta.
@@ -241,9 +230,10 @@ def _compute_nll(self, theta: ch.Tensor) -> float:
         Returns:
             Mean negative log-likelihood over training samples.
         """
-        S = self.train_loader_.dataset.S  # pylint: disable=invalid-name
-        D = self.dist(theta.detach(), self.dims)  # pylint: disable=invalid-name
-        return -D.log_prob(S).mean().item()
+        with ch.no_grad():
+            S = self.train_loader_.dataset.S  # pylint: disable=invalid-name
+            D = self.dist(theta, self.dims)  # pylint: disable=invalid-name
+            return -D.log_prob(S).mean().item()
 
     def _project_onto_sublevel_set(self, theta: ch.Tensor) -> ch.Tensor:
         """Project theta onto {θ : L(θ) ≤ nll_threshold} via bisection.
@@ -296,6 +286,12 @@ def step_post_hook(self, optimizer, args, kwargs) -> None:
         """
         if self.args.project:
             with ch.no_grad():
+                # Enforce constraints before projection so that _compute_nll
+                # inside _project_onto_sublevel_set always receives a theta in
+                # the valid domain (e.g. strictly negative for exponential /
+                # Weibull).  Constraints are applied again after projection to
+                # keep the final value valid as well.
+                self._write_theta(self._constraints(self.theta))
                 proj_theta = self._project_onto_sublevel_set(self.theta)
                 self._write_theta(self._constraints(proj_theta))
 

delphi/truncated/distributions/truncated_multivariate_normal_known_covariance.py

@@ -53,6 +53,7 @@ def __init__(
             multivariate_normal.ExponentialFamilyMultivariateNormalKnownCovariance,
             covariance_matrix,
         )
+        self.dist.calc_suff_stat = multivariate_normal.ExponentialFamilyMultivariateNormalKnownCovariance.calc_suff_stat
         super().__init__(
             args,
             phi,

delphi/truncated/distributions/truncated_poisson.py

@@ -4,12 +4,11 @@
 from collections.abc import Callable
 import logging
 
-import torch as ch
 
 from delphi import delphi_logger
 from delphi.exponential_family import poisson
 from delphi.truncated.distributions import truncated_exponential_family_distributions
-from delphi.utils import defaults, helpers
+from delphi.utils import configs
 
 
 class TruncatedPoisson(
@@ -21,25 +20,22 @@ class TruncatedPoisson(
 
     def __init__(
         self,
-        args: helpers.Parameters,
+        args: dict | configs.TruncatedExponentialFamilyDistributionConfig,
         phi: Callable,
         alpha: float,
         dims: int,
     ):
         """Initialize TruncatedPoisson.
 
         Args:
-            args: Parameter object holding hyperparameters.
+            args: Hyperparameter dict or Pydantic config.
             phi: Truncation set oracle.
             alpha: Survival probability lower bound.
             dims: Number of dimensions.
-
-        Raises:
-            TypeError: If args is not a Parameters instance.
         """
-        if not isinstance(args, helpers.Parameters):
-            raise TypeError(f"args is type {type(args).__name__}; expected Parameters.")
-        args = defaults.check_and_fill_args(args, defaults.TRUNC_POISS_DEFAULTS)
+        args = configs.make_config(
+            args, configs.TruncatedExponentialFamilyDistributionConfig
+        )
 
         logger = (
             delphi_logger.delphiLogger()
@@ -54,33 +50,25 @@ def __init__(
             logger,
         )
 
-    def _reparameterize_nat_form(self, theta):
-        """Convert canonical rate parameter to natural log form."""
-        return ch.log(theta)
-
-    def _reparameterize_canon_form(self, theta):
-        """Convert natural parameters to canonical rate parameter."""
-        return ch.exp(theta)
-
     @property
     def best_lambda_(self):
         """Return the best rate parameter estimate."""
-        return self.best_params
+        return poisson.ExponentialFamilyPoisson.to_canonical(self.best_params)
 
     @property
     def final_lambda_(self):
         """Return the final rate parameter estimate."""
-        return self.final_params
+        return poisson.ExponentialFamilyPoisson.to_canonical(self.final_params)
 
     @property
     def ema_lambda_(self):
         """Return the EMA rate parameter estimate."""
-        return self.ema_params
+        return poisson.ExponentialFamilyPoisson.to_canonical(self.ema_params)
 
     @property
     def avg_lambda_(self):
         """Return the averaged rate parameter estimate."""
-        return self.avg_params
+        return poisson.ExponentialFamilyPoisson.to_canonical(self.avg_params)
 
     def __str__(self):
         """Return a human-readable name for this distribution."""