diff --git a/dpsynth/__init__.py b/dpsynth/__init__.py index 8877780..02a37e8 100644 --- a/dpsynth/__init__.py +++ b/dpsynth/__init__.py @@ -25,6 +25,7 @@ from dpsynth import relational from dpsynth import reporting from dpsynth import serialize +from dpsynth import workload from dpsynth._calibration import calibrate from dpsynth._calibration import with_group_size from dpsynth._checkpoint import checkpoint @@ -41,6 +42,7 @@ from dpsynth.reporting import PrivacyReport from dpsynth.serialize import from_yaml from dpsynth.serialize import to_yaml +from dpsynth.workload import Workload import mbi # Route MBI callback logs (which use print-style formatting) to @@ -70,6 +72,7 @@ 'TabularConfig', 'TabularMechanism', 'TabularSynthesizer', + 'Workload', 'api', 'calibrate', 'checkpoint', @@ -82,4 +85,5 @@ 'serialize', 'to_yaml', 'with_group_size', + 'workload', ] diff --git a/dpsynth/_calibration.py b/dpsynth/_calibration.py index 8ccb963..bff8995 100644 --- a/dpsynth/_calibration.py +++ b/dpsynth/_calibration.py @@ -72,6 +72,7 @@ def calibrate( *, epsilon: float, delta: float, + workload: Any = None, delta_split: float = 0.5, poisson_sampling_prob: float = 1.0, group_size: int = 1, @@ -89,6 +90,8 @@ def calibrate( domain: Optional domain specification, forwarded to ``config.configure()``. epsilon: Target epsilon for (epsilon, delta)-DP. delta: Target delta for (epsilon, delta)-DP. + workload: Optional workload specification, forwarded to + ``config.configure()``. delta_split: Fraction of ``delta`` passed to ``config.configure()`` for sub-mechanisms that consume approximate DP budget directly (e.g. open-set partition selection). Defaults to 0.5. @@ -123,12 +126,12 @@ def calibrate( if group_size < 1: raise ValueError(f'group_size < 1: {group_size}.') + configure_kwargs: dict[str, Any] = {'delta': delta * delta_split} + if workload is not None: + configure_kwargs['workload'] = workload + def make_event_fn(rho: float) -> dp_accounting.DpEvent: - base = config.configure( - domain, - budget=rho, - delta=delta * delta_split, - ).dp_event + base = config.configure(domain, budget=rho, **configure_kwargs).dp_event base = with_group_size(base, group_size) sampled = dp_accounting.PoissonSampledDpEvent(poisson_sampling_prob, base) return base if poisson_sampling_prob == 1.0 else sampled @@ -169,8 +172,4 @@ def make_event_fn(rho: float) -> dp_accounting.DpEvent: ) optimal_rho = max(rhos.values()) - return config.configure( - domain, - budget=optimal_rho, - delta=delta * delta_split, - ) + return config.configure(domain, budget=optimal_rho, **configure_kwargs) diff --git a/dpsynth/adapters/beam.py b/dpsynth/adapters/beam.py index add3bc3..6c2b390 100644 --- a/dpsynth/adapters/beam.py +++ b/dpsynth/adapters/beam.py @@ -481,8 +481,8 @@ def _run_two_pass( column_measurements, synth.schema ).mbi_domain - assert hasattr(synth.config.discrete_mechanism, 'supporting_cliques') - workload = synth.config.discrete_mechanism.supporting_cliques(mbi_domain) + assert hasattr(synth.base_mechanism, 'supporting_cliques') + workload = synth.base_mechanism.supporting_cliques(mbi_domain) # Pass 2: compute the marginal workload. with beam.Pipeline(**pipeline_kwargs) as p: @@ -558,20 +558,18 @@ class BeamTabularConfig(api.MechanismConfig): temp_location: str | None = None pipeline_options: beam.options.pipeline_options.PipelineOptions | None = None - def __post_init__(self): - if not hasattr(self.synthesizer.discrete_mechanism, 'supporting_cliques'): - raise ValueError( - 'self.synthesizer.discrete_mechanism must have a supporting_cliques' - ' method.' - ) - - def configure(self, schema=None, *, budget, delta=0) -> BeamTabularMechanism: + def configure( + self, schema=None, *, budget, delta=0, workload=None + ) -> BeamTabularMechanism: """Returns a copy whose synthesizer is configured with the given budget.""" synthesizer = self.synthesizer.configure( schema, budget=budget, delta=delta, + workload=workload, ) + if not hasattr(synthesizer.base_mechanism, 'supporting_cliques'): + raise ValueError('base_mechanism must define supporting_cliques.') return BeamTabularMechanism( synthesizer=synthesizer, temp_location=self.temp_location, diff --git a/dpsynth/api.py b/dpsynth/api.py index e9f0eaa..39698c2 100644 --- a/dpsynth/api.py +++ b/dpsynth/api.py @@ -140,7 +140,9 @@ def get_subclass(cls, name: str) -> type[MechanismConfig] | None: return cls._registry.get(name) @abc.abstractmethod - def configure(self, domain=None, *, budget, delta=0) -> CalibratedMechanism: + def configure( + self, domain=None, *, budget, delta=0, workload=None + ) -> CalibratedMechanism: """Returns a calibrated mechanism for the given dummy budget. Converts the budget into the mechanism's natural privacy parameter @@ -166,6 +168,9 @@ def configure(self, domain=None, *, budget, delta=0) -> CalibratedMechanism: delta: Approximate DP delta consumed by the mechanism itself (e.g., for thresholding). Defaults to 0 (pure zCDP). Mechanisms that need delta should raise if it is 0. + workload: Optional workload specification (e.g., a ``dpsynth.Workload``, + sequence of attribute tuples, or mapping from attribute tuples to + weights). Mechanisms that do not use a workload ignore this argument. Returns: A calibrated, runnable mechanism. @@ -178,6 +183,7 @@ def calibrate( *, epsilon: float, delta: float, + workload: Any = None, delta_split: float = 0.5, poisson_sampling_prob: float = 1.0, group_size: int = 1, @@ -197,6 +203,7 @@ def calibrate( domain, epsilon=epsilon, delta=delta, + workload=workload, delta_split=delta_split, poisson_sampling_prob=poisson_sampling_prob, group_size=group_size, diff --git a/dpsynth/data_generation_v3.py b/dpsynth/data_generation_v3.py index bf3fa2c..477294f 100644 --- a/dpsynth/data_generation_v3.py +++ b/dpsynth/data_generation_v3.py @@ -27,6 +27,7 @@ from dpsynth import constraints from dpsynth import discrete_mechanisms from dpsynth import domain +from dpsynth import workload as workload_lib from dpsynth.discrete_mechanisms import common as dm_common from dpsynth.local_mode import initialization from dpsynth.local_mode import vectorized_transformations as vtx @@ -323,9 +324,8 @@ def from_summary( ] logging.info('[DPSynth]: Compressed discrete domain:\n%s', data.domain) - cfg = self.config.discrete_mechanism - if hasattr(cfg, 'supporting_cliques'): - cliques = cfg.supporting_cliques(data.domain) + if hasattr(self.base_mechanism, 'supporting_cliques'): + cliques = self.base_mechanism.supporting_cliques(data.domain) data = _checkpoint.get_or_compute( 'precomputed_marginals', dm_common.precompute_marginals, @@ -443,6 +443,7 @@ def configure( *, budget: float, delta: float = 0.0, + workload: workload_lib.WorkloadInput | None = None, ) -> TabularMechanism: """Returns a calibrated mechanism configured with the given privacy budget. @@ -460,12 +461,17 @@ def configure( delta: Approximate DP delta allocated to partition selection for open-set columns (split evenly across open-set columns). Must be positive when open-set categorical attributes are present. + workload: Optional workload specification (e.g. a ``dpsynth.Workload``, + sequence of attribute tuples, or mapping from attribute tuples to + weights) validated against ``schema`` and forwarded to the discrete + mechanism. Returns: A calibrated TabularMechanism ready to be run on tabular data. Raises: - ValueError: If open-set attributes exist but delta is 0. + ValueError: If open-set attributes exist but delta is 0, or if the + workload contains attributes not present in ``schema``. """ if schema is not None and isinstance(schema, domain.Schema): pass @@ -483,6 +489,7 @@ def configure( ' construction time.' ) + workload = workload_lib.Workload.from_any(workload, schema) per_col_deltas = self._compute_per_col_deltas(schema, delta) inits = create_initializers( @@ -504,7 +511,7 @@ def configure( } calibrated_discrete = self.discrete_mechanism.configure( - budget=discrete_rho, + schema, budget=discrete_rho, workload=workload ) return TabularMechanism( diff --git a/dpsynth/discrete_mechanisms/README.md b/dpsynth/discrete_mechanisms/README.md index c9380bc..f95bf45 100644 --- a/dpsynth/discrete_mechanisms/README.md +++ b/dpsynth/discrete_mechanisms/README.md @@ -53,13 +53,13 @@ does not model relationships between columns. ## `direct.py` — Caller-Defined Workload -Implements `DirectConfig`, which measures the cliques supplied through -`prespecified_marginal_queries`. It performs no data-dependent selection and -does not create its own one-way measurements, so the full budget is available -for the specified workload. Initial measurements supplied by another layer are -included when fitting the final model. +Implements `DirectConfig`, which measures the cliques supplied through the +`workload` parameter of `configure()` / `dpsynth.calibrate()`. It performs no +data-dependent selection and does not create its own one-way measurements, so +the full budget is available for the specified workload. Initial measurements +supplied by another layer are included when fitting the final model. -**Public API:** `DirectConfig(prespecified_marginal_queries=...)` +**Public API:** `DirectConfig` ## `mst.py` — Private Pairwise Spanning-Tree Selection @@ -81,7 +81,7 @@ Implements `AIMConfig`, an adaptive workload-based mechanism. Instead of selecting cliques once, it repeatedly finds a marginal that the current model approximates poorly, measures it, and updates the model. -**Public API:** `AIMConfig(workload=...)` +**Public API:** `AIMConfig` **Internal behavior:** `_one_way_cliques()` limits initial measurements to the workload; `_allocate_budget()` reserves rho for the adaptive loop; `_run()` @@ -95,7 +95,7 @@ marginals while controlling clique-tree complexity. It uses a custom junction-tree-aware estimation and sampling path rather than the standard one-pass implementation in `base.py`. -**Public API:** `SWIFTConfig(workload=...)` +**Public API:** `SWIFTConfig` **Internal behavior:** `_allocate_budget()` splits rho between selection and measurement; `_run()` compiles the workload, selects supported cliques, builds a diff --git a/dpsynth/discrete_mechanisms/aim.py b/dpsynth/discrete_mechanisms/aim.py index c6931b3..73e741f 100644 --- a/dpsynth/discrete_mechanisms/aim.py +++ b/dpsynth/discrete_mechanisms/aim.py @@ -14,12 +14,13 @@ """Implementation of the Adaptive+Iterative Mechanism (AIM).""" -from collections.abc import Iterable, Mapping +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 import workload as workload_lib from dpsynth.discrete_mechanisms import accounting from dpsynth.discrete_mechanisms import common from dpsynth.local_mode import primitives @@ -104,8 +105,6 @@ class AIMConfig(api.MechanismConfig): max_model_size >= 80. Attributes: - workload: A collection of marginal queries (and weights) the synthetic data - should be tailored to. max_rounds: The maximum number of rounds to run the mechanism. max_model_size: The maximum size of the graphical model in megabytes. Controls the utility/runtime trade-off. @@ -115,7 +114,6 @@ class AIMConfig(api.MechanismConfig): selecting two-way marginal queries. """ - workload: Mapping[mbi.Clique, float] | Iterable[mbi.Clique] | None = None max_rounds: int | None = None max_model_size: int = 80 max_marginal_size: float = 1e6 @@ -124,16 +122,19 @@ class AIMConfig(api.MechanismConfig): pgm_iters: int = 1000 marginal_oracle: mbi.MarginalOracle | None = None - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - """Returns the workload cliques filtered by max_marginal_size.""" - return common.supporting_cliques( - domain, self.workload, self.max_marginal_size - ) - - def configure(self, _=None, *, budget, delta=0): + def configure( + self, + domain=None, + *, + budget: float, + delta: float = 0.0, + workload: workload_lib.WorkloadInput | None = None, + ) -> 'AIM': + del delta return AIM( config=self, zcdp_rho=budget, + workload=workload_lib.Workload.from_any(workload, domain), ) @@ -143,6 +144,14 @@ class AIM(api.CalibratedMechanism): config: AIMConfig zcdp_rho: float + workload: workload_lib.Workload | None = None + + def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: + """Returns the workload cliques filtered by max_marginal_size.""" + workload_dict = self.workload.as_dict() if self.workload else None + return common.supporting_cliques( + domain, workload_dict, self.config.max_marginal_size + ) @property def dp_event(self) -> dp_accounting.DpEvent: @@ -170,8 +179,9 @@ def __call__( ######################################################################### # Compile workload into candidate measurements. # ######################################################################### + workload_dict = self.workload.as_dict() if self.workload else None candidates = common.compiled_workload( - data.domain, self.config.workload, self.config.max_marginal_size + data.domain, workload_dict, self.config.max_marginal_size ) estimator = mbi.estimation.MirrorDescent(self.config.marginal_oracle) diff --git a/dpsynth/discrete_mechanisms/direct.py b/dpsynth/discrete_mechanisms/direct.py index 407a650..de0755e 100644 --- a/dpsynth/discrete_mechanisms/direct.py +++ b/dpsynth/discrete_mechanisms/direct.py @@ -19,6 +19,7 @@ from absl import logging import dp_accounting from dpsynth import api +from dpsynth import workload as workload_lib from dpsynth.discrete_mechanisms import accounting from dpsynth.discrete_mechanisms import common import mbi @@ -29,23 +30,24 @@ class DirectConfig(api.MechanismConfig): """Config for the direct mechanism that measures prespecified marginals.""" - def configure(self, _=None, *, budget, delta=0): - return Direct( - config=self, - gdp_budget=accounting.zcdp_to_gdp(budget), - ) - estimator: mbi.Estimator = mbi.estimation.MirrorDescent() marginal_oracle: mbi.MarginalOracle | None = None pgm_iters: int = 5000 - prespecified_marginal_queries: list[tuple[str, ...]] = dataclasses.field( - default_factory=list - ) - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - """Returns the prespecified marginal queries.""" - del domain # Unused. - return list(self.prespecified_marginal_queries) + def configure( + self, + domain=None, + *, + budget: float, + delta: float = 0.0, + workload: workload_lib.WorkloadInput | None = None, + ) -> 'Direct': + del delta + return Direct( + config=self, + gdp_budget=accounting.zcdp_to_gdp(budget), + workload=workload_lib.Workload.from_any(workload, domain), + ) @dataclasses.dataclass(frozen=True) @@ -54,6 +56,12 @@ class Direct(api.CalibratedMechanism): config: DirectConfig gdp_budget: float + workload: workload_lib.Workload | None = None + + def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: + """Returns the prespecified marginal queries.""" + del domain # Unused. + return list(self.workload.cliques) if self.workload else [] @property def dp_event(self) -> dp_accounting.DpEvent: @@ -73,7 +81,9 @@ def __call__( """Selects, measures, estimates, and generates in the compressed domain.""" common.validate_initial_measurements(initial_measurements) phase_times = {} - selected = list(self.config.prespecified_marginal_queries) + w = self.workload + selected = list(w.cliques) if w else [] + weights = np.array(w.weights) if w and w.weights is not None else None all_cliques = [m.clique for m in initial_measurements] + list(selected) summary = mbi.summarize(data.domain, all_cliques) @@ -94,6 +104,7 @@ def __call__( data=data, # pyrefly: ignore[bad-argument-type] marginal_queries=selected, gdp_sigma=accounting.gdp_gaussian_sigma(self.gdp_budget), + weights=weights, ) measurements = list(initial_measurements) + new_measurements diff --git a/dpsynth/discrete_mechanisms/discrete.py b/dpsynth/discrete_mechanisms/discrete.py index e51ad83..2e590b6 100644 --- a/dpsynth/discrete_mechanisms/discrete.py +++ b/dpsynth/discrete_mechanisms/discrete.py @@ -29,6 +29,7 @@ from absl import logging import dp_accounting from dpsynth import api +from dpsynth import workload as workload_lib from dpsynth.discrete_mechanisms import accounting from dpsynth.discrete_mechanisms import common from dpsynth.discrete_mechanisms import mst @@ -59,13 +60,20 @@ class DiscreteConfig(api.MechanismConfig): use_jax_for_bincount: bool = False use_jax_for_generation: bool = False - def configure(self, _=None, *, budget, delta=0): + def configure( + self, + domain=None, + *, + budget: float, + delta: float = 0.0, + workload: workload_lib.WorkloadInput | None = None, + ) -> DiscreteMechanism: """Configures the synthesizer with a zCDP budget.""" + workload = workload_lib.Workload.from_any(workload, domain) one_way_rho = budget * self.one_way_budget_fraction remaining_rho = budget * (1 - self.one_way_budget_fraction) inner = self.mechanism.configure( - budget=remaining_rho, - delta=delta, + domain, budget=remaining_rho, delta=delta, workload=workload ) return DiscreteMechanism( config=self, @@ -82,6 +90,12 @@ class DiscreteMechanism(api.CalibratedMechanism): base_mechanism: api.CalibratedMechanism one_way_gdp_budget: float + def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: + """Returns the cliques needed by the inner mechanism.""" + if hasattr(self.base_mechanism, 'supporting_cliques'): + return self.base_mechanism.supporting_cliques(domain) + return [(a,) for a in domain.attributes] + @property def dp_event(self) -> dp_accounting.DpEvent: """Composes one-way measurement event with the inner mechanism's event.""" @@ -155,9 +169,10 @@ def __call__( measurements = [m.compress(mappings, data.domain) for m in measurements] # pyrefly: ignore[bad-argument-type] logging.info('[DPSynth]: Compressed discrete domain:\n%s', data.domain) - cfg = self.config.mechanism - if isinstance(data, mbi.Dataset) and hasattr(cfg, 'supporting_cliques'): - cliques = cfg.supporting_cliques(data.domain) + if isinstance(data, mbi.Dataset) and hasattr( + self.base_mechanism, 'supporting_cliques' + ): + cliques = self.base_mechanism.supporting_cliques(data.domain) data = common.precompute_marginals( data, cliques, # pyrefly: ignore[bad-argument-type] diff --git a/dpsynth/discrete_mechanisms/independent.py b/dpsynth/discrete_mechanisms/independent.py index 77da78c..80fdb43 100644 --- a/dpsynth/discrete_mechanisms/independent.py +++ b/dpsynth/discrete_mechanisms/independent.py @@ -29,7 +29,8 @@ class IndependentConfig(api.MechanismConfig): pgm_iters: int = 5000 - def configure(self, _=None, *, budget, delta=0.0): + def configure(self, _=None, *, budget, delta=0.0, workload=None): + del budget, delta, workload return Independent(config=self) def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: @@ -43,6 +44,10 @@ class Independent(api.CalibratedMechanism): config: IndependentConfig + def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: + """Returns the one-way marginals this mechanism expects to process.""" + return self.config.supporting_cliques(domain) + @property def dp_event(self) -> dp_accounting.DpEvent: """Returns a zero-cost DP event (no new measurements).""" diff --git a/dpsynth/discrete_mechanisms/mst.py b/dpsynth/discrete_mechanisms/mst.py index 36edb75..20271a2 100644 --- a/dpsynth/discrete_mechanisms/mst.py +++ b/dpsynth/discrete_mechanisms/mst.py @@ -188,7 +188,8 @@ def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: self.maximum_marginal_size, ) - def configure(self, _=None, *, budget, delta=0): + def configure(self, _=None, *, budget, delta=0, workload=None): + del delta, workload return MST( config=self, zcdp_rho=budget, @@ -202,6 +203,10 @@ class MST(api.CalibratedMechanism): config: MSTConfig zcdp_rho: float + def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: + """Returns all pairwise marginals within the size limit.""" + return self.config.supporting_cliques(domain) + @property def dp_event(self) -> dp_accounting.DpEvent: """Returns the DP event for the MST mechanism.""" diff --git a/dpsynth/discrete_mechanisms/swift.py b/dpsynth/discrete_mechanisms/swift.py index 9c1d000..6cd6eb0 100644 --- a/dpsynth/discrete_mechanisms/swift.py +++ b/dpsynth/discrete_mechanisms/swift.py @@ -37,6 +37,7 @@ import dp_accounting from dpsynth import _checkpoint from dpsynth import api +from dpsynth import workload as workload_lib from dpsynth.discrete_mechanisms import accounting from dpsynth.discrete_mechanisms import clique_tree from dpsynth.discrete_mechanisms import common @@ -56,8 +57,6 @@ class SWIFTConfig(api.MechanismConfig): prevent long compilation times. Attributes: - workload: The set of marginals to consider for the mechanism. Can be a - mapping from cliques to their weights or just an iterable of cliques. max_clique_size: The maximum size (domain product) allowed for any clique in the junction tree. This is the main knob to tune to improve utility for a given compute cost. @@ -70,7 +69,6 @@ class SWIFTConfig(api.MechanismConfig): marginals to measure. """ - 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 @@ -80,16 +78,19 @@ class SWIFTConfig(api.MechanismConfig): use_jax_for_bincount: bool = True use_jax_for_generation: bool = True - def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: - """Returns the workload cliques filtered by max_marginal_size.""" - return common.supporting_cliques( - domain, self.workload, self.max_marginal_size - ) - - def configure(self, _=None, *, budget, delta=0): + def configure( + self, + domain=None, + *, + budget: float, + delta: float = 0.0, + workload: workload_lib.WorkloadInput | None = None, + ) -> SWIFT: + del delta return SWIFT( config=self, gdp_budget=accounting.zcdp_to_gdp(budget), + workload=workload_lib.Workload.from_any(workload, domain), ) @@ -99,6 +100,14 @@ class SWIFT(api.CalibratedMechanism): config: SWIFTConfig gdp_budget: float + workload: workload_lib.Workload | None = None + + def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]: + """Returns the workload cliques filtered by max_marginal_size.""" + workload_dict = self.workload.as_dict() if self.workload else None + return common.supporting_cliques( + domain, workload_dict, self.config.max_marginal_size + ) @property def dp_event(self) -> dp_accounting.DpEvent: @@ -124,9 +133,10 @@ def select_and_measure_queries( measure_gdp_budget = self.gdp_budget - select_gdp_budget with common.timed(phase_times, 'compiled_workload'): + workload_dict = self.workload.as_dict() if self.workload else None candidates = common.compiled_workload( data.domain, - self.config.workload, + workload_dict, self.config.max_marginal_size, ) logging.info('[SWIFT] %d candidates.', len(candidates)) diff --git a/dpsynth/serialize.py b/dpsynth/serialize.py index f3ec0f0..c7ce1bc 100644 --- a/dpsynth/serialize.py +++ b/dpsynth/serialize.py @@ -26,6 +26,7 @@ from dpsynth import api from dpsynth import domain from dpsynth import reporting +from dpsynth import workload from dpsynth.relational import domain as relational_domain from etils import epath import mbi @@ -41,6 +42,8 @@ def _resolve_type(type_name: str) -> type[Any] | None: return cls if hasattr(domain, type_name): return getattr(domain, type_name) + if hasattr(workload, type_name): + return getattr(workload, type_name) if hasattr(relational_domain, type_name): return getattr(relational_domain, type_name) if hasattr(reporting, type_name): @@ -106,6 +109,14 @@ def _make_converter() -> cattrs.Converter: ) # 2. Polymorphic structuring for abstract base classes and unions + conv.register_structure_hook( + workload.Workload, + lambda data, cl: ( + cattrs.gen.make_dict_structure_fn(cl, conv)(data, cl) + if isinstance(data, Mapping) + else workload.Workload(data) + ), + ) conv.register_structure_hook( api.MechanismConfig, lambda data, _: _structure_polymorphic(data, _, conv) ) diff --git a/dpsynth/workload.py b/dpsynth/workload.py new file mode 100644 index 0000000..e052826 --- /dev/null +++ b/dpsynth/workload.py @@ -0,0 +1,116 @@ +# 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. + +"""Canonical representation and validation of marginal query workloads.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +import dataclasses +import itertools +import math +from typing import TypeAlias + +from absl import logging + + +@dataclasses.dataclass(frozen=True) +class Workload: + """Canonical representation of a marginal query workload. + + Attributes: + cliques: Sequence of attribute tuples to preserve. + weights: Optional positive weights for each clique (defaults to 1.0). + """ + + cliques: Sequence[tuple[str, ...]] = () + weights: Sequence[float] | None = None + + def __post_init__(self): + if isinstance(self.cliques, (str, Mapping)) or not isinstance( + self.cliques, Iterable + ): + raise ValueError(f'Invalid cliques: {self.cliques!r}.') + cliques: list[tuple[str, ...]] = [] + for cl in self.cliques: + if isinstance(cl, str) or not isinstance(cl, Sequence): + raise ValueError(f'Clique must be a sequence of strings, got {cl!r}.') + clique = tuple(cl) + if not clique: + raise ValueError('Workload cliques must be non-empty.') + if not all(isinstance(col, str) for col in clique): + raise ValueError(f'Clique elements must be strings, got {clique!r}.') + if len(set(clique)) != len(clique): + raise ValueError(f'Clique has duplicate attributes: {clique!r}.') + cliques.append(clique) + object.__setattr__(self, 'cliques', tuple(cliques)) + + if self.weights is not None: + weights = tuple(float(w) for w in self.weights) + if len(weights) != len(cliques): + raise ValueError(f'Length mismatch: {len(cliques)} != {len(weights)}.') + for clique, w in zip(cliques, weights): + if not math.isfinite(w) or w <= 0: + raise ValueError(f'Weight for {clique!r} must be positive: {w}.') + object.__setattr__(self, 'weights', weights) + + def as_dict(self) -> dict[tuple[str, ...], float]: + """Returns a dictionary mapping each clique to its weight.""" + if self.weights is not None: + return dict(zip(self.cliques, self.weights)) + return {cl: 1.0 for cl in self.cliques} + + @classmethod + def all_kway(cls, attributes: Iterable[str], degree: int = 3) -> Workload: + """Returns a workload containing all k-way marginals over attributes.""" + if degree < 1: + raise ValueError(f'degree must be >= 1, got {degree}.') + attrs = tuple(attributes) + k = min(degree, len(attrs)) if attrs else degree + return cls(list(itertools.combinations(attrs, k))) + + @classmethod + def from_any( + cls, + workload: WorkloadInput | None, + attributes: Iterable[str] | None = None, + ) -> Workload | None: + """Coerces a raw workload into a Workload and optionally validates it.""" + if workload is None: + return None + if isinstance(workload, Workload): + resolved = workload + elif isinstance(workload, Mapping): + resolved = cls(list(workload.keys()), list(workload.values())) + else: + resolved = cls(list(workload)) + if attributes is not None: + resolved.validate(attributes) + return resolved + + def validate(self, attributes: Iterable[str]) -> None: + """Validates that all workload attributes exist in the schema.""" + attr_set = set(attributes) + workload_attrs = set().union(*self.cliques) if self.cliques else set() + unknown = workload_attrs - attr_set + if unknown: + raise ValueError(f'Unknown workload attributes: {sorted(unknown)}.') + uncovered = attr_set - workload_attrs + if uncovered: + logging.warning('Attributes not in workload: %s', sorted(uncovered)) + + +WorkloadInput: TypeAlias = ( + Workload | Mapping[tuple[str, ...], float] | Iterable[tuple[str, ...]] +) diff --git a/tests/adapters/beam_test.py b/tests/adapters/beam_test.py index b306673..45bdb17 100644 --- a/tests/adapters/beam_test.py +++ b/tests/adapters/beam_test.py @@ -381,20 +381,19 @@ def test_end_to_end_mixed_types(self): self.assertCountEqual(result.synthetic_data.columns, ['age', 'grade']) @parameterized.named_parameters( - ('mst', discrete_mechanisms.MSTConfig(pgm_iters=250)), + ('mst', discrete_mechanisms.MSTConfig(pgm_iters=250), None), ( 'independent', discrete_mechanisms.IndependentConfig(), + None, ), ( 'direct', - discrete_mechanisms.DirectConfig( - prespecified_marginal_queries=[('a',), ('b',), ('a', 'b')], - pgm_iters=250, - ), + discrete_mechanisms.DirectConfig(pgm_iters=250), + [('a',), ('b',), ('a', 'b')], ), ) - def test_runs_across_mechanisms(self, mechanism): + def test_runs_across_mechanisms(self, mechanism, workload): """The pipeline generalizes to any mechanism via supporting_cliques.""" domains = { 'a': domain.CategoricalAttribute(possible_values=['x', 'y']), @@ -402,7 +401,7 @@ def test_runs_across_mechanisms(self, mechanism): } synth = data_generation_v3.TabularConfig(discrete_mechanism=mechanism) beam_synth = beam_adapter.BeamTabularConfig(synth).configure( - domains, budget=100.0 + domains, budget=100.0, workload=workload ) rows = [ {'a': 'x', 'b': 'p'}, diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py index 9ff26de..72d9714 100644 --- a/tests/data_generation_v3_test.py +++ b/tests/data_generation_v3_test.py @@ -73,7 +73,9 @@ def _discrete_workload_mechanism_baseline_errors( rng = np.random.default_rng(0) data = _make_discrete_data(rng) - mechanism_result = config.configure(budget=budget)(rng, data) + mechanism_result = config.configure(budget=budget, workload=workload)( + rng, data + ) baseline_result = baseline_config.configure(budget=budget)(rng, data) mechanism_error = np.mean([ @@ -101,9 +103,9 @@ def _mixed_workload_mechanism_baseline_errors( numerical_bins=numerical_bins, ) - mechanism_result = mechanism_synth.configure(domains, budget=budget)( - rng, data - ) + mechanism_result = mechanism_synth.configure( + domains, budget=budget, workload=workload + )(rng, data) baseline_result = baseline_synth.configure(domains, budget=budget)(rng, data) mechanism_error = np.mean([ @@ -314,7 +316,7 @@ def test_heterogeneous_input_dataframe_with_none(self): 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) + config = aim.AIMConfig(max_rounds=4, pgm_iters=500) baseline_config = IndependentConfig(pgm_iters=500) mechanism_error, baseline_error = ( _discrete_workload_mechanism_baseline_errors( @@ -325,13 +327,26 @@ def test_discrete_workload_regression_with_aim(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) + config = aim.AIMConfig(max_rounds=4, pgm_iters=500) baseline_config = IndependentConfig(pgm_iters=500) mechanism_error, baseline_error = _mixed_workload_mechanism_baseline_errors( config, baseline_config, workload ) self.assertLess(mechanism_error, 0.05 * baseline_error) + def test_configure_validates_workload_against_schema(self): + domains = { + 'a': domain.CategoricalAttribute(possible_values=['x', 'y']), + 'b': domain.CategoricalAttribute(possible_values=['p', 'q']), + } + config = TabularConfig(discrete_mechanism=aim.AIMConfig(pgm_iters=10)) + with self.assertRaisesRegex(ValueError, 'Unknown workload attributes'): + config.configure(domains, budget=1.0, workload=[('a', 'missing')]) + with self.assertLogs(level='WARNING') as logs: + config.configure(domains, budget=1.0, workload=[('a',)]) + matched = ['Attributes not in workload' in m for m in logs.output] + self.assertTrue(any(matched)) + def test_empty_dataset(self): """Tests that DPSynth works without crashing on empty datasets, and outputs noisy rows.""" domains = { diff --git a/tests/discrete_mechanisms/aim_test.py b/tests/discrete_mechanisms/aim_test.py index 9dd8b9e..eb6f8fa 100644 --- a/tests/discrete_mechanisms/aim_test.py +++ b/tests/discrete_mechanisms/aim_test.py @@ -59,9 +59,9 @@ class AIMTest(absltest.TestCase): def test_fits_one_way_marginals_with_aim(self): data = mbi.Dataset.synthetic(mbi.Domain(["a", "b", "c"], [3, 4, 5]), N=1000) workload = [("a",), ("b",), ("c",)] - config = aim.AIMConfig(workload=workload, max_rounds=4, pgm_iters=500) + config = aim.AIMConfig(max_rounds=4, pgm_iters=500) - calibrated = config.configure(budget=10000) + calibrated = config.configure(budget=10000, workload=workload) result = calibrated(np.random.default_rng(0), data) self.assertIsInstance(result, common.DiscreteMechanismResult) diff --git a/tests/discrete_mechanisms/direct_test.py b/tests/discrete_mechanisms/direct_test.py index 9eace30..967ac75 100644 --- a/tests/discrete_mechanisms/direct_test.py +++ b/tests/discrete_mechanisms/direct_test.py @@ -25,11 +25,9 @@ def test_fits_one_way_marginals(self): data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000) prespecified_queries = [('a', 'b'), ('a', 'c'), ('b', 'c')] - config = direct.DirectConfig( - prespecified_marginal_queries=prespecified_queries, - pgm_iters=500, - ) - result = config.configure(budget=10000)(np.random.default_rng(0), data) + config = direct.DirectConfig(pgm_iters=500) + calibrated = config.configure(budget=10000, workload=prespecified_queries) + result = calibrated(np.random.default_rng(0), data) self.assertIsInstance(result, common.DiscreteMechanismResult) self.assertLen(result.measurements, len(prespecified_queries)) @@ -46,11 +44,11 @@ def test_custom_estimator(self): data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000) prespecified_queries = [('a', 'b'), ('a', 'c'), ('b', 'c')] config = direct.DirectConfig( - prespecified_marginal_queries=prespecified_queries, estimator=mbi.estimation.InteriorGradient(), pgm_iters=500, ) - result = config.configure(budget=10000)(np.random.default_rng(0), data) + calibrated = config.configure(budget=10000, workload=prespecified_queries) + result = calibrated(np.random.default_rng(0), data) for col in data.domain: expected = data.project([col]).datavector() @@ -68,12 +66,12 @@ def tracking_oracle(potentials, total=1, constraints=()): ) config = direct.DirectConfig( - prespecified_marginal_queries=[('a', 'b')], estimator=mbi.estimation.InteriorGradient(), marginal_oracle=tracking_oracle, pgm_iters=10, ) - config.configure(budget=100)(np.random.default_rng(0), data) + calibrated = config.configure(budget=100, workload=[('a', 'b')]) + calibrated(np.random.default_rng(0), data) self.assertNotEmpty(calls) diff --git a/tests/discrete_mechanisms/discrete_mechanisms_test.py b/tests/discrete_mechanisms/discrete_mechanisms_test.py index c687b6a..2ba432e 100644 --- a/tests/discrete_mechanisms/discrete_mechanisms_test.py +++ b/tests/discrete_mechanisms/discrete_mechanisms_test.py @@ -36,13 +36,11 @@ _WORKLOAD = [('a', 'b'), ('b', 'c'), ('a',), ('b',), ('c',)] _MECHANISMS = { - 'AIM': aim.AIMConfig(workload=_WORKLOAD, max_rounds=4, pgm_iters=500), + 'AIM': aim.AIMConfig(max_rounds=4, pgm_iters=500), 'MST': mst.MSTConfig(pgm_iters=500), - 'SWIFT': swift.SWIFTConfig(workload=_WORKLOAD, pgm_iters=500), + 'SWIFT': swift.SWIFTConfig(pgm_iters=500), 'Independent': independent.IndependentConfig(), - 'Direct': direct.DirectConfig( - prespecified_marginal_queries=_WORKLOAD, pgm_iters=500 - ), + 'Direct': direct.DirectConfig(pgm_iters=500), } @@ -63,8 +61,8 @@ def test_mechanism_runs_on_precomputed_marginals(self, mechanism): data = mbi.Dataset.synthetic(domain, N=500) rng = np.random.default_rng(42) - calibrated = mechanism.configure(budget=_ZCDP_RHO) - cliques = mechanism.supporting_cliques(domain) + calibrated = mechanism.configure(budget=_ZCDP_RHO, workload=_WORKLOAD) + cliques = calibrated.supporting_cliques(domain) precomputed = common.precompute_marginals(data, cliques) @@ -86,7 +84,9 @@ def test_compression_restores_domain(self, config): data = _make_skewed_dataset(rng) original_domain = data.domain - result = synth_config.configure(budget=_ZCDP_RHO)(rng, data) + result = synth_config.configure(budget=_ZCDP_RHO, workload=_WORKLOAD)( + rng, data + ) self.assertEqual(result.synthetic_data.domain, original_domain) @@ -103,7 +103,7 @@ def test_compression_with_initial_measurements(self, config): rng, data, [('a',), ('b',)], gdp_sigma=1.0 ) - mechanism = synth_config.configure(budget=_ZCDP_RHO) + mechanism = synth_config.configure(budget=_ZCDP_RHO, workload=_WORKLOAD) result = mechanism(rng, data, initial_measurements=initial_measurements) self.assertEqual(result.synthetic_data.domain, original_domain) @@ -119,7 +119,9 @@ def test_low_epsilon_calibration(self, mechanism): return rng = np.random.default_rng(0) data = _make_skewed_dataset(rng) - result = mechanism.calibrate(epsilon=1e-3, delta=1e-5)(rng, data) + result = mechanism.calibrate(epsilon=1e-3, delta=1e-5, workload=_WORKLOAD)( + rng, data + ) self.assertIsInstance(result, common.DiscreteMechanismResult) diff --git a/tests/serialize_test.py b/tests/serialize_test.py index 53246a0..4acf290 100644 --- a/tests/serialize_test.py +++ b/tests/serialize_test.py @@ -84,10 +84,7 @@ def test_independent_config_roundtrip(self): self.assertEqual(loaded, config) def test_direct_config_roundtrip(self): - config = direct.DirectConfig( - pgm_iters=4000, - prespecified_marginal_queries=[('a', 'b'), ('c',)], - ) + config = direct.DirectConfig(pgm_iters=4000) yaml_str = serialize.to_yaml(config) raw_dict = yaml.safe_load(yaml_str) self.assertNotIn('estimator', raw_dict) @@ -97,7 +94,6 @@ def test_direct_config_roundtrip(self): def test_direct_config_with_estimator_roundtrip(self): config = direct.DirectConfig( pgm_iters=1500, - prespecified_marginal_queries=[('a', 'b')], estimator=mbi.estimation.UniversalAcceleratedMethod(linesearch=True), ) yaml_str = serialize.to_yaml(config) @@ -109,6 +105,37 @@ def test_direct_config_with_estimator_roundtrip(self): loaded = serialize.from_yaml(yaml_str) self.assertEqual(loaded, config) + def test_workload_unweighted_roundtrip(self): + workload = dpsynth.Workload([('a', 'b'), ('b', 'c')]) + yaml_str = dpsynth.to_yaml(workload) + raw_dict = yaml.safe_load(yaml_str) + self.assertEqual( + raw_dict, + {'type': 'Workload', 'cliques': [['a', 'b'], ['b', 'c']]}, + ) + loaded = dpsynth.from_yaml(yaml_str) + self.assertEqual(loaded, workload) + + def test_workload_weighted_roundtrip(self): + workload = dpsynth.Workload([('a', 'b'), ('c',)], weights=[2.5, 0.5]) + yaml_str = dpsynth.to_yaml(workload) + raw_dict = yaml.safe_load(yaml_str) + self.assertEqual( + raw_dict, + { + 'type': 'Workload', + 'cliques': [['a', 'b'], ['c']], + 'weights': [2.5, 0.5], + }, + ) + loaded = dpsynth.from_yaml(yaml_str) + self.assertEqual(loaded, workload) + + def test_workload_from_raw_list_yaml(self): + raw_yaml = '- [a, b]\n- [b, c]\n' + loaded = dpsynth.from_yaml(raw_yaml, dpsynth.Workload) + self.assertEqual(loaded, dpsynth.Workload([('a', 'b'), ('b', 'c')])) + def test_discrete_config_roundtrip(self): config = discrete.DiscreteConfig( mechanism=aim.AIMConfig(pgm_iters=400), diff --git a/tests/workload_test.py b/tests/workload_test.py new file mode 100644 index 0000000..4dcc259 --- /dev/null +++ b/tests/workload_test.py @@ -0,0 +1,95 @@ +# 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 the Workload class and schema consistency validation.""" + +import math + +from absl.testing import absltest +from absl.testing import parameterized +from dpsynth import workload + + +class WorkloadTest(parameterized.TestCase): + + def test_unweighted_workload(self): + w = workload.Workload([('a', 'b'), ('b', 'c')]) + self.assertEqual(w.cliques, (('a', 'b'), ('b', 'c'))) + self.assertIsNone(w.weights) + self.assertEqual(w.as_dict(), {('a', 'b'): 1.0, ('b', 'c'): 1.0}) + + def test_weighted_workload(self): + w = workload.Workload([('a', 'b'), ('c',)], weights=[2.5, 0.5]) + self.assertEqual(w.cliques, (('a', 'b'), ('c',))) + self.assertEqual(w.weights, (2.5, 0.5)) + self.assertEqual(w.as_dict(), {('a', 'b'): 2.5, ('c',): 0.5}) + + def test_all_kway(self): + w2 = workload.Workload.all_kway(['a', 'b', 'c'], degree=2) + self.assertEqual(w2.cliques, (('a', 'b'), ('a', 'c'), ('b', 'c'))) + self.assertIsNone(w2.weights) + + w3_clamped = workload.Workload.all_kway(['a', 'b'], degree=3) + self.assertEqual(w3_clamped.cliques, (('a', 'b'),)) + + with self.assertRaises(ValueError): + workload.Workload.all_kway(['a', 'b'], degree=0) + + def test_from_any(self): + self.assertIsNone(workload.Workload.from_any(None)) + w = workload.Workload([('a', 'b')]) + self.assertIs(workload.Workload.from_any(w), w) + self.assertEqual(workload.Workload.from_any([('a', 'b')]), w) + weighted = workload.Workload.from_any({('a', 'b'): 2.5, ('c',): 0.5}) + expected = workload.Workload([('a', 'b'), ('c',)], weights=[2.5, 0.5]) + self.assertEqual(weighted, expected) + + @parameterized.named_parameters( + dict(testcase_name='bare_string', cliques=['ab'], weights=None), + dict(testcase_name='empty_clique', cliques=[()], weights=None), + dict(testcase_name='non_string_element', cliques=[(1, 2)], weights=None), + dict( + testcase_name='duplicate_in_clique', + cliques=[('a', 'a')], + weights=None, + ), + dict( + testcase_name='mismatched_weights_len', + cliques=[('a',), ('b',)], + weights=[1.0], + ), + dict(testcase_name='zero_weight', cliques=[('a',)], weights=[0.0]), + dict(testcase_name='negative_weight', cliques=[('a',)], weights=[-1.0]), + dict(testcase_name='nan_weight', cliques=[('a',)], weights=[math.nan]), + dict(testcase_name='inf_weight', cliques=[('a',)], weights=[math.inf]), + ) + def test_invalid_workload_raises(self, cliques, weights): + with self.assertRaises(ValueError): + workload.Workload(cliques, weights=weights) + + def test_validate_unknown_attribute_raises(self): + w = workload.Workload([('a', 'b'), ('b', 'unknown')]) + with self.assertRaisesRegex(ValueError, 'Unknown workload attributes'): + w.validate(['a', 'b', 'c']) + + def test_validate_uncovered_attribute_logs_warning(self): + w = workload.Workload([('a', 'b')]) + with self.assertLogs(level='WARNING') as logs: + w.validate(['a', 'b', 'c']) + matched = ['Attributes not in workload' in m for m in logs.output] + self.assertTrue(any(matched)) + + +if __name__ == '__main__': + absltest.main()