From a7c6f91f0a6cdc4828c1e586633ca991d799b5c2 Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Mon, 24 Aug 2026 17:50:22 -0700 Subject: [PATCH] Add Schema dataclass, decouple constraints, and refactor TabularConfig into reusable preset. - Introduce `dpsynth.domain.Schema` dataclass implementing `Mapping[str, AttributeType]` to encapsulate column attribute domains and cross-attribute constraints with YAML serialization support. - Decouple `dpsynth.constraints.Constraint` so `attribute_domains` is optional, adding `bind_schema()` and `to_mbi(schema)` for automatic type resolution. - Refactor `dpsynth.TabularConfig` into a reusable algorithm preset: remove the unused `initializers` field, make `schema`/`domains` optional in the constructor, and support passing `schema` and `constraints` directly to `configure()` and `calibrate()`. - Re-export `Schema` and `Constraint` at the top level `dpsynth` package. - Update in-memory documentation and add unit tests across domain, constraints, and tabular mechanisms. PiperOrigin-RevId: 970195519 --- docs/in_memory_api.md | 53 +++-- dpsynth/__init__.py | 20 ++ dpsynth/api.py | 339 ++++++++++++++++++++++++++++++- dpsynth/constraints.py | 115 +++++++++-- dpsynth/data_generation_v3.py | 159 ++++++++------- dpsynth/domain.py | 192 +++++++++++++---- dpsynth/relational/domain.py | 4 +- pyproject.toml | 1 + tests/api_test.py | 195 ++++++++++++++++++ tests/constraints_test.py | 28 +++ tests/data_generation_v3_test.py | 48 ++++- tests/domain_test.py | 55 +++++ 12 files changed, 1060 insertions(+), 149 deletions(-) create mode 100644 tests/api_test.py diff --git a/docs/in_memory_api.md b/docs/in_memory_api.md index 5a16d3e1..c0c67e5e 100644 --- a/docs/in_memory_api.md +++ b/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 schema.to_yaml_file / 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=schema, epsilon=1.0, delta=1e-6) result = mechanism(np.random.default_rng(), sensitive_df) synthetic_df = result.synthetic_data ``` @@ -40,22 +51,23 @@ 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. +* `cross_attribute_constraints`: Optional sequence of `Constraint` objects to + enforce on generated data. +* `schema`: Optional schema specification if binding the config to a specific + dataset at construction time rather than calibration time. When calling `config.calibrate(...)`: +* `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 +76,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 +86,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=schema, epsilon=1.0, delta=1e-5) # 4. Load Data: Create sensitive input DataFrame matching the domain schema sensitive_df = pd.DataFrame({ diff --git a/dpsynth/__init__.py b/dpsynth/__init__.py index dac7c8d2..8938372d 100644 --- a/dpsynth/__init__.py +++ b/dpsynth/__init__.py @@ -21,6 +21,13 @@ from dpsynth import discrete_mechanisms from dpsynth import domain from dpsynth import relational +from dpsynth.api import CalibratedMechanism +from dpsynth.api import from_yaml +from dpsynth.api import from_yaml_file +from dpsynth.api import MechanismConfig +from dpsynth.api import to_yaml +from dpsynth.api import to_yaml_file +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,6 +37,7 @@ 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 @@ -37,18 +45,30 @@ MultiTableMechanism = relational.MultiTableMechanism __all__ = [ + 'CalibratedMechanism', 'CategoricalAttribute', + 'Constraint', + 'DiscreteConfig', + 'DiscreteMechanism', 'ForeignKeyRelation', + 'FreeFormTextAttribute', + 'MechanismConfig', 'MultiDataGenerationResult', 'MultiTableConfig', 'MultiTableMechanism', 'NumericalAttribute', 'OpenSetCategoricalAttribute', + 'Schema', 'TabularConfig', 'TabularMechanism', 'TabularSynthesizer', 'api', + 'constraints', 'discrete_mechanisms', 'domain', + 'from_yaml', + 'from_yaml_file', 'relational', + 'to_yaml', + 'to_yaml_file', ] diff --git a/dpsynth/api.py b/dpsynth/api.py index 70ce4ce3..60671da8 100644 --- a/dpsynth/api.py +++ b/dpsynth/api.py @@ -34,12 +34,21 @@ from __future__ import annotations import abc -from collections.abc import Callable +from collections.abc import Callable, Mapping +import dataclasses import functools -from typing import Any +import importlib +from typing import Any, TypeVar import warnings import dp_accounting +from dpsynth import domain +import yaml + +import pathlib +PathType = pathlib.Path | str + +SelfConfig = TypeVar('SelfConfig', bound='MechanismConfig') class CalibratedMechanism(abc.ABC): @@ -117,9 +126,58 @@ class MechanismConfig(abc.ABC): as the mechanism's own privacy characterization allows. """ + _registry: dict[str, type[MechanismConfig]] = {} + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + MechanismConfig._registry[cls.__name__] = cls + + @classmethod + def get_registered_class(cls, name: str) -> type[MechanismConfig] | None: + """Returns the registered MechanismConfig subclass for a given name.""" + return cls._registry.get(name) + + def to_dict(self) -> dict[str, Any]: + """Converts the config into a serializable dictionary.""" + return _config_to_dict(self) + + def to_yaml(self) -> str: + """Serializes the config into a YAML string.""" + return yaml.dump(self.to_dict(), default_flow_style=False, sort_keys=False) + + def to_yaml_file(self, filepath: str | PathType) -> None: + """Writes the config to a YAML file.""" + with open(filepath, 'w') as f: + f.write(self.to_yaml()) + + @classmethod + def from_dict(cls: type[SelfConfig], data: Mapping[str, Any]) -> SelfConfig: + """Instantiates a MechanismConfig from a dictionary.""" + return _config_from_dict(data, expected_cls=cls) + + @classmethod + def from_yaml(cls: type[SelfConfig], yaml_str: str) -> SelfConfig: + """Instantiates a MechanismConfig from a YAML string.""" + data = yaml.safe_load(yaml_str) + if not isinstance(data, dict): + raise ValueError(f'Expected YAML dictionary, got {type(data).__name__}.') + return cls.from_dict(data) + + @classmethod + def from_yaml_file( + cls: type[SelfConfig], filepath: str | PathType + ) -> SelfConfig: + """Reads a MechanismConfig from a YAML file.""" + with open(filepath, 'r') as f: + return cls.from_yaml(f.read()) + @abc.abstractmethod def configure( - self, *, zcdp_rho, delta=0, max_records_per_user=1 + self, + *, + zcdp_rho: float, + delta: float = 0.0, + max_records_per_user: int = 1, ) -> CalibratedMechanism: """Returns a calibrated mechanism for the given zCDP budget. @@ -209,6 +267,7 @@ def calibrate( zcdp_rho: float | None = None, poisson_sampling_prob: float = 1.0, max_records_per_user: int = 1, + **kwargs: Any, ) -> CalibratedMechanism: """Calibrate the mechanism to a target (epsilon, delta)-DP guarantee. @@ -232,6 +291,8 @@ def calibrate( scaled by this factor to provide user-level rather than record-level DP; the privacy accounting is unchanged. Soundness relies on the caller enforcing this bound. + **kwargs: Additional mechanism-specific configuration arguments (e.g. + ``schema`` or ``constraints``). Returns: A calibrated, runnable mechanism. @@ -254,6 +315,7 @@ def calibrate( return self.configure( zcdp_rho=zcdp_rho, max_records_per_user=max_records_per_user, + **kwargs, ) if epsilon is None or delta is None: @@ -264,6 +326,7 @@ def make_event_fn(rho: float) -> dp_accounting.DpEvent: zcdp_rho=rho, delta=delta, max_records_per_user=max_records_per_user, + **kwargs, ).dp_event sampled = dp_accounting.PoissonSampledDpEvent(poisson_sampling_prob, base) return base if poisson_sampling_prob == 1.0 else sampled @@ -277,6 +340,7 @@ def make_event_fn(rho: float) -> dp_accounting.DpEvent: zcdp_rho=optimal_rho, delta=delta, max_records_per_user=max_records_per_user, + **kwargs, ) @@ -298,3 +362,272 @@ def validate_max_records_per_user(value: int) -> None: """Raises ValueError if the per-user record bound is not a positive int.""" if value < 1: raise ValueError(f'max_records_per_user must be >= 1, got {value}.') + + +def _ensure_subclasses_loaded(): + """Loads standard MechanismConfig subclasses to populate registry.""" + modules = [ + 'dpsynth.data_generation_v3', + 'dpsynth.discrete_mechanisms.aim', + 'dpsynth.discrete_mechanisms.aim_gdp', + 'dpsynth.discrete_mechanisms.direct', + 'dpsynth.discrete_mechanisms.discrete', + 'dpsynth.discrete_mechanisms.independent', + 'dpsynth.discrete_mechanisms.mst', + 'dpsynth.discrete_mechanisms.swift', + 'dpsynth.relational.synthesizer', + 'dpsynth.local_mode.initialization', + ] + for mod_name in modules: + try: + importlib.import_module(mod_name) + except (ImportError, AttributeError): + pass + + +def _get_foreign_key_relation_cls() -> type[Any] | None: + """Lazily resolves ForeignKeyRelation class if available.""" + try: + mod = importlib.import_module('dpsynth.relational.domain') + return getattr(mod, 'ForeignKeyRelation', None) + except (ImportError, AttributeError): + return None + + +def _config_to_dict(obj: Any) -> Any: + """Converts an object into a YAML-serializable dictionary structure.""" + if isinstance(obj, MechanismConfig): + data = {'type': obj.__class__.__name__} + if dataclasses.is_dataclass(obj): + for f in dataclasses.fields(obj): + data[f.name] = _config_to_dict(getattr(obj, f.name)) + return data + elif isinstance( + obj, + ( + domain.CategoricalAttribute, + domain.NumericalAttribute, + domain.OpenSetCategoricalAttribute, + domain.FreeFormTextAttribute, + ), + ): + return domain.attribute_to_dict(obj) + elif isinstance(obj, domain.Schema): + return obj.to_dict() + elif hasattr(obj, 'to_dict') and callable(obj.to_dict): + return obj.to_dict() + elif dataclasses.is_dataclass(obj) and not isinstance(obj, type): + data = {} + for f in dataclasses.fields(obj): + data[f.name] = _config_to_dict(getattr(obj, f.name)) + data['type'] = obj.__class__.__name__ + return data + elif isinstance(obj, Mapping): + return {str(k): _config_to_dict(v) for k, v in obj.items()} + elif isinstance(obj, (list, tuple)): + return [_config_to_dict(item) for item in obj] + else: + return obj + + +def _instantiate_dataclass(cls: type[Any], data: dict[str, Any]) -> Any: + """Instantiates a dataclass, converting nested fields appropriately.""" + kwargs = {} + for field in dataclasses.fields(cls): + if field.name not in data: + continue + val = data[field.name] + if val is None: + kwargs[field.name] = None + continue + + if field.name in ('domains', 'domain', 'schema'): + if isinstance(val, dict): + if 'attributes' in val: + kwargs[field.name] = domain.Schema.from_dict(val) + else: + first_val = next(iter(val.values()), None) + if isinstance(first_val, dict) and any( + k in first_val + for k in ( + 'min_value', + 'possible_values', + 'default_value', + 'max_tokens', + 'type', + ) + ): + kwargs[field.name] = domain.Schema.from_dict(val) + elif isinstance(first_val, dict): + kwargs[field.name] = { + t: domain.Schema.from_dict(c) for t, c in val.items() + } + else: + kwargs[field.name] = val + else: + kwargs[field.name] = val + elif field.name in ('constraints', 'cross_attribute_constraints'): + if isinstance(val, list): + constraints_mod = importlib.import_module('dpsynth.constraints') + c_list = [ + constraints_mod.Constraint.from_dict(c) + if isinstance(c, dict) + else c + for c in val + ] + if isinstance(field.default, tuple): + kwargs[field.name] = tuple(c_list) + else: + kwargs[field.name] = c_list + else: + kwargs[field.name] = val + elif field.name in ('discrete_mechanism', 'mechanism'): + kwargs[field.name] = _config_from_dict(val, expected_cls=MechanismConfig) + elif field.name == 'initializers': + if isinstance(val, dict): + res = {} + for k, v in val.items(): + if isinstance(v, dict): + if 'type' in v or any( + f in v for f in ('target_attribute_name', 'num_bins') + ): + res[k] = _config_from_dict(v, expected_cls=MechanismConfig) + else: + res[k] = { + col: _config_from_dict(cfg, expected_cls=MechanismConfig) + for col, cfg in v.items() + } + else: + res[k] = v + kwargs[field.name] = res + else: + kwargs[field.name] = val + elif field.name in ('foreign_keys', 'relations'): + fkr_cls = _get_foreign_key_relation_cls() + if isinstance(val, (list, tuple)) and fkr_cls is not None: + fk_list = [] + for item in val: + if isinstance(item, dict): + item_dict = {k: v for k, v in item.items() if k != 'type'} + fk_list.append(fkr_cls(**item_dict)) + else: + fk_list.append(item) + if isinstance(field.default, tuple): + kwargs[field.name] = tuple(fk_list) + else: + kwargs[field.name] = fk_list + else: + kwargs[field.name] = val + elif field.name == 'workload': + if isinstance(val, list): + kwargs[field.name] = [ + tuple(c) if isinstance(c, list) else c for c in val + ] + elif isinstance(val, dict): + kwargs[field.name] = { + tuple(k) if isinstance(k, list) else k: v for k, v in val.items() + } + else: + kwargs[field.name] = val + elif field.name == 'prespecified_marginal_queries': + if isinstance(val, list): + kwargs[field.name] = [ + tuple(q) if isinstance(q, list) else q for q in val + ] + else: + kwargs[field.name] = val + elif field.name == 'attribute' and isinstance(val, dict): + kwargs[field.name] = domain.attribute_from_dict(val) + elif isinstance(val, dict) and 'type' in val: + kwargs[field.name] = _config_from_dict(val) + elif isinstance(field.default, tuple) and isinstance(val, list): + kwargs[field.name] = tuple(val) + else: + kwargs[field.name] = val + return cls(**kwargs) + + +def _config_from_dict( + data: Mapping[str, Any], expected_cls: type[Any] | None = None +) -> Any: + """Reconstructs a MechanismConfig or nested object from a dictionary.""" + if not isinstance(data, dict): + return data + + _ensure_subclasses_loaded() + + type_name = data.get('type') + if type_name: + data_without_type = {k: v for k, v in data.items() if k != 'type'} + target_cls = MechanismConfig.get_registered_class(type_name) + if target_cls is not None: + return _instantiate_dataclass(target_cls, data_without_type) + elif type_name in ( + 'CategoricalAttribute', + 'NumericalAttribute', + 'OpenSetCategoricalAttribute', + 'FreeFormTextAttribute', + ): + return domain.attribute_from_dict(data) + elif type_name == 'ForeignKeyRelation': + fkr_cls = _get_foreign_key_relation_cls() + if fkr_cls is not None: + return fkr_cls(**data_without_type) + raise ValueError(f"Unknown type in YAML: '{type_name}'") + else: + raise ValueError(f"Unknown type in YAML: '{type_name}'") + elif expected_cls is not None and expected_cls is not MechanismConfig: + if dataclasses.is_dataclass(expected_cls): + return _instantiate_dataclass(expected_cls, dict(data)) + return expected_cls(**data) + else: + raise ValueError( + "Missing 'type' field in YAML dictionary to identify MechanismConfig" + ' class.' + ) + + +def to_yaml(config: MechanismConfig) -> str: + """Serializes a MechanismConfig into a YAML string. + + Args: + config: The MechanismConfig to serialize. + + Returns: + A YAML string representation of the config. + """ + return config.to_yaml() + + +def to_yaml_file(config: MechanismConfig, filepath: str | PathType) -> None: + """Writes a MechanismConfig to a YAML file. + + Args: + config: The MechanismConfig to serialize. + filepath: Destination file path. + """ + config.to_yaml_file(filepath) + + +def from_yaml(yaml_str: str) -> MechanismConfig: + """Loads a MechanismConfig from a YAML string. + + Args: + yaml_str: YAML string encoding a MechanismConfig. + + Returns: + The reconstructed MechanismConfig instance. + """ + return MechanismConfig.from_yaml(yaml_str) + + +def from_yaml_file(filepath: str | PathType) -> MechanismConfig: + """Loads a MechanismConfig from a YAML file. + + Args: + filepath: Path to the YAML file. + + Returns: + The reconstructed MechanismConfig instance. + """ + return MechanismConfig.from_yaml_file(filepath) diff --git a/dpsynth/constraints.py b/dpsynth/constraints.py index f6bf7b00..8f1c91a5 100644 --- a/dpsynth/constraints.py +++ b/dpsynth/constraints.py @@ -28,7 +28,9 @@ def _validate(c: Constraint) -> None: """Validate a Constraint's fields.""" - if len(c.attribute_names) != len(c.attribute_domains): + if c.attribute_domains is not None and len(c.attribute_names) != len( + c.attribute_domains + ): raise ValueError( 'attribute_names and attribute_domains must have the same length, got' f' {len(c.attribute_names)} != {len(c.attribute_domains)}.' @@ -87,7 +89,8 @@ class Constraint: Attributes: attribute_names: Names of the constrained attributes. - attribute_domains: Categorical domain for each attribute. + attribute_domains: Categorical domain for each attribute. Optional when + associated with a Schema or provided to ``to_mbi(schema)``. possible_combinations: Allowed value combinations. impossible_combinations: Forbidden value combinations. functional_dependency: Dict mapping fine attribute values to coarse @@ -95,7 +98,7 @@ class Constraint: """ attribute_names: tuple[str, ...] - attribute_domains: tuple[domain.CategoricalAttribute, ...] + attribute_domains: tuple[domain.CategoricalAttribute, ...] | None = None possible_combinations: Sequence[tuple[Any, ...]] | None = None impossible_combinations: Sequence[tuple[Any, ...]] | None = None functional_dependency: Mapping[Any, Any] | None = None @@ -103,27 +106,115 @@ class Constraint: def __post_init__(self): _validate(self) - def to_mbi(self) -> mbi.Constraint: + def bind_schema( + self, schema: Mapping[str, domain.AttributeType] + ) -> Constraint: + """Returns a new Constraint with attribute_domains resolved from schema.""" + if self.attribute_domains is not None: + return self + resolved_domains = [] + for name in self.attribute_names: + if name not in schema: + raise KeyError( + f"Attribute '{name}' from constraint not found in schema." + ) + attr = schema[name] + if not isinstance(attr, domain.CategoricalAttribute): + raise TypeError( + f"Constraint attribute '{name}' must be CategoricalAttribute, got" + f' {type(attr).__name__}.' + ) + resolved_domains.append(attr) + return dataclasses.replace(self, attribute_domains=tuple(resolved_domains)) + + def to_mbi( + self, schema: Mapping[str, domain.AttributeType] | None = None + ) -> mbi.Constraint: """Convert to an mbi.Constraint.""" - shape = tuple(d.size for d in self.attribute_domains) - mbi_domain = mbi.Domain(self.attribute_names, shape) + bound = self.bind_schema(schema) if schema is not None else self + if bound.attribute_domains is None: + raise ValueError( + 'Constraint has no attribute_domains; provide a schema to to_mbi() or' + ' construct with attribute_domains.' + ) + + shape = tuple(d.size for d in bound.attribute_domains) + mbi_domain = mbi.Domain(bound.attribute_names, shape) encoders = [ - transformations.discrete_encoder(d) for d in self.attribute_domains + transformations.discrete_encoder(d) for d in bound.attribute_domains ] - if self.functional_dependency is not None: + if bound.functional_dependency is not None: _, coarse_enc = encoders - fine_values = self.attribute_domains[0].possible_values + fine_values = bound.attribute_domains[0].possible_values coarse_indices = [ - coarse_enc(self.functional_dependency[v]) for v in fine_values + coarse_enc(bound.functional_dependency[v]) for v in fine_values ] return mbi.Constraint( domain=mbi_domain, mapping=np.array(coarse_indices, dtype=np.int32) ) - combos = self.possible_combinations or self.impossible_combinations + combos = bound.possible_combinations or bound.impossible_combinations encoded = [[enc(v) for enc, v in zip(encoders, c)] for c in combos] # pyrefly: ignore[not-iterable] indices = np.array(encoded, dtype=np.int32) - if self.possible_combinations is not None: + if bound.possible_combinations is not None: return mbi.Constraint(domain=mbi_domain, valid=indices) return mbi.Constraint(domain=mbi_domain, invalid=indices) + + def to_dict(self) -> dict[str, Any]: + """Converts the Constraint into a serializable dictionary.""" + data: dict[str, Any] = {'attribute_names': list(self.attribute_names)} + if self.possible_combinations is not None: + data['possible_combinations'] = [ + list(c) for c in self.possible_combinations + ] + if self.impossible_combinations is not None: + data['impossible_combinations'] = [ + list(c) for c in self.impossible_combinations + ] + if self.functional_dependency is not None: + data['functional_dependency'] = dict(self.functional_dependency) + if self.attribute_domains is not None: + data['attribute_domains'] = [ + domain.attribute_to_dict(d) for d in self.attribute_domains + ] + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> Constraint: + """Instantiates a Constraint from a dictionary.""" + attr_names = tuple(data['attribute_names']) + attr_domains = None + if 'attribute_domains' in data and data['attribute_domains'] is not None: + parsed_domains = [] + for d in data['attribute_domains']: + parsed = domain.attribute_from_dict(d) + if not isinstance(parsed, domain.CategoricalAttribute): + raise TypeError( + 'Constraint attribute_domains must be CategoricalAttribute, got' + f' {type(parsed).__name__}.' + ) + parsed_domains.append(parsed) + attr_domains = tuple(parsed_domains) + possible_combinations = None + if ( + 'possible_combinations' in data + and data['possible_combinations'] is not None + ): + possible_combinations = [tuple(c) for c in data['possible_combinations']] + impossible_combinations = None + if ( + 'impossible_combinations' in data + and data['impossible_combinations'] is not None + ): + impossible_combinations = [ + tuple(c) for c in data['impossible_combinations'] + ] + functional_dependency = data.get('functional_dependency') + return cls( + attribute_names=attr_names, + attribute_domains=attr_domains, + possible_combinations=possible_combinations, + impossible_combinations=impossible_combinations, + functional_dependency=functional_dependency, + ) diff --git a/dpsynth/data_generation_v3.py b/dpsynth/data_generation_v3.py index e760a070..9f543ba4 100644 --- a/dpsynth/data_generation_v3.py +++ b/dpsynth/data_generation_v3.py @@ -18,15 +18,17 @@ from collections.abc import Mapping, Sequence import dataclasses +import typing import warnings from absl import logging import dp_accounting from dpsynth import api -from dpsynth import constraints +from dpsynth import constraints as constraints_mod from dpsynth import discrete_mechanisms -from dpsynth import domain +from dpsynth import domain as domain_mod from dpsynth.discrete_mechanisms import common as dm_common +from dpsynth.domain import Schema from dpsynth.local_mode import initialization from dpsynth.local_mode import primitives from dpsynth.local_mode import vectorized_transformations as vtx @@ -36,7 +38,7 @@ def create_initializers( - domains: Mapping[str, domain.AttributeType], + domains: Mapping[str, domain_mod.AttributeType], numerical_bins: int, ) -> dict[str, api.MechanismConfig]: """Creates per-column initializers from the domain specification. @@ -53,18 +55,18 @@ def create_initializers( """ initializers = {} for col, attr in domains.items(): - if isinstance(attr, domain.NumericalAttribute): + if isinstance(attr, domain_mod.NumericalAttribute): initializers[col] = initialization.NumericalInitializerConfig( name=col, num_partitions=numerical_bins, attribute=attr, ) - elif isinstance(attr, domain.CategoricalAttribute): + elif isinstance(attr, domain_mod.CategoricalAttribute): initializers[col] = initialization.CategoricalInitializerConfig( name=col, attribute=attr, ) - elif isinstance(attr, domain.OpenSetCategoricalAttribute): + elif isinstance(attr, domain_mod.OpenSetCategoricalAttribute): initializers[col] = initialization.OpenSetInitializerConfig( name=col, attribute=attr, @@ -88,14 +90,15 @@ class ColumnCodec: """ column_measurement: initialization.ColumnMeasurement - attribute: domain.AttributeType + attribute: domain_mod.AttributeType def encode(self, values: np.ndarray) -> np.ndarray: """Encodes raw column values to discrete integer ids.""" if self.column_measurement.bin_edges is not None: return vtx.discretize( - # pyrefly: ignore[bad-argument-type] - values, self.column_measurement.bin_edges, self.attribute + values, + self.column_measurement.bin_edges, + typing.cast(domain_mod.NumericalAttribute, self.attribute), ) return vtx.discrete_encode( values, self.column_measurement.categorical_attribute @@ -105,8 +108,10 @@ def decode(self, ids: np.ndarray, rng: np.random.Generator) -> np.ndarray: """Decodes synthetic discrete ids back to the original domain.""" if self.column_measurement.bin_edges is not None: return vtx.undiscretize( - # pyrefly: ignore[bad-argument-type] - ids, self.column_measurement.bin_edges, self.attribute, rng=rng + ids, + self.column_measurement.bin_edges, + typing.cast(domain_mod.NumericalAttribute, self.attribute), + rng=rng, ) return vtx.discrete_decode( ids, self.column_measurement.categorical_attribute @@ -127,7 +132,7 @@ class TabularCodec: def from_measurements( cls, results: Mapping[str, initialization.ColumnMeasurement], - domains: Mapping[str, domain.AttributeType], + domains: Mapping[str, domain_mod.AttributeType], ) -> TabularCodec: """Builds a codec from initialization results and the original domains.""" columns = {col: ColumnCodec(m, domains[col]) for col, m in results.items()} @@ -198,11 +203,11 @@ class TabularMechanism(api.CalibratedMechanism): """ config: TabularConfig - domains: Mapping[str, domain.AttributeType] + domains: Mapping[str, domain_mod.AttributeType] base_mechanism: discrete_mechanisms.CalibratedMechanism initializers: dict[str, api.CalibratedMechanism] total_count_sigma: float = dataclasses.field(repr=False) - cross_attribute_constraints: Sequence[constraints.Constraint] = () + cross_attribute_constraints: Sequence[constraints_mod.Constraint] = () max_records_per_user: int = 1 @property @@ -226,7 +231,7 @@ def __call__( rng: np.random.Generator, data: pd.DataFrame, *, - cross_attribute_constraints: Sequence[constraints.Constraint] = (), + cross_attribute_constraints: Sequence[constraints_mod.Constraint] = (), ) -> DataGenerationResult: """Generates differentially private synthetic data. @@ -248,9 +253,11 @@ def __call__( f'{col=} not found in dataset. Available: {list(data.columns)}' ) if not cross_attribute_constraints: - cross_attribute_constraints = self.config.cross_attribute_constraints + cross_attribute_constraints = self.cross_attribute_constraints - mbi_constraints = tuple(c.to_mbi() for c in cross_attribute_constraints) + mbi_constraints = tuple( + c.to_mbi(self.domains) for c in cross_attribute_constraints + ) # Phase 1: Per-column initialization. # Measure total count first, then run per-column initializers. @@ -316,35 +323,42 @@ class TabularConfig(api.MechanismConfig): Usage:: - config = TabularConfig(domains=domains) - calibrated = config.configure(zcdp_rho=1.0) + config = TabularConfig() + calibrated = config.calibrate(schema=schema, epsilon=1.0, delta=1e-5) result = calibrated(rng, df) synthetic_df = result.synthetic_data Attributes: - domains: Mapping from column names to attribute domain specifications. discrete_mechanism: The mechanism to run on the discretized data. numerical_bins: Number of bins for numerical attribute discretization. init_budget_fraction: Fraction of total zCDP budget allocated to per-column initialization (the rest goes to the discrete mechanism). cross_attribute_constraints: Constraints to enforce on generated data. + schema: Schema or mapping from column names to attribute domain + specifications. Optional; can be supplied at calibration time via + ``configure(schema=...)`` or ``calibrate(schema=...)``. + domains: Alias for ``schema`` for backward compatibility. """ - domains: Mapping[str, domain.AttributeType] discrete_mechanism: api.MechanismConfig = discrete_mechanisms.MSTConfig() numerical_bins: int = 32 init_budget_fraction: float = 0.1 - initializers: dict[str, api.MechanismConfig] | None = None - cross_attribute_constraints: Sequence[constraints.Constraint] = () - - def _compute_per_col_deltas(self, delta): + cross_attribute_constraints: Sequence[constraints_mod.Constraint] = () + schema: domain_mod.Schema | Mapping[str, domain_mod.AttributeType] | None = ( + None + ) + domains: Mapping[str, domain_mod.AttributeType] | None = None + + def _compute_per_col_deltas( + self, schema: Mapping[str, domain_mod.AttributeType], delta: float + ) -> dict[str, float]: # Split delta across open-set columns, analogous to splitting zcdp_rho. # Under calibrate(), any delta not consumed here is automatically # available for the zCDP-to-(epsilon, delta) conversion, so this # simple additive split is tight. num_open_set = sum( - isinstance(attr, domain.OpenSetCategoricalAttribute) - for attr in self.domains.values() + isinstance(attr, domain_mod.OpenSetCategoricalAttribute) + for attr in schema.values() ) if num_open_set > 0 and delta <= 0: raise ValueError( @@ -355,8 +369,8 @@ def _compute_per_col_deltas(self, delta): thresholding_delta = self.init_budget_fraction * delta per_col_deltas = {} - for col in self.domains: - if isinstance(self.domains[col], domain.OpenSetCategoricalAttribute): + for col in schema: + if isinstance(schema[col], domain_mod.OpenSetCategoricalAttribute): per_col_deltas[col] = thresholding_delta / num_open_set else: per_col_deltas[col] = 0.0 @@ -366,60 +380,56 @@ def configure( self, *, zcdp_rho: float, + schema: ( + domain_mod.Schema | Mapping[str, domain_mod.AttributeType] | None + ) = None, + domain: ( + domain_mod.Schema | Mapping[str, domain_mod.AttributeType] | None + ) = None, + constraints: Sequence[constraints_mod.Constraint] | None = None, delta: float = 0.0, max_records_per_user: int = 1, ) -> TabularMechanism: - """Returns a calibrated mechanism configured with the given privacy budget. - - Splits the budget additively, just as it does for ``zcdp_rho``: - - - ``init_budget_fraction`` of ``zcdp_rho`` goes to per-column initializers - (split evenly, including a total-count mechanism); the remainder goes to - the discrete mechanism. - - ``init_budget_fraction`` of ``delta`` is reserved for open-set partition - selection (split evenly across open-set columns); the remaining delta is - unused by pure-zCDP sub-mechanisms. - - When ``calibrate(epsilon, delta)`` is called, the base class binary search - passes the guarantee delta here. Because the thresholding delta is honestly - reported in the composite ``dp_event``, the binary search automatically - ensures the overall (epsilon, delta) guarantee is tight. - - Args: - zcdp_rho: The zCDP privacy budget. - delta: Overall approximate DP delta for the mechanism. A fraction - (``init_budget_fraction``) is allocated to partition selection for - open-set columns. Must be positive when open-set categorical attributes - are present. - max_records_per_user: Assumed upper bound on the number of records a - single user contributes. Values greater than 1 scale the added noise - (and mechanism sensitivity) to provide user-level rather than - record-level DP; the privacy accounting is unchanged. This bound is NOT - enforced -- soundness relies on the caller guaranteeing it via - preprocessing. - - Returns: - A calibrated TabularMechanism ready to be run on tabular data. - - Raises: - ValueError: If open-set attributes exist but delta is 0. - """ + """Returns a calibrated mechanism configured with the given privacy budget.""" api.validate_max_records_per_user(max_records_per_user) - per_col_deltas = self._compute_per_col_deltas(delta) + if schema is not None: + resolved_schema = schema + elif domain is not None: + resolved_schema = domain + elif self.schema is not None: + resolved_schema = self.schema + elif self.domains is not None: + resolved_schema = self.domains + else: + raise ValueError( + "Must provide 'schema' to configure() or in TabularConfig" + ' constructor.' + ) - inits = ( - self.initializers - if self.initializers is not None - else create_initializers(self.domains, self.numerical_bins) + attr_schema = ( + resolved_schema.attributes + if isinstance(resolved_schema, Schema) + else resolved_schema ) + schema_constraints = ( + resolved_schema.constraints + if isinstance(resolved_schema, Schema) + else () + ) + + resolved_constraints = ( + constraints + if constraints is not None + else (self.cross_attribute_constraints or schema_constraints) + ) + + per_col_deltas = self._compute_per_col_deltas(attr_schema, delta) + inits = create_initializers(attr_schema, self.numerical_bins) init_rho = self.init_budget_fraction * zcdp_rho - # +1 for the DPGaussianCount that always measures the total. per_col_rho = init_rho / (len(inits) + 1) discrete_rho = zcdp_rho - init_rho - calibrated_inits: dict[str, api.CalibratedMechanism] - - calibrated_inits = { + calibrated_inits: dict[str, api.CalibratedMechanism] = { col: init.configure( zcdp_rho=per_col_rho, delta=per_col_deltas[col], @@ -436,10 +446,11 @@ def configure( return TabularMechanism( config=self, - domains=self.domains, + domains=attr_schema, base_mechanism=calibrated_discrete, initializers=calibrated_inits, total_count_sigma=total_count_sigma, + cross_attribute_constraints=resolved_constraints, max_records_per_user=max_records_per_user, ) diff --git a/dpsynth/domain.py b/dpsynth/domain.py index d485cd1d..1070a216 100644 --- a/dpsynth/domain.py +++ b/dpsynth/domain.py @@ -51,7 +51,6 @@ from typing import Any, Literal, TypeAlias -from absl import logging import numpy as np import yaml @@ -304,42 +303,167 @@ class FreeFormTextAttribute: | FreeFormTextAttribute ) -Schema: TypeAlias = Mapping[str, AttributeType] +@dataclasses.dataclass(frozen=True, eq=False) +class Schema(Mapping[str, AttributeType]): + """Schema defining attribute domains and optional cross-attribute constraints. -def to_yaml_file(domain: Mapping[str, AttributeType], filepath: str | PathType): - """Writes a dictionary of Attribute objects to a YAML file.""" - yaml_data = {} - for name, attr_obj in domain.items(): - attr_data = dataclasses.asdict(attr_obj) - attr_data['type'] = attr_obj.__class__.__name__ - yaml_data[name] = attr_data - with open(filepath, 'w') as f: - yaml.dump(yaml_data, f, default_flow_style=False) + Implements ``collections.abc.Mapping[str, AttributeType]`` so it can be + indexed like a dictionary (e.g. ``schema['col']``, ``'col' in schema``, + ``len(schema)``, ``for col in schema``). + Attributes: + attributes: Mapping from column names to attribute domain specifications. + constraints: Cross-attribute constraints associated with this schema. + """ -def from_yaml_file(filepath: str | PathType) -> Mapping[str, AttributeType]: - """Reads a dictionary of Attribute objects from a YAML file.""" - with open(filepath, 'r') as f: - yaml_data = yaml.safe_load(f) - domain = {} - - for name, attr_data in yaml_data.items(): - attr_type = attr_data.pop('type', None) - if attr_type is None: - logging.warning( - 'Field "type" missing in domain YAML; re-save using `to_yaml_file`.' - 'In the future, missing this field will raise an error.' + attributes: Mapping[str, AttributeType] + constraints: Sequence[Any] = () + + def __getitem__(self, key: str) -> AttributeType: + return self.attributes[key] + + def __contains__(self, key: object) -> bool: + return key in self.attributes + + def __iter__(self) -> Any: + return iter(self.attributes) + + def __len__(self) -> int: + return len(self.attributes) + + def __eq__(self, other: object) -> bool: + if isinstance(other, Schema): + return ( + self.attributes == other.attributes + and self.constraints == other.constraints ) - if 'possible_values' in attr_data: - domain[name] = CategoricalAttribute(**attr_data) - elif 'min_value' in attr_data: - domain[name] = NumericalAttribute(**attr_data) - elif 'max_tokens' in attr_data: - domain[name] = FreeFormTextAttribute(**attr_data) - elif 'default_value' in attr_data or not attr_data: - domain[name] = OpenSetCategoricalAttribute(**attr_data) - else: - raise ValueError(f'Invalid YAML data for attribute: {name}') + if isinstance(other, Mapping): + return not self.constraints and self.attributes == other + return False + + def to_dict(self) -> dict[str, Any]: + """Converts the Schema into a serializable dictionary.""" + attrs_dict = schema_to_dict(self.attributes) + if not self.constraints: + return attrs_dict + return { + 'attributes': attrs_dict, + 'constraints': [ + c.to_dict() if hasattr(c, 'to_dict') else c + for c in self.constraints + ], + } + + def to_yaml(self) -> str: + """Serializes the Schema to a YAML string.""" + return yaml.dump(self.to_dict(), default_flow_style=False, sort_keys=False) + + def to_yaml_file(self, filepath: str | PathType) -> None: + """Writes the Schema to a YAML file.""" + with open(filepath, 'w') as f: + f.write(self.to_yaml()) + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> Schema: + """Constructs a Schema from a dictionary.""" + if isinstance(data, Schema): + return data + if 'attributes' in data and isinstance(data['attributes'], dict): + attrs = schema_from_dict(data['attributes']) + constraints_list = [] + if 'constraints' in data and data['constraints']: + import importlib # pylint: disable=g-import-not-at-top + + constraints_mod = importlib.import_module('dpsynth.constraints') + for c in data['constraints']: + if isinstance(c, dict): + constraints_list.append(constraints_mod.Constraint.from_dict(c)) + else: + constraints_list.append(c) + return cls(attributes=attrs, constraints=tuple(constraints_list)) + attrs = schema_from_dict(data) + return cls(attributes=attrs) + + @classmethod + def from_yaml(cls, yaml_str: str) -> Schema: + """Deserializes a Schema from a YAML string.""" + yaml_data = yaml.safe_load(yaml_str) + return cls.from_dict(yaml_data) + + @classmethod + def from_yaml_file(cls, filepath: str | PathType) -> Schema: + """Reads a Schema from a YAML file.""" + with open(filepath, 'r') as f: + return cls.from_yaml(f.read()) + + +def attribute_to_dict(attr: AttributeType) -> dict[str, Any]: + """Converts an AttributeType object to a dictionary.""" + attr_data = dataclasses.asdict(attr) + attr_data['type'] = attr.__class__.__name__ + return attr_data + + +def attribute_from_dict(attr_data: Mapping[str, Any]) -> AttributeType: + """Instantiates an AttributeType from a dictionary.""" + data = dict(attr_data) + attr_type = data.pop('type', None) + if attr_type == 'CategoricalAttribute' or ( + attr_type is None and 'possible_values' in data + ): + return CategoricalAttribute(**data) + elif attr_type == 'NumericalAttribute' or ( + attr_type is None and 'min_value' in data + ): + return NumericalAttribute(**data) + elif attr_type == 'FreeFormTextAttribute' or ( + attr_type is None and 'max_tokens' in data + ): + return FreeFormTextAttribute(**data) + elif attr_type == 'OpenSetCategoricalAttribute' or ( + attr_type is None and ('default_value' in data or not data) + ): + return OpenSetCategoricalAttribute(**data) + else: + raise ValueError(f'Invalid data for attribute: {attr_data}') + + +def schema_to_dict( + domain: Mapping[str, AttributeType], +) -> dict[str, dict[str, Any]]: + """Converts a mapping of column names to AttributeType to a dictionary.""" + return {name: attribute_to_dict(attr) for name, attr in domain.items()} + + +def schema_from_dict(data: Mapping[str, Any]) -> dict[str, AttributeType]: + """Constructs a mapping of column names to AttributeType from a dictionary.""" + return { + name: attribute_from_dict(attr_data) for name, attr_data in data.items() + } + + +def to_yaml(domain: Schema | Mapping[str, AttributeType]) -> str: + """Serializes a domain dictionary or Schema to a YAML string.""" + if isinstance(domain, Schema): + return domain.to_yaml() + return yaml.dump(schema_to_dict(domain), default_flow_style=False) + + +def from_yaml(yaml_str: str) -> Schema: + """Deserializes a domain dictionary or Schema from a YAML string.""" + return Schema.from_yaml(yaml_str) + + +def to_yaml_file( + domain: Schema | Mapping[str, AttributeType], filepath: str | PathType +): + """Writes a dictionary of Attribute objects or Schema to a YAML file.""" + with open(filepath, 'w') as f: + f.write(to_yaml(domain)) + - return domain +def from_yaml_file(filepath: str | PathType) -> Schema: + """Reads a dictionary of Attribute objects or Schema from a YAML file.""" + with open(filepath, 'r') as f: + return from_yaml(f.read()) diff --git a/dpsynth/relational/domain.py b/dpsynth/relational/domain.py index e4235ccc..023c45d9 100644 --- a/dpsynth/relational/domain.py +++ b/dpsynth/relational/domain.py @@ -189,10 +189,10 @@ def from_dict( for table_name, table_schema in config['tables'].items(): if not isinstance(table_schema, Mapping): raise ValueError(f'Table schema for {table_name!r} must be a mapping.') - table_domains[table_name] = { + table_domains[table_name] = domain.Schema({ col_name: _parse_attribute(table_name, col_name, spec) for col_name, spec in table_schema.items() - } + }) foreign_keys: list[ForeignKeyRelation] = [] for fk in config.get('foreign_keys', []): diff --git a/pyproject.toml b/pyproject.toml index 08ae5868..08998c05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ keywords = [] dependencies = [ "absl-py", "attrs", + "etils[epath]", "numpy", "pandas", "pydantic>=2.0", diff --git a/tests/api_test.py b/tests/api_test.py new file mode 100644 index 00000000..003ccba5 --- /dev/null +++ b/tests/api_test.py @@ -0,0 +1,195 @@ +# 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 MechanismConfig serialization in api.py.""" + +import os +from absl.testing import absltest +from absl.testing import parameterized +import dpsynth +from dpsynth import api +from dpsynth import constraints +from dpsynth import data_generation_v3 +from dpsynth import domain +from dpsynth import relational +from dpsynth.discrete_mechanisms import aim +from dpsynth.discrete_mechanisms import aim_gdp +from dpsynth.discrete_mechanisms import direct +from dpsynth.discrete_mechanisms import discrete +from dpsynth.discrete_mechanisms import independent +from dpsynth.discrete_mechanisms import mst +from dpsynth.discrete_mechanisms import swift + + +class ApiYamlSerializationTest(parameterized.TestCase): + + def test_mst_config_roundtrip(self): + config = mst.MSTConfig( + pgm_iters=2500, + select_budget_fraction=0.75, + maximum_marginal_size=5_000_000, + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + self.assertEqual(mst.MSTConfig.from_yaml(yaml_str), config) + + def test_aim_config_roundtrip(self): + config = aim.AIMConfig( + max_rounds=25, + max_model_size=50, + max_marginal_size=2e6, + anneal_factor=3.0, + select_budget_fraction=0.2, + pgm_iters=500, + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_swift_config_roundtrip(self): + config = swift.SWIFTConfig( + max_clique_size=5e6, + max_marginal_size=2e6, + pgm_iters=8000, + select_budget_frac=0.15, + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_independent_config_roundtrip(self): + config = independent.IndependentConfig(pgm_iters=3000) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_direct_config_roundtrip(self): + config = direct.DirectConfig( + pgm_iters=4000, + prespecified_marginal_queries=[('a', 'b'), ('c',)], + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_aim_gdp_config_roundtrip(self): + config = aim_gdp.AIMGDPConfig( + pgm_iters=500, + max_rounds=10, + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_discrete_config_roundtrip(self): + config = discrete.DiscreteConfig( + mechanism=aim.AIMConfig(pgm_iters=400), + compress_columns=['col1', 'col2'], + one_way_budget_fraction=0.2, + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_tabular_config_pure_preset_roundtrip(self): + config = data_generation_v3.TabularConfig( + discrete_mechanism=mst.MSTConfig(pgm_iters=1500), + numerical_bins=64, + init_budget_fraction=0.15, + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_tabular_config_with_schema_roundtrip(self): + schema = domain.Schema( + attributes={ + 'age': domain.NumericalAttribute(min_value=0, max_value=120), + 'gender': domain.CategoricalAttribute(possible_values=['M', 'F']), + 'notes': domain.FreeFormTextAttribute(max_tokens=100), + 'state': domain.OpenSetCategoricalAttribute(), + }, + constraints=( + constraints.Constraint( + attribute_names=('gender',), + impossible_combinations=[('X',)], + ), + ), + ) + config = data_generation_v3.TabularConfig( + schema=schema, + discrete_mechanism=mst.MSTConfig(pgm_iters=1000), + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_multitable_config_roundtrip(self): + config = relational.MultiTableConfig( + domains={ + 'users': { + 'age': domain.NumericalAttribute(min_value=18, max_value=80), + }, + 'orders': { + 'amount': domain.NumericalAttribute(min_value=0, max_value=100), + }, + }, + foreign_keys=[ + relational.ForeignKeyRelation( + parent_table='users', + parent_primary_key='user_id', + child_table='orders', + child_foreign_key='user_id', + max_children_per_parent=5, + ), + ], + discrete_mechanism=mst.MSTConfig(pgm_iters=500), + ) + yaml_str = config.to_yaml() + loaded = api.MechanismConfig.from_yaml(yaml_str) + self.assertEqual(loaded, config) + + def test_file_io_roundtrip(self): + config = data_generation_v3.TabularConfig( + domains={'age': domain.NumericalAttribute(min_value=0, max_value=100)}, + discrete_mechanism=mst.MSTConfig(pgm_iters=1000), + ) + tmp_path = os.path.join(self.create_tempdir().full_path, 'config.yaml') + config.to_yaml_file(tmp_path) + loaded = api.MechanismConfig.from_yaml_file(tmp_path) + self.assertEqual(loaded, config) + + def test_module_level_helpers(self): + config = mst.MSTConfig(pgm_iters=1234) + yaml_str = api.to_yaml(config) + self.assertEqual(api.from_yaml(yaml_str), config) + self.assertEqual(dpsynth.to_yaml(config), yaml_str) + self.assertEqual(dpsynth.from_yaml(yaml_str), config) + + tmp_path = os.path.join(self.create_tempdir().full_path, 'helper.yaml') + dpsynth.to_yaml_file(config, tmp_path) + self.assertEqual(dpsynth.from_yaml_file(tmp_path), config) + + def test_missing_type_raises(self): + with self.assertRaises(ValueError): + api.MechanismConfig.from_yaml('pgm_iters: 1000\n') + + def test_unknown_type_raises(self): + with self.assertRaises(ValueError): + api.MechanismConfig.from_yaml('type: NonExistentConfig\n') + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/constraints_test.py b/tests/constraints_test.py index 1f71c8a0..77a81d7f 100644 --- a/tests/constraints_test.py +++ b/tests/constraints_test.py @@ -138,6 +138,34 @@ def test_functional_dependency(self): self.assertEqual(vals[0, 1], -np.inf) self.assertEqual(vals[2, 0], -np.inf) + def test_bind_schema_and_to_mbi_with_schema(self): + schema = domain.Schema({ + 'Software': self.software, + 'OS': self.os, + }) + c = constraints.Constraint( + attribute_names=('Software', 'OS'), + possible_combinations=[ + ('GameSuite', 'Windows'), + ('DevTool', 'Linux'), + ], + ) + bound_c = c.bind_schema(schema) + self.assertEqual(bound_c.attribute_domains, (self.software, self.os)) + + mbi_c = c.to_mbi(schema) + self.assertIsInstance(mbi_c, mbi.Constraint) + + def test_dict_roundtrip(self): + c = constraints.Constraint( + attribute_names=('Software', 'OS'), + attribute_domains=(self.software, self.os), + impossible_combinations=[('GameSuite', 'Linux')], + ) + d = c.to_dict() + loaded = constraints.Constraint.from_dict(d) + self.assertEqual(loaded, c) + if __name__ == '__main__': absltest.main() diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py index 9bb04912..888e225b 100644 --- a/tests/data_generation_v3_test.py +++ b/tests/data_generation_v3_test.py @@ -17,6 +17,7 @@ from absl.testing import absltest from absl.testing import parameterized import dp_accounting +from dpsynth import constraints from dpsynth import data_generation_v3 from dpsynth import discrete_mechanisms from dpsynth import domain @@ -403,14 +404,55 @@ def test_open_set_with_k_supported(self): synthetic_df = mech(np.random.default_rng(0), df).synthetic_data self.assertListEqual(synthetic_df.columns.tolist(), ['A']) - def test_custom_initializers_inherit_k(self): + def test_initializers_inherit_k(self): domains = self._categorical_domains() - inits = data_generation_v3.create_initializers(domains, 32) - config = TabularConfig(domains=domains, initializers=inits) + config = TabularConfig(domains=domains) calibrated = config.configure(zcdp_rho=100.0, max_records_per_user=2) for init in calibrated.initializers.values(): self.assertEqual(init.max_records_per_user, 2) + def test_pure_preset_configure_and_calibrate_with_schema(self): + schema = domain.Schema({ + 'A': domain.CategoricalAttribute(possible_values=['a', 'b', 'c']), + 'B': domain.NumericalAttribute(min_value=0, max_value=10), + }) + df = pd.DataFrame({'A': ['a', 'b', 'c'], 'B': [1.0, 5.0, 10.0]}) + rng = np.random.default_rng(0) + + # Pure preset with no schema in constructor + preset = TabularConfig(numerical_bins=16) + + # 1. configure(schema=...) + calibrated = preset.configure(schema=schema, zcdp_rho=100.0) + result = calibrated(rng, df).synthetic_data + self.assertListEqual(result.columns.tolist(), ['A', 'B']) + + # 2. calibrate(schema=...) + calibrated2 = preset.calibrate(schema=schema, epsilon=1.0, delta=1e-5) + result2 = calibrated2(rng, df).synthetic_data + self.assertListEqual(result2.columns.tolist(), ['A', 'B']) + + def test_configure_with_schema_constraints(self): + c = constraints.Constraint( + attribute_names=('A', 'B'), + possible_combinations=[('a', 'x'), ('b', 'y')], + ) + schema = domain.Schema( + attributes={ + 'A': domain.CategoricalAttribute(possible_values=['a', 'b']), + 'B': domain.CategoricalAttribute(possible_values=['x', 'y']), + }, + constraints=(c,), + ) + df = pd.DataFrame({'A': ['a', 'b'], 'B': ['x', 'y']}) + rng = np.random.default_rng(0) + + preset = TabularConfig() + calibrated = preset.configure(schema=schema, zcdp_rho=100.0) + synthetic_df = calibrated(rng, df).synthetic_data + for _, row in synthetic_df.iterrows(): + self.assertIn((row['A'], row['B']), [('a', 'x'), ('b', 'y')]) + @parameterized.named_parameters(('zero', 0), ('negative', -3)) def test_invalid_k_raises(self, k): config = TabularConfig(domains=self._categorical_domains()) diff --git a/tests/domain_test.py b/tests/domain_test.py index 6e98455c..be8b0a84 100644 --- a/tests/domain_test.py +++ b/tests/domain_test.py @@ -15,6 +15,7 @@ import math from absl.testing import absltest +from dpsynth import constraints from dpsynth import domain import numpy as np @@ -211,6 +212,60 @@ def test_open_set_yaml_roundtrip(self): loaded_domain = domain.from_yaml_file(temp_file.full_path) self.assertEqual(loaded_domain, original_domain) + def test_to_from_yaml_string_roundtrip(self): + original_domain = { + 'cat': domain.CategoricalAttribute(possible_values=['A', 'B']), + 'num': domain.NumericalAttribute(min_value=0, max_value=10), + } + yaml_str = domain.to_yaml(original_domain) + loaded = domain.from_yaml(yaml_str) + self.assertEqual(loaded, original_domain) + + def test_schema_mapping_interface(self): + attrs = { + 'age': domain.NumericalAttribute(min_value=0, max_value=100), + 'gender': domain.CategoricalAttribute(possible_values=['M', 'F']), + } + schema = domain.Schema(attributes=attrs) + self.assertEqual(schema['age'], attrs['age']) + self.assertIn('gender', schema) + self.assertNotIn('unknown', schema) + self.assertLen(schema, 2) + self.assertEqual(list(schema), ['age', 'gender']) + self.assertEqual(schema.get('age'), attrs['age']) + self.assertIsNone(schema.get('unknown')) + + def test_schema_with_constraints_yaml_roundtrip(self): + attrs = { + 'age': domain.NumericalAttribute(min_value=0, max_value=100), + 'gender': domain.CategoricalAttribute(possible_values=['M', 'F']), + } + c = constraints.Constraint( + attribute_names=('gender',), + impossible_combinations=[('X',)], + ) + schema = domain.Schema(attributes=attrs, constraints=(c,)) + yaml_str = schema.to_yaml() + loaded = domain.Schema.from_yaml(yaml_str) + self.assertEqual(loaded.attributes, schema.attributes) + self.assertLen(loaded.constraints, 1) + self.assertEqual(loaded.constraints[0].attribute_names, c.attribute_names) + self.assertEqual( + loaded.constraints[0].impossible_combinations, c.impossible_combinations + ) + + def test_schema_from_legacy_yaml(self): + yaml_str = ( + 'age:\n type: NumericalAttribute\n min_value: 0\n max_value: 100\n' + ) + schema = domain.Schema.from_yaml(yaml_str) + self.assertIsInstance(schema, domain.Schema) + self.assertIn('age', schema) + self.assertEqual( + schema['age'], domain.NumericalAttribute(min_value=0, max_value=100) + ) + self.assertEqual(schema.constraints, ()) + if __name__ == '__main__': absltest.main()