Refactor DiscreteMechanism to use Composition instead of Inheritance.

dpsynth/data_generation_v2.py

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

dpsynth/data_generation_v3.py

@@ -188,7 +188,7 @@ class DataGenerationResult:
   """Result of end-to-end DP synthetic data generation."""
 
   synthetic_data: pd.DataFrame
-  discrete_mechanism_result: dm_common.DiscreteMechanismResult
+  discrete_mechanism_result: dm_common.DiscreteSynthesizerResult
   codec: TabularCodec
 
 
@@ -227,13 +227,14 @@ class TabularSynthesizer(api.DPMechanism):
   """
 
   domains: Mapping[str, domain.AttributeType]
-  discrete_mechanism: discrete_mechanisms.DiscreteMechanism = dataclasses.field(
+  discrete_mechanism: dm_common.DiscreteSynthesizerProtocol = dataclasses.field(
       default_factory=discrete_mechanisms.MSTMechanism
   )
   numerical_bins: int = 32
   init_budget_fraction: float = 0.1
   initializers: dict[str, api.DPMechanism] | None = None
   total_count_sigma: float | None = dataclasses.field(default=None, repr=False)
+  compress_columns: bool | Sequence[str] = False
   cross_attribute_constraints: Sequence[constraints.Constraint] = ()
   experimental_max_records_per_user: int = 1
 
@@ -324,6 +325,19 @@ def configure(  # pyrefly: ignore[bad-override]
         self.discrete_mechanism,
         max_records_per_user=self.experimental_max_records_per_user,
     ).configure(zcdp_rho=discrete_rho)
+
+    if hasattr(self.discrete_mechanism, '__dataclass_fields__'):
+      calibrated_discrete = dataclasses.replace(  # pytype: disable=wrong-arg-types
+          self.discrete_mechanism,
+          max_records_per_user=self.experimental_max_records_per_user,
+      ).configure(
+          zcdp_rho=discrete_rho
+      )
+    else:
+      # Some mechanisms (or mocks) might not be dataclasses. Just configure them.
+      calibrated_discrete = self.discrete_mechanism.configure(
+          zcdp_rho=discrete_rho
+      )
     return dataclasses.replace(
         self,
         initializers=calibrated_inits,
@@ -416,6 +430,16 @@ def __call__(
     mbi_constraints = tuple(
         c.to_mbi() for c in self.cross_attribute_constraints
     )
+
+    mappings = dm_common.compression_mappings(
+        initial_measurements, self.compress_columns, constraints=mbi_constraints
+    )
+    if mappings and hasattr(discrete, 'compress'):
+      discrete = discrete.compress(mappings)  # pyrefly: ignore[bad-argument-type]
+      initial_measurements = [
+          m.compress(mappings, discrete.domain) for m in initial_measurements  # pyrefly: ignore[bad-argument-type]
+      ]
+
     mechanism_result = self.discrete_mechanism(
         rng,
         data=discrete,
@@ -424,6 +448,13 @@ def __call__(
     )
     logging.info('[DPSynth]: Generated discrete synthetic data.')
 
+    if mappings:
+      mechanism_result = dataclasses.replace(
+          mechanism_result,
+          synthetic_data=mechanism_result.synthetic_data.decompress(mappings),
+          mappings=mappings,
+      )
+
     synthetic_data = codec.decode(
         mechanism_result.synthetic_data, rng, column_order
     )

dpsynth/discrete_mechanisms/__init__.py

@@ -18,7 +18,7 @@
 
 from dpsynth.discrete_mechanisms.aim import AIMMechanism
 from dpsynth.discrete_mechanisms.aim_gdp import AIMGDPMechanism
-from dpsynth.discrete_mechanisms.base import DiscreteMechanism
+from dpsynth.discrete_mechanisms.base import DiscreteSynthesizer
 from dpsynth.discrete_mechanisms.common import DiscreteMechanismResult
 from dpsynth.discrete_mechanisms.common import MechanismDiagnostics
 from dpsynth.discrete_mechanisms.direct import DirectMechanism

dpsynth/discrete_mechanisms/aim.py

@@ -19,6 +19,7 @@
 
 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
@@ -87,7 +88,16 @@ def _worst_approximated(
 
 
 @dataclasses.dataclass
-class AIMMechanism(base.DiscreteMechanism):
+class AIMMechanism(api.DPMechanism):
+  """AIM mechanism."""
+
+  marginal_oracle: mbi.MarginalOracle | None = None
+  zcdp_rho: float | None = None
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
+
   """Configuration for the AIM mechanism.
 
   Details are described in the paper:
@@ -124,6 +134,9 @@ class AIMMechanism(base.DiscreteMechanism):
   pgm_iters: int = 1000
   _loop_rho: float | None = dataclasses.field(default=None, repr=False)
 
+  def configure(self, *, zcdp_rho: float, **kwargs) -> AIMMechanism:
+    return dataclasses.replace(self, zcdp_rho=zcdp_rho, _loop_rho=zcdp_rho)
+
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
     """Returns the workload cliques filtered by max_marginal_size."""
     return common.supporting_cliques(
@@ -141,12 +154,17 @@ def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the DP event for the AIM mechanism."""
-    self._check_calibration()
-    events = self._one_way_dp_event()
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    events = []
     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):
+  def __call__(self, rng, data, *, initial_measurements=None, constraints=()):
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    phase_times = {}
+    measurements = list(initial_measurements or [])
     """Adaptively selects, measures, and estimates in an annealed loop."""
     logging.info('[AIM]: Starting Mechanism.')
     zcdp_rho = self.zcdp_rho
@@ -266,4 +284,9 @@ def _run(self, rng, data, measurements, constraints, phase_times):
         logging.info('[AIM] Reducing sigma: %.1f', sigma)
 
     synthetic_data = model.synthetic_data()
-    return model, synthetic_data, measurements
+    return common.DiscreteMechanismResult(
+        model=model,
+        synthetic_data=synthetic_data,
+        measurements=measurements,
+        diagnostics=common.clique_stats(model, phase_times),
+    )

dpsynth/discrete_mechanisms/aim_gdp.py

@@ -20,6 +20,7 @@
 
 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
@@ -144,7 +145,15 @@ def _worst_approximated(
 
 # select loop, injecting the budgeting strategy (zCDP vs. GDP) as configuration.
 @dataclasses.dataclass
-class AIMGDPMechanism(base.DiscreteMechanism):
+class AIMGDPMechanism(api.DPMechanism):
+
+  marginal_oracle: mbi.MarginalOracle | None = None
+  zcdp_rho: float | None = None
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
+
   """Configuration for the AIM mechanism with Gaussian DP.
 
   Details are described in the paper:
@@ -198,20 +207,28 @@ def _one_way_cliques(self, data):
     """Returns only the workload-specified one-way cliques."""
     return common.one_way_cliques(self.workload, data.domain)
 
+  def configure(self, *, zcdp_rho: float, **kwargs) -> AIMGDPMechanism:
+    return dataclasses.replace(self, zcdp_rho=zcdp_rho, _loop_rho=zcdp_rho)
+
   def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
     """Allocates the entire remaining budget to the adaptive loop."""
     return {'_loop_rho': remaining_rho}
 
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the DP event for the AIM-GDP mechanism."""
-    self._check_calibration()
-    events = self._one_way_dp_event()
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    events = []
     # The loop's privacy cost in zCDP terms.
     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):
+  def __call__(self, rng, data, *, initial_measurements=None, constraints=()):
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    phase_times = {}
+    measurements = list(initial_measurements or [])
     """Adaptively selects, measures, and estimates in an annealed loop (GDP)."""
     logging.info('[AIM] Starting Mechanism.')
 
@@ -348,4 +365,9 @@ def _run(self, rng, data, measurements, constraints, phase_times):
         )
 
     synthetic_data = model.synthetic_data()
-    return model, synthetic_data, measurements
+    return common.DiscreteMechanismResult(
+        model=model,
+        synthetic_data=synthetic_data,
+        measurements=measurements,
+        diagnostics=common.clique_stats(model, phase_times),
+    )