Move truncated distribution configs to delphi/utils/configs.py.

TruncatedExponentialFamilyDistributionConfig and
TruncatedMultivariateNormalConfig now live alongside TrainerConfig and
OptimizerConfig so any module can import them without pulling in the
full distribution implementation.

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

Author: pstefanou12

delphi/truncated/distributions/truncated_exponential_family_distributions.py

@@ -6,7 +6,6 @@
 from __future__ import annotations
 from collections.abc import Callable
 
-import pydantic
 import torch as ch
 from torch import nn
 from torch.distributions import exp_family
@@ -16,77 +15,6 @@
 from delphi.utils import configs, datasets
 
 
-class TruncatedExponentialFamilyDistributionConfig(
-    configs.TrainerConfig, configs.OptimizerConfig
-):
-    """Configuration for truncated exponential family distribution algorithms.
-
-    Attributes:
-        val: Fraction of data held out for validation.
-        eps: Numerical stability constant for the NLL criterion.
-        min_radius: Initial NLL budget above the empirical initialization
-            for the sublevel-set projection (phase 1).
-        max_radius: Maximum NLL budget; the procedure stops when reached.
-        rate: Multiplicative budget expansion factor per phase.
-        batch_size: Mini-batch size for training.
-        num_samples: Monte Carlo samples drawn per NLL evaluation.
-        max_phases: Maximum number of radius-expansion phases.
-        loss_convergence_tol: Absolute loss improvement threshold for
-            stopping between phases.
-        relative_loss_tol: Relative loss improvement threshold between phases.
-        loss_increase_tol: Loss increase threshold for detecting overshoot.
-        project: Enable per-step sublevel-set projection.
-    """
-
-    model_config = pydantic.ConfigDict(extra="ignore")
-
-    # Override parent defaults for distribution training.
-    tol: float = pydantic.Field(default=1e-1, ge=0.0)
-    record_params_every: int = pydantic.Field(default=1, ge=1)
-    epochs: int | None = pydantic.Field(default=1, ge=1)
-
-    # Distribution-specific fields.
-    val: float = pydantic.Field(default=0.2, ge=0.0, le=1.0)
-    eps: float = pydantic.Field(default=1e-5, gt=0.0)
-    min_radius: float = pydantic.Field(default=3.0, ge=0.0)
-    max_radius: float = pydantic.Field(default=10.0, ge=0.0)
-    rate: float = pydantic.Field(default=1.1, gt=1.0)
-    batch_size: int = pydantic.Field(default=10, ge=1)
-    num_samples: int = pydantic.Field(default=10000, ge=1)
-    max_phases: int = pydantic.Field(default=1, ge=1)
-    loss_convergence_tol: float = pydantic.Field(default=1e-3, ge=0.0)
-    relative_loss_tol: float = pydantic.Field(default=float("inf"), ge=0.0)
-    loss_increase_tol: float = pydantic.Field(default=float("inf"), ge=0.0)
-    project: bool = True
-
-    @pydantic.model_validator(mode="after")
-    def check_radius(self) -> TruncatedExponentialFamilyDistributionConfig:
-        """Validate that min_radius does not exceed max_radius."""
-        if self.min_radius > self.max_radius:
-            raise ValueError(
-                f"min_radius ({self.min_radius}) must be <= "
-                f"max_radius ({self.max_radius})."
-            )
-        return self
-
-    @pydantic.model_validator(mode="after")
-    def resolve_epochs_iterations(
-        self,
-    ) -> TruncatedExponentialFamilyDistributionConfig:
-        """Clear the default epochs when iterations is explicitly provided.
-
-        The trainer uses exactly one stopping criterion.  When the user
-        supplies iterations, the per-phase epochs default is cleared so
-        the trainer stops on iterations instead.
-        """
-        if (
-            "iterations" in self.model_fields_set
-            and "epochs" not in self.model_fields_set
-        ):
-            object.__setattr__(self, "epochs", None)
-        return self
-
-
 class TruncatedExponentialFamilyDistribution(distributions.distributions):  # pylint: disable=too-many-instance-attributes
     """Base class for truncated exponential family distribution models.
 
@@ -112,7 +40,7 @@ class TruncatedExponentialFamilyDistribution(distributions.distributions):  # py
 
     def __init__(
         self,
-        args: TruncatedExponentialFamilyDistributionConfig,
+        args: configs.TruncatedExponentialFamilyDistributionConfig,
         phi: Callable,
         alpha: float,
         dims: int,

delphi/truncated/distributions/truncated_multivariate_normal.py

@@ -4,7 +4,6 @@
 from collections.abc import Callable
 import logging
 
-from pydantic import Field
 import torch as ch
 from torch import nn
 
@@ -14,22 +13,6 @@
 from delphi.utils import configs
 
 
-class TruncatedMultivariateNormalConfig(
-    truncated_exponential_family_distributions.TruncatedExponentialFamilyDistributionConfig
-):
-    """Configuration for truncated multivariate normal distributions.
-
-    Attributes:
-        eigenvalue_lower_bound: Minimum eigenvalue enforced during the
-            negative-definite cone projection of the precision matrix T.
-        covariance_matrix_lr: Optional separate learning rate for the
-            covariance matrix parameter; falls back to lr when None.
-    """
-
-    eigenvalue_lower_bound: float = Field(default=1e-2, gt=0.0)
-    covariance_matrix_lr: float | None = Field(default=None, gt=0.0)
-
-
 class TruncatedMultivariateNormal(
     truncated_exponential_family_distributions.TruncatedExponentialFamilyDistribution
 ):
@@ -44,7 +27,7 @@ class TruncatedMultivariateNormal(
 
     def __init__(
         self,
-        args: dict | TruncatedMultivariateNormalConfig,
+        args: dict | configs.TruncatedMultivariateNormalConfig,
         phi: Callable,
         alpha: float,
         dims: int,
@@ -59,7 +42,7 @@ def __init__(
             dims: Number of dimensions.
             sampler: Optional sampler override.
         """
-        args = configs.make_config(args, TruncatedMultivariateNormalConfig)
+        args = configs.make_config(args, configs.TruncatedMultivariateNormalConfig)
         self.eigenvalue_lower_bound = args.eigenvalue_lower_bound
 
         logger = (

delphi/truncated/distributions/truncated_multivariate_normal_known_covariance.py

@@ -11,7 +11,6 @@
 from delphi import delphi_logger
 from delphi.distributions import multivariate_normal
 from delphi.truncated.distributions import truncated_exponential_family_distributions
-from delphi.truncated.distributions import truncated_multivariate_normal
 from delphi.utils import configs
 
 
@@ -26,7 +25,7 @@ class TruncatedMultivariateNormalKnownCovariance(
 
     def __init__(
         self,
-        args: dict | truncated_multivariate_normal.TruncatedMultivariateNormalConfig,
+        args: dict | configs.TruncatedMultivariateNormalConfig,
         phi: Callable,
         alpha: float,
         dims: int,
@@ -43,9 +42,7 @@ def __init__(
             covariance_matrix: Known covariance matrix.
             sampler: Optional sampler override.
         """
-        args = configs.make_config(
-            args, truncated_multivariate_normal.TruncatedMultivariateNormalConfig
-        )
+        args = configs.make_config(args, configs.TruncatedMultivariateNormalConfig)
         logger = (
             delphi_logger.delphiLogger()
             if args.verbose

delphi/truncated/distributions/truncated_normal.py

@@ -4,14 +4,15 @@
 from collections.abc import Callable
 
 from delphi.truncated.distributions import truncated_multivariate_normal
+from delphi.utils import configs
 
 
 class TruncatedNormal(truncated_multivariate_normal.TruncatedMultivariateNormal):
     """Truncated normal distribution with unknown variance."""
 
     def __init__(
         self,
-        args: dict | truncated_multivariate_normal.TruncatedMultivariateNormalConfig,
+        args: dict | configs.TruncatedMultivariateNormalConfig,
         phi: Callable,
         alpha: float,
         sampler: Callable = None,

delphi/truncated/distributions/truncated_normal_known_variance.py

@@ -6,9 +6,9 @@
 import torch as ch
 
 from delphi.truncated.distributions import (
-    truncated_multivariate_normal,
     truncated_multivariate_normal_known_covariance,
 )
+from delphi.utils import configs
 
 
 class TruncatedNormalKnownVariance(
@@ -18,7 +18,7 @@ class TruncatedNormalKnownVariance(
 
     def __init__(
         self,
-        args: dict | truncated_multivariate_normal.TruncatedMultivariateNormalConfig,
+        args: dict | configs.TruncatedMultivariateNormalConfig,
         phi: Callable,
         alpha: float,
         covariance_matrix: ch.Tensor | None,

delphi/utils/configs.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 
 def make_config(args: dict | BaseModel, config_class: type[BaseModel]) -> BaseModel:
@@ -113,3 +113,86 @@ class OptimizerConfig(BaseModel):
     differentiable: bool = False
     fused: bool | None = None
     scheduler: str | None = None
+
+
+class TruncatedExponentialFamilyDistributionConfig(TrainerConfig, OptimizerConfig):
+    """Configuration for truncated exponential family distribution algorithms.
+
+    Attributes:
+        val: Fraction of data held out for validation.
+        eps: Numerical stability constant for the NLL criterion.
+        min_radius: Initial NLL budget above the empirical initialization
+            for the sublevel-set projection (phase 1).
+        max_radius: Maximum NLL budget; the procedure stops when reached.
+        rate: Multiplicative budget expansion factor per phase.
+        batch_size: Mini-batch size for training.
+        num_samples: Monte Carlo samples drawn per NLL evaluation.
+        max_phases: Maximum number of radius-expansion phases.
+        loss_convergence_tol: Absolute loss improvement threshold for
+            stopping between phases.
+        relative_loss_tol: Relative loss improvement threshold between phases.
+        loss_increase_tol: Loss increase threshold for detecting overshoot.
+        project: Enable per-step sublevel-set projection.
+    """
+
+    model_config = ConfigDict(extra="ignore")
+
+    # Override parent defaults for distribution training.
+    tol: float = Field(default=1e-1, ge=0.0)
+    record_params_every: int = Field(default=1, ge=1)
+    epochs: int | None = Field(default=1, ge=1)
+
+    # Distribution-specific fields.
+    val: float = Field(default=0.2, ge=0.0, le=1.0)
+    eps: float = Field(default=1e-5, gt=0.0)
+    min_radius: float = Field(default=3.0, ge=0.0)
+    max_radius: float = Field(default=10.0, ge=0.0)
+    rate: float = Field(default=1.1, gt=1.0)
+    batch_size: int = Field(default=10, ge=1)
+    num_samples: int = Field(default=10000, ge=1)
+    max_phases: int = Field(default=1, ge=1)
+    loss_convergence_tol: float = Field(default=1e-3, ge=0.0)
+    relative_loss_tol: float = Field(default=float("inf"), ge=0.0)
+    loss_increase_tol: float = Field(default=float("inf"), ge=0.0)
+    project: bool = True
+
+    @model_validator(mode="after")
+    def check_radius(self) -> TruncatedExponentialFamilyDistributionConfig:
+        """Validate that min_radius does not exceed max_radius."""
+        if self.min_radius > self.max_radius:
+            raise ValueError(
+                f"min_radius ({self.min_radius}) must be <= "
+                f"max_radius ({self.max_radius})."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def resolve_epochs_iterations(
+        self,
+    ) -> TruncatedExponentialFamilyDistributionConfig:
+        """Clear the default epochs when iterations is explicitly provided.
+
+        The trainer uses exactly one stopping criterion. When the user
+        supplies iterations, the per-phase epochs default is cleared so
+        the trainer stops on iterations instead.
+        """
+        if (
+            "iterations" in self.model_fields_set
+            and "epochs" not in self.model_fields_set
+        ):
+            object.__setattr__(self, "epochs", None)
+        return self
+
+
+class TruncatedMultivariateNormalConfig(TruncatedExponentialFamilyDistributionConfig):
+    """Configuration for truncated multivariate normal distributions.
+
+    Attributes:
+        eigenvalue_lower_bound: Minimum eigenvalue enforced during the
+            negative-definite cone projection of the precision matrix T.
+        covariance_matrix_lr: Optional separate learning rate for the
+            covariance matrix parameter; falls back to lr when None.
+    """
+
+    eigenvalue_lower_bound: float = Field(default=1e-2, gt=0.0)
+    covariance_matrix_lr: float | None = Field(default=None, gt=0.0)