diff --git a/dpsynth/__init__.py b/dpsynth/__init__.py index 5d87272..6e24f6f 100644 --- a/dpsynth/__init__.py +++ b/dpsynth/__init__.py @@ -20,9 +20,7 @@ from dpsynth import constraints from dpsynth import discrete_mechanisms from dpsynth import domain -from dpsynth.api import CalibratedMechanism -from dpsynth.api import MechanismConfig -from dpsynth.constraints import Constraint +from dpsynth import relational from dpsynth.data_generation_v3 import TabularConfig from dpsynth.data_generation_v3 import TabularMechanism from dpsynth.data_generation_v3 import TabularSynthesizer @@ -30,3 +28,25 @@ from dpsynth.domain import FreeFormTextAttribute from dpsynth.domain import NumericalAttribute from dpsynth.domain import OpenSetCategoricalAttribute + +ForeignKeyRelation = relational.ForeignKeyRelation +MultiDataGenerationResult = relational.MultiDataGenerationResult +MultiTableConfig = relational.MultiTableConfig +MultiTableMechanism = relational.MultiTableMechanism + +__all__ = [ + 'CategoricalAttribute', + 'ForeignKeyRelation', + 'MultiDataGenerationResult', + 'MultiTableConfig', + 'MultiTableMechanism', + 'NumericalAttribute', + 'OpenSetCategoricalAttribute', + 'TabularConfig', + 'TabularMechanism', + 'TabularSynthesizer', + 'api', + 'discrete_mechanisms', + 'domain', + 'relational', +] diff --git a/dpsynth/relational/__init__.py b/dpsynth/relational/__init__.py index e692cbe..efe45d2 100644 --- a/dpsynth/relational/__init__.py +++ b/dpsynth/relational/__init__.py @@ -17,9 +17,6 @@ # pylint: disable=g-importing-member from dpsynth.relational.domain import ForeignKeyRelation -from dpsynth.relational.domain import from_dict -from dpsynth.relational.domain import from_yaml_file -from dpsynth.relational.domain import topological_sort_hierarchy from dpsynth.relational.synthesizer import MultiDataGenerationResult from dpsynth.relational.synthesizer import MultiTableConfig from dpsynth.relational.synthesizer import MultiTableMechanism @@ -29,7 +26,4 @@ 'MultiDataGenerationResult', 'MultiTableConfig', 'MultiTableMechanism', - 'from_dict', - 'from_yaml_file', - 'topological_sort_hierarchy', ] diff --git a/dpsynth/relational/synthesizer.py b/dpsynth/relational/synthesizer.py index a86ed1c..8ecbfb1 100644 --- a/dpsynth/relational/synthesizer.py +++ b/dpsynth/relational/synthesizer.py @@ -675,11 +675,11 @@ def _synthesize_relational_link( ) assert hasattr(mech_res, 'measurements') - wide_measurements = post_processing.symmetrize_to_wide_domain( - measurements=mech_res.measurements, - max_children_per_parent=fk_relation.max_children_per_parent, - num_permutation_slots=num_permutation_slots, - ) + # wide_measurements = post_processing.symmetrize_to_wide_domain( + # measurements=mech_res.measurements, + # max_children_per_parent=fk_relation.max_children_per_parent, + # num_permutation_slots=num_permutation_slots, + # ) wide_domain = transformations.build_exploration_domain( parent_domain=parent_dataset.domain, @@ -703,11 +703,16 @@ def _synthesize_relational_link( else max(1, int(round(noisy_root_total))) ) + pgm_iters = getattr( + getattr(discrete_mechanism, 'config', None), 'pgm_iters', 5000 + ) + synth_wide_records = _fit_and_sample_wide_link_mrf( wide_domain=wide_domain, - wide_measurements=wide_measurements, + wide_measurements=mech_res.measurements, # wide_measurements wide_constraints=wide_constraints, num_rows=num_rows, + iters=pgm_iters, ) synth_parent_dataset: mbi.Dataset | None = None diff --git a/examples/relational/california_census/domain.yaml b/examples/relational/california_census/domain.yaml new file mode 100644 index 0000000..bde2aa2 --- /dev/null +++ b/examples/relational/california_census/domain.yaml @@ -0,0 +1,449 @@ +tables: + household: + FARM: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + OWNERSHP: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + ACREHOUS: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + TAXINCL: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + PROPINSR: + type: NumericalAttribute + clip_to_range: true + dtype: int + min_value: 0.0 + max_value: 59.0 + COSTELEC: + type: NumericalAttribute + clip_to_range: true + dtype: int + min_value: 0.0 + max_value: 99.0 + VALUEH: + type: NumericalAttribute + clip_to_range: true + dtype: int + min_value: 0.0 + max_value: 99.0 + ROOMS: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + - 8 + - 9 + PLUMBING: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + PUMA: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + - 8 + - 9 + - 10 + - 11 + - 12 + - 13 + - 14 + - 15 + - 16 + - 17 + - 18 + - 19 + - 20 + - 21 + - 22 + - 23 + - 24 + - 25 + - 26 + - 27 + - 28 + - 29 + - 30 + - 31 + - 32 + - 33 + - 34 + - 35 + - 36 + - 37 + - 38 + - 39 + - 40 + - 41 + - 42 + - 43 + - 44 + - 45 + - 46 + - 47 + - 48 + - 49 + - 50 + - 51 + - 52 + - 53 + - 54 + - 55 + - 56 + - 57 + - 58 + - 59 + - 60 + - 61 + - 62 + - 63 + - 64 + - 65 + - 66 + - 67 + - 68 + - 69 + - 70 + - 71 + - 72 + - 73 + - 74 + - 75 + - 76 + - 77 + - 78 + - 79 + - 80 + - 81 + - 82 + - 83 + - 84 + - 85 + - 86 + - 87 + - 88 + - 89 + - 90 + - 91 + - 92 + - 93 + - 94 + - 95 + - 96 + - 97 + - 98 + - 99 + - 100 + - 101 + - 102 + - 103 + - 104 + - 105 + - 106 + - 107 + - 108 + - 109 + - 110 + - 111 + - 112 + - 113 + - 114 + - 115 + - 116 + - 117 + - 118 + - 119 + - 120 + - 121 + - 122 + - 123 + - 124 + - 125 + - 126 + - 127 + - 128 + - 129 + - 130 + - 131 + - 132 + - 133 + - 134 + - 135 + - 136 + - 137 + - 138 + - 139 + - 140 + - 141 + - 142 + - 143 + - 144 + - 145 + - 146 + - 147 + - 148 + - 149 + - 150 + - 151 + - 152 + - 153 + - 154 + - 155 + - 156 + - 157 + - 158 + - 159 + - 160 + - 161 + - 162 + - 163 + - 164 + - 165 + - 166 + - 167 + - 168 + - 169 + - 170 + - 171 + - 172 + - 173 + - 174 + - 175 + - 176 + - 177 + - 178 + - 179 + - 180 + - 181 + - 182 + - 183 + - 184 + - 185 + - 186 + - 187 + - 188 + - 189 + - 190 + - 191 + - 192 + - 193 + - 194 + - 195 + - 196 + - 197 + - 198 + - 199 + - 200 + - 201 + - 202 + - 203 + - 204 + - 205 + - 206 + - 207 + - 208 + - 209 + - 210 + - 211 + - 212 + - 213 + - 214 + - 215 + - 216 + - 217 + - 218 + - 219 + - 220 + - 221 + - 222 + - 223 + - 224 + - 225 + - 226 + - 227 + - 228 + - 229 + - 230 + - 231 + - 232 + + individual: + RELATE: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + SEX: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + AGE: + type: NumericalAttribute + clip_to_range: true + dtype: int + min_value: 0.0 + max_value: 85.0 + MARST: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + RACE: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + - 8 + CITIZEN: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + SPEAKENG: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + SCHOOL: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + EDUC: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + - 8 + - 9 + - 10 + GRADEATT: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + SCHLTYPE: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + EMPSTAT: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + - 3 + CLASSWKR: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + INCTOT: + type: NumericalAttribute + clip_to_range: true + dtype: int + min_value: 0.0 + max_value: 99.0 + DISABWRK: + type: CategoricalAttribute + out_of_domain_index: 0 + possible_values: + - 0 + - 1 + - 2 + +foreign_keys: + - parent_table: household + parent_primary_key: HOUSEHOLD + child_table: individual + child_foreign_key: HOUSEHOLD + max_children_per_parent: 8 diff --git a/examples/relational/california_census/example_california_census.py b/examples/relational/california_census/example_california_census.py new file mode 100644 index 0000000..3edb5bc --- /dev/null +++ b/examples/relational/california_census/example_california_census.py @@ -0,0 +1,232 @@ +# 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. + +"""Example demonstrating Differentially Private synthesis on California Census data. + +This script demonstrates how to synthesize multi-table relational data under +Differential Privacy (DP) using DPSynth on the California Census (PUMS) dataset. + +Pipeline steps: +1. Load relational schema and domain constraints from domain.yaml. +2. Load parent (household) and child (individual) tables from CSV storage. +3. Configure and calibrate MultiTableConfig with Differential Privacy +parameters. +4. Execute relational synthesis mechanism. +5. Validate relational integrity (primary key uniqueness, foreign key linkage). +6. Evaluate statistical utility (Total Variation Distance) and ML utility +(TSTR). +7. Generate SDMetrics multi-table quality and diagnostic reports. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import os +from typing import Literal + +from absl import app +from absl import logging +from dpsynth import discrete_mechanisms +from dpsynth import domain +from dpsynth.examples.relational.california_census import example_eval_california_census as eval_utils +from dpsynth.relational import domain as rel_domain +from dpsynth.relational import synthesizer as rel_synth +from etils import epath +import numpy as np +import pandas as pd + +# ============================================================================== +# Configuration Constants +# ============================================================================== +DATA_DIR: str = './data/California' +OUTPUT_DIR: str = './data/California/synthetic' +DOMAIN_PATH: epath.Path = ( + epath.resource_path('dpsynth') + / 'examples/relational/california_census/domain.yaml' +) +EPSILON: float = 3.2 +DELTA: float = 1e-6 +NUM_PERMUTATION_SLOTS: int = 3 +EXPLORATION_STRATEGY: Literal['empty_token', 'size_sliced'] = 'size_sliced' +RANDOM_SEED: int = 42 +PGM_ITERS: int = 100 # Increase for higher fidelity, lower for faster runtime. + + +# ============================================================================== +# Step 1: Load Domain Schema & Datasets +# ============================================================================== +def load_domain_and_data( + domain_path: epath.Path, + data_dir: str, +) -> tuple[ + dict[str, domain.Schema], + list[rel_domain.ForeignKeyRelation], + dict[str, pd.DataFrame], +]: + """Loads relational schema from YAML and California Census tables from storage. + + Args: + domain_path: Path to the relational domain YAML file. + data_dir: Directory containing household.csv and individual.csv. + + Returns: + A tuple of (table_domains, foreign_keys, tables_dict). + """ + logging.info('Loading relational domain schema from: %s', domain_path) + table_domains, foreign_keys = rel_domain.from_yaml_file(str(domain_path)) + + logging.info('Loading dataset tables from: %s', data_dir) + with epath.Path(f'{data_dir}/household.csv').open('r') as f_h: + household_df = pd.read_csv(f_h) + with epath.Path(f'{data_dir}/individual.csv').open('r') as f_i: + individual_df = pd.read_csv(f_i) + + tables = { + 'household': household_df, + 'individual': individual_df, + } + logging.info( + 'Loaded %d households and %d individuals.', + len(tables['household']), + len(tables['individual']), + ) + return table_domains, foreign_keys, tables + + +# ============================================================================== +# Step 2: Configure & Calibrate Multi-Table DP Mechanism +# ============================================================================== +def create_calibrated_mechanism( + table_domains: Mapping[str, domain.Schema], + foreign_keys: Sequence[rel_domain.ForeignKeyRelation], + epsilon: float, + delta: float, + num_permutation_slots: int, + exploration_strategy: Literal['empty_token', 'size_sliced'], +) -> rel_synth.MultiTableMechanism: + """Configures and calibrates the differential privacy multi-table synthesizer. + + Args: + table_domains: Dictionary mapping table name to column domain dict. + foreign_keys: Sequence of ForeignKeyRelation constraints. + epsilon: Differential privacy epsilon parameter. + delta: Differential privacy delta parameter. + num_permutation_slots: Permutation order for exploration. + exploration_strategy: 'empty_token' or 'size_sliced'. + + Returns: + A calibrated MultiTableMechanism ready to synthesize data. + """ + logging.info( + 'Configuring MultiTableConfig with epsilon=%.2f, delta=%e, slots=%d,' + ' strategy=%s', + epsilon, + delta, + num_permutation_slots, + exploration_strategy, + ) + config = rel_synth.MultiTableConfig( + domains=table_domains, + foreign_keys=foreign_keys, + discrete_mechanism=discrete_mechanisms.MSTConfig(pgm_iters=PGM_ITERS), + num_permutation_slots=num_permutation_slots, + exploration_strategy=exploration_strategy, + ) + mechanism = config.calibrate(epsilon=epsilon, delta=delta) + assert isinstance(mechanism, rel_synth.MultiTableMechanism) + logging.info('Mechanism calibrated successfully.') + return mechanism + + +# ============================================================================== +# Step 3: Save Synthetic Data to Storage +# ============================================================================== +def save_synthetic_tables( + synthetic_tables: Mapping[str, pd.DataFrame], + output_dir: str, +) -> None: + """Saves synthetic tables as CSV files to the specified output directory. + + Args: + synthetic_tables: Mapping of table names to synthesized DataFrames. + output_dir: Output directory path on disk or local storage. + """ + logging.info('Saving synthetic tables to: %s', output_dir) + output_path = epath.Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + for table_name, df in synthetic_tables.items(): + file_path = output_path / f'{table_name}.csv' + logging.info('Writing %s (%d rows) to %s', table_name, len(df), file_path) + with file_path.open('w') as f: + df.to_csv(f, index=False) + logging.info('Synthetic tables successfully saved.') + + +# ============================================================================== +# Main Orchestration Loop +# ============================================================================== +def main(argv: list[str]) -> None: + if len(argv) > 1: + raise app.UsageError('Too many command-line arguments.') + + logging.info('=== Relational DPSynth: California Census Example ===') + + # 1. Load schema and real data + table_domains, foreign_keys, real_tables = load_domain_and_data( + domain_path=DOMAIN_PATH, + data_dir=DATA_DIR, + ) + + # 2. Configure and calibrate DP mechanism + mechanism = create_calibrated_mechanism( + table_domains=table_domains, + foreign_keys=foreign_keys, + epsilon=EPSILON, + delta=DELTA, + num_permutation_slots=NUM_PERMUTATION_SLOTS, + exploration_strategy=EXPLORATION_STRATEGY, + ) + + # 3. Execute DP synthesis + logging.info('Starting relational synthesis...') + rng = np.random.default_rng(RANDOM_SEED) + result = mechanism(rng=rng, data=real_tables) + synthetic_tables = dict(result.synthetic_tables) + logging.info( + 'Synthesis complete: %d households, %d individuals.', + len(synthetic_tables['household']), + len(synthetic_tables['individual']), + ) + + # 4. Save synthetic datasets to storage + save_synthetic_tables(synthetic_tables, output_dir=OUTPUT_DIR) + + # 5. Validate relational integrity (PK uniqueness, 0 orphans, max capacity) + eval_utils.validate_relational_integrity(synthetic_tables) + + # 6. Statistical & ML utility evaluations + # eval_utils.evaluate_statistical_fidelity(real_tables, synthetic_tables) + # eval_utils.evaluate_downstream_ml_utility( + # real_tables, synthetic_tables, random_state=RANDOM_SEED + # ) + + # 7. SDMetrics multi-table diagnostic and quality reports + eval_utils.generate_sdmetrics_reports(real_tables, synthetic_tables) + logging.info('=== Relational DPSynth Pipeline Completed Successfully ===') + os._exit(0) # pylint: disable=protected-access + + +if __name__ == '__main__': + app.run(main) diff --git a/examples/relational/california_census/example_eval_california_census.py b/examples/relational/california_census/example_eval_california_census.py new file mode 100644 index 0000000..43f95ae --- /dev/null +++ b/examples/relational/california_census/example_eval_california_census.py @@ -0,0 +1,260 @@ +# 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. + +"""Evaluation and validation utilities for California Census relational synthesis. + +Provides modular evaluation routines for: +1. Relational Integrity: Validating PK uniqueness, referential integrity (0 +orphans), and group size capacity bounds. +2. Statistical Utility: 1-way marginal Total Variation Distance (TVD). +3. Machine Learning Utility: Train-on-Synthetic, Test-on-Real (TSTR) classifier +for EMPSTAT. +4. Multi-Table Quality & Diagnostics: SDMetrics Multi-Table Diagnostic and +Quality reports. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from absl import logging +import pandas as pd +import sdmetrics +from sklearn import ensemble +from sklearn import metrics +from sklearn import model_selection + + +# ============================================================================== +# 1. Validate Relational Integrity +# ============================================================================== +def validate_relational_integrity( + synthetic_tables: Mapping[str, pd.DataFrame], + max_children_per_parent: int = 8, +) -> None: + """Validates relational consistency and foreign key integrity. + + Args: + synthetic_tables: Dictionary of synthesized DataFrames. + max_children_per_parent: Upper bound on child records per parent. + + Raises: + ValueError: If primary key uniqueness is violated, orphaned records exist, + or max child capacity bound is exceeded. + """ + synth_h = synthetic_tables['household'] + synth_i = synthetic_tables['individual'] + + # 1. Primary key uniqueness in parent table + if synth_h['HOUSEHOLD'].nunique() != len(synth_h): + raise ValueError('Primary key uniqueness violation in household table!') + logging.info('Household primary key uniqueness: PASSED') + + # 2. Referential integrity (no orphaned children) + orphan_mask = ~synth_i['HOUSEHOLD'].isin(synth_h['HOUSEHOLD']) + num_orphans = int(orphan_mask.sum()) + if num_orphans > 0: + raise ValueError(f'Found {num_orphans} orphaned individual records!') + logging.info('Referential integrity (0 orphans): PASSED') + + # 3. Capacity constraints + counts_per_h = synth_i['HOUSEHOLD'].value_counts() + max_found = int(counts_per_h.max()) if not counts_per_h.empty else 0 + if max_found > max_children_per_parent: + raise ValueError( + f'Max children capacity exceeded: {max_found} >' + f' {max_children_per_parent}' + ) + mean_found = float(counts_per_h.mean()) if not counts_per_h.empty else 0.0 + logging.info( + 'Max children capacity check (max found = %d <= %d, mean = %.2f): PASSED', + max_found, + max_children_per_parent, + mean_found, + ) + + +# ============================================================================== +# 2. Statistical Fidelity (Total Variation Distance) +# ============================================================================== +def compute_tvd(real_series: pd.Series, synth_series: pd.Series) -> float: + """Computes Total Variation Distance (TVD in [0, 1]) between two series.""" + real_dist = real_series.value_counts(normalize=True) + synth_dist = synth_series.value_counts(normalize=True) + all_cats = real_dist.index.union(synth_dist.index) + real_dist = real_dist.reindex(all_cats, fill_value=0.0) + synth_dist = synth_dist.reindex(all_cats, fill_value=0.0) + return float(0.5 * (real_dist - synth_dist).abs().sum()) + + +def evaluate_statistical_fidelity( + real_tables: Mapping[str, pd.DataFrame], + synthetic_tables: Mapping[str, pd.DataFrame], +) -> dict[str, float]: + """Evaluates 1-way marginal TVD across key household and individual columns.""" + scores = { + 'household.OWNERSHP': compute_tvd( + real_tables['household']['OWNERSHP'], + synthetic_tables['household']['OWNERSHP'], + ), + 'household.ROOMS': compute_tvd( + real_tables['household']['ROOMS'], + synthetic_tables['household']['ROOMS'], + ), + 'individual.EDUC': compute_tvd( + real_tables['individual']['EDUC'], + synthetic_tables['individual']['EDUC'], + ), + 'individual.MARST': compute_tvd( + real_tables['individual']['MARST'], + synthetic_tables['individual']['MARST'], + ), + } + for col_name, score in scores.items(): + logging.info('TVD [%s]: %.4f', col_name, score) + return scores + + +# ============================================================================== +# 3. Machine Learning Utility (TSTR) +# ============================================================================== +def evaluate_downstream_ml_utility( + real_tables: Mapping[str, pd.DataFrame], + synthetic_tables: Mapping[str, pd.DataFrame], + random_state: int = 42, +) -> tuple[float, float, float, float]: + """Evaluates multi-table cross-feature correlations using TSTR benchmark. + + Trains a classifier to predict individual employment status (EMPSTAT) + using features from both individual and household tables. + + Args: + real_tables: Real input tables dictionary. + synthetic_tables: Synthesized tables dictionary. + random_state: Random seed for model reproducibility. + + Returns: + Tuple of (real_acc, real_fscore, synth_acc, synth_fscore). + """ + real_joined = real_tables['individual'].merge( + real_tables['household'], on='HOUSEHOLD', how='inner' + ) + synth_joined = synthetic_tables['individual'].merge( + synthetic_tables['household'], on='HOUSEHOLD', how='inner' + ) + + features = ['AGE', 'SEX', 'EDUC', 'MARST', 'OWNERSHP', 'ROOMS', 'PUMA'] + target = 'EMPSTAT' + + train_real, test_real = model_selection.train_test_split( + real_joined, test_size=0.3, random_state=random_state + ) + + # 1. Baseline: Train on Real, Test on Real + clf_real = ensemble.RandomForestClassifier( + n_estimators=50, random_state=random_state + ) + clf_real.fit(train_real[features], train_real[target]) + pred_real = clf_real.predict(test_real[features]) + real_acc = float(metrics.accuracy_score(test_real[target], pred_real)) + real_fscore = float( + metrics.f1_score(test_real[target], pred_real, average='macro') + ) + + # 2. Synthetic: Train on Synthetic, Test on Real (TSTR) + clf_synth = ensemble.RandomForestClassifier( + n_estimators=50, random_state=random_state + ) + clf_synth.fit(synth_joined[features], synth_joined[target]) + pred_synth = clf_synth.predict(test_real[features]) + synth_acc = float(metrics.accuracy_score(test_real[target], pred_synth)) + synth_fscore = float( + metrics.f1_score(test_real[target], pred_synth, average='macro') + ) + + logging.info( + 'Downstream ML (EMPSTAT) - Real Baseline (TRTR) -> Acc: %.4f,' + ' Macro-FScore: %.4f', + real_acc, + real_fscore, + ) + logging.info( + 'Downstream ML (EMPSTAT) - Synthetic (TSTR) -> Acc: %.4f,' + ' Macro-FScore: %.4f', + synth_acc, + synth_fscore, + ) + return real_acc, real_fscore, synth_acc, synth_fscore + + +# ============================================================================== +# 4. SDMetrics Multi-Table Reports +# ============================================================================== +def generate_sdmetrics_reports( + real_tables: Mapping[str, pd.DataFrame], + synthetic_tables: Mapping[str, pd.DataFrame], +) -> tuple[ + sdmetrics.reports.multi_table.DiagnosticReport, + sdmetrics.reports.multi_table.QualityReport, +]: + """Generates SDMetrics Multi-Table Diagnostic and Quality reports. + + Args: + real_tables: Real input tables dictionary. + synthetic_tables: Synthesized tables dictionary. + + Returns: + A tuple of (diagnostic_report, quality_report). + """ + synth_h = synthetic_tables['household'] + synth_i = synthetic_tables['individual'] + + # Construct multi-table metadata for SDMetrics + metadata = { + 'tables': { + 'household': { + 'primary_key': 'HOUSEHOLD', + 'columns': { + col: {'sdtype': 'categorical'} + for col in synth_h.columns + if col != 'HOUSEHOLD' + }, + }, + 'individual': { + 'columns': { + col: {'sdtype': 'categorical'} + for col in synth_i.columns + if col != 'HOUSEHOLD' + }, + }, + }, + 'relationships': [{ + 'parent_table_name': 'household', + 'parent_primary_key': 'HOUSEHOLD', + 'child_table_name': 'individual', + 'child_foreign_key': 'HOUSEHOLD', + }], + } + + logging.info('Generating SDMetrics DiagnosticReport...') + diag_report = sdmetrics.reports.multi_table.DiagnosticReport() + diag_report.generate(real_tables, synthetic_tables, metadata) + logging.info('SDMetrics Diagnostic Score: %.4f', diag_report.get_score()) + + logging.info('Generating SDMetrics QualityReport...') + qual_report = sdmetrics.reports.multi_table.QualityReport() + qual_report.generate(real_tables, synthetic_tables, metadata) + logging.info('SDMetrics Quality Score: %.4f', qual_report.get_score()) + + return diag_report, qual_report diff --git a/tests/examples/relational/california_census/domain_test.py b/tests/examples/relational/california_census/domain_test.py new file mode 100644 index 0000000..b2f460d --- /dev/null +++ b/tests/examples/relational/california_census/domain_test.py @@ -0,0 +1,84 @@ +# 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 California Census relational domain configuration.""" + +from __future__ import annotations +from absl.testing import absltest +from dpsynth import domain +from dpsynth.relational import domain as rel_domain +from etils import epath + +_DOMAIN_PATH = ( + epath.resource_path('dpsynth') + / 'examples/relational/california_census/domain.yaml' +) + + +class CaliforniaCensusDomainTest(absltest.TestCase): + + def test_load_california_census_domain(self): + table_domains, foreign_keys = rel_domain.from_yaml_file(str(_DOMAIN_PATH)) + + # Validate table set + self.assertCountEqual( + list(table_domains.keys()), ['household', 'individual'] + ) + + # Validate household attributes + household_schema = table_domains['household'] + self.assertLen(household_schema, 10) + self.assertIsInstance(household_schema['FARM'], domain.CategoricalAttribute) + self.assertEqual(household_schema['FARM'].size, 2) + self.assertIsInstance( + household_schema['PROPINSR'], domain.NumericalAttribute + ) + self.assertEqual(household_schema['PROPINSR'].min_value, 0.0) + self.assertEqual(household_schema['PROPINSR'].max_value, 59.0) + + # Validate individual attributes + individual_schema = table_domains['individual'] + self.assertLen(individual_schema, 15) + self.assertIsInstance( + individual_schema['RELATE'], domain.CategoricalAttribute + ) + self.assertEqual(individual_schema['RELATE'].size, 8) + self.assertIsInstance(individual_schema['AGE'], domain.NumericalAttribute) + self.assertEqual(individual_schema['AGE'].min_value, 0.0) + self.assertEqual(individual_schema['AGE'].max_value, 85.0) + + # Validate foreign keys + self.assertLen(foreign_keys, 1) + fk = foreign_keys[0] + self.assertEqual(fk.parent_table, 'household') + self.assertEqual(fk.parent_primary_key, 'HOUSEHOLD') + self.assertEqual(fk.child_table, 'individual') + self.assertEqual(fk.child_foreign_key, 'HOUSEHOLD') + self.assertEqual(fk.max_children_per_parent, 8) + + # Validate topological hierarchy + hierarchy = rel_domain.topological_sort_hierarchy( + list(table_domains.keys()), foreign_keys + ) + self.assertEqual( + hierarchy, + [ + (0, 'household', None), + (1, 'individual', fk), + ], + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/examples/relational/california_census/synthesis_integration_test.py b/tests/examples/relational/california_census/synthesis_integration_test.py new file mode 100644 index 0000000..1f98137 --- /dev/null +++ b/tests/examples/relational/california_census/synthesis_integration_test.py @@ -0,0 +1,180 @@ +# 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. + +"""Integration test for California Census relational differential privacy synthesis.""" + +from __future__ import annotations + +from absl.testing import absltest +from absl.testing import parameterized +from dpsynth import discrete_mechanisms +from dpsynth.relational import domain as rel_domain +from dpsynth.relational import synthesizer as rel_synth +from etils import epath +import jax +import numpy as np +import pandas as pd + +jax.config.update('jax_enable_compilation_cache', False) + +_DOMAIN_PATH = ( + epath.resource_path('dpsynth') + / 'examples/relational/california_census/domain.yaml' +) + + +def _generate_mock_california_data( + num_households: int = 50, + rng: np.random.Generator | None = None, +) -> dict[str, pd.DataFrame]: + """Generates mock California census DataFrames adhering to domain.yaml.""" + if rng is None: + rng = np.random.default_rng(42) + + # 1. Generate Household records + h_ids = np.arange(1, num_households + 1) + household_df = pd.DataFrame({ + 'HOUSEHOLD': h_ids, + 'FARM': rng.integers(0, 2, size=num_households), + 'OWNERSHP': rng.integers(0, 3, size=num_households), + 'ACREHOUS': rng.integers(0, 3, size=num_households), + 'TAXINCL': rng.integers(0, 3, size=num_households), + 'PROPINSR': rng.integers(0, 60, size=num_households), + 'COSTELEC': rng.integers(0, 100, size=num_households), + 'VALUEH': rng.integers(0, 100, size=num_households), + 'ROOMS': rng.integers(0, 10, size=num_households), + 'PLUMBING': rng.integers(0, 3, size=num_households), + 'PUMA': rng.integers(0, 233, size=num_households), + }) + + # 2. Generate variable-sized child individual records + # (1 to 4 individuals per household) + group_sizes = rng.integers(1, 5, size=num_households) + ind_households = np.repeat(h_ids, group_sizes) + num_individuals = len(ind_households) + + individual_df = pd.DataFrame({ + 'HOUSEHOLD': ind_households, + 'RELATE': rng.integers(0, 8, size=num_individuals), + 'SEX': rng.integers(0, 2, size=num_individuals), + 'AGE': rng.integers(0, 86, size=num_individuals), + 'MARST': rng.integers(0, 6, size=num_individuals), + 'RACE': rng.integers(0, 9, size=num_individuals), + 'CITIZEN': rng.integers(0, 4, size=num_individuals), + 'SPEAKENG': rng.integers(0, 6, size=num_individuals), + 'SCHOOL': rng.integers(0, 3, size=num_individuals), + 'EDUC': rng.integers(0, 11, size=num_individuals), + 'GRADEATT': rng.integers(0, 8, size=num_individuals), + 'SCHLTYPE': rng.integers(0, 4, size=num_individuals), + 'EMPSTAT': rng.integers(0, 4, size=num_individuals), + 'CLASSWKR': rng.integers(0, 3, size=num_individuals), + 'INCTOT': rng.integers(0, 100, size=num_individuals), + 'DISABWRK': rng.integers(0, 3, size=num_individuals), + }) + + return {'household': household_df, 'individual': individual_df} + + +class CaliforniaCensusSynthesisIntegrationTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.table_domains, self.foreign_keys = rel_domain.from_yaml_file( + str(_DOMAIN_PATH) + ) + + def test_california_census_pipeline_e2e_mst(self): + rng = np.random.default_rng(12345) + data = _generate_mock_california_data(num_households=50, rng=rng) + + config = rel_synth.MultiTableConfig( + domains=self.table_domains, + foreign_keys=self.foreign_keys, + discrete_mechanism=discrete_mechanisms.MSTConfig(pgm_iters=10), + num_permutation_slots=1, + exploration_strategy='empty_token', + numerical_bins=2, + ) + calibrated_mechanism = config.calibrate(epsilon=3.2, delta=1e-6) + + # Synthesize + result = calibrated_mechanism(rng=rng, data=data) + + self.assertIsInstance(result, rel_synth.MultiDataGenerationResult) + self.assertIn('household', result.synthetic_tables) + self.assertIn('individual', result.synthetic_tables) + + synth_h = result.synthetic_tables['household'] + synth_i = result.synthetic_tables['individual'] + + # Non-empty tables + self.assertNotEmpty(synth_h) + self.assertNotEmpty(synth_i) + + # Verify column presence + self.assertCountEqual( + synth_h.columns, + [ + 'HOUSEHOLD', + 'FARM', + 'OWNERSHP', + 'ACREHOUS', + 'TAXINCL', + 'PROPINSR', + 'COSTELEC', + 'VALUEH', + 'ROOMS', + 'PLUMBING', + 'PUMA', + ], + ) + self.assertCountEqual( + synth_i.columns, + [ + 'HOUSEHOLD', + 'RELATE', + 'SEX', + 'AGE', + 'MARST', + 'RACE', + 'CITIZEN', + 'SPEAKENG', + 'SCHOOL', + 'EDUC', + 'GRADEATT', + 'SCHLTYPE', + 'EMPSTAT', + 'CLASSWKR', + 'INCTOT', + 'DISABWRK', + ], + ) + + # Verify relational integrity: No orphaned children + h_pks = set(synth_h['HOUSEHOLD']) + i_fks = set(synth_i['HOUSEHOLD']) + self.assertTrue( + i_fks.issubset(h_pks), 'Found orphaned individual foreign keys!' + ) + + # Verify group size capacity: max children per parent <= 8 + counts_per_h = synth_i['HOUSEHOLD'].value_counts() + self.assertTrue( + (counts_per_h <= 8).all(), + f'Max children exceeded: {counts_per_h.max()}', + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/relational/synthesizer_test.py b/tests/relational/synthesizer_test.py index 54f20d0..0923bc5 100644 --- a/tests/relational/synthesizer_test.py +++ b/tests/relational/synthesizer_test.py @@ -16,11 +16,13 @@ import math import unittest.mock + from absl.testing import absltest import dp_accounting from dpsynth import api from dpsynth import discrete_mechanisms from dpsynth import domain +from dpsynth import relational as rel_pkg from dpsynth.local_mode import initialization from dpsynth.relational import domain as rel_domain from dpsynth.relational import post_processing @@ -1718,6 +1720,14 @@ def test_multi_table_mechanism_call_size_sliced(self): self.assertTrue(set(synth_p['hid']).issubset(set(synth_h['hid']))) self.assertIn('Household->Person', result.discrete_mechanism_results) + def test_public_api_exports(self): + self.assertIs(rel_pkg.MultiTableConfig, synthesizer.MultiTableConfig) + self.assertIs(rel_pkg.MultiTableMechanism, synthesizer.MultiTableMechanism) + self.assertIs( + rel_pkg.MultiDataGenerationResult, synthesizer.MultiDataGenerationResult + ) + self.assertIs(rel_pkg.ForeignKeyRelation, rel_domain.ForeignKeyRelation) + if __name__ == '__main__': absltest.main()