From e51cafd4b9976c45ed64d8947e2e91108b23d465 Mon Sep 17 00:00:00 2001 From: DPSynth Team Date: Thu, 20 Aug 2026 02:59:51 -0700 Subject: [PATCH] Implement postprocessing, MultiTableMechanism.__call__ and output materialization in the relational synthesizer: 1. Decompression (_decompress_synthetic_datasets): - Probabilistically reverses categorical domain compression by sampling original category preimages from composite `` bins via `mbi.Dataset.decompress`. 2. Native TabularCodec Decoding (_decode_synthetic_tables): - Converts discrete integer bucket tokens back into continuous numerical floats (uniform interval dequantization) and categorical strings via `TabularCodec.decode()`. 3. Relational Key Assignment (_assign_relational_keys): - Generates synthetic integer surrogate primary keys (0..N-1) for parent tables. - Vectorizes child foreign key assignment by indexing parent primary keys via `parent_mappings` from Phase 2 copula unstacking, guaranteeing 100% referential integrity without hash joins. 4. End-to-End Orchestration (MultiTableMechanism.__call__): - Connectsdiscretization & initializers, cascading discrete relational hierarchy synthesis & copula matching, and postprocessing & key linkage into `MultiDataGenerationResult`. 5. Testing & Optimization: - Added unit tests for decompression, decoding, multi-tier PK/FK assignment, and empty child edge cases. - Added end-to-end integration test verifying calibrated `MultiTableMechanism` execution with differential privacy guarantees and referential integrity. - Add parameterized subtests covering all eager `MultiTableConfig.__post_init__` validation failures. 6. Eager MultiTableConfig.__post_init__ Validation: * Make `foreign_keys` a required argument with no default (`= ()` removed). * Enforce naming constraints PiperOrigin-RevId: 967697951 --- dpsynth/relational/synthesizer.py | 276 +++++++++++++- tests/relational/synthesizer_test.py | 522 ++++++++++++++++++++++++++- 2 files changed, 781 insertions(+), 17 deletions(-) diff --git a/dpsynth/relational/synthesizer.py b/dpsynth/relational/synthesizer.py index 8a8f215..a86ed1c 100644 --- a/dpsynth/relational/synthesizer.py +++ b/dpsynth/relational/synthesizer.py @@ -333,7 +333,7 @@ def _encode_and_compress_tables( @dataclasses.dataclass(frozen=True) class PreprocessedTables: - """Engine-agnostic container for preprocessed relational tables from Phase 1. + """Engine-agnostic container for preprocessed relational tables. Attributes: compressed_datasets: Mapping from table name to discrete compressed @@ -443,7 +443,7 @@ def _compute_table_col_deltas( Args: domains: Mapping from table names to per-column AttributeType schemas. delta: Total DP delta for partition selection thresholding. - init_budget_fraction: Fraction of delta allocated to Phase 1 initialization. + init_budget_fraction: Fraction of delta allocated to column initialization. Returns: A nested mapping from table name and column name to its allocated delta. @@ -591,7 +591,7 @@ class SynthesizedLinkResult: unstacked_child_dataset: mbi.Dataset parent_row_indices: np.ndarray synth_parent_dataset: mbi.Dataset | None = None # Needed for the root table. - discrete_mechanism_result: dm_common.DiscreteMechanismResult | None = None + discrete_mechanism_result: Any | None = None def _synthesize_relational_link( @@ -668,7 +668,7 @@ def _synthesize_relational_link( else: exploration_constraints = () - mech_res: dm_common.DiscreteMechanismResult = discrete_mechanism( + mech_res: Any = discrete_mechanism( rng=rng, data=exploration_dataset, constraints=exploration_constraints, @@ -754,7 +754,7 @@ def _synthesize_relational_hierarchy( Args: mechanism: Calibrated MultiTableMechanism. - preprocessed: PreprocessedTables container from Phase 1. + preprocessed: PreprocessedTables container holding encoded data. hierarchy: Ordered topological synthesis levels. rng: NumPy random generator. @@ -811,6 +811,106 @@ def _synthesize_relational_hierarchy( return synth_datasets, parent_mappings, discrete_mechanism_results +def _decompress_synthetic_datasets( + synth_datasets: Mapping[str, mbi.Dataset], + compression_mappings: Mapping[str, Mapping[str, np.ndarray]], +) -> dict[str, mbi.Dataset]: + """Probabilistically decompresses synthetic discrete datasets for each table. + + Reverses domain category compression mappings by sampling original category + preimages uniformly from composite bins. + + Args: + synth_datasets: Mapping from table name to synthesized mbi.Dataset. + compression_mappings: Mapping from table name to per-column compression + mapping arrays. + + Returns: + A mapping from table name to decompressed mbi.Dataset. + """ + decompressed: dict[str, mbi.Dataset] = {} + for table_name, dataset in synth_datasets.items(): + mappings = compression_mappings.get(table_name, {}) + if mappings: + raw_mappings = {str(col): arr for col, arr in mappings.items()} + decompressed[table_name] = dataset.decompress(raw_mappings) + else: + decompressed[table_name] = dataset + return decompressed + + +def _decode_synthetic_tables( + decompressed_datasets: Mapping[str, mbi.Dataset], + column_codecs: Mapping[str, data_generation_v3.TabularCodec], + domains: Mapping[str, domain.Schema], + rng: np.random.Generator, +) -> dict[str, pd.DataFrame]: + """Decodes decompressed discrete datasets into continuous/categorical DataFrames. + + Converts discrete integer bucket tokens back to native strings and continuous + floats (uniformly dequantized within bucket intervals) per table schema. + + Args: + decompressed_datasets: Mapping from table name to decompressed mbi.Dataset. + column_codecs: Mapping from table name to TabularCodec. + domains: Mapping from table name to per-column AttributeType schemas. + rng: NumPy random generator for uniform interval dequantization. + + Returns: + A mapping from table name to decoded pandas DataFrame. + """ + decoded_tables: dict[str, pd.DataFrame] = {} + for table_name, dataset in decompressed_datasets.items(): + codec = column_codecs[table_name] + column_order = list(domains[table_name].keys()) + decoded_tables[table_name] = codec.decode( + synthetic=dataset, rng=rng, column_order=column_order + ) + return decoded_tables + + +def _assign_relational_keys( + tables: Mapping[str, pd.DataFrame], + foreign_keys: Sequence[rel_domain.ForeignKeyRelation], + parent_mappings: Mapping[str, np.ndarray], + hierarchy: Sequence[tuple[int, str, rel_domain.ForeignKeyRelation | None]], +) -> dict[str, pd.DataFrame]: + """Assigns synthetic primary and foreign keys to link relational tables. + + For each table in topological order, generates integer surrogate primary keys + (0..N-1) and assigns child foreign keys by indexing parent primary keys via + the parent-row mappings from copula matching and unstacking. + + Args: + tables: Mapping from table name to decoded DataFrame. + foreign_keys: Sequence of foreign key relationships. + parent_mappings: Mapping from child table name to 1D int array of parent row + indices. + hierarchy: Ordered topological synthesis levels. + + Returns: + A mapping from table name to DataFrame with assigned PK and FK columns. + """ + linked_tables = {name: df.copy() for name, df in tables.items()} + primary_keys: dict[str, dict[str, np.ndarray]] = {} + + for fk in foreign_keys: + if fk.parent_table not in primary_keys: + primary_keys[fk.parent_table] = {} + if fk.parent_primary_key not in primary_keys[fk.parent_table]: + pk_arr = np.arange(len(linked_tables[fk.parent_table]), dtype=np.int64) + primary_keys[fk.parent_table][fk.parent_primary_key] = pk_arr + linked_tables[fk.parent_table][fk.parent_primary_key] = pk_arr + + for _, table_name, fk in hierarchy: + if fk is not None: + p_keys = primary_keys[fk.parent_table][fk.parent_primary_key] + parent_indices = parent_mappings[table_name] + linked_tables[table_name][fk.child_foreign_key] = p_keys[parent_indices] + + return linked_tables + + @dataclasses.dataclass(frozen=True) class MultiDataGenerationResult: """Results of multi-table relational DP synthetic data generation. @@ -883,10 +983,64 @@ def __call__( rng: np.random.Generator, data: Mapping[str, pd.DataFrame], ) -> MultiDataGenerationResult: - """Generates synthetic multi-table relational data.""" - del rng, data - raise NotImplementedError( - 'MultiTableMechanism.__call__ is not yet implemented.' + """Generates differentially private synthetic multi-table relational data. + + Executes the multi-table synthesis process across three steps: + 1. Preprocessing: Weights children tables to guarantee sensitivity = 1 for + each parent record. Then individually for each table, discretizes + numerical columns into bins, encodes & compresses categories under DP. + 2. Candidate Exploration: Explores set of candidate marginals from + permutations of children records and their parent, representing + cross-table correlations. + 3. Relational Synthesis: Generates synthetic records top-down from root + parent to leaf child tables, fitting the MRF per parent->child pair + preserving the chosen cross-table marginal correlation marginals. + 4. Output Decoding & Key Assignment: Converts synthetic tokens back to + original data types (continuous numbers & category strings) and assigns + matching primary and foreign keys to ensure referential integrity. + + Args: + rng: NumPy random number generator. + data: Mapping from table name to input pandas DataFrame. + + Returns: + A MultiDataGenerationResult containing synthetic DataFrames and + diagnostics. + """ + preprocessed = _run_table_preprocessing(self, rng=rng, data=data) + + hierarchy = rel_domain.topological_sort_hierarchy( + list(self.domains.keys()), self.foreign_keys + ) + synth_datasets, parent_mappings, discrete_results = ( + _synthesize_relational_hierarchy( + mechanism=self, + preprocessed=preprocessed, + hierarchy=hierarchy, + rng=rng, + ) + ) + + decompressed = _decompress_synthetic_datasets( + synth_datasets=synth_datasets, + compression_mappings=preprocessed.compression_mappings, + ) + decoded_tables = _decode_synthetic_tables( + decompressed_datasets=decompressed, + column_codecs=preprocessed.column_codecs, + domains=self.domains, + rng=rng, + ) + final_tables = _assign_relational_keys( + tables=decoded_tables, + foreign_keys=self.foreign_keys, + parent_mappings=parent_mappings, + hierarchy=hierarchy, + ) + + return MultiDataGenerationResult( + synthetic_tables=final_tables, + discrete_mechanism_results=discrete_results, ) @@ -901,14 +1055,15 @@ class MultiTableConfig(api.MechanismConfig): discrete_mechanism: Discrete mechanism config (e.g. AIM, MST) for relational links. numerical_bins: Number of bins for numerical attribute discretization. - init_budget_fraction: Fraction of total zCDP budget allocated to Phase 1. + init_budget_fraction: Fraction of total privacy budget allocated to column + initialization. num_permutation_slots: Permutation exploration slot count (o), default 2. exploration_strategy: Exploration strategy ('empty_token' or 'size_sliced'). """ domains: Mapping[str, domain.Schema] - foreign_keys: Sequence[rel_domain.ForeignKeyRelation] = () - discrete_mechanism: discrete_mechanisms.MechanismConfig = dataclasses.field( + foreign_keys: Sequence[rel_domain.ForeignKeyRelation] + discrete_mechanism: api.MechanismConfig = dataclasses.field( default_factory=discrete_mechanisms.AIMConfig ) numerical_bins: int = 32 @@ -929,9 +1084,9 @@ def __post_init__(self): 'MultiTableConfig requires at least one foreign key relationship in' ' foreign_keys. For single-table synthesis, use TabularConfig.' ) - if not 0.0 <= self.init_budget_fraction <= 1.0: + if not (0.0 < self.init_budget_fraction < 1.0): raise ValueError( - 'init_budget_fraction must be in [0.0, 1.0], got' + 'init_budget_fraction must be strictly in (0.0, 1.0), got' f' {self.init_budget_fraction}.' ) if self.numerical_bins < 1: @@ -947,6 +1102,97 @@ def __post_init__(self): raise ValueError( f'Unsupported exploration_strategy {self.exploration_strategy!r}.' ) + if not isinstance(self.discrete_mechanism, api.MechanismConfig): + raise ValueError( + 'discrete_mechanism must be an instance of MechanismConfig, got' + f' {type(self.discrete_mechanism).__name__}.' + ) + + # 1. Validate table names, column names, and attribute types. + for table_name, schema in self.domains.items(): + if '.' in table_name: + raise ValueError( + f"Table name {table_name!r} must not contain '.' characters." + ) + if not schema: + raise ValueError( + f'Table {table_name!r} schema in domains cannot be empty.' + ) + for col_name, attr in schema.items(): + if '.' in col_name: + raise ValueError( + f"Table {table_name!r} column {col_name!r} must not contain '.'" + ' (reserved for wide relational slot prefixes).' + ) + if col_name == 'group_size': + raise ValueError( + f"Table {table_name!r} column name 'group_size' is reserved for" + ' relational exploration.' + ) + if col_name.startswith('slot_'): + raise ValueError( + f'Table {table_name!r} column {col_name!r} cannot start with' + " 'slot_' (reserved for permutation slots)." + ) + if not isinstance( + attr, + ( + domain.NumericalAttribute, + domain.CategoricalAttribute, + domain.OpenSetCategoricalAttribute, + ), + ): + raise ValueError( + f'Table {table_name!r} column {col_name!r} has unsupported' + f' attribute type {type(attr).__name__}.' + ) + + # 2. Validate DAG hierarchy (acyclicity, known tables, + # in-degree <= 1, single root). + hierarchy = rel_domain.topological_sort_hierarchy( + list(self.domains.keys()), self.foreign_keys + ) + roots = [t for _, t, fk in hierarchy if fk is None] + if len(roots) > 1: + raise ValueError( + 'MultiTableConfig expects a single root table, but found' + f' {len(roots)}: {roots}. All tables must be connected in a single' + ' tree hierarchy.' + ) + + # 3. Ensure PK and FK columns are not present in domain schemas. + for fk in self.foreign_keys: + if fk.parent_primary_key in self.domains[fk.parent_table]: + raise ValueError( + f'Primary key column {fk.parent_primary_key!r} of table' + f' {fk.parent_table!r} must not be in domains[{fk.parent_table!r}].' + ) + if fk.child_foreign_key in self.domains[fk.child_table]: + raise ValueError( + f'Foreign key column {fk.child_foreign_key!r} of table' + f' {fk.child_table!r} must not be in domains[{fk.child_table!r}].' + ) + + # 4. Validate custom initializers structure if provided. + if self.initializers is not None: + if set(self.initializers.keys()) != set(self.domains.keys()): + raise ValueError( + f'Custom initializers tables {set(self.initializers.keys())} do not' + f' match domains tables {set(self.domains.keys())}.' + ) + for table_name, table_inits in self.initializers.items(): + if set(table_inits.keys()) != set(self.domains[table_name].keys()): + raise ValueError( + f'Custom initializers for table {table_name!r}' + f' columns {set(table_inits.keys())} do not match' + f' domains columns {set(self.domains[table_name].keys())}.' + ) + for col_name, init_cfg in table_inits.items(): + if not isinstance(init_cfg, api.MechanismConfig): + raise ValueError( + f'Custom initializer for {table_name}.{col_name} must be an' + f' api.MechanismConfig, got {type(init_cfg).__name__}.' + ) def configure( self, @@ -955,7 +1201,7 @@ def configure( delta: float = 0.0, max_records_per_user: int = 1, ) -> MultiTableMechanism: - """Configures privacy budgets across Phase 1 initializers and Phase 2 links. + """Configures privacy budgets across column initializers and relational links. Formal Guarantees: - Additive zCDP Partitioning: The total zCDP budget zcdp_rho is diff --git a/tests/relational/synthesizer_test.py b/tests/relational/synthesizer_test.py index 070e022..54f20d0 100644 --- a/tests/relational/synthesizer_test.py +++ b/tests/relational/synthesizer_test.py @@ -15,6 +15,7 @@ """Unit tests for dpsynth.relational.synthesizer.""" import math +import unittest.mock from absl.testing import absltest import dp_accounting from dpsynth import api @@ -382,9 +383,9 @@ def test_configure_hyperparameter_validations(self): with self.assertRaisesRegex(ValueError, 'zcdp_rho must be positive'): config.configure(zcdp_rho=-0.1) - with self.subTest('invalid_init_budget_fraction'): + with self.subTest('invalid_init_budget_fraction_above_one'): with self.assertRaisesRegex( - ValueError, 'init_budget_fraction must be in' + ValueError, 'init_budget_fraction must be strictly in' ): synthesizer.MultiTableConfig( domains=domains, @@ -392,6 +393,16 @@ def test_configure_hyperparameter_validations(self): init_budget_fraction=1.5, ) + with self.subTest('invalid_init_budget_fraction_zero'): + with self.assertRaisesRegex( + ValueError, 'init_budget_fraction must be strictly in' + ): + synthesizer.MultiTableConfig( + domains=domains, + foreign_keys=foreign_keys, + init_budget_fraction=0.0, + ) + with self.subTest('invalid_numerical_bins'): with self.assertRaisesRegex(ValueError, 'numerical_bins must be >= 1'): synthesizer.MultiTableConfig( @@ -420,6 +431,146 @@ def test_configure_hyperparameter_validations(self): exploration_strategy='unsupported_strategy', ) + with self.subTest('invalid_discrete_mechanism'): + with self.assertRaisesRegex( + ValueError, + 'discrete_mechanism must be an instance of MechanismConfig', + ): + synthesizer.MultiTableConfig( + domains=domains, + foreign_keys=foreign_keys, + discrete_mechanism='not_a_config', # pyrefly: ignore[bad-argument-type] + ) + + with self.subTest('dot_in_table_name'): + with self.assertRaisesRegex(ValueError, "must not contain '.'"): + synthesizer.MultiTableConfig( + domains={ + 'House.hold': {'income': domain.NumericalAttribute(0, 100)}, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + }, + foreign_keys=[ + rel_domain.ForeignKeyRelation( + parent_table='House.hold', + parent_primary_key='hid', + child_table='Person', + child_foreign_key='hid', + max_children_per_parent=2, + ) + ], + ) + + with self.subTest('dot_in_column_name'): + with self.assertRaisesRegex(ValueError, "must not contain '.'"): + synthesizer.MultiTableConfig( + domains={ + 'Household': {'inc.ome': domain.NumericalAttribute(0, 100)}, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('reserved_column_name_group_size'): + with self.assertRaisesRegex( + ValueError, 'reserved for relational exploration' + ): + synthesizer.MultiTableConfig( + domains={ + 'Household': {'group_size': domain.NumericalAttribute(0, 100)}, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('reserved_column_name_slot_prefix'): + with self.assertRaisesRegex(ValueError, 'reserved for permutation slots'): + synthesizer.MultiTableConfig( + domains={ + 'Household': {'slot_1': domain.NumericalAttribute(0, 100)}, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('empty_table_schema'): + with self.assertRaisesRegex( + ValueError, 'schema in domains cannot be empty' + ): + synthesizer.MultiTableConfig( + domains={ + 'Household': {}, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('unsupported_attribute_type'): + with self.assertRaisesRegex(ValueError, 'unsupported attribute type'): + synthesizer.MultiTableConfig( + domains={ + 'Household': {'text': domain.FreeFormTextAttribute()}, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('multi_root_forest_raises'): + with self.assertRaisesRegex(ValueError, 'expects a single root table'): + synthesizer.MultiTableConfig( + domains={ + 'Household': {'income': domain.NumericalAttribute(0, 100)}, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + 'Unlinked': {'type': domain.CategoricalAttribute(['A', 'B'])}, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('pk_in_domain_schema_raises'): + with self.assertRaisesRegex(ValueError, 'must not be in domains'): + synthesizer.MultiTableConfig( + domains={ + 'Household': { + 'income': domain.NumericalAttribute(0, 100), + 'hid': domain.CategoricalAttribute(['H1', 'H2']), + }, + 'Person': {'age': domain.NumericalAttribute(0, 100)}, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('fk_in_domain_schema_raises'): + with self.assertRaisesRegex(ValueError, 'must not be in domains'): + synthesizer.MultiTableConfig( + domains={ + 'Household': {'income': domain.NumericalAttribute(0, 100)}, + 'Person': { + 'age': domain.NumericalAttribute(0, 100), + 'hid': domain.CategoricalAttribute(['H1', 'H2']), + }, + }, + foreign_keys=foreign_keys, + ) + + with self.subTest('custom_initializers_mismatched_tables'): + with self.assertRaisesRegex(ValueError, 'do not match domains tables'): + synthesizer.MultiTableConfig( + domains=domains, + foreign_keys=foreign_keys, + initializers={'Household': {}}, + ) + + with self.subTest('custom_initializers_mismatched_columns'): + mock_cfg = unittest.mock.MagicMock(spec=api.MechanismConfig) + with self.assertRaisesRegex(ValueError, 'do not match domains columns'): + synthesizer.MultiTableConfig( + domains=domains, + foreign_keys=foreign_keys, + initializers={ + 'Household': {'wrong_col': mock_cfg}, + 'Person': {'age': mock_cfg}, + }, + ) + def test_validate_input_table_columns_success(self): domains = { 'Household': { @@ -1200,6 +1351,373 @@ def test_synthesize_relational_hierarchy_branching(self): self.assertIn('Household->Person', discrete_results) self.assertIn('Household->Vehicle', discrete_results) + def test_decompress_synthetic_datasets(self): + # Table 1: compressed domain (original size 4, compressed to size 2 + # via [0, 0, 1, 1]) + orig_dom_h = mbi.Domain(('region',), (4,)) + comp_dom_h = mbi.Domain(('region',), (2,)) + comp_ds_h = mbi.Dataset({'region': np.array([0, 1, 0, 1])}, comp_dom_h) + + # Table 2: uncompressed domain + dom_p = mbi.Domain(('age',), (5,)) + ds_p = mbi.Dataset({'age': np.array([0, 2, 4])}, dom_p) + + synth_datasets = {'Household': comp_ds_h, 'Person': ds_p} + compression_mappings = { + 'Household': {'region': np.array([0, 0, 1, 1], dtype=np.int64)}, + 'Person': {}, + } + + decompressed = synthesizer._decompress_synthetic_datasets( + synth_datasets, compression_mappings + ) + + self.assertIn('Household', decompressed) + self.assertIn('Person', decompressed) + # Household should have restored original domain + self.assertEqual(decompressed['Household'].domain, orig_dom_h) + self.assertEqual(decompressed['Household'].records, 4) + # Values mapped from 0 should be in {0, 1}, and from 1 should be in {2, 3} + h_vals = decompressed['Household'].data['region'] + self.assertIn(h_vals[0], (0, 1)) + self.assertIn(h_vals[1], (2, 3)) + self.assertIn(h_vals[2], (0, 1)) + self.assertIn(h_vals[3], (2, 3)) + # Person should be unchanged + self.assertEqual(decompressed['Person'].domain['age'], 5) + np.testing.assert_array_equal( + decompressed['Person'].data['age'], np.array([0, 2, 4]) + ) + + def test_decode_synthetic_tables(self): + rng = np.random.default_rng(42) + domains = { + 'Household': { + 'income': domain.NumericalAttribute(min_value=0.0, max_value=100.0), + 'region': domain.CategoricalAttribute( + possible_values=['Urban', 'Rural'] + ), + }, + 'Person': { + 'age': domain.NumericalAttribute(min_value=0, max_value=100), + 'gender': domain.CategoricalAttribute(possible_values=['M', 'F']), + }, + } + # Mock column measurements for initializers + h_measurements = { + 'income': initialization.ColumnMeasurement( + categorical_attribute=domain.CategoricalAttribute( + possible_values=['[0.0, 50.0)', '[50.0, 100.0]'] + ), + bin_edges=np.array([50.0]), + ), + 'region': initialization.ColumnMeasurement( + categorical_attribute=domain.CategoricalAttribute( + possible_values=['Urban', 'Rural'] + ), + ), + } + p_measurements = { + 'age': initialization.ColumnMeasurement( + categorical_attribute=domain.CategoricalAttribute( + possible_values=['[0, 50)', '[50, 100]'] + ), + bin_edges=np.array([50.0]), + ), + 'gender': initialization.ColumnMeasurement( + categorical_attribute=domain.CategoricalAttribute( + possible_values=['M', 'F'] + ), + ), + } + codecs = { + 'Household': ( + synthesizer.data_generation_v3.TabularCodec.from_measurements( + h_measurements, domains['Household'] + ) + ), + 'Person': synthesizer.data_generation_v3.TabularCodec.from_measurements( + p_measurements, domains['Person'] + ), + } + decompressed_datasets = { + 'Household': mbi.Dataset( + {'income': np.array([0, 1]), 'region': np.array([0, 1])}, + codecs['Household'].mbi_domain, + ), + 'Person': mbi.Dataset( + { + 'age': np.array([0, 1, 0]), + 'gender': np.array([1, 0, 1]), + }, + codecs['Person'].mbi_domain, + ), + } + + decoded_tables = synthesizer._decode_synthetic_tables( + decompressed_datasets=decompressed_datasets, + column_codecs=codecs, + domains=domains, + rng=rng, + ) + + self.assertIn('Household', decoded_tables) + self.assertIn('Person', decoded_tables) + h_df = decoded_tables['Household'] + p_df = decoded_tables['Person'] + self.assertEqual(list(h_df.columns), ['income', 'region']) + self.assertEqual(list(p_df.columns), ['age', 'gender']) + self.assertLen(h_df, 2) + self.assertLen(p_df, 3) + # Check numerical dequantization bounds + self.assertTrue(np.all((h_df['income'] >= 0.0) & (h_df['income'] <= 100.0))) + self.assertTrue(np.all((p_df['age'] >= 0) & (p_df['age'] <= 100))) + # Check categorical decoding strings + self.assertEqual(list(h_df['region']), ['Urban', 'Rural']) + self.assertEqual(list(p_df['gender']), ['F', 'M', 'F']) + + def test_assign_relational_keys_3_tier(self): + tables = { + 'Household': pd.DataFrame({'income': [50.0, 80.0]}), + 'Person': pd.DataFrame({'age': [25, 30, 45]}), + 'Activity': pd.DataFrame({'amt': [10.0, 20.0, 30.0, 40.0]}), + } + fks = [ + rel_domain.ForeignKeyRelation( + parent_table='Household', + parent_primary_key='hid', + child_table='Person', + child_foreign_key='hid', + max_children_per_parent=2, + ), + rel_domain.ForeignKeyRelation( + parent_table='Person', + parent_primary_key='pid', + child_table='Activity', + child_foreign_key='pid', + max_children_per_parent=2, + ), + ] + parent_mappings = { + 'Person': np.array([0, 0, 1], dtype=np.int64), + 'Activity': np.array([0, 1, 1, 2], dtype=np.int64), + } + hierarchy = [ + (0, 'Household', None), + (1, 'Person', fks[0]), + (2, 'Activity', fks[1]), + ] + + linked = synthesizer._assign_relational_keys( + tables=tables, + foreign_keys=fks, + parent_mappings=parent_mappings, + hierarchy=hierarchy, + ) + + self.assertIn('hid', linked['Household'].columns) + np.testing.assert_array_equal(linked['Household']['hid'], np.array([0, 1])) + + self.assertIn('hid', linked['Person'].columns) + self.assertIn('pid', linked['Person'].columns) + np.testing.assert_array_equal(linked['Person']['hid'], np.array([0, 0, 1])) + np.testing.assert_array_equal(linked['Person']['pid'], np.array([0, 1, 2])) + + self.assertIn('pid', linked['Activity'].columns) + np.testing.assert_array_equal( + linked['Activity']['pid'], np.array([0, 1, 1, 2]) + ) + + def test_assign_relational_keys_branching_and_empty_child(self): + tables = { + 'Household': pd.DataFrame({'income': [50.0, 80.0]}), + 'Person': pd.DataFrame({'age': [25, 30]}), + 'Vehicle': pd.DataFrame({'type': []}), + } + fks = [ + rel_domain.ForeignKeyRelation( + parent_table='Household', + parent_primary_key='hid', + child_table='Person', + child_foreign_key='hid', + max_children_per_parent=2, + ), + rel_domain.ForeignKeyRelation( + parent_table='Household', + parent_primary_key='hid', + child_table='Vehicle', + child_foreign_key='hid', + max_children_per_parent=2, + ), + ] + parent_mappings = { + 'Person': np.array([0, 1], dtype=np.int64), + 'Vehicle': np.empty(0, dtype=np.int64), + } + hierarchy = [ + (0, 'Household', None), + (1, 'Person', fks[0]), + (1, 'Vehicle', fks[1]), + ] + + linked = synthesizer._assign_relational_keys( + tables=tables, + foreign_keys=fks, + parent_mappings=parent_mappings, + hierarchy=hierarchy, + ) + + np.testing.assert_array_equal(linked['Household']['hid'], np.array([0, 1])) + np.testing.assert_array_equal(linked['Person']['hid'], np.array([0, 1])) + self.assertEmpty(linked['Vehicle']) + self.assertIn('hid', linked['Vehicle'].columns) + self.assertEmpty(linked['Vehicle']['hid']) + + def test_multi_table_mechanism_call_end_to_end(self): + rng = np.random.default_rng(42) + h_df = pd.DataFrame({ + 'hid': [0, 1, 2, 3], + 'income': [20.0, 50.0, 80.0, 100.0], + 'region': ['Urban', 'Rural', 'Urban', 'Rural'], + }) + p_df = pd.DataFrame({ + 'hid': [0, 0, 1, 2, 2, 3], + 'pid': [10, 11, 12, 13, 14, 15], + 'age': [15, 45, 30, 20, 55, 60], + 'gender': ['M', 'F', 'M', 'F', 'M', 'F'], + }) + data = {'Household': h_df, 'Person': p_df} + + domains = { + 'Household': { + 'income': domain.NumericalAttribute(min_value=0.0, max_value=100.0), + 'region': domain.CategoricalAttribute( + possible_values=['Urban', 'Rural'] + ), + }, + 'Person': { + 'age': domain.NumericalAttribute(min_value=0, max_value=100), + 'gender': domain.CategoricalAttribute(possible_values=['M', 'F']), + }, + } + fks = [ + rel_domain.ForeignKeyRelation( + parent_table='Household', + parent_primary_key='hid', + child_table='Person', + child_foreign_key='hid', + max_children_per_parent=2, + ) + ] + + cfg = synthesizer.MultiTableConfig( + domains=domains, + foreign_keys=fks, + discrete_mechanism=discrete_mechanisms.AIMConfig( + max_rounds=2, + pgm_iters=50, + ), + numerical_bins=4, + init_budget_fraction=0.2, + ) + mechanism = cfg.configure( + zcdp_rho=1.0, + ) + + result = mechanism(rng=rng, data=data) + + self.assertIsInstance(result, synthesizer.MultiDataGenerationResult) + self.assertIn('Household', result.synthetic_tables) + self.assertIn('Person', result.synthetic_tables) + + synth_h = result.synthetic_tables['Household'] + synth_p = result.synthetic_tables['Person'] + + # Referential integrity check + self.assertIn('hid', synth_h.columns) + self.assertIn('hid', synth_p.columns) + self.assertTrue(set(synth_p['hid']).issubset(set(synth_h['hid']))) + + # Column existence and bounds + self.assertIn('income', synth_h.columns) + self.assertIn('region', synth_h.columns) + self.assertIn('age', synth_p.columns) + self.assertIn('gender', synth_p.columns) + + self.assertTrue( + np.all((synth_h['income'] >= 0.0) & (synth_h['income'] <= 100.0)) + ) + self.assertTrue(np.all((synth_p['age'] >= 0) & (synth_p['age'] <= 100))) + self.assertTrue(set(synth_h['region']).issubset({'Urban', 'Rural'})) + self.assertTrue(set(synth_p['gender']).issubset({'M', 'F'})) + self.assertIn('Household->Person', result.discrete_mechanism_results) + + def test_multi_table_mechanism_call_size_sliced(self): + rng = np.random.default_rng(42) + h_df = pd.DataFrame({ + 'hid': [0, 1, 2, 3], + 'income': [20.0, 50.0, 80.0, 100.0], + 'region': ['Urban', 'Rural', 'Urban', 'Rural'], + }) + p_df = pd.DataFrame({ + 'hid': [0, 0, 1, 2, 2, 3], + 'pid': [10, 11, 12, 13, 14, 15], + 'age': [15, 45, 30, 20, 55, 60], + 'gender': ['M', 'F', 'M', 'F', 'M', 'F'], + }) + data = {'Household': h_df, 'Person': p_df} + + domains = { + 'Household': { + 'income': domain.NumericalAttribute(min_value=0.0, max_value=100.0), + 'region': domain.CategoricalAttribute( + possible_values=['Urban', 'Rural'] + ), + }, + 'Person': { + 'age': domain.NumericalAttribute(min_value=0, max_value=100), + 'gender': domain.CategoricalAttribute(possible_values=['M', 'F']), + }, + } + fks = [ + rel_domain.ForeignKeyRelation( + parent_table='Household', + parent_primary_key='hid', + child_table='Person', + child_foreign_key='hid', + max_children_per_parent=2, + ) + ] + + cfg = synthesizer.MultiTableConfig( + domains=domains, + foreign_keys=fks, + discrete_mechanism=discrete_mechanisms.AIMConfig( + max_rounds=2, + pgm_iters=50, + ), + numerical_bins=4, + init_budget_fraction=0.2, + exploration_strategy='size_sliced', + ) + mechanism = cfg.configure( + zcdp_rho=1.0, + ) + + result = mechanism(rng=rng, data=data) + + self.assertIsInstance(result, synthesizer.MultiDataGenerationResult) + self.assertIn('Household', result.synthetic_tables) + self.assertIn('Person', result.synthetic_tables) + + synth_h = result.synthetic_tables['Household'] + synth_p = result.synthetic_tables['Person'] + + self.assertIn('hid', synth_h.columns) + self.assertIn('hid', synth_p.columns) + self.assertTrue(set(synth_p['hid']).issubset(set(synth_h['hid']))) + self.assertIn('Household->Person', result.discrete_mechanism_results) + if __name__ == '__main__': absltest.main()