Refactor discrete_mechanisms to Select-Measure-Estimate base class

dpsynth/discrete_mechanisms/__init__.py

@@ -16,9 +16,10 @@
 
 # pylint: disable=g-importing-member
 
-from dpsynth.api import DPMechanism as DiscreteMechanism
 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 GaussianMarginalMeasurement
 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,8 +19,8 @@
 
 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
@@ -86,7 +86,7 @@ def _worst_approximated(
 
 
 @dataclasses.dataclass
-class AIMMechanism(api.DPMechanism):
+class AIMMechanism(base.DiscreteMechanism):
   """Configuration for the AIM mechanism.
 
   Details are described in the paper:
@@ -106,102 +106,66 @@ class AIMMechanism(api.DPMechanism):
     workload: A collection of marginal queries (and weights) the synthetic data
       should be tailored to.
     max_rounds: The maximum number of rounds to run the mechanism.
-    pgm_iters: The number of iterations for the mirror descent algorithm.
     max_model_size: The maximum size of the graphical model in megabytes.
       Controls the utility/runtime trade-off.
     max_marginal_size: The maximum size of a marginal query to consider.
-    marginal_oracle: The marginal oracle to use for the mirror descent
-      algorithm.
     anneal_factor: The factor by which to anneal the privacy.
-    one_way_budget_fraction: The fraction of the total budget to use for one-way
-      marginal queries.
-    compress_columns: Controls domain compression. True compresses all columns,
-      False disables compression, or a list of column names to compress.
     select_budget_fraction: The fraction of the total budget to use for
       selecting two-way marginal queries.
   """
 
   workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None
   max_rounds: int | None = None
-  pgm_iters: int = 1000
   max_model_size: int = 80
   max_marginal_size: float = 1e6
-  marginal_oracle: mbi.MarginalOracle | None = None
   anneal_factor: float = 4.0
-  one_way_budget_fraction: float = 0.1
-  compress_columns: bool | Sequence[str] = False
   select_budget_fraction: float = 0.1
-  zcdp_rho: float | None = None
+  _loop_rho: float | None = dataclasses.field(default=None, repr=False)
 
   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 configure(self, *, zcdp_rho: float, delta: float = 0.0) -> 'AIMMechanism':
-    """Returns a new instance calibrated to the given zCDP budget."""
-    return dataclasses.replace(self, zcdp_rho=zcdp_rho)
-
-  @property
-  def dp_event(self) -> dp_accounting.DpEvent:
-    """Returns the DP event for the AIM mechanism."""
-    if self.zcdp_rho is None:
-      raise ValueError('Must call calibrate() before using the mechanism.')
-    return dp_accounting.ZCDpEvent(self.zcdp_rho)
+  def _one_way_cliques(self, data):
+    """Returns only the workload-specified one-way cliques."""
+    return common.one_way_cliques(self.workload, data.domain)
 
-  def __call__(
+  def configure(
       self,
-      rng: np.random.Generator,
-      data: mbi.Dataset | mbi.CliqueVector,
       *,
+      zcdp_rho: float,
+      delta: float = 0.0,
       initial_measurements: list[mbi.LinearMeasurement] | None = None,
-      constraints: tuple[mbi.Constraint, ...] = (),
-  ) -> common.DiscreteMechanismResult:
-    """Runs the AIM mechanism on the given data.
-
-    Args:
-      rng: A numpy random number generator.
-      data: The input data. Must be an mbi.Dataset for domain compression;
-        mbi.CliqueVector is supported but compression will be skipped.
-      initial_measurements: Optional initial measurements to start from.
-      constraints: Structural constraints for the estimation.
-
-    Returns:
-      A DiscreteMechanismResult containing the estimated data distribution.
-    """
-    if self.zcdp_rho is None:
-      raise ValueError('Must call calibrate() before using the mechanism.')
+      **kwargs,
+  ) -> 'AIMMechanism':
+    """Returns a new instance calibrated to the given zCDP budget."""
+    result = super().configure(
+        zcdp_rho=zcdp_rho,
+        delta=delta,
+        initial_measurements=initial_measurements,
+    )
+    return dataclasses.replace(result, _loop_rho=result.remaining_rho)
 
+  @property
+  def dp_event(self) -> dp_accounting.DpEvent:
+    """Returns the DP event for the AIM mechanism."""
+    self._check_calibration()
+    events = []
+    if self.one_way_mechanism is not None:
+      events.append(self.one_way_mechanism.dp_event)
+    events.append(dp_accounting.ZCDpEvent(self._loop_rho))
+    return dp_accounting.ComposedDpEvent(events)
+
+  def _run(self, rng, data, measurements, constraints, phase_times, mappings):
+    """Adaptively selects, measures, and estimates in an annealed loop."""
     logging.info('[AIM]: Starting Mechanism.')
-    phase_times = {}
-
     zcdp_rho = self.zcdp_rho
     terminate = False
-    rho_remaining = zcdp_rho
+    rho_remaining = self._loop_rho
     max_rounds = self.max_rounds or 16 * len(data.domain)
-    rho_per_round = zcdp_rho / max_rounds
-
-    if initial_measurements is None:
-      rho_remaining -= self.one_way_budget_fraction * zcdp_rho
-      marginal_queries = common.one_way_cliques(self.workload, data.domain)
-      measurements = common.measure_marginals_with_noise(
-          rng,
-          data,
-          marginal_queries=marginal_queries,  # pyrefly: ignore[bad-argument-type]
-          gdp_sigma=accounting.zcdp_gaussian_sigma(
-              zcdp_rho * self.one_way_budget_fraction
-          ),
-      )
-    else:
-      measurements = list(initial_measurements)
-
-    mappings = common.compression_mappings(
-        measurements, self.compress_columns, constraints
-    )
-    if mappings:
-      data = data.compress(mappings)
-      measurements = [m.compress(mappings, data.domain) for m in measurements]
+    rho_per_round = self._loop_rho / max_rounds
 
     #########################################################################
     # Compile workload into candidate measurements, and precompute answers. #
@@ -305,16 +269,7 @@ def __call__(
         sigma = accounting.zcdp_gaussian_sigma((1 - fraction) * rho_per_round)
         logging.info('[AIM] Reducing sigma: %.1f', sigma)
 
-    diagnostics = common.clique_stats(model)
-    diagnostics.phase_times = phase_times
-    diagnostics.num_rounds = t
     synthetic_data = model.synthetic_data()
     if mappings:
       synthetic_data = synthetic_data.decompress(mappings)
-    return common.DiscreteMechanismResult(
-        model=model,
-        synthetic_data=synthetic_data,
-        measurements=measurements,
-        diagnostics=diagnostics,
-        mappings=mappings,
-    )
+    return model, synthetic_data, measurements

dpsynth/discrete_mechanisms/aim_gdp.py

@@ -20,8 +20,8 @@
 
 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
@@ -124,7 +124,7 @@ def _worst_approximated(
 
 
 @dataclasses.dataclass
-class AIMGDPMechanism(api.DPMechanism):
+class AIMGDPMechanism(base.DiscreteMechanism):
   """Configuration for the AIM mechanism with Gaussian DP.
 
   Details are described in the paper:
@@ -146,117 +146,76 @@ class AIMGDPMechanism(api.DPMechanism):
       each marginal query. A default value of 1.0 will be assigned if the
       workload is provided as a list.
     max_rounds: The maximum number of rounds to run the mechanism.
-    pgm_iters: The number of iterations for the mirror descent algorithm.
     max_model_size: The maximum size of the graphical model in megabytes.
       Controls the utility/runtime trade-off.
     max_marginal_size: The maximum size of a marginal query to consider.
     max_candidates_per_round: The maximum number of candidates to consider per
       round. This can improve privacy budget utilization as well as speed up the
       "Select" step, which in some settings is the main bottelneck of the
       mechanism.
-    marginal_oracle: The marginal oracle to use for the mirror descent
-      algorithm.
     anneal_factor: The factor by which to anneal the privacy budget.
-    one_way_budget_fraction: The fraction of the privacy budget to use for
-      one-way marginals.
-    compress_columns: Controls domain compression. True compresses all columns,
-      False disables compression, or a list of column names to compress.
     select_budget_fraction: The fraction of the privacy budget to use for the
       "Select" step.
-    gdp_sigma: The GDP sigma of the end-to-end mechanism. Privacy budget is
-      split across rounds internally.
   """
 
   workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None
   max_rounds: int | None = None
-  pgm_iters: int = 1000
   max_model_size: int = 80
   max_marginal_size: float = 1e6
   max_candidates_per_round: int = 16
-  marginal_oracle: mbi.MarginalOracle | None = None
   anneal_factor: float = 4.0
-  one_way_budget_fraction: float = 0.1
-  compress_columns: bool | Sequence[str] = False
   select_budget_fraction: float = 0.1
-  gdp_sigma: float | None = None
+  _loop_rho: float | None = dataclasses.field(default=None, repr=False)
 
   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 _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, delta: float = 0.0
+      self,
+      *,
+      zcdp_rho: float,
+      delta: float = 0.0,
+      initial_measurements: list[mbi.LinearMeasurement] | None = None,
+      **kwargs,
   ) -> 'AIMGDPMechanism':
     """Returns a new instance calibrated to the given zCDP budget."""
-    return dataclasses.replace(
-        self, gdp_sigma=accounting.zcdp_gaussian_sigma(zcdp_rho)
+    result = super().configure(
+        zcdp_rho=zcdp_rho,
+        delta=delta,
+        initial_measurements=initial_measurements,
     )
+    return dataclasses.replace(result, _loop_rho=result.remaining_rho)
 
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the DP event for the AIM-GDP mechanism."""
-    if self.gdp_sigma is None:
-      raise ValueError('Must call calibrate() before using the mechanism.')
-    return dp_accounting.GaussianDpEvent(noise_multiplier=self.gdp_sigma)
-
-  def __call__(
-      self,
-      rng: np.random.Generator,
-      data: mbi.Dataset | mbi.CliqueVector,
-      *,
-      initial_measurements: list[mbi.LinearMeasurement] | None = None,
-      constraints: tuple[mbi.Constraint, ...] = (),
-  ) -> common.DiscreteMechanismResult:
-    """Runs the AIM-GDP mechanism on the given data.
-
-    Args:
-      rng: A numpy random number generator.
-      data: The input data. Must be an mbi.Dataset for domain compression;
-        mbi.CliqueVector is supported but compression will be skipped.
-      initial_measurements: Optional initial measurements to start from.
-      constraints: Structural constraints for the estimation.
-
-    Returns:
-      A DiscreteMechanismResult containing the estimated data distribution.
-    """
-    if self.gdp_sigma is None:
-      raise ValueError('Must call calibrate() before using the mechanism.')
-
+    self._check_calibration()
+    events = []
+    if self.one_way_mechanism is not None:
+      events.append(self.one_way_mechanism.dp_event)
+    # The loop's privacy cost in zCDP terms.
+    events.append(dp_accounting.ZCDpEvent(self._loop_rho))
+    return dp_accounting.ComposedDpEvent(events)
+
+  def _run(self, rng, data, measurements, constraints, phase_times, mappings):
+    """Adaptively selects, measures, and estimates in an annealed loop (GDP)."""
     logging.info('[AIM] Starting Mechanism.')
-    phase_times = {}
 
-    # Convert end-to-end GDP sigma to budget for internal allocation.
-    gdp_budget = 1.0 / self.gdp_sigma**2
+    # Convert loop's zCDP budget to GDP budget for internal allocation.
+    gdp_budget = accounting.zcdp_to_gdp(self._loop_rho)
 
     terminate = False
     budget_remaining = gdp_budget
     max_rounds = self.max_rounds or 16 * len(data.domain)
     budget_per_round = budget_remaining / max_rounds
 
-    if initial_measurements is None:
-      one_way_budget = self.one_way_budget_fraction * gdp_budget
-      one_way_gdp_sigma = accounting.gdp_gaussian_sigma(one_way_budget)
-      budget_remaining -= one_way_budget
-      marginal_queries = common.one_way_cliques(self.workload, data.domain)
-      # measure_marginals_with_noise splits one_way_gdp_sigma across queries.
-      measurements = common.measure_marginals_with_noise(
-          rng,
-          data,
-          marginal_queries=marginal_queries,
-          gdp_sigma=one_way_gdp_sigma,
-      )
-    else:
-      measurements = list(initial_measurements)
-
-    mappings = common.compression_mappings(
-        measurements, self.compress_columns, constraints
-    )
-    if mappings:
-      data = data.compress(mappings)
-      measurements = [m.compress(mappings, data.domain) for m in measurements]
-
     #########################################################################
     # Compile workload into candidate measurements, and precompute answers. #
     #########################################################################
@@ -376,16 +335,7 @@ def __call__(
             '[AIM] Increasing budget per round: %.5f', budget_per_round
         )
 
-    diagnostics = common.clique_stats(model)
-    diagnostics.phase_times = phase_times
-    diagnostics.num_rounds = t
     synthetic_data = model.synthetic_data()
     if mappings:
       synthetic_data = synthetic_data.decompress(mappings)
-    return common.DiscreteMechanismResult(
-        model=model,
-        synthetic_data=synthetic_data,
-        measurements=measurements,
-        diagnostics=diagnostics,
-        mappings=mappings,
-    )
+    return model, synthetic_data, measurements