Define domain.Schema, relational.RelationalSchema, and support presets with backwards compatibility.

bin/main.py

@@ -105,10 +105,9 @@ def main(_):
     case _:
       raise ValueError(f'Unknown mechanism: {_MECHANISM.value}')
 
-  mechanism = dpsynth.TabularSynthesizer(
-      domains=attribute_domains,
+  mechanism = dpsynth.TabularConfig(
       discrete_mechanism=mechanism_config,
-  ).calibrate(epsilon=_EPSILON.value, delta=_DELTA.value)
+  ).calibrate(attribute_domains, epsilon=_EPSILON.value, delta=_DELTA.value)
   result = mechanism(np.random.default_rng(_SEED.value), df)  # pyrefly: ignore[bad-argument-type]
   result.synthetic_data.to_csv(_OUTPUT_PATH.value, index=False)
 

docs/in_memory_api.md

@@ -14,8 +14,9 @@ within a single machine's RAM.
 ## Python API: `dpsynth.TabularConfig`
 
 The primary entry point for in-memory synthesis is
-`dpsynth.TabularConfig`. It accepts a dictionary of attribute domains and
-mechanism options, is calibrated with a privacy budget to produce a
+`dpsynth.TabularConfig`. It configures the algorithm hyperparameters (e.g.
+discrete mechanism, numerical bin count, budget allocation), is calibrated
+with a dataset `dpsynth.Schema` and privacy budget to produce a
 `dpsynth.TabularMechanism`, and generates a fully synthetic, differentially
 private DataFrame matching the exact schema and data types of your input.
 
@@ -24,14 +25,24 @@ private DataFrame matching the exact schema and data types of your input.
 ```python
 import dpsynth
 from dpsynth import discrete_mechanisms
+from dpsynth import domain
 import numpy as np
 import pandas as pd
 
+# Define schema (or load from domain.from_yaml_file)
+schema = dpsynth.Schema({
+    "age": domain.NumericalAttribute(min_value=18, max_value=90),
+    "workclass": domain.CategoricalAttribute(possible_values=["Private", "Gov", "Other"]),
+})
+
+# Reusable algorithm preset
 config = dpsynth.TabularConfig(
-    domains=domains,
     discrete_mechanism=discrete_mechanisms.MSTConfig(),
+    numerical_bins=32,
 )
-mechanism = config.calibrate(epsilon=1.0, delta=1e-6)
+
+# Calibrate with schema and privacy budget
+mechanism = config.calibrate(schema, epsilon=1.0, delta=1e-6)
 result = mechanism(np.random.default_rng(), sensitive_df)
 synthetic_df = result.synthetic_data
 ```
@@ -40,22 +51,19 @@ synthetic_df = result.synthetic_data
 
 When initializing `dpsynth.TabularConfig`:
 
-*   `domains`: Mapping of column names to domain specifications
-    ([`CategoricalAttribute`, `NumericalAttribute`, or `OpenSetCategoricalAttribute`](data_and_terminology.md)).
-    Every key must exist in `data.columns`.
 *   `discrete_mechanism`: Configuration object specifying which DP synthesis
-    mechanism to run (e.g., `MSTConfig()`, `AIMConfig()`,
+    mechanism to run (e.g., `MSTConfig()`, `AIMConfig()`, `SWIFTConfig()`,
     `IndependentConfig()`).
 *   `numerical_bins`: Number of equal-frequency quantile buckets used to
     discretize continuous numerical columns (default: `32`).
 *   `init_budget_fraction`: Fraction of total `(epsilon, delta)` budget
     allocated for per-column initialization such as bounds computation and
     partition selection (default: `0.1`).
-*   `cross_attribute_constraints`: Optional sequence of constraints to enforce
-    on generated data.
 
-When calling `config.calibrate(...)`:
+When calling `config.calibrate(schema, ...)` or `config.configure(schema, ...)`:
 
+*   `schema`: The `dpsynth.Schema` (or mapping of column names to attribute
+    domains) defining the dataset columns and optional constraints.
 *   `epsilon`, `delta`: Total differential privacy budget parameters. Returns a
     runnable `TabularMechanism`.
 
@@ -64,7 +72,7 @@ When calling `config.calibrate(...)`:
 ## Standalone End-to-End Python Example
 
 Here is a complete, self-contained Python script demonstrating how to specify a
-domain, set up a `TabularConfig`, calibrate the mechanism with a privacy budget,
+schema, set up a `TabularConfig`, calibrate the mechanism with a privacy budget,
 load sensitive data, synthesize records, and print the first few rows.
 
 ```python
@@ -74,26 +82,25 @@ from dpsynth import domain
 import numpy as np
 import pandas as pd
 
-# 1. Domain Specification: Define the schema of the tabular dataset
-attribute_domains = {
-    "age": domain.NumericalAttribute(lower_bound=18, upper_bound=90),
+# 1. Schema Specification: Define the schema of the tabular dataset
+schema = dpsynth.Schema({
+    "age": domain.NumericalAttribute(min_value=18, max_value=90),
     "workclass": domain.CategoricalAttribute(
-        allowed_values=["Private", "Self-emp", "Gov", "Other"]
+        possible_values=["Private", "Self-emp", "Gov", "Other"]
     ),
     "education": domain.CategoricalAttribute(
-        allowed_values=["HS-grad", "Bachelors", "Masters", "PhD"]
+        possible_values=["HS-grad", "Bachelors", "Masters", "PhD"]
     ),
-}
+})
 
-# 2. Setup Config: Configure synthesizer with domain and mechanism choices
+# 2. Setup Config: Configure synthesizer hyperparameter preset
 config = dpsynth.TabularConfig(
-    domains=attribute_domains,
     discrete_mechanism=discrete_mechanisms.MSTConfig(),
     numerical_bins=16,
 )
 
-# 3. Calibrate Mechanism: Allocate privacy budget to get runnable mechanism
-mechanism = config.calibrate(epsilon=1.0, delta=1e-5)
+# 3. Calibrate Mechanism: Allocate privacy budget with schema to get runnable mechanism
+mechanism = config.calibrate(schema, epsilon=1.0, delta=1e-5)
 
 # 4. Load Data: Create sensitive input DataFrame matching the domain schema
 sensitive_df = pd.DataFrame({

dpsynth/__init__.py

@@ -21,6 +21,9 @@
 from dpsynth import discrete_mechanisms
 from dpsynth import domain
 from dpsynth import relational
+from dpsynth.api import CalibratedMechanism
+from dpsynth.api import MechanismConfig
+from dpsynth.constraints import Constraint
 from dpsynth.data_generation_v3 import TabularConfig
 from dpsynth.data_generation_v3 import TabularMechanism
 from dpsynth.data_generation_v3 import TabularSynthesizer
@@ -30,24 +33,35 @@
 from dpsynth.domain import FreeFormTextAttribute
 from dpsynth.domain import NumericalAttribute
 from dpsynth.domain import OpenSetCategoricalAttribute
+from dpsynth.domain import Schema
 
 ForeignKeyRelation = relational.ForeignKeyRelation
 MultiDataGenerationResult = relational.MultiDataGenerationResult
 MultiTableConfig = relational.MultiTableConfig
 MultiTableMechanism = relational.MultiTableMechanism
+RelationalSchema = relational.RelationalSchema
 
 __all__ = [
+    'CalibratedMechanism',
     'CategoricalAttribute',
+    'Constraint',
+    'DiscreteConfig',
+    'DiscreteMechanism',
     'ForeignKeyRelation',
+    'FreeFormTextAttribute',
+    'MechanismConfig',
     'MultiDataGenerationResult',
     'MultiTableConfig',
     'MultiTableMechanism',
     'NumericalAttribute',
     'OpenSetCategoricalAttribute',
+    'RelationalSchema',
+    'Schema',
     'TabularConfig',
     'TabularMechanism',
     'TabularSynthesizer',
     'api',
+    'constraints',
     'discrete_mechanisms',
     'domain',
     'relational',

dpsynth/adapters/beam.py

@@ -24,7 +24,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 import dataclasses
 import io
 import math
@@ -74,12 +74,14 @@ def __init__(self, initializers: dict[str, Initializer]):
     for column, init in initializers.items():
       if isinstance(init, initialization.NumericalInitializerConfig):
         attr = init.attribute
-        lower, upper, gs = init.grid_spec
+        assert attr is not None
+        lower, upper, gs = init.grid_spec(attr)
         delta = (upper - lower) / (gs - 1)
         meta = dict(attribute=attr, lower=lower, upper=upper, delta=delta)
         self._specs.append((column, 'numerical', meta))
 
       elif isinstance(init, initialization.CategoricalInitializerConfig):
+        assert init.attribute is not None
         meta = {
             'lookup': init.attribute.lookup,
             'default': init.attribute.out_of_domain_index,
@@ -211,10 +213,10 @@ def run_from_summary(
   for column, init in initializers.items():
     sparse = sparse_stats[column]
     if isinstance(init, initialization.NumericalInitializer):
-      counts = _sparse_to_dense_numerical(sparse, init.config.grid_spec[2])
+      counts = _sparse_to_dense_numerical(sparse, init.grid_size)
       results[column] = init.from_summary(rng, counts)
     elif isinstance(init, initialization.CategoricalInitializer):
-      counts = _sparse_to_dense_categorical(sparse, init.config.attribute.size)
+      counts = _sparse_to_dense_categorical(sparse, init.attribute.size)
       results[column] = init.from_summary(rng, counts)
     elif isinstance(init, initialization.OpenSetInitializer):
       unique_values, value_counts = _sparse_to_openset(sparse)
@@ -530,21 +532,23 @@ class BeamTabularConfig(api.MechanismConfig):
 
   Usage::
 
-      config = data_generation_v3.TabularConfig(domains=domains)
-      beam_synth = BeamTabularConfig(config).configure(zcdp_rho=1.0)
+      config = data_generation_v3.TabularConfig()
+      beam_synth = BeamTabularConfig(config).configure(schema, zcdp_rho=1.0)
       result = beam_synth(rng, create_rows_fn)
 
   Attributes:
-    synthesizer: The wrapped local-mode TabularConfig. Supplies the domain,
-      sub-mechanisms, constraints, and privacy calibration.
+    synthesizer: The wrapped local-mode TabularConfig. Supplies the
+      sub-mechanisms and initialization parameters.
     temp_location: Directory used to shuttle small singleton results between the
       pipeline and the driver. Must be readable and writable by all workers --
       i.e. a shared distributed filesystem for distributed runners. Defaults to
       a local temp directory, which is only valid for in-process runners.
     pipeline_options: Optional Beam pipeline options applied to both passes.
   """
 
-  synthesizer: data_generation_v3.TabularConfig
+  synthesizer: data_generation_v3.TabularConfig = dataclasses.field(
+      default_factory=data_generation_v3.TabularConfig
+  )
   temp_location: str | None = None
   pipeline_options: beam.options.pipeline_options.PipelineOptions | None = None
 
@@ -555,11 +559,25 @@ def __post_init__(self):
           ' method.'
       )
 
+  @property
+  def domains(
+      self,
+  ) -> domain.Schema | Mapping[str, domain.AttributeType] | None:
+    return self.synthesizer.domains
+
   def configure(
-      self, *, zcdp_rho, delta=0, max_records_per_user=1
+      self,
+      schema: domain.Schema | Mapping[str, domain.AttributeType] | None = None,
+      *,
+      zcdp_rho: float,
+      delta: float = 0.0,
+      max_records_per_user: int = 1,
   ) -> BeamTabularMechanism:
     """Returns a copy whose synthesizer is configured with the given budget."""
+    if schema is None:
+      schema = self.domains
     synthesizer = self.synthesizer.configure(
+        schema,
         zcdp_rho=zcdp_rho,
         delta=delta,
         max_records_per_user=max_records_per_user,

dpsynth/api.py

@@ -119,7 +119,12 @@ class MechanismConfig(abc.ABC):
 
   @abc.abstractmethod
   def configure(
-      self, *, zcdp_rho, delta=0, max_records_per_user=1
+      self,
+      domain: Any = None,
+      *,
+      zcdp_rho: float,
+      delta: float = 0.0,
+      max_records_per_user: int = 1,
   ) -> CalibratedMechanism:
     """Returns a calibrated mechanism for the given zCDP budget.
 
@@ -133,6 +138,8 @@ def configure(
     not provided (i.e., is 0).
 
     Args:
+      domain: Mechanism-specific domain or schema specification (e.g.
+        ``domain.Schema``, ``domain.AttributeType``, or ``mbi.Domain``).
       zcdp_rho: The zCDP privacy budget (rho).
       delta: Approximate DP delta consumed by the mechanism itself (e.g., for
         thresholding). Defaults to 0 (pure zCDP). Mechanisms that need delta
@@ -203,6 +210,7 @@ def _find_optimal_rho(
 
   def calibrate(
       self,
+      domain: Any = None,
       *,
       epsilon: float | None = None,
       delta: float | None = None,
@@ -221,6 +229,7 @@ def calibrate(
       ``configure(zcdp_rho=...)`` directly instead.
 
     Args:
+      domain: Mechanism-specific domain or schema specification.
       epsilon: Target epsilon for (epsilon, delta)-DP.
       delta: Target delta for (epsilon, delta)-DP.
       zcdp_rho: Deprecated. Direct zCDP budget. Use ``configure()`` instead.
@@ -252,6 +261,7 @@ def calibrate(
           stacklevel=2,
       )
       return self.configure(
+          domain,
           zcdp_rho=zcdp_rho,
           max_records_per_user=max_records_per_user,
       )
@@ -261,6 +271,7 @@ def calibrate(
 
     def make_event_fn(rho: float) -> dp_accounting.DpEvent:
       base = self.configure(
+          domain,
           zcdp_rho=rho,
           delta=delta,
           max_records_per_user=max_records_per_user,
@@ -274,6 +285,7 @@ def make_event_fn(rho: float) -> dp_accounting.DpEvent:
         target_delta=delta,
     )
     return self.configure(
+        domain,
         zcdp_rho=optimal_rho,
         delta=delta,
         max_records_per_user=max_records_per_user,

dpsynth/data_generation_v2.py

@@ -23,6 +23,7 @@
 from dpsynth import constraints
 from dpsynth import discrete_mechanisms
 from dpsynth import domain
+from dpsynth.data_generation_v3 import TabularConfig
 from dpsynth.data_generation_v3 import TabularSynthesizer
 import numpy as np
 import pandas as pd
@@ -51,14 +52,17 @@ def generate(
       stacklevel=2,
   )
   del skip_compression  # Not supported by TabularSynthesizer.
-  synth = TabularSynthesizer(
-      domains=domains,
+  config = TabularConfig(
       discrete_mechanism=discrete_config,  # pyrefly: ignore[bad-argument-type]
-      cross_attribute_constraints=cross_attribute_constraints,
       numerical_bins=numerical_bins,
       init_budget_fraction=one_way_marginal_budget_fraction,
   )
-  result = synth.calibrate(
+  schema = domain.Schema(
+      attributes=domains,
+      constraints=cross_attribute_constraints or (),
+  )
+  result = config.calibrate(
+      schema,
       epsilon=epsilon,
       delta=delta,
   )(np.random.default_rng(), data)