diff --git a/dpsynth/adapters/beam.py b/dpsynth/adapters/beam.py index 2685570..ae1e9ff 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 70ce4ce..90ea540 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 5ebe289..cb3db12 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 4b5a29c..3984811 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 9796440..1910c1f 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 778ce5e..a06164b 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 e58bf58..0000000 --- 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 60d51de..0c16e4f 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 5a4d350..89a2f95 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 f512662..60e0a1b 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 0000000..d438194 --- /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 9d2e803..631560f 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 b56ad23..094b4b3 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 b020f08..88ce64c 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 d66899d..7c339f3 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 8583bc6..b5090b1 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 3b75385..5e9dbed 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 00a717e..88fdd3e 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 a6d50da..7370d85 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 743d20e..0000000 --- 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 6b2fac2..9a98c09 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 0000000..7fed53d --- /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 369cca2..74dba1b 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 2788030..4ddb172 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 365c8fb..b11afb4 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()