Move workload from MechanismConfig subclasses (DirectConfig, AIMConfig, SWIFTConfig) to configure(..., workload=...) and calibrate(..., workload=...) across TabularConfig, DiscreteConfig, BeamTabularConfig, DirectConfig, AIMConfig, and SWIFTConfig, validating workloads against the schema/domain at configuration time via Workload.from_any(workload, domain).

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',
 ]

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)

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,

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,

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(

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

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)