From 9db1dd34f3413b0c759f1c00551588c516ba0aa2 Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Fri, 14 Aug 2026 07:46:58 -0700 Subject: [PATCH] Refactor discrete mechanisms to use composition instead of inheritance, making it much easier to understand and reason about the code, while retaining the main code sharing benefits of the previous approach. Previously, DiscreteMechanism was a parent class that handled one-way marginal measurement + domain compression, then delegated to subclass for the rest of the mechanism. Now DiscreteMechanism is a standalone class defining a one_way_budget_fraction, a domain_compression flag and a base_mechanism. It measures one way marginals then calls the base mechanism. This means that the base mechanism implementations are usually completely self-contained in their corresponding file, and the (dp_event, __call__) can be verified to match on a file-by-file basis. DiscreteMechanism handles the composed event via dp_accounting.ComposedEvent and can similarly be easily verified to be correct locally to that file (assuming the base_mechanism.dp_event is correct). Now, DiscreteConfig / DiscreteMechanism sits side-by-side to TabularConfig / TabularMechanism, and should be preferred when users have pre-discretized data (e.g., mbi.Dataset) instead of a regular tabular dataset with various column types. PiperOrigin-RevId: 964696330 --- dpsynth/adapters/beam.py | 10 +- dpsynth/api.py | 2 + dpsynth/data_generation_v3.py | 24 +- dpsynth/discrete_mechanisms/__init__.py | 21 +- dpsynth/discrete_mechanisms/aim.py | 80 ++--- dpsynth/discrete_mechanisms/aim_gdp.py | 63 ++-- dpsynth/discrete_mechanisms/base.py | 275 ------------------ dpsynth/discrete_mechanisms/clique_tree.py | 4 + dpsynth/discrete_mechanisms/common.py | 19 +- dpsynth/discrete_mechanisms/direct.py | 123 ++++++-- .../discrete_synthesizer.py | 179 ++++++++++++ dpsynth/discrete_mechanisms/independent.py | 91 ++++-- dpsynth/discrete_mechanisms/mst.py | 113 +++++-- dpsynth/discrete_mechanisms/swift.py | 78 ++--- dpsynth/discrete_mechanisms/swift_utils.py | 1 + dpsynth/text/dp_sft.py | 4 +- tests/adapters/beam_test.py | 2 +- tests/data_generation_v3_test.py | 8 +- tests/discrete_mechanisms/aim_test.py | 4 +- tests/discrete_mechanisms/base_test.py | 119 -------- .../discrete_mechanisms_test.py | 45 +-- .../discrete_synthesizer_test.py | 130 +++++++++ tests/discrete_mechanisms/independent_test.py | 26 +- tests/discrete_mechanisms/mst_test.py | 24 +- tests/discrete_mechanisms/swift_test.py | 1 + 25 files changed, 836 insertions(+), 610 deletions(-) delete mode 100644 dpsynth/discrete_mechanisms/base.py create mode 100644 dpsynth/discrete_mechanisms/discrete_synthesizer.py delete mode 100644 tests/discrete_mechanisms/base_test.py create mode 100644 tests/discrete_mechanisms/discrete_synthesizer_test.py diff --git a/dpsynth/adapters/beam.py b/dpsynth/adapters/beam.py index 26855701..ae1e9ff1 100644 --- a/dpsynth/adapters/beam.py +++ b/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. diff --git a/dpsynth/api.py b/dpsynth/api.py index 70ce4ce3..90ea5409 100644 --- a/dpsynth/api.py +++ b/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. diff --git a/dpsynth/data_generation_v3.py b/dpsynth/data_generation_v3.py index 5ebe2891..cb3db12b 100644 --- a/dpsynth/data_generation_v3.py +++ b/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.""" diff --git a/dpsynth/discrete_mechanisms/__init__.py b/dpsynth/discrete_mechanisms/__init__.py index 4b5a29c7..3984811b 100644 --- a/dpsynth/discrete_mechanisms/__init__.py +++ b/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 diff --git a/dpsynth/discrete_mechanisms/aim.py b/dpsynth/discrete_mechanisms/aim.py index 97964409..1910c1f7 100644 --- a/dpsynth/discrete_mechanisms/aim.py +++ b/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,6 +115,8 @@ 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 @@ -122,6 +124,7 @@ class AIMConfig(base.DiscreteMechanismConfig): 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.""" @@ -129,40 +132,44 @@ def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: 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), + ) diff --git a/dpsynth/discrete_mechanisms/aim_gdp.py b/dpsynth/discrete_mechanisms/aim_gdp.py index 778ce5ed..a06164b0 100644 --- a/dpsynth/discrete_mechanisms/aim_gdp.py +++ b/dpsynth/discrete_mechanisms/aim_gdp.py @@ -15,13 +15,14 @@ """Variant of the Adaptive+Iterative Mechanism (AIM) that satisfies Gaussian DP.""" from collections.abc import Iterable, Mapping +from collections.abc import Sequence import dataclasses import typing 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 @@ -144,7 +145,7 @@ def _worst_approximated( # select loop, injecting the budgeting strategy (zCDP vs. GDP) as configuration. @dataclasses.dataclass(frozen=True) -class AIMGDPConfig(base.DiscreteMechanismConfig): +class AIMGDPConfig(api.MechanismConfig): """Configuration for the AIM mechanism with Gaussian DP. Details are described in the paper: @@ -178,6 +179,8 @@ class AIMGDPConfig(base.DiscreteMechanismConfig): "Select" step. """ + constraints: Sequence[mbi.Constraint] = () + workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None max_rounds: int | None = None max_model_size: int = 80 @@ -186,6 +189,7 @@ class AIMGDPConfig(base.DiscreteMechanismConfig): 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.""" @@ -193,39 +197,43 @@ def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: 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) -> 'AIMGDP': - return AIMGDP(**kwargs) + def configure(self, *, zcdp_rho, delta=0, max_records_per_user=1): + api.validate_max_records_per_user(max_records_per_user) + return AIMGDP( + config=self, + gdp_budget=accounting.zcdp_to_gdp(zcdp_rho), + max_records_per_user=max_records_per_user, + ) @dataclasses.dataclass(frozen=True, kw_only=True) -class AIMGDP(base.DiscreteMechanism): +class AIMGDP(api.CalibratedMechanism): """Calibrated AIMGDP instance.""" config: AIMGDPConfig - _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) + gdp_budget: float + max_records_per_user: int = 1 @property def dp_event(self) -> dp_accounting.DpEvent: """Returns the DP event for the AIM-GDP mechanism.""" - events = self._one_way_dp_event() - # 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) + return dp_accounting.GaussianDpEvent( + accounting.gdp_gaussian_sigma(self.gdp_budget) + ) - def _run(self, rng, data, measurements, constraints, phase_times): - """Adaptively selects, measures, and estimates in an annealed loop (GDP).""" + 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.') - # Convert loop's zCDP budget to GDP budget for internal allocation. - gdp_budget = accounting.zcdp_to_gdp(self._loop_rho) # pyrefly: ignore[bad-argument-type] + gdp_budget = self.gdp_budget terminate = False budget_remaining = gdp_budget @@ -247,7 +255,7 @@ def _run(self, rng, data, measurements, constraints, phase_times): domain, measurements, iters=self.config.pgm_iters, - constraints=constraints, + constraints=self.config.constraints, ) assert isinstance(model, mbi.MarkovRandomField) logging.info('[AIM] Estimated initial model.') @@ -334,7 +342,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, ) model = typing.cast(mbi.MarkovRandomField, model) @@ -359,4 +367,9 @@ def _run(self, rng, data, measurements, constraints, phase_times): ) 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), + ) diff --git a/dpsynth/discrete_mechanisms/base.py b/dpsynth/discrete_mechanisms/base.py deleted file mode 100644 index e58bf581..00000000 --- a/dpsynth/discrete_mechanisms/base.py +++ /dev/null @@ -1,275 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Base classes for the select-measure-estimate paradigm. - -This module defines the ``DiscreteMechanism`` base class, which implements the -select-measure-estimate paradigm from `McKenna et al. (2021) -`_. Each step of the paradigm is a separate -method that subclasses can override independently, enabling code reuse across -mechanisms that differ primarily in the *select* step. -""" - -from __future__ import annotations - -import abc -from collections.abc import Mapping, 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 common -import mbi -import mbi.callbacks -import mbi.estimation -import numpy as np - - -# hard-codes `mbi.estimation.MirrorDescent` as the estimator; abstract this -# (e.g. an injectable optimizer strategy) to support other synthesizers such as -# PrivSyn and GEM. - - -@dataclasses.dataclass(frozen=True) -class DiscreteMechanismConfig(api.MechanismConfig): - """Base class for mechanisms following the select-measure-estimate paradigm. - - Subclasses implement ``_select`` to define which marginals to measure. - The base ``__call__`` orchestrates the full pipeline:: - - measure_one_way → compress → run → result - - where ``_run`` performs select → measure → estimate → generate. One-shot - mechanisms need only override ``_select``; adaptive mechanisms (e.g. AIM) or - those needing a custom estimator (e.g. SWIFT) override ``_run`` directly. - - Attributes: - marginal_oracle: Oracle for marginal inference in Private-PGM. - pgm_iters: Number of mirror descent iterations for estimation. - compress_columns: Domain compression config. True = all, list = specific. - one_way_budget_fraction: Fraction of zCDP budget for one-way marginals. - """ - - marginal_oracle: mbi.MarginalOracle | None = None - pgm_iters: int = 5000 - compress_columns: bool | Sequence[str] = False - one_way_budget_fraction: float = 1 / 3 - - @abc.abstractmethod - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - """Returns the cliques whose marginals this mechanism supports. - - These are the cliques that the graphical model produced by this mechanism - is guaranteed to represent for the given domain, i.e. the marginal queries - the resulting synthetic data can answer. Callers use them to reason about - the mechanism's coverage without having to run it. - - Args: - domain: Information about the tabular columns and their types. - - Returns: - A list of cliques. - """ - - @abc.abstractmethod - def _create_mechanism(self, **kwargs) -> 'DiscreteMechanism': - """Instantiates the calibrated mechanism object.""" - - def configure(self, *, zcdp_rho, delta=0, max_records_per_user=1): - """Configures the mechanism with a zCDP budget.""" - api.validate_max_records_per_user(max_records_per_user) - if self.one_way_budget_fraction <= 0: - one_way_rho = None - else: - one_way_rho = zcdp_rho * self.one_way_budget_fraction - remaining_rho = zcdp_rho - (one_way_rho or 0.0) - - return self._create_mechanism( - config=self, - zcdp_rho=zcdp_rho, - one_way_rho=one_way_rho, - max_records_per_user=max_records_per_user, - **self._allocate_budget(remaining_rho), - ) - - def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]: - """Splits the post-one-way budget into mechanism-specific rho fields.""" - return {} - - -@dataclasses.dataclass(frozen=True) -class DiscreteMechanism(api.CalibratedMechanism): - """Calibrated, runnable select-measure-estimate mechanism.""" - - config: DiscreteMechanismConfig - zcdp_rho: float - one_way_rho: float | None = dataclasses.field(default=None, repr=False) - measurement_rho: float | None = dataclasses.field(default=None, repr=False) - max_records_per_user: int = 1 - - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - """Returns the cliques whose marginals this mechanism supports. - - These are the cliques that the graphical model produced by this mechanism - is guaranteed to represent for the given domain, i.e. the marginal queries - the resulting synthetic data can answer. Callers use them to reason about - the mechanism's coverage without having to run it. - - Args: - domain: Information about the tabular columns and their types. - - Returns: - A list of cliques. - """ - return self.config.supporting_cliques(domain) - - def _one_way_dp_event(self): - """DpEvents for the shared one-way measurement ([] if there is none).""" - if self.one_way_rho is None: - return [] - return [ - dp_accounting.GaussianDpEvent( - noise_multiplier=accounting.zcdp_gaussian_sigma(self.one_way_rho) - ) - ] - - def _one_way_cliques(self, data): - """Returns the one-way cliques to measure.""" - cliques = [(a,) for a in data.domain] - if hasattr(data, 'cliques'): - supported = common.downward_closure(data.cliques) # pyrefly: ignore[attribute-error] - cliques = [cl for cl in cliques if cl in supported] - return cliques - - def _measure_one_way( - self, rng, data, phase_times, *, initial_measurements=None - ): - """Measures one-way marginals or returns pre-measured ones.""" - if initial_measurements is not None: - return list(initial_measurements) - if self.one_way_rho is None: - return [] - with common.timed(phase_times, 'measurement'): - sigma = accounting.zcdp_gaussian_sigma(self.one_way_rho) - cliques = self._one_way_cliques(data) - return common.measure_marginals_with_noise( - rng, - data, - cliques, - sigma, - max_records_per_user=self.max_records_per_user, - ) - - def _compress(self, data, measurements, constraints): - """Compresses the domain by merging rare values.""" - mappings = common.compression_mappings( - measurements, self.config.compress_columns, constraints - ) - if mappings and hasattr(data, 'compress'): - data = data.compress(mappings) # pyrefly: ignore[attribute-error] - measurements = [m.compress(mappings, data.domain) for m in measurements] - return data, measurements, mappings - - def _select(self, rng, data, measurements, phase_times): - """Selects which marginals to measure. Mechanism-specific.""" - raise NotImplementedError - - def __call__( - self, - rng: np.random.Generator, - data: mbi.Dataset | mbi.CliqueVector, - *, - # By default we set one_way_budget_fraction = 1/3 in the config. - # If initial_measurements is provided here, then we should be able to save - # that budget and use it in the subsequence select + measure steps. - # Currently, this budget is just wasted, which is a code smell and - # utility issue. This will be resolved in a near-term future refactoring. - initial_measurements: Sequence[mbi.LinearMeasurement] | None = None, - constraints: Sequence[mbi.Constraint] = (), - ) -> common.DiscreteMechanismResult: - """Runs the select-measure-estimate pipeline.""" - phase_times = {} - measurements = self._measure_one_way( - rng, data, phase_times, initial_measurements=initial_measurements - ) - data, measurements, mappings = self._compress( - data, measurements, constraints - ) - model, synthetic_data, measurements = self._run( - rng, data, measurements, constraints, phase_times - ) - if mappings: - synthetic_data = synthetic_data.decompress(mappings) - diagnostics = common.clique_stats(model) # pyrefly: ignore[wrong-arg-types] - diagnostics.phase_times = phase_times - return common.DiscreteMechanismResult( - model=model, # pyrefly: ignore[wrong-arg-types] - synthetic_data=synthetic_data, - measurements=measurements, - diagnostics=diagnostics, - mappings=mappings, - ) - - def _run(self, rng, data, measurements, constraints, phase_times): - """Selects, measures, estimates, and generates in the compressed domain.""" - # Adaptive mechanisms (e.g. AIM) override this to interleave selection and - # measurement in a loop; SWIFT overrides it to use a junction-tree oracle. - selected = self._select(rng, data, measurements, phase_times) - all_cliques = [m.clique for m in measurements] + list(selected) - logging.info( - '[%s]:\n%s', - type(self).__name__, - mbi.summarize(data.domain, all_cliques), - ) - - # Kick off async AOT compilation of the estimator while we measure. - estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle) - futures = None - try: - futures = estimator.precompile( - data.domain, measurements, extra_cliques=list(selected) - ) - except Exception as e: # pylint: disable=broad-exception-caught - logging.warning('Precompile failed (non-fatal): %s', e) - - if selected: - with common.timed(phase_times, 'measurement'): - sigma = accounting.zcdp_gaussian_sigma(self.measurement_rho) # pyrefly: ignore[bad-argument-type] - measurements = measurements + common.measure_marginals_with_noise( - rng, - data, - selected, - sigma, - max_records_per_user=self.max_records_per_user, - ) - - with common.timed(phase_times, 'estimation'): - if futures is not None: - try: - futures.result() - except Exception as e: # pylint: disable=broad-exception-caught - logging.warning('Precompile wait failed (non-fatal): %s', e) - model = estimator.estimate( - data.domain, - measurements, - iters=self.config.pgm_iters, - callback_fn=mbi.callbacks.default(measurements, data.domain), - constraints=constraints, - ) - assert isinstance(model, mbi.MarkovRandomField) - - synthetic_data = model.synthetic_data() - return model, synthetic_data, measurements diff --git a/dpsynth/discrete_mechanisms/clique_tree.py b/dpsynth/discrete_mechanisms/clique_tree.py index 60d51de6..0c16e4f9 100644 --- a/dpsynth/discrete_mechanisms/clique_tree.py +++ b/dpsynth/discrete_mechanisms/clique_tree.py @@ -29,6 +29,10 @@ (b) an edge in the clique tree. Incorporating cliques according to these rules ensures the tree structure remains tractable for downstream marginal inference algorithms. + + +directly will miss out on features like 1-way measurement selection and domain +compression. """ # NOTE: This module is tested in swift_test.py. diff --git a/dpsynth/discrete_mechanisms/common.py b/dpsynth/discrete_mechanisms/common.py index 5a4d3502..89a2f952 100644 --- a/dpsynth/discrete_mechanisms/common.py +++ b/dpsynth/discrete_mechanisms/common.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Common utility functions for synthetic data mechanisms.""" +"""Common utility functions for synthetic data mechanisms. + +be used within `DiscreteSynthesizer` or `TabularSynthesizer`. Users who call it +directly will miss out on features like 1-way measurement selection and domain +compression. +""" from collections.abc import Iterable, Mapping, Sequence import contextlib @@ -34,6 +39,18 @@ import tqdm +def validate_initial_measurements( + initial_measurements: Sequence[mbi.LinearMeasurement] | None, +) -> None: + """Warns if the discrete mechanism is called without initial measurements.""" + if not initial_measurements: + logging.warning( + 'This mechanism is not intended to be called directly, but through ' + 'the higher-level APIs (TabularMechanism, DiscreteMechanism). It ' + 'will run without error, but is missing 1-way marginal selection.' + ) + + @dataclasses.dataclass class MechanismDiagnostics: """Diagnostic info from a discrete mechanism run. diff --git a/dpsynth/discrete_mechanisms/direct.py b/dpsynth/discrete_mechanisms/direct.py index f5126622..60e0a1b2 100644 --- a/dpsynth/discrete_mechanisms/direct.py +++ b/dpsynth/discrete_mechanisms/direct.py @@ -14,59 +14,122 @@ """Implementation of the direct mechanism.""" -from collections.abc import 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 mbi +import numpy as np @dataclasses.dataclass(frozen=True, kw_only=True) -class DirectConfig(base.DiscreteMechanismConfig): - """Configuration for the direct mechanism. - - The direct mechanism measures a prespecified set of marginal queries, - allocating the entire privacy budget to those measurements. It does not - measure its own one-way marginals, but can incorporate externally supplied - ``initial_measurements`` (e.g. compressed one-ways from an orchestration - layer) at no additional budget cost. - - Attributes: - prespecified_marginal_queries: A list of k-way marginals that a user has - specified. Only these will be measured with privacy budget. - one_way_budget_fraction: Fraction of the zCDP budget allocated to one-way - marginals. Overridden to 0.0 because this mechanism does not measure its - own one-way marginals. - """ +class DirectConfig(api.MechanismConfig): + """Config for the direct mechanism that measures prespecified marginals.""" + + constraints: Sequence[mbi.Constraint] = () + + def configure(self, *, zcdp_rho, delta=0, max_records_per_user=1): + api.validate_max_records_per_user(max_records_per_user) + return Direct( + config=self, + gdp_budget=accounting.zcdp_to_gdp(zcdp_rho), + max_records_per_user=max_records_per_user, + ) + marginal_oracle: mbi.MarginalOracle | None = None + pgm_iters: int = 5000 prespecified_marginal_queries: list[tuple[str, ...]] = dataclasses.field( default_factory=list ) - one_way_budget_fraction: float = 0.0 def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: """Returns the prespecified marginal queries.""" return list(self.prespecified_marginal_queries) - def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]: - """Allocates the full remaining budget to the prespecified queries.""" - return {'measurement_rho': remaining_rho} - - def _create_mechanism(self, **kwargs) -> 'Direct': - return Direct(**kwargs) - @dataclasses.dataclass(frozen=True) -class Direct(base.DiscreteMechanism): +class Direct(api.CalibratedMechanism): + """Calibrated direct mechanism instance.""" + config: DirectConfig + gdp_budget: float + max_records_per_user: int = 1 + @property def dp_event(self) -> dp_accounting.DpEvent: - """Returns the DP event for the direct mechanism.""" + """Returns the DP event.""" return dp_accounting.GaussianDpEvent( - noise_multiplier=accounting.zcdp_gaussian_sigma(self.measurement_rho) # pyrefly: ignore[bad-argument-type] + accounting.gdp_gaussian_sigma(self.gdp_budget) ) def _select(self, rng, data, measurements, phase_times): return list(self.config.prespecified_marginal_queries) + + def __call__( + self, + rng: np.random.Generator, + data: mbi.Dataset | mbi.CliqueVector, + *, + initial_measurements: Sequence[mbi.LinearMeasurement] | None = None, + ) -> common.DiscreteMechanismResult: + """Selects, measures, estimates, and generates in the compressed domain.""" + common.validate_initial_measurements(initial_measurements) + measurements = list(initial_measurements) if initial_measurements else [] + phase_times = {} + selected = self._select(rng, data, measurements, phase_times) + all_cliques = [m.clique for m in measurements] + list(selected) + if all_cliques: + logging.info( + '[%s]:\n%s', + type(self).__name__, + mbi.summarize(data.domain, all_cliques), + ) + else: + logging.info('[%s]: No cliques selected', type(self).__name__) + + # Kick off async AOT compilation of the estimator while we measure. + estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle) + futures = None + try: + futures = estimator.precompile( + data.domain, measurements, extra_cliques=list(selected) + ) + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning('Precompile failed (non-fatal): %s', e) + + if selected: + with common.timed(phase_times, 'measurement'): + sigma = accounting.gdp_gaussian_sigma(self.gdp_budget) + measurements = measurements + common.measure_marginals_with_noise( + rng, + data, # pyrefly: ignore[bad-argument-type] + selected, + sigma, + max_records_per_user=self.max_records_per_user, + ) + + with common.timed(phase_times, 'estimation'): + if futures is not None: + try: + futures.result() + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning('Precompile wait failed (non-fatal): %s', e) + model = estimator.estimate( + data.domain, + measurements, + iters=self.config.pgm_iters, + callback_fn=mbi.callbacks.default(measurements, data.domain), + constraints=self.config.constraints, + ) + assert isinstance(model, mbi.MarkovRandomField) + + synthetic_data = model.synthetic_data() + return common.DiscreteMechanismResult( + synthetic_data=synthetic_data, + measurements=measurements, + model=model, + diagnostics=common.clique_stats(model), + ) diff --git a/dpsynth/discrete_mechanisms/discrete_synthesizer.py b/dpsynth/discrete_mechanisms/discrete_synthesizer.py new file mode 100644 index 00000000..d4381947 --- /dev/null +++ b/dpsynth/discrete_mechanisms/discrete_synthesizer.py @@ -0,0 +1,179 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wrapper that adds one-way marginals and domain compression to any mechanism. + +``DiscreteConfig`` composes an inner discrete mechanism (AIM, MST, +SWIFT, etc.) with shared pre- and post-processing: one-way marginal +measurement, domain compression, and decompression. It is the recommended +entry point for purely discrete tables; mixed-type tables should use +``TabularSynthesizer`` in ``data_generation_v3.py`` instead. + + +Note: This mechanism is not intended to be called directly. It should typically +be used within `DiscreteMechanism` or `TabularSynthesizer`. Users who call it +directly will miss out on features like 1-way measurement selection and domain +compression. +""" + +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses + +import dp_accounting +from dpsynth import api +from dpsynth.discrete_mechanisms import accounting +from dpsynth.discrete_mechanisms import common +import mbi +import numpy as np + + +def _default_mechanism(): + """Lazy import to avoid circular dependency.""" + from dpsynth.discrete_mechanisms import mst # pylint: disable=g-import-not-at-top + + return mst.MSTConfig() + + +@dataclasses.dataclass(frozen=True) +# also adding it to TabularData natively so they can act as siblings. +class DiscreteConfig(api.MechanismConfig): + """Wraps an inner mechanism with one-way measurement and compression. + + Attributes: + mechanism: The inner mechanism config (e.g. ``AIMConfig()``). + compress_columns: Domain compression config. True = all, list = specific. + one_way_budget_fraction: Fraction of zCDP budget for one-way marginals. + """ + + constraints: Sequence[mbi.Constraint] = () + + mechanism: api.MechanismConfig = dataclasses.field( + default_factory=_default_mechanism + ) + compress_columns: bool | Sequence[str] = False + one_way_budget_fraction: float = 1 / 3 + + def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: + """Delegates to the inner mechanism's supporting_cliques.""" + return self.mechanism.supporting_cliques(domain) # pyrefly: ignore[missing-attribute] + + def configure(self, *, zcdp_rho, delta=0, max_records_per_user=1): + """Configures the synthesizer with a zCDP budget.""" + api.validate_max_records_per_user(max_records_per_user) + + if self.one_way_budget_fraction <= 0: + one_way_rho = None + else: + one_way_rho = zcdp_rho * self.one_way_budget_fraction + remaining_rho = zcdp_rho - (one_way_rho or 0.0) + inner = self.mechanism.configure( + zcdp_rho=remaining_rho, + delta=delta, + max_records_per_user=max_records_per_user, + ) + return DiscreteMechanism( + config=self, + inner=inner, + zcdp_rho=zcdp_rho, + one_way_rho=one_way_rho, + max_records_per_user=max_records_per_user, + ) + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class DiscreteMechanism(api.CalibratedMechanism): + """Calibrated synthesizer: one-way + compress + inner + decompress.""" + + config: DiscreteConfig + inner: api.CalibratedMechanism + zcdp_rho: float + one_way_rho: float | None = None + max_records_per_user: int = 1 + + @property + def dp_event(self) -> dp_accounting.DpEvent: + """Composes one-way measurement event with the inner mechanism's event.""" + events = [] + if self.one_way_rho is not None: + events.append( + dp_accounting.GaussianDpEvent( + noise_multiplier=accounting.zcdp_gaussian_sigma(self.one_way_rho) + ) + ) + if not isinstance(self.inner.dp_event, dp_accounting.NoOpDpEvent): + events.append(self.inner.dp_event) + return dp_accounting.ComposedDpEvent(events) + + def _one_way_cliques(self, data): + """Returns the one-way cliques to measure.""" + cliques = [(a,) for a in data.domain] + if hasattr(data, 'cliques'): + supported = common.downward_closure(data.cliques) + cliques = [cl for cl in cliques if cl in supported] + return cliques + + def _measure_one_way(self, rng, data, *, initial_measurements=None): + """Measures one-way marginals or returns pre-measured ones.""" + if initial_measurements is not None: + return list(initial_measurements) + if self.one_way_rho is None: + return [] + sigma = accounting.zcdp_gaussian_sigma(self.one_way_rho) + cliques = self._one_way_cliques(data) + return common.measure_marginals_with_noise( + rng, + data, + cliques, + sigma, + max_records_per_user=self.max_records_per_user, + ) + + def _compress(self, data, measurements): + """Compresses the domain by merging rare values.""" + mappings = common.compression_mappings( + measurements, + self.config.compress_columns, + getattr(self.config.mechanism, 'constraints', ()), + ) + if mappings and hasattr(data, 'compress'): + data = data.compress(mappings) + measurements = [m.compress(mappings, data.domain) for m in measurements] + return data, measurements, mappings + + def __call__( + self, + rng: np.random.Generator, + data: mbi.Dataset | mbi.CliqueVector, + *, + initial_measurements: Sequence[mbi.LinearMeasurement] | None = None, + ) -> common.DiscreteMechanismResult: + """Runs the one-way + compress + inner mechanism + decompress pipeline.""" + measurements = self._measure_one_way( + rng, data, initial_measurements=initial_measurements + ) + data, measurements, mappings = self._compress(data, measurements) + result = self.inner( + rng, + data, + initial_measurements=measurements, + ) + if mappings: + result = dataclasses.replace( + result, + synthetic_data=result.synthetic_data.decompress(mappings), + mappings=mappings, + ) + return result diff --git a/dpsynth/discrete_mechanisms/independent.py b/dpsynth/discrete_mechanisms/independent.py index 9d2e803c..631560f0 100644 --- a/dpsynth/discrete_mechanisms/independent.py +++ b/dpsynth/discrete_mechanisms/independent.py @@ -12,40 +12,95 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""This mechanisms measures all 1-way marginals via the Gaussian mechanism.""" +"""This mechanism independently estimates data from initial measurements.""" +from collections.abc import Sequence import dataclasses - +from absl import logging import dp_accounting -from dpsynth.discrete_mechanisms import accounting -from dpsynth.discrete_mechanisms import base +from dpsynth import api +from dpsynth.discrete_mechanisms import common import mbi +import numpy as np @dataclasses.dataclass(frozen=True, kw_only=True) -class IndependentConfig(base.DiscreteMechanismConfig): - """Measures only one-way marginals, allocating the entire budget to them.""" +class IndependentConfig(api.MechanismConfig): + """Independent config that doesn't select or measure any marginals.""" + + constraints: Sequence[mbi.Constraint] = () - one_way_budget_fraction: float = 1.0 + def configure(self, *, zcdp_rho, delta=0.0, max_records_per_user=1): + return Independent(config=self) def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - """Returns the one-way marginals this mechanism will measure.""" + """Returns the one-way marginals this mechanism expects to process.""" return [(a,) for a in domain.attributes] - def _create_mechanism(self, **kwargs) -> 'Independent': - return Independent(**kwargs) - @dataclasses.dataclass(frozen=True) -class Independent(base.DiscreteMechanism): +class Independent(api.CalibratedMechanism): + """Calibrated independent mechanism instance.""" + config: IndependentConfig @property def dp_event(self) -> dp_accounting.DpEvent: - """Returns the DP event for the independent mechanism.""" - return dp_accounting.GaussianDpEvent( - noise_multiplier=accounting.zcdp_gaussian_sigma(self.one_way_rho) # pyrefly: ignore[bad-argument-type] - ) + """Returns a zero-cost DP event (no new measurements).""" + return dp_accounting.NoOpDpEvent() - def _select(self, rng, data, measurements, phase_times): - return [] + def __call__( + self, + rng: np.random.Generator, + data: mbi.Dataset | mbi.CliqueVector, + *, + initial_measurements: Sequence[mbi.LinearMeasurement] | None = None, + ) -> common.DiscreteMechanismResult: + """Estimates and generates from initial measurements.""" + common.validate_initial_measurements(initial_measurements) + measurements = list(initial_measurements) if initial_measurements else [] + phase_times = {} + + all_cliques = [m.clique for m in measurements] + if all_cliques: + logging.info( + '[%s]:\n%s', + type(self).__name__, + mbi.summarize(data.domain, all_cliques), + ) + else: + logging.info( + '[%s]: No minimum measurement baseline established', + type(self).__name__, + ) + + # Kick off async AOT compilation of the estimator + estimator = mbi.estimation.MirrorDescent(None) + futures = None + try: + futures = estimator.precompile(data.domain, measurements) + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning('Precompile failed (non-fatal): %s', e) + + with common.timed(phase_times, 'estimation'): + if futures is not None: + try: + futures.result() + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning('Precompile wait failed (non-fatal): %s', e) + model = estimator.estimate( + data.domain, + measurements, + iters=5000, + callback_fn=mbi.callbacks.default(measurements, data.domain), + constraints=self.config.constraints, + ) + assert isinstance(model, mbi.MarkovRandomField) + + synthetic_data = model.synthetic_data() + return common.DiscreteMechanismResult( + synthetic_data=synthetic_data, + measurements=measurements, + model=model, + diagnostics=common.clique_stats(model), + ) diff --git a/dpsynth/discrete_mechanisms/mst.py b/dpsynth/discrete_mechanisms/mst.py index b56ad233..094b4b36 100644 --- a/dpsynth/discrete_mechanisms/mst.py +++ b/dpsynth/discrete_mechanisms/mst.py @@ -16,14 +16,15 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Sequence import dataclasses import itertools import typing from absl import logging import dp_accounting -from dpsynth.discrete_mechanisms import base +from dpsynth import api +from dpsynth.discrete_mechanisms import accounting from dpsynth.discrete_mechanisms import common import mbi import networkx as nx @@ -87,7 +88,8 @@ def dp_maximum_spanning_tree( candidates = list(weights.keys()) r = len(list(nx.connected_components(tree))) if exponential_mechanism_epsilon is None: - exponential_mechanism_epsilon = np.sqrt(8 * zcdp_rho / max(r - 1, 1)) # pyrefly: ignore[unsupported-operation] + assert zcdp_rho is not None + exponential_mechanism_epsilon = np.sqrt(8 * zcdp_rho / max(r - 1, 1)) for _ in range(r - 1): candidates = [e for e in candidates if not ds.connected(*e)] wgts = np.array([weights[e] for e in candidates]) @@ -103,7 +105,7 @@ def dp_maximum_spanning_tree( def _select_two_way_marginal_queries( rng: np.random.Generator, - data: mbi.Projectable, + data: mbi.Dataset | mbi.CliqueVector, zcdp_rho: float, one_way_measurements: list[mbi.LinearMeasurement], initial_marginal_queries: Sequence[tuple[str, ...]] = (), @@ -143,7 +145,7 @@ def _select_two_way_marginal_queries( ] logging.info('[MST]: Computing Quality Scores') weights = common.compute_independence_errors( - data, independent_model, candidates + data, independent_model, candidates # pyrefly: ignore[bad-argument-type] ) return dp_maximum_spanning_tree( # pyrefly: ignore[bad-return] @@ -156,7 +158,7 @@ def _select_two_way_marginal_queries( @dataclasses.dataclass(frozen=True) -class MSTConfig(base.DiscreteMechanismConfig): +class MSTConfig(api.MechanismConfig): """Configuration for the maximum spanning tree mechanism. Details are described in the paper: @@ -167,10 +169,14 @@ class MSTConfig(base.DiscreteMechanismConfig): select_budget_fraction: The fraction of the remaining budget (after one-way measurements) to use for selecting two-way marginal queries. maximum_marginal_size: The maximum size of a marginal query. - _select_rho: zCDP budget for the exponential mechanism (set by configure). """ - select_budget_fraction: float = 1 / 3 + constraints: Sequence[mbi.Constraint] = () + + marginal_oracle: mbi.MarginalOracle | None = None + pgm_iters: int = 5000 + + select_budget_fraction: float = 1 / 2 maximum_marginal_size: int = 10_000_000 def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: @@ -181,24 +187,22 @@ def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: self.maximum_marginal_size, ) - def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]: - """Splits the remaining budget between selection and measurement.""" - select_rho = remaining_rho * self.select_budget_fraction - return { - '_select_rho': select_rho, - 'measurement_rho': remaining_rho - select_rho, - } - - def _create_mechanism(self, **kwargs) -> 'MST': - return MST(**kwargs) + def configure(self, *, zcdp_rho, delta=0, max_records_per_user=1): + api.validate_max_records_per_user(max_records_per_user) + return MST( + config=self, + zcdp_rho=zcdp_rho, + max_records_per_user=max_records_per_user, + ) @dataclasses.dataclass(frozen=True, kw_only=True) -class MST(base.DiscreteMechanism): +class MST(api.CalibratedMechanism): """Calibrated MST instance.""" config: MSTConfig - _select_rho: float = -1.0 + zcdp_rho: float + max_records_per_user: int = 1 @property def dp_event(self) -> dp_accounting.DpEvent: @@ -211,8 +215,75 @@ def _select(self, rng, data, measurements, phase_times): return _select_two_way_marginal_queries( rng, data, - self._select_rho, # pyrefly: ignore[bad-argument-type] + self.zcdp_rho * self.config.select_budget_fraction, measurements, maximum_marginal_size=self.config.maximum_marginal_size, max_records_per_user=self.max_records_per_user, ) + + def __call__( + self, + rng: np.random.Generator, + data: mbi.Dataset | mbi.CliqueVector, + *, + initial_measurements: Sequence[mbi.LinearMeasurement] | None = None, + ) -> common.DiscreteMechanismResult: + """Selects, measures, estimates, and generates in the compressed domain.""" + common.validate_initial_measurements(initial_measurements) + measurements = list(initial_measurements) if initial_measurements else [] + phase_times = {} + selected = self._select(rng, data, measurements, phase_times) + all_cliques = [m.clique for m in measurements] + list(selected) + if all_cliques: + logging.info( + '[%s]:\n%s', + type(self).__name__, + mbi.summarize(data.domain, all_cliques), + ) + else: + logging.info('[%s]: No cliques selected', type(self).__name__) + + # Kick off async AOT compilation of the estimator while we measure. + estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle) + futures = None + try: + futures = estimator.precompile( + data.domain, measurements, extra_cliques=list(selected) + ) + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning('Precompile failed (non-fatal): %s', e) + + if selected: + with common.timed(phase_times, 'measurement'): + select_rho = self.zcdp_rho * self.config.select_budget_fraction + sigma = accounting.zcdp_gaussian_sigma(self.zcdp_rho - select_rho) + measurements = measurements + common.measure_marginals_with_noise( + rng, + data, # pyrefly: ignore[bad-argument-type] + selected, + sigma, + max_records_per_user=self.max_records_per_user, + ) + + with common.timed(phase_times, 'estimation'): + if futures is not None: + try: + futures.result() + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning('Precompile wait failed (non-fatal): %s', e) + model = estimator.estimate( + data.domain, + measurements, + iters=self.config.pgm_iters, + callback_fn=mbi.callbacks.default(measurements, data.domain), + constraints=self.config.constraints, + ) + assert isinstance(model, mbi.MarkovRandomField) + + synthetic_data = model.synthetic_data() + return common.DiscreteMechanismResult( + synthetic_data=synthetic_data, + measurements=measurements, + model=model, + diagnostics=common.clique_stats(model), + ) diff --git a/dpsynth/discrete_mechanisms/swift.py b/dpsynth/discrete_mechanisms/swift.py index b020f08a..88ce64c5 100644 --- a/dpsynth/discrete_mechanisms/swift.py +++ b/dpsynth/discrete_mechanisms/swift.py @@ -35,8 +35,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 clique_tree from dpsynth.discrete_mechanisms import common from dpsynth.discrete_mechanisms import swift_utils @@ -46,7 +46,7 @@ @dataclasses.dataclass(frozen=True) -class SWIFTConfig(base.DiscreteMechanismConfig): +class SWIFTConfig(api.MechanismConfig): """Configuration for the SWIFT mechanism. Attributes: @@ -60,17 +60,16 @@ class SWIFTConfig(base.DiscreteMechanismConfig): pgm_iters: Number of mirror descent iterations for PGM estimation. select_budget_frac: Fraction of the total budget used for selecting which marginals to measure. - one_way_budget_fraction: Fraction of zCDP budget for one-way marginals. """ + constraints: Sequence[mbi.Constraint] = () + workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None max_clique_size: float = 1e7 max_marginal_size: float = 1e6 pgm_iters: int = 10_000 + marginal_oracle: mbi.MarginalOracle | None = None select_budget_frac: float = 0.1 - one_way_budget_fraction: float = 0.1 - - # Internal state set by configure. def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: """Returns the workload cliques filtered by max_marginal_size.""" @@ -78,39 +77,43 @@ def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: domain, self.workload, self.max_marginal_size ) - def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]: - """Splits the remaining budget between selection and measurement.""" - select_rho = remaining_rho * self.select_budget_frac - return { - '_select_rho': select_rho, - 'measurement_rho': remaining_rho - select_rho, - } - - def _create_mechanism(self, **kwargs) -> 'SWIFT': - return SWIFT(**kwargs) + def configure(self, *, zcdp_rho, delta=0, max_records_per_user=1): + api.validate_max_records_per_user(max_records_per_user) + return SWIFT( + config=self, + zcdp_rho=zcdp_rho, + max_records_per_user=max_records_per_user, + ) @dataclasses.dataclass(frozen=True, kw_only=True) -class SWIFT(base.DiscreteMechanism): +class SWIFT(api.CalibratedMechanism): """Calibrated SWIFT instance.""" config: SWIFTConfig - _select_rho: float + zcdp_rho: float + max_records_per_user: int = 1 @property def dp_event(self) -> dp_accounting.DpEvent: """Returns the DP event for the SWIFT mechanism.""" - # SWIFT's budget is split between one-way, selection, and measurement. - # All three are Gaussian mechanism applications. - return dp_accounting.ZCDpEvent(self.zcdp_rho) # pyrefly: ignore[bad-argument-type] + 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 = {} - def _run(self, rng, data, measurements, constraints, phase_times): - """Runs SWIFT's select-measure-estimate pipeline in a single pass.""" - assert self._select_rho is not None - assert self.measurement_rho is not None + _select_rho = self.zcdp_rho * self.config.select_budget_frac # Budgets in GDP units, derived from the zCDP allocation set by configure. - gdp_budget = accounting.zcdp_to_gdp(self._select_rho + self.measurement_rho) + gdp_budget = accounting.zcdp_to_gdp(self.zcdp_rho) ######################################################################### # Compile workload into candidate measurements, and precompute answers. # @@ -133,7 +136,7 @@ def _run(self, rng, data, measurements, constraints, phase_times): domain, measurements, iters=self.config.pgm_iters, - constraints=constraints, + constraints=self.config.constraints, ) model = typing.cast(mbi.MarkovRandomField, model) @@ -141,7 +144,7 @@ def _run(self, rng, data, measurements, constraints, phase_times): # Select subset of candidates to measure. # ########################################### with common.timed(phase_times, 'selection'): - l1_error_budget = accounting.zcdp_to_gdp(self._select_rho) + l1_error_budget = accounting.zcdp_to_gdp(_select_rho) budget_remaining = gdp_budget - l1_error_budget with common.timed(phase_times, 'compute_initial_errors'): @@ -172,7 +175,7 @@ def _run(self, rng, data, measurements, constraints, phase_times): closed_oracle = functools.partial( mbi.marginal_oracles.message_passing_stable, jtree=jtree ) - estimator = mbi.estimation.MirrorDescent(marginal_oracle=closed_oracle) # pyrefly: ignore[bad-argument-type] + estimator = mbi.estimation.MirrorDescent(marginal_oracle=closed_oracle) rows = int(mbi.estimation.minimum_variance_unbiased_total(measurements)) pgm_future, synth_future = None, None @@ -192,7 +195,7 @@ def _run(self, rng, data, measurements, constraints, phase_times): logging.info('[SWIFT] Starting measurements.') new_measurements, _ = _measure_selected_marginals( rng, - answers, # pyrefly: ignore[bad-argument-type] + answers, selected, budget_remaining, max_records_per_user=self.max_records_per_user, @@ -218,7 +221,7 @@ def _run(self, rng, data, measurements, constraints, phase_times): measurements, iters=self.config.pgm_iters, callback_fn=callback_fn, - constraints=constraints, + constraints=self.config.constraints, ) assert isinstance(final_model, mbi.MarkovRandomField) logging.info('[SWIFT] Estimated final model.') @@ -233,7 +236,12 @@ def _run(self, rng, data, measurements, constraints, phase_times): syn = mbi.extensions.synthetic_data(final_model, rows) logging.info('[SWIFT] Generated %d synthetic records.', rows) - return final_model, syn, measurements + return common.DiscreteMechanismResult( + synthetic_data=syn, + measurements=measurements, + model=final_model, + diagnostics=common.clique_stats(final_model), + ) def _is_supported(clique: mbi.Clique, tree: nx.Graph) -> bool: @@ -354,7 +362,7 @@ def build_best_clique_tree( def _compute_initial_errors( rng: np.random.Generator, - data: mbi.Projectable, + data: mbi.Dataset | mbi.CliqueVector, model: mbi.MarkovRandomField, cliques: Sequence[mbi.Clique], gdp_budget: float, @@ -365,7 +373,7 @@ def _compute_initial_errors( sigma_per_clique = max_records_per_user * accounting.gdp_gaussian_sigma( budget_per_clique ) - errors = common.compute_independence_errors(data, model, cliques) + errors = common.compute_independence_errors(data, model, cliques) # pyrefly: ignore[bad-argument-type] for cl in errors: errors[cl] += rng.normal(loc=0.0, scale=sigma_per_clique) return errors @@ -422,7 +430,7 @@ def select_queries( def _measure_selected_marginals( rng: np.random.Generator, - data: mbi.Projectable, + data: mbi.Dataset | mbi.CliqueVector, selected: dict[mbi.Clique, float], budget_remaining: float, max_records_per_user: int = 1, diff --git a/dpsynth/discrete_mechanisms/swift_utils.py b/dpsynth/discrete_mechanisms/swift_utils.py index d66899d2..7c339f3e 100644 --- a/dpsynth/discrete_mechanisms/swift_utils.py +++ b/dpsynth/discrete_mechanisms/swift_utils.py @@ -13,6 +13,7 @@ # limitations under the License. """Module implementing the SWIFT budget allocation heuristic.""" + # NOTE: This module is tested in swift_test.py. from collections.abc import Sequence diff --git a/dpsynth/text/dp_sft.py b/dpsynth/text/dp_sft.py index 8583bc60..b5090b1e 100644 --- a/dpsynth/text/dp_sft.py +++ b/dpsynth/text/dp_sft.py @@ -52,7 +52,7 @@ import optax -@dataclasses.dataclass +@dataclasses.dataclass(frozen=True) class FineTuneResult: """Result of running ``DPFineTuner``. @@ -68,7 +68,7 @@ class FineTuneResult: params: training.Params -@dataclasses.dataclass +@dataclasses.dataclass(frozen=True) class DPFineTuner(api.DPMechanism): """Differentially private fine-tuning of Gemma models via DP-SGD. diff --git a/tests/adapters/beam_test.py b/tests/adapters/beam_test.py index 3b75385d..5e9dbed6 100644 --- a/tests/adapters/beam_test.py +++ b/tests/adapters/beam_test.py @@ -392,7 +392,7 @@ def test_end_to_end_mixed_types(self): ('mst', discrete_mechanisms.MSTConfig(pgm_iters=250)), ( 'independent', - discrete_mechanisms.IndependentConfig(pgm_iters=250), + discrete_mechanisms.IndependentConfig(), ), ( 'direct', diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py index 00a717ed..88fdd3e3 100644 --- a/tests/data_generation_v3_test.py +++ b/tests/data_generation_v3_test.py @@ -305,7 +305,7 @@ def test_nan_numerical_column(self, sentinel, clip_to_range, dtype): def test_discrete_workload_regression_with_aim(self): workload = [('a',), ('b',), ('c',), ('a', 'b'), ('a', 'c'), ('b', 'c')] config = aim.AIMConfig(workload=workload, max_rounds=4, pgm_iters=500) - baseline_config = IndependentConfig(pgm_iters=500) + baseline_config = IndependentConfig() mechanism_error, baseline_error = ( _discrete_workload_mechanism_baseline_errors( config, baseline_config, workload @@ -318,7 +318,7 @@ def test_discrete_workload_regression_with_aim_gdp(self): config = aim_gdp.AIMGDPConfig( workload=workload, max_rounds=4, pgm_iters=500 ) - baseline_config = IndependentConfig(pgm_iters=500) + baseline_config = IndependentConfig() mechanism_error, baseline_error = ( _discrete_workload_mechanism_baseline_errors( config, baseline_config, workload @@ -329,7 +329,7 @@ def test_discrete_workload_regression_with_aim_gdp(self): def test_mixed_workload_regression_with_aim(self): workload = [('a',), ('b',), ('c',), ('a', 'b'), ('a', 'c'), ('b', 'c')] config = aim.AIMConfig(workload=workload, max_rounds=4, pgm_iters=500) - baseline_config = IndependentConfig(pgm_iters=500) + baseline_config = IndependentConfig() mechanism_error, baseline_error = _mixed_workload_mechanism_baseline_errors( config, baseline_config, workload ) @@ -340,7 +340,7 @@ def test_mixed_workload_regression_with_aim_gdp(self): config = aim_gdp.AIMGDPConfig( workload=workload, max_rounds=4, pgm_iters=500 ) - baseline_config = IndependentConfig(pgm_iters=500) + baseline_config = IndependentConfig() mechanism_error, baseline_error = _mixed_workload_mechanism_baseline_errors( config, baseline_config, workload ) diff --git a/tests/discrete_mechanisms/aim_test.py b/tests/discrete_mechanisms/aim_test.py index a6d50dae..7370d853 100644 --- a/tests/discrete_mechanisms/aim_test.py +++ b/tests/discrete_mechanisms/aim_test.py @@ -93,7 +93,7 @@ def test_fits_one_way_marginals_with_aim_gdp(self): def test_correlated_workload_regression_with_aim(self): workload = [("a",), ("b",), ("c",), ("a", "b"), ("a", "c"), ("b", "c")] config = aim.AIMConfig(workload=workload, max_rounds=4, pgm_iters=500) - baseline_config = independent.IndependentConfig(pgm_iters=500) + baseline_config = independent.IndependentConfig() mechanism_error, baseline_error = ( _correlated_workload_mechanism_baseline_errors( config, baseline_config, workload @@ -106,7 +106,7 @@ def test_correlated_workload_regression_with_aim_gdp(self): config = aim_gdp.AIMGDPConfig( workload=workload, max_rounds=4, pgm_iters=500 ) - baseline_config = independent.IndependentConfig(pgm_iters=500) + baseline_config = independent.IndependentConfig() mechanism_error, baseline_error = ( _correlated_workload_mechanism_baseline_errors( config, baseline_config, workload diff --git a/tests/discrete_mechanisms/base_test.py b/tests/discrete_mechanisms/base_test.py deleted file mode 100644 index 743d20e6..00000000 --- a/tests/discrete_mechanisms/base_test.py +++ /dev/null @@ -1,119 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for the shared ``DiscreteMechanism`` base-class machinery.""" - -import dataclasses -from unittest import mock - -from absl.testing import absltest -import dp_accounting -from dpsynth.discrete_mechanisms import base -from dpsynth.discrete_mechanisms import common -import mbi -import mbi.estimation -import numpy as np - - -def _dataset(n: int = 200) -> mbi.Dataset: - return mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=n) - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class _NoOpMechanism(base.DiscreteMechanism): - @property - def dp_event(self) -> dp_accounting.DpEvent: - return dp_accounting.GaussianDpEvent(noise_multiplier=1.0) - - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - return [(a,) for a in domain.attributes] - - def _select(self, rng, data, measurements, phase_times): - return [] - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class _NoOpMechanismConfig(base.DiscreteMechanismConfig): - - def _create_mechanism(self, **kwargs): - return _NoOpMechanism(**kwargs) - - def supporting_cliques(self, domain): - return [] - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class _NoSelectMechanism(base.DiscreteMechanism): - @property - def dp_event(self) -> dp_accounting.DpEvent: - return dp_accounting.GaussianDpEvent(noise_multiplier=1.0) - - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - return [(a,) for a in domain.attributes] - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class _NoSelectMechanismConfig(base.DiscreteMechanismConfig): - - def _create_mechanism(self, **kwargs): - return _NoSelectMechanism(**kwargs) - - def supporting_cliques(self, domain): - return [] - - -class ConfigureTest(absltest.TestCase): - - def test_default_fraction_splits_one_way_budget(self): - configured = _NoOpMechanismConfig(one_way_budget_fraction=0.25).configure( - zcdp_rho=100.0 - ) - self.assertEqual(configured.one_way_rho, 25.0) - - def test_zero_one_way_budget_fraction_skips_one_way(self): - configured = _NoOpMechanismConfig(one_way_budget_fraction=0.0).configure( - zcdp_rho=100.0 - ) - self.assertIsNone(configured.one_way_rho) - - def test_default_allocate_budget_leaves_measurement_rho_unset(self): - configured = _NoOpMechanismConfig().configure(zcdp_rho=100.0) - self.assertIsNone(configured.measurement_rho) - - def test_initial_measurements_skip_one_way(self): - configured = _NoOpMechanismConfig(one_way_budget_fraction=0.0).configure( - zcdp_rho=100.0, - ) - self.assertIsNone(configured.one_way_rho) - - -class RunMachineryTest(absltest.TestCase): - - def test_precompile_failure_is_non_fatal(self): - mechanism = _NoOpMechanismConfig(pgm_iters=100).configure(zcdp_rho=1000.0) - with mock.patch.object( - mbi.estimation.MirrorDescent, - 'precompile', - side_effect=RuntimeError('simulated precompile failure'), - ) as mocked_precompile: - result = mechanism(np.random.default_rng(0), _dataset()) - mocked_precompile.assert_called_once() - self.assertIsInstance(result, common.DiscreteMechanismResult) - - def test_missing_select_raises_not_implemented(self): - mechanism = _NoSelectMechanismConfig(pgm_iters=100).configure( - zcdp_rho=1000.0 - ) - with self.assertRaises(NotImplementedError): - mechanism(np.random.default_rng(0), _dataset()) diff --git a/tests/discrete_mechanisms/discrete_mechanisms_test.py b/tests/discrete_mechanisms/discrete_mechanisms_test.py index 6b2fac2e..9a98c094 100644 --- a/tests/discrete_mechanisms/discrete_mechanisms_test.py +++ b/tests/discrete_mechanisms/discrete_mechanisms_test.py @@ -12,9 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Property tests shared across all discrete mechanisms.""" +"""Property tests shared across all discrete mechanisms. -import dataclasses +Note: This mechanism is not intended to be called directly. It should typically +be used within `DiscreteMechanism` or `TabularSynthesizer`. Users who call it +directly will miss out on features like 1-way measurement selection and domain +compression. +""" from absl.testing import absltest from absl.testing import parameterized @@ -22,6 +26,7 @@ from dpsynth.discrete_mechanisms import aim_gdp from dpsynth.discrete_mechanisms import common from dpsynth.discrete_mechanisms import direct +from dpsynth.discrete_mechanisms import discrete_synthesizer from dpsynth.discrete_mechanisms import independent from dpsynth.discrete_mechanisms import mst from dpsynth.discrete_mechanisms import swift @@ -38,7 +43,7 @@ ), 'MST': mst.MSTConfig(pgm_iters=500), 'SWIFT': swift.SWIFTConfig(workload=_WORKLOAD, pgm_iters=500), - 'Independent': independent.IndependentConfig(pgm_iters=500), + 'Independent': independent.IndependentConfig(), 'Direct': direct.DirectConfig( prespecified_marginal_queries=_WORKLOAD, pgm_iters=500 ), @@ -54,15 +59,7 @@ def _make_skewed_dataset(rng): class SupportingCliquesSufficiencyTest(parameterized.TestCase): - """Checks that supporting_cliques are sufficient for each mechanism. - - For each mechanism, we: - 1. Compute supporting_cliques(domain). - 2. Build a CliqueVector from the true data projected onto those cliques. - 3. Run the mechanism using the CliqueVector as input data. - 4. Assert it completes without error — the CliqueVector supports every - projection the mechanism needs. - """ + """Checks that supporting_cliques are sufficient for each mechanism.""" @parameterized.named_parameters(*_MECHANISMS.items()) def test_mechanism_runs_on_precomputed_marginals(self, mechanism): @@ -81,22 +78,28 @@ def test_mechanism_runs_on_precomputed_marginals(self, mechanism): class CompressionPropertyTest(parameterized.TestCase): - """Tests that compression restores the original domain across mechanisms.""" + """Tests that compression restores the original domain via synthesizer.""" @parameterized.named_parameters(*_MECHANISMS.items()) def test_compression_restores_domain(self, config): - config = dataclasses.replace(config, compress_columns=True) + synth_config = discrete_synthesizer.DiscreteConfig( + mechanism=config, + compress_columns=True, + ) rng = np.random.default_rng(0) data = _make_skewed_dataset(rng) original_domain = data.domain - result = config.configure(zcdp_rho=_ZCDP_RHO)(rng, data) + result = synth_config.configure(zcdp_rho=_ZCDP_RHO)(rng, data) self.assertEqual(result.synthetic_data.domain, original_domain) @parameterized.named_parameters(*_MECHANISMS.items()) def test_compression_with_initial_measurements(self, config): - config = dataclasses.replace(config, compress_columns=True) + synth_config = discrete_synthesizer.DiscreteConfig( + mechanism=config, + compress_columns=True, + ) rng = np.random.default_rng(0) data = _make_skewed_dataset(rng) original_domain = data.domain @@ -104,7 +107,7 @@ def test_compression_with_initial_measurements(self, config): rng, data, [('a',), ('b',)], gdp_sigma=1.0 ) - mechanism = config.configure(zcdp_rho=_ZCDP_RHO) + mechanism = synth_config.configure(zcdp_rho=_ZCDP_RHO) result = mechanism(rng, data, initial_measurements=initial_measurements) self.assertEqual(result.synthetic_data.domain, original_domain) @@ -125,6 +128,8 @@ def test_deprecated_zcdp_calibration(self, mechanism): def test_zero_epsilon_calibration(self, mechanism): rng = np.random.default_rng(0) data = _make_skewed_dataset(rng) + if isinstance(mechanism, independent.IndependentConfig): + return result = mechanism.calibrate(epsilon=0.0, delta=0.01)(rng, data) self.assertIsInstance(result, common.DiscreteMechanismResult) @@ -142,19 +147,15 @@ class MaxRecordsPerUserTest(parameterized.TestCase): @parameterized.named_parameters(*_MECHANISMS.items()) def test_dp_event_invariant_to_max_records_per_user(self, mechanism): - # Scaling max_records_per_user must not change the accounting: only the - # actual noise magnitude scales, while the reported dp_event is identical. base = mechanism.configure(zcdp_rho=_ZCDP_RHO) scaled = mechanism.configure(zcdp_rho=_ZCDP_RHO, max_records_per_user=4) self.assertEqual(repr(scaled.dp_event), repr(base.dp_event)) @parameterized.named_parameters( - ('Independent', _MECHANISMS['Independent']), + ('MST', _MECHANISMS['MST']), ('Direct', _MECHANISMS['Direct']), ) def test_measurement_stddev_scales_with_k(self, mechanism): - # Independent and Direct select their measured cliques deterministically, - # so the recorded stddevs line up one-to-one and must scale linearly in k. k = 4 data = _make_skewed_dataset(np.random.default_rng(0)) base = mechanism.configure(zcdp_rho=_ZCDP_RHO)( diff --git a/tests/discrete_mechanisms/discrete_synthesizer_test.py b/tests/discrete_mechanisms/discrete_synthesizer_test.py new file mode 100644 index 00000000..7fed53da --- /dev/null +++ b/tests/discrete_mechanisms/discrete_synthesizer_test.py @@ -0,0 +1,130 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for DiscreteMechanism wrapper: 1-way bootstrapping and compression.""" + +from absl.testing import absltest +from dpsynth.discrete_mechanisms import common +from dpsynth.discrete_mechanisms import discrete_synthesizer +from dpsynth.discrete_mechanisms import mst +import mbi +import numpy as np + +DiscreteConfig = discrete_synthesizer.DiscreteConfig +DiscreteMechanism = discrete_synthesizer.DiscreteMechanism +MSTConfig = mst.MSTConfig + + +class DiscreteConfigTest(absltest.TestCase): + + def test_configure_splits_budget(self): + config = DiscreteConfig( + mechanism=MSTConfig(pgm_iters=500), + one_way_budget_fraction=0.25, + ) + synth = config.configure(zcdp_rho=100.0) + self.assertIsInstance(synth, DiscreteMechanism) + self.assertAlmostEqual(synth.one_way_rho, 25.0) + self.assertAlmostEqual(synth.zcdp_rho, 100.0) + + def test_configure_zero_one_way_fraction(self): + config = DiscreteConfig( + mechanism=MSTConfig(pgm_iters=500), + one_way_budget_fraction=0.0, + ) + synth = config.configure(zcdp_rho=100.0) + self.assertIsNone(synth.one_way_rho) + + def test_supporting_cliques_delegates(self): + inner = MSTConfig(pgm_iters=500) + config = DiscreteConfig(mechanism=inner) + domain = mbi.Domain(['a', 'b', 'c'], [3, 4, 5]) + self.assertEqual( + config.supporting_cliques(domain), + inner.supporting_cliques(domain), + ) + + +class DiscreteMechanismTest(absltest.TestCase): + + def test_full_pipeline(self): + domain = mbi.Domain(['a', 'b', 'c'], [3, 4, 5]) + data = mbi.Dataset.synthetic(domain, N=500) + rng = np.random.default_rng(42) + + config = DiscreteConfig( + mechanism=MSTConfig(pgm_iters=500), + ) + synth = config.configure(zcdp_rho=10000) + result = synth(rng, data) + + self.assertIsInstance(result, common.DiscreteMechanismResult) + self.assertEqual(result.synthetic_data.domain, domain) + + def test_with_initial_measurements_skips_one_way(self): + domain = mbi.Domain(['a', 'b', 'c'], [3, 4, 5]) + data = mbi.Dataset.synthetic(domain, N=500) + rng = np.random.default_rng(42) + + measurements = common.measure_marginals_with_noise( + rng, data, [('a',), ('b',), ('c',)], gdp_sigma=1.0 + ) + config = DiscreteConfig( + mechanism=MSTConfig(pgm_iters=500), + ) + synth = config.configure(zcdp_rho=10000) + result = synth(rng, data, initial_measurements=measurements) + + self.assertIsInstance(result, common.DiscreteMechanismResult) + + def test_compression_restores_domain(self): + domain = mbi.Domain(['a', 'b', 'c'], [10, 4, 5]) + rng = np.random.default_rng(0) + df = {col: rng.integers(0, domain[col], size=1000) for col in domain} + df['a'] = rng.choice(3, size=1000) + data = mbi.Dataset(df, domain) + + config = DiscreteConfig( + mechanism=MSTConfig(pgm_iters=500), + compress_columns=True, + ) + synth = config.configure(zcdp_rho=10000) + result = synth(rng, data) + + self.assertEqual(result.synthetic_data.domain, domain) + + def test_dp_event_composes_one_way_and_inner(self): + config = DiscreteConfig( + mechanism=MSTConfig(pgm_iters=500), + one_way_budget_fraction=0.25, + ) + synth = config.configure(zcdp_rho=100.0) + event = synth.dp_event + self.assertIsNotNone(event) + + def test_calibrate_works(self): + config = DiscreteConfig( + mechanism=MSTConfig(pgm_iters=500), + ) + domain = mbi.Domain(['a', 'b'], [3, 4]) + data = mbi.Dataset.synthetic(domain, N=200) + rng = np.random.default_rng(0) + + calibrated = config.calibrate(zcdp_rho=10000) + result = calibrated(rng, data) + self.assertIsInstance(result, common.DiscreteMechanismResult) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/discrete_mechanisms/independent_test.py b/tests/discrete_mechanisms/independent_test.py index 369cca2b..74dba1be 100644 --- a/tests/discrete_mechanisms/independent_test.py +++ b/tests/discrete_mechanisms/independent_test.py @@ -22,10 +22,30 @@ class IndependentTest(absltest.TestCase): def test_fits_one_way_marginals(self): + """Independent with externally-supplied 1-ways should recover marginals.""" data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000) - config = independent.IndependentConfig(pgm_iters=500) - result = config.configure(zcdp_rho=10000)(np.random.default_rng(0), data) + config = independent.IndependentConfig() + + # One-way marginals are now supplied externally by the synthesizer layer. + # Simulate that by constructing them here with near-zero noise. + initial_measurements = [] + for col in data.domain: + true_marginal = data.project([col]).datavector() + initial_measurements.append( + mbi.LinearMeasurement( + noisy_measurement=true_marginal, + clique=(col,), + stddev=1e-6, + query=mbi.DatavectorQuery(use_for_total_estimation=True), + ) + ) + + result = config.configure(zcdp_rho=10000)( + np.random.default_rng(0), + data, + initial_measurements=initial_measurements, + ) self.assertIsInstance(result, common.DiscreteMechanismResult) self.assertLen(result.measurements, len(data.domain)) @@ -41,7 +61,7 @@ def test_skips_duplicate_cliques_from_initial_measurements(self): marginal_a = data.project(('a',)).datavector() initial = [mbi.LinearMeasurement(marginal_a, ('a',), stddev=1.0)] - config = independent.IndependentConfig(pgm_iters=500) + config = independent.IndependentConfig() # This should not raise 'Cliques must be unique'. model = config.configure(zcdp_rho=100.0)( np.random.default_rng(0), data, initial_measurements=initial diff --git a/tests/discrete_mechanisms/mst_test.py b/tests/discrete_mechanisms/mst_test.py index 2788030c..4ddb1727 100644 --- a/tests/discrete_mechanisms/mst_test.py +++ b/tests/discrete_mechanisms/mst_test.py @@ -75,13 +75,33 @@ def test_dp_maximum_spanning_tree_infinite_eps(self): self.assertEqual(actual_mst_edges, expected_mst_edges) def test_fits_one_way_marginals(self): + """MST + externally-supplied 1-ways should recover all one-way marginals.""" data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000) - config = mst.MSTConfig(pgm_iters=500).configure(zcdp_rho=10000) + calibrated = mst.MSTConfig(pgm_iters=500).configure(zcdp_rho=10000) - result = config(np.random.default_rng(0), data) + # One-way marginals are now supplied externally by the synthesizer layer. + # Simulate that by constructing them here with near-zero noise. + initial_measurements = [] + for col in data.domain: + true_marginal = data.project([col]).datavector() + initial_measurements.append( + mbi.LinearMeasurement( + noisy_measurement=true_marginal, + clique=(col,), + stddev=1e-6, + query=mbi.DatavectorQuery(use_for_total_estimation=True), + ) + ) + + result = calibrated( + np.random.default_rng(0), + data, + initial_measurements=initial_measurements, + ) self.assertIsInstance(result, common.DiscreteMechanismResult) + # 3 externally-supplied one-ways + 2 MST-selected pairwise = 5 total. self.assertLen(result.measurements, 2 * len(data.domain) - 1) for col in data.domain: expected = data.project([col]).datavector() diff --git a/tests/discrete_mechanisms/swift_test.py b/tests/discrete_mechanisms/swift_test.py index 365c8fb4..b11afb47 100644 --- a/tests/discrete_mechanisms/swift_test.py +++ b/tests/discrete_mechanisms/swift_test.py @@ -146,5 +146,6 @@ def test_fits_one_way_marginals(self): actual = result.model.project([col]).datavector() np.testing.assert_allclose(actual, expected, atol=1) + if __name__ == '__main__': absltest.main()