Refactor discrete mechanisms from inheritance to composition, simplifying Independent, Direct, and SWIFT configurations.

dpsynth/adapters/beam.py

@@ -402,14 +402,12 @@ def generate_from_marginals(
   )
 
   initial_measurements = [total_measurement, *codec.one_way_measurements()]
-  mbi_constraints = tuple(c.to_mbi() for c in synth.cross_attribute_constraints)
   logging.info('[DPSynth/Beam]: Running discrete mechanism.')
   # pyrefly: ignore[missing-attribute,not-callable]
   mechanism_result = synth.calibrated_discrete_mechanism(
       rng,
       data=marginals,
       initial_measurements=initial_measurements,
-      constraints=mbi_constraints,
   )
   synthetic_data = codec.decode(
       mechanism_result.synthetic_data, rng, column_order
@@ -473,9 +471,7 @@ def _run_two_pass(
         column_measurements, synth.domains
     ).mbi_domain
     # pyrefly: ignore[missing-attribute]
-    workload = synth.calibrated_discrete_mechanism.supporting_cliques(
-        mbi_domain
-    )
+    workload = synth.config.discrete_mechanism.supporting_cliques(mbi_domain)
 
     # Pass 2: compute the marginal workload.
     with beam.Pipeline(**pipeline_kwargs) as p:
@@ -501,7 +497,7 @@ def _run_two_pass(
       shutil.rmtree(temp_dir, ignore_errors=True)
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class BeamTabularMechanism(api.CalibratedMechanism):
   """Beam-backed DPMechanism with the TabularMechanism calibrate->run API."""
 
@@ -527,7 +523,7 @@ def __call__(
     )
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class BeamTabularConfig(api.MechanismConfig):
   """Beam-backed DPMechanism with the TabularConfig calibrate->run API.
 

dpsynth/api.py

@@ -35,6 +35,7 @@
 
 import abc
 from collections.abc import Callable
+import dataclasses
 import functools
 from typing import Any
 import warnings
@@ -74,6 +75,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
     """
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
 class MechanismConfig(abc.ABC):
   """A recipe that produces a calibrated, runnable mechanism.
 

dpsynth/data_generation_v3.py

@@ -198,6 +198,7 @@ class TabularMechanism(api.CalibratedMechanism):
       user contributes.
   """
 
+  config: TabularConfig
   domains: Mapping[str, domain.AttributeType]
   calibrated_discrete_mechanism: discrete_mechanisms.DiscreteMechanism
   calibrated_initializers: dict[str, api.CalibratedMechanism]
@@ -217,6 +218,9 @@ def dp_event(self) -> dp_accounting.DpEvent:
         dp_accounting.GaussianDpEvent(noise_multiplier=self.total_count_sigma)
     )
     events.append(self.calibrated_discrete_mechanism.dp_event)
+    events = [e for e in events if not isinstance(e, dp_accounting.NoOpDpEvent)]
+    if not events:
+      return dp_accounting.NoOpDpEvent()
     return dp_accounting.ComposedDpEvent(events)
 
   def __call__(
@@ -274,14 +278,10 @@ def __call__(
     # initial measurements so the mechanism does not re-measure them.
     column_order = [col for col in data.columns if col in self.domains]
     initial_measurements = [total_measurement, *codec.one_way_measurements()]
-    mbi_constraints = tuple(
-        c.to_mbi() for c in self.cross_attribute_constraints
-    )
     mechanism_result = self.calibrated_discrete_mechanism(
         rng,
         data=discrete,
         initial_measurements=initial_measurements,
-        constraints=mbi_constraints,
     )
     logging.info('[DPSynth]: Generated discrete synthetic data.')
 
@@ -297,7 +297,9 @@ def __call__(
     )
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True, kw_only=True)
+# be a no-op for numerical / open set attributes and group low counts for
+# categorical.
 class TabularConfig(api.MechanismConfig):
   """Configures end-to-end DP synthetic data generation.
 
@@ -423,12 +425,20 @@ def configure(
     }
     total_count_sigma = math.sqrt(0.5 / per_col_rho)
 
-    calibrated_discrete = self.discrete_mechanism.configure(
+    mbi_constraints = tuple(
+        c.to_mbi() for c in self.cross_attribute_constraints
+    )
+    discrete_with_constraints = dataclasses.replace(
+        self.discrete_mechanism,
+        constraints=mbi_constraints,  # pyrefly: ignore[unexpected-keyword]
+    )
+    calibrated_discrete = discrete_with_constraints.configure(
         max_records_per_user=max_records_per_user,
         zcdp_rho=discrete_rho,
     )
 
     return TabularMechanism(
+        config=self,
         domains=self.domains,
         calibrated_discrete_mechanism=calibrated_discrete,
         calibrated_initializers=calibrated_inits,
@@ -438,7 +448,7 @@ def configure(
     )
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True, kw_only=True)
 class TabularSynthesizer(TabularConfig):
   """Deprecated. Use TabularConfig and TabularMechanism instead."""
 

dpsynth/discrete_mechanisms/__init__.py

@@ -12,20 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Implementations of mechanisms that operate over discrete data."""
+"""Implementations of mechanisms that operate over discrete data.
+
+Note: This mechanism is not intended to be called directly. It should typically
+be used within `DiscreteSynthesizer` or `TabularSynthesizer`. Users who call it
+directly will miss out on features like 1-way measurement selection and domain
+compression.
+"""
 
 # pylint: disable=g-importing-member
 
+from dpsynth.api import CalibratedMechanism
+from dpsynth.api import MechanismConfig
 from dpsynth.discrete_mechanisms.aim import AIM
 from dpsynth.discrete_mechanisms.aim import AIMConfig
 from dpsynth.discrete_mechanisms.aim_gdp import AIMGDP
 from dpsynth.discrete_mechanisms.aim_gdp import AIMGDPConfig
-from dpsynth.discrete_mechanisms.base import DiscreteMechanism
-from dpsynth.discrete_mechanisms.base import DiscreteMechanismConfig
 from dpsynth.discrete_mechanisms.common import DiscreteMechanismResult
 from dpsynth.discrete_mechanisms.common import MechanismDiagnostics
 from dpsynth.discrete_mechanisms.direct import Direct
 from dpsynth.discrete_mechanisms.direct import DirectConfig
+from dpsynth.discrete_mechanisms.discrete_synthesizer import DiscreteConfig
+from dpsynth.discrete_mechanisms.discrete_synthesizer import DiscreteMechanism
 from dpsynth.discrete_mechanisms.independent import Independent
 from dpsynth.discrete_mechanisms.independent import IndependentConfig
 from dpsynth.discrete_mechanisms.mst import MST
@@ -40,3 +48,10 @@
 IndependentMechanism = IndependentConfig
 MSTMechanism = MSTConfig
 SWIFTMechanism = SWIFTConfig
+
+# Legacy aliases. Downstream code (nested.py, data_generation_v2.py) use these
+# as type annotations that accept any mechanism config (MSTConfig, etc.).
+
+
+DiscreteMechanismConfig = MechanismConfig
+DiscreteMechanism = CalibratedMechanism

dpsynth/discrete_mechanisms/aim.py

@@ -15,12 +15,12 @@
 """Implementation of the Adaptive+Iterative Mechanism (AIM)."""
 
 from collections.abc import Iterable, Mapping
+from collections.abc import Sequence
 import dataclasses
-
 from absl import logging
 import dp_accounting
+from dpsynth import api
 from dpsynth.discrete_mechanisms import accounting
-from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 import jax.numpy as jnp
 import mbi
@@ -87,7 +87,7 @@ def _worst_approximated(
 
 
 @dataclasses.dataclass(frozen=True)
-class AIMConfig(base.DiscreteMechanismConfig):
+class AIMConfig(api.MechanismConfig):
   """Configuration for the AIM mechanism.
 
   Details are described in the paper:
@@ -115,54 +115,61 @@ class AIMConfig(base.DiscreteMechanismConfig):
       selecting two-way marginal queries.
   """
 
+  constraints: Sequence[mbi.Constraint] = ()
+
   workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None
   max_rounds: int | None = None
   max_model_size: int = 80
   max_marginal_size: float = 1e6
   anneal_factor: float = 4.0
   select_budget_fraction: float = 0.1
   pgm_iters: int = 1000
+  marginal_oracle: mbi.MarginalOracle | None = None
 
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
     """Returns the workload cliques filtered by max_marginal_size."""
     return common.supporting_cliques(
         domain, self.workload, self.max_marginal_size
     )
 
-  def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
-    """Allocates the entire remaining budget to the adaptive loop."""
-    return {'_loop_rho': remaining_rho}
-
-  def _create_mechanism(self, **kwargs) -> 'AIM':
-    return AIM(**kwargs)
+  def configure(self, *, zcdp_rho, delta=0, max_records_per_user=1):
+    api.validate_max_records_per_user(max_records_per_user)
+    return AIM(
+        config=self,
+        zcdp_rho=zcdp_rho,
+        max_records_per_user=max_records_per_user,
+    )
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
-class AIM(base.DiscreteMechanism):
+class AIM(api.CalibratedMechanism):
   """Calibrated AIM instance."""
 
   config: AIMConfig
-  _loop_rho: float
-
-  def _one_way_cliques(self, data):
-    """Returns only the workload-specified one-way cliques."""
-    return common.one_way_cliques(self.config.workload, data.domain)
+  zcdp_rho: float
+  max_records_per_user: int = 1
 
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the DP event for the AIM mechanism."""
-    events = self._one_way_dp_event()
-    events.append(dp_accounting.ZCDpEvent(self._loop_rho))  # pyrefly: ignore[bad-argument-type]
-    return dp_accounting.ComposedDpEvent(events)
-
-  def _run(self, rng, data, measurements, constraints, phase_times):
-    """Adaptively selects, measures, and estimates in an annealed loop."""
+    return dp_accounting.ZCDpEvent(self.zcdp_rho)
+
+  def __call__(
+      self,
+      rng: np.random.Generator,
+      data: mbi.Dataset | mbi.CliqueVector,
+      *,
+      initial_measurements: Sequence[mbi.LinearMeasurement] | None = None,
+  ) -> common.DiscreteMechanismResult:
+    common.validate_initial_measurements(initial_measurements)
+    measurements = list(initial_measurements) if initial_measurements else []
+    phase_times = {}
     logging.info('[AIM]: Starting Mechanism.')
     zcdp_rho = self.zcdp_rho
     terminate = False
-    rho_remaining = self._loop_rho
+    rho_remaining = self.zcdp_rho
     max_rounds = self.config.max_rounds or 16 * len(data.domain)
-    rho_per_round = self._loop_rho / max_rounds  # pyrefly: ignore[unsupported-operation]
+    rho_per_round = self.zcdp_rho / max_rounds
 
     #########################################################################
     # Compile workload into candidate measurements, and precompute answers. #
@@ -178,14 +185,14 @@ def _run(self, rng, data, measurements, constraints, phase_times):
         data.domain,
         measurements,
         iters=self.config.pgm_iters,
-        constraints=constraints,
+        constraints=self.config.constraints,
     )
     assert isinstance(model, mbi.MarkovRandomField)
 
     t = 0
     while not terminate:
       t += 1
-      if rho_remaining < 2 * rho_per_round:  # pyrefly: ignore[unsupported-operation]
+      if rho_remaining < 2 * rho_per_round:
         logging.info('[AIM] Final round, Using all remaining privacy budget.')
         rho_per_round = rho_remaining
         terminate = True
@@ -194,11 +201,13 @@ def _run(self, rng, data, measurements, constraints, phase_times):
       # Select a marginal query worst approximated by the current model.     #
       ########################################################################
       with common.timed(phase_times, 'selection'):
-        rho_remaining -= rho_per_round  # pyrefly: ignore[unsupported-operation]
+        rho_remaining -= rho_per_round
         fraction = self.config.select_budget_fraction
-        sigma = accounting.zcdp_gaussian_sigma((1 - fraction) * rho_per_round)  # pyrefly: ignore[unsupported-operation]
-        epsilon = accounting.zcdp_exponential_eps(fraction * rho_per_round)  # pyrefly: ignore[unsupported-operation]
-        size_limit = self.config.max_model_size * (zcdp_rho - rho_remaining) / zcdp_rho  # pyrefly: ignore[unsupported-operation]
+        sigma = accounting.zcdp_gaussian_sigma((1 - fraction) * rho_per_round)
+        epsilon = accounting.zcdp_exponential_eps(fraction * rho_per_round)
+        size_limit = (
+            self.config.max_model_size * (zcdp_rho - rho_remaining) / zcdp_rho
+        )
         small_candidates = _filter_candidates(candidates, model, size_limit)
 
         estimates = mbi.marginal_oracles.bulk_variable_elimination(
@@ -222,7 +231,7 @@ def _run(self, rng, data, measurements, constraints, phase_times):
           '[AIM] Round %d, Budget used: %.4f, Measuring: %s, Candidates: %d,'
           ' cliques: %d, treewidth: %d, memory: %d bytes',
           t,
-          (zcdp_rho - rho_remaining) / zcdp_rho,  # pyrefly: ignore[unsupported-operation]
+          (zcdp_rho - rho_remaining) / zcdp_rho,
           marginal_query,
           len(small_candidates),
           summary.num_cliques,
@@ -255,7 +264,7 @@ def _run(self, rng, data, measurements, constraints, phase_times):
             warm_start=model,
             iters=self.config.pgm_iters,
             callback_fn=callback_fn,
-            constraints=constraints,
+            constraints=self.config.constraints,
         )
         assert isinstance(model, mbi.MarkovRandomField)
 
@@ -272,10 +281,15 @@ def _run(self, rng, data, measurements, constraints, phase_times):
       )
       if np.linalg.norm(new_estimate - old_estimate, ord=1) <= threshold:
         # No useful information at this noise level, increase budget per round.
-        rho_per_round *= self.config.anneal_factor  # pyrefly: ignore[unsupported-operation]
+        rho_per_round *= self.config.anneal_factor
         fraction = self.config.select_budget_fraction
         sigma = accounting.zcdp_gaussian_sigma((1 - fraction) * rho_per_round)
         logging.info('[AIM] Reducing sigma: %.1f', sigma)
 
     synthetic_data = model.synthetic_data()
-    return model, synthetic_data, measurements
+    return common.DiscreteMechanismResult(
+        synthetic_data=synthetic_data,
+        measurements=measurements,
+        model=model,
+        diagnostics=common.clique_stats(model),
+    )