Split discrete mechanisms and initializers into distinct configuration classes and calibrated mechanisms.

docs/in_memory_api.md

@@ -29,7 +29,7 @@ import pandas as pd
 
 synth = dpsynth.TabularSynthesizer(
     domains=domains,
-    discrete_mechanism=discrete_mechanisms.MSTMechanism(),
+    discrete_mechanism=discrete_mechanisms.MSTConfig(),
 )
 result = synth.calibrate(
     epsilon=1.0,

dpsynth/api.py

@@ -19,7 +19,7 @@
 
 Example usage::
 
-  mechanism = dpsynth.discrete_mechanisms.AIMMechanism(pgm_iters=500)
+  mechanism = dpsynth.discrete_mechanisms.AIMConfig(pgm_iters=500)
 
   # Option 1: Calibrate to (epsilon, delta)-DP (tight PLD accounting).
   calibrated = mechanism.calibrate(epsilon=1.0, delta=1e-5)
@@ -42,24 +42,58 @@
 import dp_accounting
 
 
-class DPMechanism(abc.ABC):
-  """Abstract base class for differentially private mechanisms.
+class CalibratedMechanism(abc.ABC):
+  """A privacy-calibrated, runnable differentially private mechanism.
 
-  A DPMechanism encapsulates a randomized algorithm that satisfies differential
-  privacy. Usage follows a three-phase pattern:
+  Produced by ``MechanismConfig.configure()`` / ``.calibrate()``: its natural
+  privacy parameter (e.g. Gaussian sigma) is populated and it is ready to run.
+  It exposes the exact ``DpEvent`` characterizing its privacy cost and is
+  directly callable on data.
 
-  1. **Construct**: Create the mechanism with algorithm-specific parameters
-     (e.g., ``AIMMechanism(pgm_iters=500)``).
+  Subclasses must implement:
+
+  - ``dp_event``: return the exact ``DpEvent`` characterizing the mechanism.
+  - ``__call__``: run the mechanism on data.
+  """
+
+  @property
+  @abc.abstractmethod
+  def dp_event(self) -> dp_accounting.DpEvent:
+    """The DpEvent characterizing the privacy cost of this mechanism."""
+
+  @abc.abstractmethod
+  def __call__(self, *args: Any, **kwargs: Any) -> Any:
+    """Runs the mechanism on the given data.
+
+    Subclass signatures vary, but typically accept at least the data to operate
+    on and a source of randomness.
+
+    Args:
+      *args: Positional arguments (subclass-specific).
+      **kwargs: Keyword arguments (subclass-specific).
+    """
+
+
+class MechanismConfig(abc.ABC):
+  """A recipe that produces a calibrated, runnable mechanism.
+
+  A config holds the mechanism's structural and hyperparameter fields (the
+  parts a user writes, logs, and serializes) and knows how to turn a privacy
+  budget into a runnable ``CalibratedMechanism``. Usage follows a three-phase
+  pattern:
+
+  1. **Construct**: Create the config with algorithm-specific parameters
+     (e.g., ``AIMConfig(pgm_iters=500)``).
   2. **Calibrate**: Call ``calibrate(epsilon=..., delta=...)`` or
-     ``configure(zcdp_rho=...)`` to bind a privacy budget, returning a new
-     frozen instance with the mechanism's natural privacy parameter set.
+     ``configure(zcdp_rho=...)`` to bind a privacy budget, returning a
+     ``CalibratedMechanism`` with the mechanism's natural privacy parameter set.
   3. **Run**: Call the calibrated mechanism on data via ``__call__``.
 
   **Design: configure vs calibrate.**  The API separates two concerns:
 
   - ``configure(zcdp_rho, **kwargs)`` is the low-level primitive that each
-    mechanism must implement. It maps a zCDP budget to the mechanism's natural
-    privacy parameter (e.g., Gaussian sigma) and returns a new frozen instance.
+    config must implement. It maps a zCDP budget to the mechanism's natural
+    privacy parameter (e.g., Gaussian sigma) and returns a runnable mechanism.
     This is lightweight — just arithmetic — and produces reasonably tight
     parameter settings for most mechanisms.
 
@@ -82,24 +116,17 @@ class DPMechanism(abc.ABC):
   search exploits this: it evaluates each candidate's raw ``dp_event`` rather
   than relying on the zCDP conversion, so the final calibration is as tight
   as the mechanism's own privacy characterization allows.
-
-  Subclasses must implement:
-
-  - ``configure(zcdp_rho, **kwargs)``: set the mechanism's natural privacy
-    parameter (e.g., Gaussian sigma) from a zCDP budget.
-  - ``dp_event``: return the exact ``DpEvent`` characterizing the mechanism.
-  - ``__call__``: run the mechanism on data.
   """
 
   @abc.abstractmethod
   def configure(
-      self, *, zcdp_rho: float, delta: float = 0.0, **kwargs: Any
-  ) -> DPMechanism:
-    """Returns a new mechanism configured with the given zCDP budget.
+      self, *, zcdp_rho, delta=0.0, max_records_per_user=1
+  ) -> CalibratedMechanism:
+    """Returns a calibrated mechanism for the given zCDP budget.
 
     Converts the zCDP budget into the mechanism's natural privacy parameter
-    (e.g., Gaussian sigma) and returns a new frozen instance with that
-    parameter set.
+    (e.g., Gaussian sigma) and returns a runnable ``CalibratedMechanism`` with
+    that parameter set.
 
     Most mechanisms are pure zCDP and ignore ``delta``. Mechanisms that
     consume approximate DP budget (e.g., partition selection with Gaussian
@@ -111,27 +138,11 @@ def configure(
       delta: Approximate DP delta consumed by the mechanism itself (e.g., for
         thresholding). Defaults to 0 (pure zCDP). Mechanisms that need delta
         should raise if it is 0.
-      **kwargs: Mechanism-specific hyperparameters.
+      **kwargs: Mechanism-specific hyperparameters (e.g.
+        ``max_records_per_user`` for mechanisms that support user-level DP).
 
     Returns:
-      A new DPMechanism instance configured for the given budget.
-    """
-
-  @property
-  @abc.abstractmethod
-  def dp_event(self) -> dp_accounting.DpEvent:
-    """The DpEvent characterizing the privacy cost of this mechanism."""
-
-  @abc.abstractmethod
-  def __call__(self, *args: Any, **kwargs: Any) -> Any:
-    """Runs the mechanism on the given data.
-
-    Subclass signatures vary, but typically accept at least the data to operate
-    on and a source of randomness.
-
-    Args:
-      *args: Positional arguments (subclass-specific).
-      **kwargs: Keyword arguments (subclass-specific).
+      A calibrated, runnable mechanism.
     """
 
   def _find_optimal_rho(
@@ -193,7 +204,7 @@ def calibrate(
       delta: float | None = None,
       zcdp_rho: float | None = None,
       **kwargs: Any,
-  ) -> DPMechanism:
+  ) -> CalibratedMechanism:
     """Calibrate the mechanism to a target (epsilon, delta)-DP guarantee.
 
     Performs a binary search over zCDP budgets, calling ``configure`` at each
@@ -211,7 +222,7 @@ def calibrate(
       **kwargs: Forwarded to ``configure()``.
 
     Returns:
-      A new calibrated DPMechanism instance.
+      A calibrated, runnable mechanism.
 
     Raises:
       ValueError: If neither (epsilon, delta) nor zcdp_rho is specified, or
@@ -243,6 +254,20 @@ def calibrate(
     return self.configure(zcdp_rho=optimal_rho, delta=delta, **kwargs)
 
 
+class DPMechanism(MechanismConfig, CalibratedMechanism, abc.ABC):
+  """Transitional monolithic base: both a config and a runnable mechanism.
+
+  Historically every mechanism was a single class that was constructed, then
+  ``configure``d in place, then run. That design is being split into a
+  ``MechanismConfig`` recipe and a ``CalibratedMechanism`` runnable (so a
+  calibrated instance can never be un-calibrated and needs no nullable
+  privacy-parameter fields or guards). Mechanisms are migrated one layer at a
+  time; those not yet split still subclass this monolith, which exposes exactly
+  the historical abstract surface (``configure`` + ``dp_event`` + ``__call__``,
+  with ``calibrate`` inherited). Remove once every mechanism is split.
+  """
+
+
 def validate_max_records_per_user(value: int) -> None:
   """Raises ValueError if the per-user record bound is not a positive int."""
   if value < 1:

dpsynth/data_generation_v2.py

@@ -35,8 +35,9 @@ def generate(
     delta: float,
     *,
     discrete_config: (
-        discrete_mechanisms.DiscreteMechanism
-    ) = discrete_mechanisms.MSTMechanism(),
+        discrete_mechanisms.DiscreteMechanismConfig
+        | discrete_mechanisms.DiscreteMechanism
+    ) = discrete_mechanisms.MSTConfig(),
     numerical_bins: int = 32,
     one_way_marginal_budget_fraction: float = 0.1,
     cross_attribute_constraints: Sequence[constraints.Constraint] = (),

dpsynth/data_generation_v3.py

@@ -21,7 +21,6 @@
 
 from absl import logging
 import dp_accounting
-from dpsynth import api
 from dpsynth import constraints
 from dpsynth import discrete_mechanisms
 from dpsynth import domain
@@ -38,18 +37,13 @@ def _create_initializers(
     domains: Mapping[str, domain.AttributeType],
     numerical_bins: int,
     init_delta: float,
-    max_records_per_user: int = 1,
 ) -> dict[str, primitives.DPMechanism]:
   """Creates per-column initializers from the domain specification.
 
   Args:
     domains: Mapping from column names to attribute domain specifications.
     numerical_bins: Number of bins for numerical discretization.
     init_delta: Delta for open-set categorical partition selection.
-    max_records_per_user: Assumed upper bound on the number of records a single
-      user contributes. Sensitivity (and hence added noise) is scaled by this
-      factor to provide user-level rather than record-level DP; the privacy
-      analysis is unchanged. Soundness relies on the caller enforcing the bound.
 
   Returns:
     A dictionary mapping column names to uncalibrated initializer instances.
@@ -60,22 +54,20 @@ def _create_initializers(
   initializers = {}
   for col, attr in domains.items():
     if isinstance(attr, domain.NumericalAttribute):
-      initializers[col] = initialization.NumericalInitializer(
+      initializers[col] = initialization.NumericalInitializerConfig(
           name=col,
           num_partitions=numerical_bins,
           attribute=attr,
-          max_records_per_user=max_records_per_user,
       )
     elif isinstance(attr, domain.CategoricalAttribute):
-      initializers[col] = initialization.CategoricalInitializer(
-          name=col, attribute=attr, max_records_per_user=max_records_per_user
+      initializers[col] = initialization.CategoricalInitializerConfig(
+          name=col, attribute=attr
       )
     elif isinstance(attr, domain.OpenSetCategoricalAttribute):
-      initializers[col] = initialization.OpenSetCategoricalInitializer(
+      initializers[col] = initialization.OpenSetCategoricalInitializerConfig(
           name=col,
           attribute=attr,
           delta=init_delta,
-          max_records_per_user=max_records_per_user,
       )
     else:
       raise ValueError(
@@ -225,9 +217,10 @@ class TabularSynthesizer(primitives.DPMechanism):
   """
 
   domains: Mapping[str, domain.AttributeType]
-  discrete_mechanism: discrete_mechanisms.DiscreteMechanism = dataclasses.field(
-      default_factory=discrete_mechanisms.MSTMechanism
-  )
+  discrete_mechanism: (
+      discrete_mechanisms.DiscreteMechanismConfig
+      | discrete_mechanisms.DiscreteMechanism
+  ) = dataclasses.field(default_factory=discrete_mechanisms.MSTConfig)
   numerical_bins: int = 32
   init_budget_fraction: float = 0.1
   initializers: dict[str, primitives.DPMechanism] | None = None
@@ -236,13 +229,15 @@ class TabularSynthesizer(primitives.DPMechanism):
   experimental_max_records_per_user: int = 1
 
   def __post_init__(self):
-    api.validate_max_records_per_user(self.experimental_max_records_per_user)
+    if self.experimental_max_records_per_user <= 0:
+      raise ValueError('experimental_max_records_per_user must be >= 1')
 
   def configure(  # pyrefly: ignore[bad-override]
       self,
       *,
       zcdp_rho: float,
       delta: float = 0.0,
+      max_records_per_user: int | None = None,
   ) -> TabularSynthesizer:
     """Returns a copy configured with the given privacy budget.
 
@@ -266,6 +261,11 @@ def configure(  # pyrefly: ignore[bad-override]
         (``init_budget_fraction``) is allocated to partition selection for
         open-set columns. Must be positive when open-set categorical attributes
         are present.
+      max_records_per_user: Upper bound on the number of records a single user
+        contributes. Values greater than 1 scale the added noise (and mechanism
+        sensitivity) to provide user-level rather than record-level DP; the
+        privacy accounting is unchanged. This bound is NOT enforced -- soundness
+        relies on the caller guaranteeing it via preprocessing.
 
     Returns:
       A new TabularSynthesizer with calibrated sub-mechanisms.
@@ -282,6 +282,9 @@ def configure(  # pyrefly: ignore[bad-override]
           'delta must be positive when open-set categorical attributes are'
           ' present. It is used for Gaussian partition selection.'
       )
+    if max_records_per_user is None:
+      max_records_per_user = self.experimental_max_records_per_user
+
     # Split delta across open-set columns, analogous to splitting zcdp_rho.
     # Under calibrate(), any delta not consumed here is automatically
     # available for the zCDP-to-(epsilon, delta) conversion, so this
@@ -297,38 +300,33 @@ def configure(  # pyrefly: ignore[bad-override]
           self.domains,
           self.numerical_bins,
           per_col_delta,
-          self.experimental_max_records_per_user,
       )
-    elif self.experimental_max_records_per_user > 1:
-      # The synthesizer's experimental_max_records_per_user is the single
-      # source of truth: the total-count and discrete mechanisms already use it,
-      # so propagate it to caller-supplied initializers too.
-      propagated = {}
-      for col, init in inits.items():
-        propagated[col] = dataclasses.replace(  # pytype: disable=wrong-arg-types
-            init, max_records_per_user=self.experimental_max_records_per_user
-        )
-      inits = propagated
     init_rho = self.init_budget_fraction * zcdp_rho
     # +1 for the DPGaussianCount that always measures the total.
     per_col_rho = init_rho / (len(inits) + 1)
     discrete_rho = zcdp_rho - init_rho
 
     calibrated_inits = {
-        col: init.configure(zcdp_rho=per_col_rho) for col, init in inits.items()
+        col: init.configure(
+            zcdp_rho=per_col_rho,
+            max_records_per_user=max_records_per_user,
+        )
+        for col, init in inits.items()
     }
-    calibrated_total = primitives.DPGaussianCount(
-        max_records_per_user=self.experimental_max_records_per_user
-    ).configure(zcdp_rho=per_col_rho)
-    calibrated_discrete = dataclasses.replace(
-        self.discrete_mechanism,
-        max_records_per_user=self.experimental_max_records_per_user,
-    ).configure(zcdp_rho=discrete_rho)
+    calibrated_total = primitives.DPGaussianCountConfig().configure(
+        zcdp_rho=per_col_rho,
+        max_records_per_user=max_records_per_user,
+    )
+    calibrated_discrete = self.discrete_mechanism.configure(  # pyrefly: ignore[missing-attribute]
+        zcdp_rho=discrete_rho,
+        max_records_per_user=max_records_per_user,
+    )
     return dataclasses.replace(
         self,
         initializers=calibrated_inits,
         discrete_mechanism=calibrated_discrete,
         total_count_mechanism=calibrated_total,
+        experimental_max_records_per_user=max_records_per_user,
     )
 
   @property
@@ -347,7 +345,7 @@ def dp_event(self) -> dp_accounting.DpEvent:
       )
     events = [init.dp_event for init in self.initializers.values()]
     events.append(self.total_count_mechanism.dp_event)
-    events.append(self.discrete_mechanism.dp_event)
+    events.append(self.discrete_mechanism.dp_event)  # pyrefly: ignore[missing-attribute]
     return dp_accounting.ComposedDpEvent(events)
 
   def __call__(
@@ -385,7 +383,7 @@ def __call__(
     total_measurement = mbi.LinearMeasurement(
         np.array([total]),  # pyrefly: ignore[bad-argument-type]
         (),
-        stddev=k * self.total_count_mechanism.sigma,  # pyrefly: ignore[unsupported-operation]
+        stddev=k * self.total_count_mechanism.sigma,
     )
 
     results: dict[str, initialization.ColumnMeasurement] = {}
@@ -408,7 +406,7 @@ def __call__(
     mbi_constraints = tuple(
         c.to_mbi() for c in self.cross_attribute_constraints
     )
-    mechanism_result = self.discrete_mechanism(
+    mechanism_result = self.discrete_mechanism(  # pyrefly: ignore[not-callable]
         rng,
         data=discrete,
         initial_measurements=initial_measurements,