diff --git a/dpsynth/relational/README.md b/dpsynth/relational/README.md new file mode 100644 index 0000000..6484417 --- /dev/null +++ b/dpsynth/relational/README.md @@ -0,0 +1,65 @@ +# `dpsynth.relational` (Multi-Table Relational DP Synthesis) + +> [!WARNING] +> **UNDER ACTIVE DEVELOPMENT / EXPERIMENTAL**: This module is in an early +> experimental stage and is under active development. The APIs, internal +> interfaces, and algorithms are subject to breaking changes. It is not yet +> recommended for production use. + +--- + +## Overview + +The `dpsynth.relational` package extends DP Synth to support **hierarchical +multi-table relational databases** under differential privacy. + +In relational databases (e.g. `Household -> Person -> Activity`), privacy is +protected at the **root parent record** (e.g. the household), while valid +foreign-key linkages and cross-table statistical correlations are preserved +across child tables without generating orphaned records or requiring +intractable flat Cartesian joins. + +### Core Approach + +- **Cascading Relational Synthesis**: Synthesizes databases table-by-table + down foreign-key hierarchy trees, ensuring root-entity differential privacy + without materializing expensive Cartesian joins. +- **Permutation Modeling**: Captures rich parent-child and sibling-to-sibling + correlations through slot permutation and exchangeability, inspired by + PrivPetal ([Cai et al., 2025](https://arxiv.org/abs/2503.22970)). +- **Built on `dpsynth` & `mbi` Core**: Reuses existing discrete mechanisms (e.g. + AIM, MST) and the `mbi` (Private-PGM) graphical model engine under the hood. +- **Unified Mechanism API**: Integrates directly with the `MechanismConfig` / + `CalibratedMechanism` paradigm used throughout `dpsynth`. + +--- + +## Planned API + +```python +import dpsynth +from dpsynth.relational import ForeignKeyRelation, MultiTableConfig + +# Define schemas and foreign key relations +foreign_keys = [ + ForeignKeyRelation( + parent_table="households", + parent_primary_key="household_id", + child_table="persons", + child_foreign_key="household_id", + max_children_per_parent=5, + ), +] + +# Configure relational synthesizer +config = MultiTableConfig( + domains=table_domains, + foreign_keys=foreign_keys, + discrete_mechanism=dpsynth.discrete_mechanisms.AIMConfig(), +) + +# Calibrate and run +calibrated = config.calibrate(epsilon=1.0, delta=1e-5) +result = calibrated(rng, {"households": df_h, "persons": df_p}) +synth_tables = result.synthetic_tables +``` diff --git a/dpsynth/relational/__init__.py b/dpsynth/relational/__init__.py new file mode 100644 index 0000000..e692cbe --- /dev/null +++ b/dpsynth/relational/__init__.py @@ -0,0 +1,35 @@ +# 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. + +"""Public API for multi-table relational differential privacy synthesis.""" + +# 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 + +__all__ = [ + 'ForeignKeyRelation', + 'MultiDataGenerationResult', + 'MultiTableConfig', + 'MultiTableMechanism', + 'from_dict', + 'from_yaml_file', + 'topological_sort_hierarchy', +] diff --git a/dpsynth/relational/domain.py b/dpsynth/relational/domain.py new file mode 100644 index 0000000..d2cac55 --- /dev/null +++ b/dpsynth/relational/domain.py @@ -0,0 +1,275 @@ +# 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. + +"""Domain representations, schema definitions, and DAG validators for relational data.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import dataclasses +from typing import Any + +from absl import logging +from dpsynth import domain +from etils import epath +import networkx as nx +import yaml + +PathType = epath.PathLike + + +@dataclasses.dataclass(frozen=True) +class ForeignKeyRelation: + """Defines a directed foreign key relationship between parent and child tables. + + Attributes: + parent_table: Name of the parent table (e.g. 'households'). + parent_primary_key: Name of the parent primary key column (e.g. + 'household_id'). + child_table: Name of the child table (e.g. 'persons'). + child_foreign_key: Name of the child foreign key column referencing parent. + max_children_per_parent: Maximum number of children associated with a single + parent record (group size capacity bound s). Must be >= 1. Determines the + wide MRF generation slot count (s) and directly scales cascading DP + sensitivity (Delta_k = prod s_ancestors) for downstream child tables. + """ + + parent_table: str + parent_primary_key: str + child_table: str + child_foreign_key: str + max_children_per_parent: int + + def __post_init__(self): + if self.max_children_per_parent < 1: + raise ValueError( + 'max_children_per_parent must be >= 1, got' + f' {self.max_children_per_parent}.' + ) + + +def topological_sort_hierarchy( + tables: Sequence[str], + foreign_keys: Sequence[ForeignKeyRelation], +) -> list[tuple[int, str, ForeignKeyRelation | None]]: + """Validates DAG tree structure and computes topological synthesis levels. + + Args: + tables: Sequence of all table names in the database. + foreign_keys: Sequence of foreign key relationships between tables. + + Returns: + An ordered list of (depth, table_name, foreign_key_relation) tuples, where + depth is 0 for root tables (foreign_key_relation is None) and depth >= 1 for + child tables (foreign_key_relation links the table to its immediate parent). + + Raises: + ValueError: If foreign keys contain cycles, missing tables, or if a child + table references more than one parent table (in-degree > 1). + """ + logging.debug( + 'Computing topological sort for %d tables with %d foreign keys.', + len(tables), + len(foreign_keys), + ) + table_set = set(tables) + graph = nx.DiGraph() + graph.add_nodes_from(tables) + + # Maps child_table -> ForeignKeyRelation (incoming edge, e.g. + # 'persons' -> fk_household_person). + incoming_fk_map: dict[str, ForeignKeyRelation] = {} + for fk in foreign_keys: + if fk.parent_table not in table_set or fk.child_table not in table_set: + raise ValueError(f'Foreign key references unknown table in {fk}.') + if fk.child_table in incoming_fk_map: + raise ValueError( + f'Child table {fk.child_table!r} has multiple parents;' + ' in-degree must be <= 1.' + ) + if fk.parent_table == fk.child_table: + raise ValueError(f'Self-referential cycle in table {fk.parent_table!r}.') + incoming_fk_map[fk.child_table] = fk + graph.add_edge(fk.parent_table, fk.child_table) + + if not nx.is_directed_acyclic_graph(graph): + raise ValueError('Cycle detected in foreign keys.') + + roots = [t for t in tables if t not in incoming_fk_map] + logging.debug('Identified %d root privacy unit tables: %s', len(roots), roots) + + result: list[tuple[int, str, ForeignKeyRelation | None]] = [] + for depth, generation in enumerate(nx.topological_generations(graph)): + for table in generation: + result.append((depth, table, incoming_fk_map.get(table))) + + logging.info( + 'Computed topological synthesis levels: %s', + [(d, t) for d, t, _ in result], + ) + return result + + +def _parse_attribute( + table_name: str, col_name: str, spec: Any +) -> domain.AttributeType: + """Parses a single attribute specification.""" + if isinstance( + spec, + ( + domain.CategoricalAttribute, + domain.NumericalAttribute, + domain.OpenSetCategoricalAttribute, + domain.FreeFormTextAttribute, + ), + ): + return spec + if not isinstance(spec, Mapping): + raise ValueError( + f'Invalid attribute specification for {table_name}.{col_name}: {spec}' + ) + attr_data = dict(spec) + attr_data.pop('type', None) + if 'possible_values' in attr_data: + return domain.CategoricalAttribute(**attr_data) + elif 'min_value' in attr_data: + return domain.NumericalAttribute(**attr_data) + elif 'max_tokens' in attr_data: + return domain.FreeFormTextAttribute(**attr_data) + elif 'default_value' in attr_data or not attr_data: + return domain.OpenSetCategoricalAttribute(**attr_data) + else: + raise ValueError( + f'Invalid attribute specification for {table_name}.{col_name}:' + f' {attr_data}' + ) + + +def from_dict( + config: Mapping[str, Any], +) -> tuple[dict[str, domain.Schema], list[ForeignKeyRelation]]: + """Parses multi-table schema and foreign keys from a dictionary. + + Args: + config: Dictionary with 'tables' and optional 'foreign_keys' blocks. + + Returns: + A tuple of (table_domains, foreign_keys). + + Raises: + ValueError: If configuration format or attribute specifications are invalid. + """ + if 'tables' not in config or not isinstance(config['tables'], Mapping): + raise ValueError("'tables' block missing or invalid in config.") + + logging.debug( + 'Parsing multi-table schema dictionary for %d tables.', + len(config['tables']), + ) + table_domains: dict[str, domain.Schema] = {} + 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] = { + 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', []): + if isinstance(fk, ForeignKeyRelation): + foreign_keys.append(fk) + elif isinstance(fk, Mapping): + foreign_keys.append(ForeignKeyRelation(**fk)) + else: + raise ValueError(f'Invalid foreign key specification: {fk}') + + logging.info( + 'Successfully parsed multi-table schema: %d tables, %d foreign keys.', + len(table_domains), + len(foreign_keys), + ) + return table_domains, foreign_keys + + +def from_yaml_file( + filepath: str | PathType, +) -> tuple[dict[str, domain.Schema], list[ForeignKeyRelation]]: + """Reads multi-table schema and foreign keys from a YAML file. + + Args: + filepath: Path to the YAML schema file. + + Returns: + A tuple of (table_domains, foreign_keys). + """ + logging.info('Loading relational domain schema from YAML file: %s', filepath) + path = epath.Path(filepath) + with path.open('r') as f: + config = yaml.safe_load(f) + if not isinstance(config, Mapping): + raise ValueError(f'YAML root in {filepath} must be a mapping.') + return from_dict(config) + + +def to_dict( + table_domains: Mapping[str, domain.Schema], + foreign_keys: Sequence[ForeignKeyRelation] = (), +) -> dict[str, Any]: + """Converts multi-table schemas and foreign keys to a dictionary. + + Args: + table_domains: Mapping from table name to per-column AttributeType schemas. + foreign_keys: Optional sequence of ForeignKeyRelation objects. + + Returns: + A dictionary with 'tables' and optional 'foreign_keys' blocks. + """ + tables_dict: dict[str, dict[str, Any]] = {} + for table_name, schema in table_domains.items(): + table_dict: dict[str, Any] = {} + for col_name, attr in schema.items(): + attr_dict = dataclasses.asdict(attr) + attr_dict['type'] = attr.__class__.__name__ + table_dict[col_name] = attr_dict + tables_dict[table_name] = table_dict + + result: dict[str, Any] = {'tables': tables_dict} + if foreign_keys: + result['foreign_keys'] = [dataclasses.asdict(fk) for fk in foreign_keys] + return result + + +def to_yaml_file( + table_domains: Mapping[str, domain.Schema], + foreign_keys: Sequence[ForeignKeyRelation], + filepath: str | PathType, +) -> None: + """Writes multi-table schema and foreign keys to a YAML file. + + Args: + table_domains: Mapping from table name to per-column AttributeType schemas. + foreign_keys: Sequence of ForeignKeyRelation objects. + filepath: Destination path for the YAML schema file. + """ + logging.info( + 'Saving relational domain schema (%d tables, %d foreign keys) to: %s', + len(table_domains), + len(foreign_keys), + filepath, + ) + data = to_dict(table_domains, foreign_keys) + path = epath.Path(filepath) + with path.open('w') as f: + yaml.dump(data, f, default_flow_style=False, sort_keys=False) diff --git a/dpsynth/relational/synthesizer.py b/dpsynth/relational/synthesizer.py new file mode 100644 index 0000000..fded2b4 --- /dev/null +++ b/dpsynth/relational/synthesizer.py @@ -0,0 +1,134 @@ +# 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. + +"""Multi-table relational differential privacy synthesizer.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import dataclasses +from typing import Any, Literal + +from absl import logging +import dp_accounting +from dpsynth import api +from dpsynth import discrete_mechanisms +from dpsynth import domain +from dpsynth.relational import domain as rel_domain +import numpy as np +import pandas as pd + +# pylint: disable=unused-import +_LOGGING_UNUSED = logging +# pylint: enable=unused-import + + +@dataclasses.dataclass(frozen=True) +class MultiDataGenerationResult: + """Results of multi-table relational DP synthetic data generation. + + Attributes: + synthetic_tables: Mapping from table names to synthetic DataFrames. + discrete_mechanism_results: Mapping from link/table names to mechanism + diagnostics. + """ + + synthetic_tables: Mapping[str, pd.DataFrame] + discrete_mechanism_results: Mapping[str, Any] = dataclasses.field( + default_factory=dict + ) + + +@dataclasses.dataclass +class MultiTableMechanism(api.CalibratedMechanism): + """Calibrated, runnable multi-table relational differential privacy mechanism. + + Attributes: + domains: Mapping from table name to per-column attribute specifications. + foreign_keys: Sequence of foreign key relationships defining the hierarchy. + calibrated_discrete_mechanisms: Mapping from link names to calibrated + discrete mechanisms. + calibrated_initializers: Mapping from table and column to calibrated + initializers. + total_count_sigma: Sigma for the root table total-count mechanism. + num_permutation_slots: Permutation exploration slot count (o), default 2. + exploration_strategy: Exploration strategy ('empty_token' or 'size_sliced'). + max_records_per_user: Assumed upper bound on records a single root user + contributes. + + Note: For simplicity, user-defined contraints are not supported yet. + """ + + domains: Mapping[str, domain.Schema] + foreign_keys: Sequence[rel_domain.ForeignKeyRelation] + calibrated_discrete_mechanisms: Mapping[str, api.CalibratedMechanism] + calibrated_initializers: Mapping[str, Mapping[str, api.CalibratedMechanism]] + total_count_sigma: float = dataclasses.field(repr=False) + num_permutation_slots: int = 2 + exploration_strategy: Literal['empty_token', 'size_sliced'] = 'empty_token' + max_records_per_user: int = 1 + + @property + def dp_event(self) -> dp_accounting.DpEvent: + """Returns the composed DpEvent for all sub-mechanisms.""" + raise NotImplementedError('dp_event is not yet implemented.') + + def __call__( + self, + 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.' + ) + + +@dataclasses.dataclass +class MultiTableConfig(api.MechanismConfig): + """Configuration recipe for multi-table relational differential privacy synthesis. + + Attributes: + domains: Mapping from table name to per-column attribute domain + specifications. + foreign_keys: Sequence of foreign key relationships defining the hierarchy. + 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. + 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.DiscreteMechanismConfig = ( + dataclasses.field(default_factory=discrete_mechanisms.AIMConfig) + ) + numerical_bins: int = 32 + init_budget_fraction: float = 0.1 + num_permutation_slots: int = 2 + exploration_strategy: Literal['empty_token', 'size_sliced'] = 'empty_token' + + def configure( + self, + *, + zcdp_rho: float, + delta: float = 0.0, + max_records_per_user: int = 1, + ) -> MultiTableMechanism: + """Configures privacy budgets across Phase 1 and Phase 2 relational links.""" + del zcdp_rho, delta, max_records_per_user + raise NotImplementedError('configure is not yet implemented.') diff --git a/dpsynth/relational/transformations.py b/dpsynth/relational/transformations.py new file mode 100644 index 0000000..5b9d81b --- /dev/null +++ b/dpsynth/relational/transformations.py @@ -0,0 +1,606 @@ +# 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. + +"""Pure, deterministic relational data transformers and mathematical helpers.""" + +from __future__ import annotations + +from collections.abc import Hashable, Mapping, Sequence +import itertools +from typing import Literal + +from dpsynth.relational import domain as rel_domain +import mbi +import numpy as np +import pandas as pd + + +def _compute_row_root_mappings( + tables: Mapping[str, pd.DataFrame], + hierarchy: Sequence[tuple[int, str, rel_domain.ForeignKeyRelation | None]], + rng: np.random.Generator | None = None, +) -> dict[str, pd.Series]: + """Maps each row in every table to its root parent (root_table, root_row_idx). + + Traverses the relational hierarchy top-down. For each child table, maps + foreign key references to parent row positions and checks capacity limits. + When a parent record exceeds its `max_children_per_parent` bound (s), exactly + `s` children are selected uniformly at random without replacement. + Truncated, orphaned, or descendant records under dropped parents map to None. + + Example: + households: [H0, H1] + persons (s1=2): [P0(H0), P1(H0), P2(H0), P3(H1)] -> P2 truncated (None) + activities (s2=2): [A0(P0), A1(P1), A2(P2), A3(P3), A4(orphan)] + -> A0 maps to ('households', 0) + -> A1 maps to ('households', 0) + -> A2 maps to None (parent P2 was truncated) + -> A3 maps to ('households', 1) + -> A4 maps to None (orphan foreign key) + + Result: + { + 'households': pd.Series([('households', 0), ('households', 1)]), + 'persons': pd.Series([('households', 0), ('households', 0), None, + ('households', 1)]), + 'activities': pd.Series([('households', 0), ('households', 0), None, + ('households', 1), None]), + } + + Args: + + Args: + tables: Mapping from table name to input DataFrame. + hierarchy: Ordered topological synthesis levels from + `topological_sort_hierarchy()`. + rng: Random number generator for uniform child record truncation. + + Returns: + A dictionary mapping table name to a pd.Series of root identifier tuples or + None, aligned 1-to-1 with DataFrame rows. + + Raises: + ValueError: If required primary or foreign key columns are missing from + schemas. + """ + if rng is None: + rng = np.random.default_rng() + + row_to_root: dict[str, pd.Series] = {} + for depth, table_name, fk in hierarchy: + child_df = tables[table_name] + + # Depth 0: Root privacy unit table (no incoming foreign key). + # Each root record is its own root ancestor: (root_table, row_idx). + if depth == 0 or fk is None: + row_to_root[table_name] = pd.Series( + pd.MultiIndex.from_product( + [[table_name], range(len(child_df))] + ).to_numpy(), + index=child_df.index, + dtype=object, + ) + continue + + # Schema integrity validation (public schema check; safe to raise errors). + if fk.parent_primary_key not in tables[fk.parent_table].columns: + raise ValueError( + f'Parent primary key column {fk.parent_primary_key!r} not in table' + f' {fk.parent_table!r}.' + ) + if fk.child_foreign_key not in child_df.columns: + raise ValueError( + f'Child foreign key column {fk.child_foreign_key!r} not in table' + f' {table_name!r}.' + ) + + parent_df = tables[fk.parent_table] + parent_roots = row_to_root[fk.parent_table] + + # Fast path for empty tables: returns all None without failing. + if child_df.empty or parent_df.empty: + row_to_root[table_name] = pd.Series( + [None] * len(child_df), index=child_df.index, dtype=object + ) + continue + + # 1. Parent lookup maps parent primary keys to row numbers in parent_df. + # Ignores NaN primary keys and deduplicates repeated keys (keeping first). + # parent_lookup: (Index = parent_pk, Value = parent_df row index). + parent_pos = pd.Series(range(len(parent_df)), index=parent_df.index) + try: + parent_valid_mask = ( + parent_df[fk.parent_primary_key].notna() # No NaN keys. + & ~parent_df[fk.parent_primary_key].duplicated() # Keep only first. + ) + parent_keys = parent_df.loc[parent_valid_mask, fk.parent_primary_key] + parent_lookup = pd.Series( + parent_pos.loc[parent_valid_mask].values, index=parent_keys + ) + except TypeError: + # Safe DP fallback if primary key column contains unhashable objects. + hashable_mask = [ + pd.notna(v) and isinstance(v, Hashable) + for v in parent_df[fk.parent_primary_key] + ] + filtered_parent_df = parent_df[hashable_mask] + filtered_parent_pos = parent_pos[hashable_mask] + parent_valid_mask = ~filtered_parent_df[ + fk.parent_primary_key + ].duplicated() + parent_keys = filtered_parent_df.loc[ + parent_valid_mask, fk.parent_primary_key + ] + parent_lookup = pd.Series( + filtered_parent_pos.loc[parent_valid_mask].values, index=parent_keys + ) + + # 2. Vectorized translation of child foreign keys to parent_df row indices. + # Non-matching keys (orphans), NaNs, and malformed entries evaluate to NaN. + # child_p_idx : (Index = child row, Value = parent row | NaN). + try: + child_p_idx = child_df[fk.child_foreign_key].map(parent_lookup) + except TypeError: + # Safe DP fallback: handles unhashable or unsupported object types. + lookup_dict = parent_lookup.to_dict() + child_p_idx = pd.Series( + [ + lookup_dict.get(v, np.nan) + if (pd.notna(v) and isinstance(v, Hashable)) + else np.nan + for v in child_df[fk.child_foreign_key] + ], + index=child_df.index, + ) + + # 3. Discard unlinked children + # valid_children: (Index = child row, Value = parent row). + valid_children = child_p_idx.dropna().astype(int) + if valid_children.empty: + row_to_root[table_name] = pd.Series( + [None] * len(child_df), index=child_df.index, dtype=object + ) + continue + + # 4. Enforce cascading truncation: drop children whose parent root is None. + # Map valid parent indices back to parent_roots positions + parent_active_mask = parent_roots.notna().iloc[valid_children.values].values + valid_children = valid_children[parent_active_mask] + if valid_children.empty: + row_to_root[table_name] = pd.Series( + [None] * len(child_df), index=child_df.index, dtype=object + ) + continue + + # 5. Intra-group uniform random ranking via Pandas, for uniform truncation. + # Assigns random float to each child; ranks within each parent group. + # random_scores: (Index = child row, Value = random float). + # group_ranks: (Index = child row, Value = rank 1-to-n within parent group). + random_scores = pd.Series( + rng.random(len(valid_children)), index=valid_children.index + ) + group_ranks = random_scores.groupby(valid_children.values).rank( + method='first' + ) + selected_children = valid_children[ + group_ranks <= fk.max_children_per_parent + ] + + # 6. Assign root lineages to selected children in Pandas. + # Initialize full child table to None; update only selected active rows. + child_roots_series = pd.Series( + [None] * len(child_df), index=child_df.index, dtype=object + ) + child_roots_series.loc[selected_children.index] = parent_roots.iloc[ + selected_children.values + ].values + + row_to_root[table_name] = child_roots_series + return row_to_root + + +def compute_hierarchical_weights( + tables: Mapping[str, pd.DataFrame], + hierarchy: Sequence[tuple[int, str, rel_domain.ForeignKeyRelation | None]], + rng: np.random.Generator | None = None, +) -> dict[str, np.ndarray]: + """Computes standalone sensitivity weights (w = 1/k_eff) for Phase 1 initializers. + + Calculates a 1D weight array for each table such that the sum of weights + associated with any single root entity (e.g. household) equals 1.0, ensuring + global unit sensitivity (Delta = 1.0) without Cartesian joins or noise + scaling. + + For each table, active records belonging to a root with k_eff active rows + receive weight 1.0 / k_eff. Inactive rows (truncated or orphaned) receive 0.0. + + Args: + tables: Mapping from table name to input DataFrame. + hierarchy: Ordered topological synthesis levels from + `topological_sort_hierarchy()`. + rng: Random number generator for child record truncation. + + Returns: + A dictionary mapping table name to a 1D float64 array of row weights. + """ + row_to_root = _compute_row_root_mappings(tables, hierarchy, rng=rng) + + weights: dict[str, np.ndarray] = {} + for table_name, roots in row_to_root.items(): + root_counts = roots.value_counts() + table_weights = ( + roots.map(1.0 / root_counts).fillna(0.0).to_numpy(dtype=np.float64) + ) + weights[table_name] = table_weights + + return weights + + +def _build_exploration_domain( + parent_domain: mbi.Domain, + child_domain: mbi.Domain, + max_group_size: int, + num_permutation_slots: int, + strategy: Literal['empty_token', 'size_sliced'], +) -> mbi.Domain: + """Constructs the discrete mbi.Domain for the permuted exploration table. + + Args: + parent_domain: Domain of parent table attributes. + child_domain: Domain of single-child attributes. + max_group_size: Maximum observed or clipped child group size. + num_permutation_slots: Permutation exploration slot count (o). + strategy: 'empty_token' (extends domain by +1 for ) or 'size_sliced'. + + Returns: + An mbi.Domain encompassing parent columns, group_size, and o child slots. + """ + + # Add group_size to parent domain (cardinality = max_group_size + 1). + attrs = list(parent_domain.attributes) + ['group_size'] + shapes = list(parent_domain.shape) + [max_group_size + 1] + + # Add o slots, each with size+1 for empty_token or size for size_sliced. + for i in range(1, num_permutation_slots + 1): + for attr, size in zip(child_domain.attributes, child_domain.shape): + slot_size = size + 1 if strategy == 'empty_token' else size + attrs.append(f'slot_{i}.{attr}') + shapes.append(slot_size) + return mbi.Domain(tuple(attrs), tuple(shapes)) + + +def _get_slot_permutation_patterns( + k: int, + num_permutation_slots: int, + strategy: Literal['empty_token', 'size_sliced'], +) -> tuple[list[tuple[int, ...]], float]: + """Generates slot index permutation patterns and row weight for (k, o). + + Args: + k: Number of children in the household. + num_permutation_slots: Permutation exploration slot count (o). + strategy: 'empty_token' (permutes real and ) or 'size_sliced' (clone + tiling). + + Returns: + A tuple of (patterns, weight) where patterns is a list of o-tuples with + child relative indices [0, k-1] (or -1 for ), and weight is the float + weight for each emitted row such that len(patterns) * weight == 1.0. + """ + if k == 0: + empty_val = -1 if strategy == 'empty_token' else 0 + return [tuple(empty_val for _ in range(num_permutation_slots))], 1.0 + + if k < num_permutation_slots: + if strategy == 'size_sliced': + return [tuple(i % k for i in range(num_permutation_slots))], 1.0 + items = list(range(k)) + [-1] * (num_permutation_slots - k) + patterns = list(dict.fromkeys(itertools.permutations(items))) + return patterns, 1.0 / len(patterns) + + patterns = list(itertools.permutations(range(k), num_permutation_slots)) + return patterns, 1.0 / len(patterns) + + +def build_permuted_exploration_dataset( + parent_dataset: mbi.Dataset, + child_dataset: mbi.Dataset, + parent_primary_keys: Sequence[Hashable] | np.ndarray | pd.Series, + child_foreign_keys: Sequence[Hashable] | np.ndarray | pd.Series, + max_group_size: int, + num_permutation_slots: int = 2, + strategy: Literal['empty_token', 'size_sliced'] = 'empty_token', +) -> mbi.Dataset: + """Constructs the permuted multi-slot exploration dataset for candidate selection. + + Args: + parent_dataset: Encoded discrete mbi.Dataset for the parent table. + child_dataset: Encoded discrete mbi.Dataset for the child table. + parent_primary_keys: Sequence or array of parent primary key identifiers. + child_foreign_keys: Sequence or array of child foreign key references. + max_group_size: Public upper bound for child group capacity (s >= 1). + num_permutation_slots: Number of permutation slots (o) in exploration table, + default 2. + strategy: Exploration strategy ('empty_token' with or + 'size_sliced'). + + Returns: + An mbi.Dataset instance representing the permuted exploration table. + + Raises: + ValueError: If strategy is unsupported, num_permutation_slots < 1, + max_group_size < 1, or key lengths do not match dataset record counts. + """ + + # This input validation can probably later be moved earlier in the pipeline. + if max_group_size < 1: + raise ValueError(f'max_group_size must be >= 1, got {max_group_size}') + if num_permutation_slots < 1: + raise ValueError( + f'num_permutation_slots must be >= 1, got {num_permutation_slots}' + ) + if strategy not in ('empty_token', 'size_sliced'): + raise ValueError( + f"strategy must be 'empty_token' or 'size_sliced', got {strategy!r}" + ) + + num_parents = parent_dataset.records + if len(parent_primary_keys) != num_parents: + raise ValueError( + f'parent_primary_keys length ({len(parent_primary_keys)}) does not' + f' match parent_dataset records ({num_parents})' + ) + if len(child_foreign_keys) != child_dataset.records: + raise ValueError( + f'child_foreign_keys length ({len(child_foreign_keys)}) does not' + f' match child_dataset records ({child_dataset.records})' + ) + + # Vectorized translation of child foreign keys to parent row indices. + # parent_lookup: (Index = parent_pk, Value = parent row index [0, N_p-1]). + parent_lookup = pd.Series( + np.arange(num_parents), index=pd.Series(parent_primary_keys).values + ) + parent_lookup = parent_lookup[~parent_lookup.index.duplicated(keep='first')] + + # child_parent_idx: (Index = child row, Value = parent row index [0, N_p-1]). + child_parent_idx = ( + pd.Series(child_foreign_keys).map(parent_lookup).dropna().astype(int) + ) + + # Vectorized child counting and intra-parent ranking. + # parent_group_sizes: (Index = parent row, Value = child count k). + # child_ranks: 1D array of N_c intra-parent ranks (0, 1, ... k-1) per child. + parent_group_sizes = pd.Series(0, index=np.arange(num_parents)) + if not child_parent_idx.empty: + counts = child_parent_idx.value_counts() + parent_group_sizes.loc[counts.index] = counts.values + child_ranks = ( + child_parent_idx.groupby(child_parent_idx).cumcount().to_numpy() + ) + else: + child_ranks = np.empty(0, dtype=int) + + # Construct exploration domain fixed strictly by public max_group_size. + exploration_domain = _build_exploration_domain( + parent_domain=parent_dataset.domain, + child_domain=child_dataset.domain, + max_group_size=max_group_size, + num_permutation_slots=num_permutation_slots, + strategy=strategy, + ) + + parent_cols = list(parent_dataset.domain.attributes) + child_cols = list(child_dataset.domain.attributes) + + # Example: real children have ages [0,9] -> empty_token: {'age': 10 = |[0,9]|} + empty_tokens = dict( + zip(child_dataset.domain.attributes, child_dataset.domain.shape) + ) + + # Vectorized block assembly grouped by unique family size k. + # Avoids iterating over all N_p parent rows by processing all parents of the + # same family size in bulk NumPy operations (at most s iterations total). + block_arrays: dict[str | int, list[np.ndarray]] = { + attr: [] for attr in exploration_domain.attributes + } + weights_blocks: list[np.ndarray] = [] + + valid_child_parents = child_parent_idx.to_numpy() + valid_child_rows = child_parent_idx.index.to_numpy() + + for k in np.unique(parent_group_sizes.values): + parent_indices_k = np.where(parent_group_sizes.values == k)[0] + n_k = len(parent_indices_k) + if n_k == 0: + continue + + # Get permutation template matrix (p_k, o) and row weight (w = 1 / p_k). + # Examples for o=2 under strategy='empty_token': + # k=0: pattern_matrix = [[-1, -1]] (p_k=1, weight=1.0) + # k=1: pattern_matrix = [[ 0, -1], [-1, 0]] (p_k=2, weight=0.5) + # k=2: pattern_matrix = [[ 0, 1], [ 1, 0]] (p_k=2, weight=0.5) + # k=3: pattern_matrix = [[0,1],[0,2],[1,0],[1,2]...] (p_k=6, weight=1/6) + patterns, weight = _get_slot_permutation_patterns( + int(k), num_permutation_slots, strategy + ) + p_k = len(patterns) + pattern_matrix = np.array(patterns, dtype=np.int64) + + # Broadcast parent features, group_size, and weights (length = n_k * p_k). + # Example: parent_indices_k=[0, 1], p_k=2 -> parent_rep=[0, 0, 1, 1] + parent_rep = np.repeat(parent_indices_k, p_k) + weights_blocks.append(np.full(n_k * p_k, weight, dtype=np.float64)) + block_arrays['group_size'].append(np.full(n_k * p_k, k, dtype=np.int64)) + for p_col in parent_cols: + block_arrays[p_col].append(parent_dataset.data[p_col][parent_rep]) + + # Broadcast child slot features across repeated parents. + if k == 0: + # Childless: fill all o slots with (or 0 for size_sliced). + # Example for o=2: empty_tokens={'age': 10} -> slot_1=[10], slot_2=[10] + for slot_idx in range(1, num_permutation_slots + 1): + for c_col in child_cols: + val = empty_tokens[c_col] if strategy == 'empty_token' else 0 + block_arrays[f'slot_{slot_idx}.{c_col}'].append( + np.full(n_k * p_k, val, dtype=np.int64) + ) + else: + # Multi-child: construct 2D index grid (n_k, k) mapping + # (local_parent_idx, intra_group_rank) -> global child row index. + # Example: H_A has children [10, 11], H_B has [20, 21] + # -> child_grid = [[10, 11], [20, 21]] (row=household, col=sibling rank) + child_mask_k = np.isin(valid_child_parents, parent_indices_k) + local_parent_pos = ( + pd.Series(np.arange(n_k), index=parent_indices_k) + .loc[valid_child_parents[child_mask_k]] + .to_numpy() + ) + + child_grid = np.empty((n_k, k), dtype=np.int64) + child_grid[local_parent_pos, child_ranks[child_mask_k]] = ( + valid_child_rows[child_mask_k] + ) + + # Broadcast relative permutation slot indices into global child rows. + # Example for o=2, pattern_matrix=[[0, 1], [1, 0]]: + # Slot 1: rel_indices=tile([0, 1], 2) -> [0, 1, 0, 1] + # picks: (0,0)->10, (0,1)->11, (1,0)->20, (1,1)->21 + # Slot 2: rel_indices=tile([1, 0], 2) -> [1, 0, 1, 0] + # picks: (0,1)->11, (0,0)->10, (1,1)->21, (1,0)->20 + local_parent_rep = np.repeat(np.arange(n_k), p_k) + for slot_idx in range(1, num_permutation_slots + 1): + rel_indices = np.tile(pattern_matrix[:, slot_idx - 1], n_k) + is_empty = rel_indices == -1 + safe_rel_indices = np.maximum(rel_indices, 0) + active_child_rows = child_grid[local_parent_rep, safe_rel_indices] + + for c_col in child_cols: + slot_col_name = f'slot_{slot_idx}.{c_col}' + child_vals = child_dataset.data[c_col][active_child_rows] + if strategy == 'empty_token': + # For empty slot (-1), replace dummy child row with token. + child_vals = np.where(is_empty, empty_tokens[c_col], child_vals) + block_arrays[slot_col_name].append(child_vals) + + # Concatenate block arrays into final mbi.Dataset. + data_arrays: dict[str | int, np.ndarray] = { + attr: np.concatenate(arrs) if arrs else np.empty(0, dtype=np.int64) + for attr, arrs in block_arrays.items() + } + weights_array = ( + np.concatenate(weights_blocks) + if weights_blocks + else np.empty(0, dtype=np.float64) + ) + return mbi.Dataset(data_arrays, exploration_domain, weights=weights_array) + + +def create_slot_linear_chain_constraints( + child_domain: mbi.Domain, + num_permutation_slots: int = 2, +) -> list[mbi.Constraint]: + """Creates adjacent pairwise mbi.Constraint objects for monolithic slot validity. + + For each slot, generates D-1 pairwise adjacent constraints ((S_i.A_1, + S_i.A_2), + (S_i.A_2, S_i.A_3), ...) setting log-potential to -inf on mixed states, + ensuring + sampled slots are 100% Real or 100% with bounded treewidth <= 2. + + Args: + child_domain: Sub-domain representing attributes of a single child record. + num_permutation_slots: Number of permutation slots (o), default 2. + + Returns: + A list of mbi.Constraint instances enforcing monolithic slot locking. + """ + del child_domain, num_permutation_slots + raise NotImplementedError( + 'create_slot_linear_chain_constraints is not yet implemented.' + ) + + +def symmetrize_to_wide_domain( + measurements: Sequence[mbi.LinearMeasurement], + max_children_per_parent: int, + num_permutation_slots: int = 2, +) -> list[mbi.LinearMeasurement]: + """Replicates selected exploration measurements across all generation slots. + + Equivariantly replicates candidate measurements from (S_1) and (S_1, S_2) + to all s slots and all comb(s, 2) sibling pairs in the wide generation MRF. + + Args: + measurements: Noisy marginal measurements from exploration candidate + selection. + max_children_per_parent: Maximum group capacity bound (s). + num_permutation_slots: Number of permutation exploration slots (o), default + 2. + + Returns: + A list of expanded LinearMeasurement objects for the wide generation MRF. + """ + del measurements, max_children_per_parent, num_permutation_slots + raise NotImplementedError('symmetrize_to_wide_domain is not yet implemented.') + + +def quantile_copula_coupling( + synth_parents: mbi.Dataset, + synth_wide_children: mbi.Dataset, + parent_columns: Sequence[str], + rng: np.random.Generator | None = None, +) -> mbi.Dataset: + """Couples synthetic parents and wide child records via Quantile Copula Matching. + + Applies randomized within-bin tie-breaking and lexicographical sorting along + parent feature coordinates to align parent records with wide child records. + + Args: + synth_parents: Discrete mbi.Dataset of synthesized parent records. + synth_wide_children: Discrete mbi.Dataset of synthesized wide child records. + parent_columns: Parent feature columns used as the coupling anchor. + rng: Random number generator for within-bin tie-breaking permutation. + + Returns: + The coupled wide child discrete mbi.Dataset aligned with synthetic parents. + """ + del synth_parents, synth_wide_children, parent_columns, rng + raise NotImplementedError('quantile_copula_coupling is not yet implemented.') + + +def unstack_wide_family_records( + synth_wide_dataset: mbi.Dataset, + child_domain: mbi.Domain, + max_children_per_parent: int, +) -> tuple[mbi.Dataset, np.ndarray]: + """Unstacks wide family records into a standard normalized child mbi.Dataset. + + Reads group_size = k on each wide row, emits the active child records, and + returns the unstacked child dataset along with a 1D mapping array of parent + row indices. + + Args: + synth_wide_dataset: Discrete mbi.Dataset of wide family records. + child_domain: Single-child mbi.Domain defining attribute sizes. + max_children_per_parent: Maximum group capacity bound (s). + + Returns: + A tuple of (unstacked_child_dataset, parent_row_indices) where + parent_row_indices maps each unstacked child record to its parent row. + """ + del synth_wide_dataset, child_domain, max_children_per_parent + raise NotImplementedError( + 'unstack_wide_family_records is not yet implemented.' + ) diff --git a/tests/relational/domain_test.py b/tests/relational/domain_test.py new file mode 100644 index 0000000..6308322 --- /dev/null +++ b/tests/relational/domain_test.py @@ -0,0 +1,315 @@ +# 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 dpsynth.relational.domain.""" + +import textwrap + +from absl.testing import absltest +from dpsynth import domain as base_domain +from dpsynth.relational import domain + + +class DomainTest(absltest.TestCase): + + def test_foreign_key_relation_initialization(self): + fk = domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='household_id', + child_table='persons', + child_foreign_key='household_id', + max_children_per_parent=5, + ) + self.assertEqual(fk.parent_table, 'households') + self.assertEqual(fk.parent_primary_key, 'household_id') + self.assertEqual(fk.child_table, 'persons') + self.assertEqual(fk.child_foreign_key, 'household_id') + self.assertEqual(fk.max_children_per_parent, 5) + + def test_foreign_key_relation_invalid_capacity(self): + with self.assertRaisesRegex( + ValueError, 'max_children_per_parent must be >= 1' + ): + domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='household_id', + child_table='persons', + child_foreign_key='household_id', + max_children_per_parent=0, + ) + + def test_topological_sort_linear_chain(self): + tables = ['households', 'persons', 'activities'] + fk1 = domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 3) + fk2 = domain.ForeignKeyRelation('persons', 'pid', 'activities', 'pid', 2) + order = domain.topological_sort_hierarchy(tables, [fk1, fk2]) + self.assertEqual( + order, + [ + (0, 'households', None), + (1, 'persons', fk1), + (2, 'activities', fk2), + ], + ) + + def test_topological_sort_branching(self): + tables = ['households', 'persons', 'vehicles'] + fk_p = domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 3) + fk_v = domain.ForeignKeyRelation('households', 'hid', 'vehicles', 'hid', 2) + order = domain.topological_sort_hierarchy(tables, [fk_p, fk_v]) + self.assertEqual(order[0], (0, 'households', None)) + self.assertCountEqual( + order[1:], [(1, 'persons', fk_p), (1, 'vehicles', fk_v)] + ) + + def test_topological_sort_forest(self): + tables = ['households', 'persons', 'companies', 'departments'] + fk_h = domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 3) + fk_c = domain.ForeignKeyRelation( + 'companies', 'cid', 'departments', 'cid', 10 + ) + order = domain.topological_sort_hierarchy(tables, [fk_h, fk_c]) + depths = {t: d for d, t, _ in order} + self.assertEqual(depths['households'], 0) + self.assertEqual(depths['companies'], 0) + self.assertEqual(depths['persons'], 1) + self.assertEqual(depths['departments'], 1) + + def test_topological_sort_single_table_and_empty(self): + self.assertEqual( + domain.topological_sort_hierarchy(['households'], []), + [(0, 'households', None)], + ) + self.assertEqual(domain.topological_sort_hierarchy([], []), []) + + def test_topological_sort_cycle_raises_error(self): + tables = ['a', 'b'] + fk1 = domain.ForeignKeyRelation('a', 'id', 'b', 'id', 1) + fk2 = domain.ForeignKeyRelation('b', 'id', 'a', 'id', 1) + with self.assertRaisesRegex(ValueError, 'Cycle detected'): + domain.topological_sort_hierarchy(tables, [fk1, fk2]) + + def test_topological_sort_self_cycle_raises_error(self): + tables = ['a'] + fk = domain.ForeignKeyRelation('a', 'id', 'a', 'id', 1) + with self.assertRaisesRegex(ValueError, 'Self-referential cycle'): + domain.topological_sort_hierarchy(tables, [fk]) + + def test_topological_sort_multiple_parents_raises_error(self): + tables = ['p1', 'p2', 'c'] + fk1 = domain.ForeignKeyRelation('p1', 'id', 'c', 'p1_id', 1) + fk2 = domain.ForeignKeyRelation('p2', 'id', 'c', 'p2_id', 1) + with self.assertRaisesRegex(ValueError, 'multiple parents'): + domain.topological_sort_hierarchy(tables, [fk1, fk2]) + + def test_topological_sort_unknown_table_raises_error(self): + tables = ['households'] + fk = domain.ForeignKeyRelation('households', 'hid', 'unknown', 'hid', 1) + with self.assertRaisesRegex(ValueError, 'unknown table'): + domain.topological_sort_hierarchy(tables, [fk]) + + def test_from_dict_valid_3tier_schema(self): + config = { + 'tables': { + 'households': { + 'income': {'min_value': 0.0, 'max_value': 200000.0}, + 'region': {'possible_values': ['Urban', 'Rural']}, + }, + 'persons': { + 'age': {'min_value': 0, 'max_value': 100, 'dtype': 'int'}, + 'gender': {'possible_values': ['M', 'F']}, + }, + }, + 'foreign_keys': [{ + 'parent_table': 'households', + 'parent_primary_key': 'household_id', + 'child_table': 'persons', + 'child_foreign_key': 'household_id', + 'max_children_per_parent': 3, + }], + } + table_domains, fks = domain.from_dict(config) + self.assertIn('households', table_domains) + self.assertIn('persons', table_domains) + self.assertIsInstance( + table_domains['households']['income'], base_domain.NumericalAttribute + ) + self.assertIsInstance( + table_domains['households']['region'], + base_domain.CategoricalAttribute, + ) + self.assertLen(fks, 1) + self.assertEqual(fks[0].parent_table, 'households') + self.assertEqual(fks[0].parent_primary_key, 'household_id') + self.assertEqual(fks[0].child_table, 'persons') + self.assertEqual(fks[0].child_foreign_key, 'household_id') + self.assertEqual(fks[0].max_children_per_parent, 3) + + def test_from_dict_missing_tables_block_raises_error(self): + with self.assertRaisesRegex(ValueError, "'tables' block missing"): + domain.from_dict({}) + + def test_from_dict_invalid_table_schema_raises_error(self): + with self.assertRaisesRegex(ValueError, 'must be a mapping'): + domain.from_dict({'tables': {'households': 'invalid'}}) + + def test_from_dict_invalid_attribute_spec_raises_error(self): + with self.assertRaisesRegex(ValueError, 'Invalid attribute specification'): + domain.from_dict({'tables': {'households': {'income': 123}}}) + + def test_from_dict_invalid_foreign_key_spec_raises_error(self): + with self.assertRaisesRegex(ValueError, 'Invalid foreign key'): + domain.from_dict({'tables': {'h': {}}, 'foreign_keys': [123]}) + + def test_from_yaml_file_roundtrip(self): + yaml_content = textwrap.dedent("""\ + tables: + households: + income: + min_value: 0.0 + max_value: 100000.0 + region: + possible_values: ["Urban", "Rural"] + persons: + age: + min_value: 0 + max_value: 100 + dtype: int + gender: + possible_values: ["M", "F"] + foreign_keys: + - parent_table: households + parent_primary_key: hid + child_table: persons + child_foreign_key: hid + max_children_per_parent: 4 + """) + tmp_path = self.create_tempfile(content=yaml_content).full_path + table_domains, fks = domain.from_yaml_file(tmp_path) + self.assertIn('households', table_domains) + self.assertIn('persons', table_domains) + self.assertIsInstance( + table_domains['persons']['age'], base_domain.NumericalAttribute + ) + self.assertIsInstance( + table_domains['persons']['gender'], base_domain.CategoricalAttribute + ) + self.assertLen(fks, 1) + self.assertEqual(fks[0].parent_table, 'households') + self.assertEqual(fks[0].parent_primary_key, 'hid') + self.assertEqual(fks[0].child_table, 'persons') + self.assertEqual(fks[0].child_foreign_key, 'hid') + self.assertEqual(fks[0].max_children_per_parent, 4) + + def test_to_dict_and_roundtrip(self): + table_domains = { + 'households': { + 'income': base_domain.NumericalAttribute( + min_value=0.0, max_value=200000.0, dtype='float' + ), + 'region': base_domain.CategoricalAttribute( + possible_values=['Urban', 'Rural'] + ), + }, + 'persons': { + 'age': base_domain.NumericalAttribute( + min_value=0, max_value=100, dtype='int' + ), + 'gender': base_domain.CategoricalAttribute( + possible_values=['M', 'F'] + ), + }, + } + fks = [ + domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=3, + ) + ] + serialized = domain.to_dict(table_domains, fks) + self.assertIn('tables', serialized) + self.assertIn('foreign_keys', serialized) + self.assertLen(serialized['foreign_keys'], 1) + self.assertEqual( + serialized['tables']['households']['income']['type'], + 'NumericalAttribute', + ) + + # Roundtrip verification + rt_domains, rt_fks = domain.from_dict(serialized) + self.assertEqual( + rt_domains['households']['income'].min_value, + table_domains['households']['income'].min_value, + ) + self.assertEqual( + rt_domains['persons']['gender'].possible_values, + table_domains['persons']['gender'].possible_values, + ) + self.assertEqual(rt_fks, fks) + + def test_to_yaml_file_and_roundtrip(self): + table_domains = { + 'households': { + 'income': base_domain.NumericalAttribute( + min_value=0.0, max_value=150000.0 + ), + 'region': base_domain.CategoricalAttribute( + possible_values=['Urban', 'Suburban'] + ), + }, + 'persons': { + 'age': base_domain.NumericalAttribute( + min_value=0, max_value=120, dtype='int' + ), + }, + } + fks = [ + domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='household_id', + child_table='persons', + child_foreign_key='household_id', + max_children_per_parent=5, + ) + ] + tmp_path = self.create_tempfile().full_path + domain.to_yaml_file(table_domains, fks, tmp_path) + + rt_domains, rt_fks = domain.from_yaml_file(tmp_path) + self.assertIn('households', rt_domains) + self.assertIn('persons', rt_domains) + self.assertEqual(rt_domains['households']['income'].max_value, 150000.0) + self.assertEqual(rt_fks, fks) + + def test_to_dict_without_foreign_keys(self): + table_domains = { + 'single_table': { + 'col_a': base_domain.CategoricalAttribute( + possible_values=['x', 'y'] + ) + } + } + serialized = domain.to_dict(table_domains) + self.assertIn('tables', serialized) + self.assertNotIn('foreign_keys', serialized) + rt_domains, rt_fks = domain.from_dict(serialized) + self.assertIn('single_table', rt_domains) + self.assertEmpty(rt_fks) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/relational/synthesizer_test.py b/tests/relational/synthesizer_test.py new file mode 100644 index 0000000..29178ca --- /dev/null +++ b/tests/relational/synthesizer_test.py @@ -0,0 +1,34 @@ +# 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 dpsynth.relational.synthesizer.""" + +from absl.testing import absltest +from dpsynth import api +from dpsynth.relational import synthesizer + + +class SynthesizerTest(absltest.TestCase): + + def test_classes_inherit_correct_api_abstractions(self): + self.assertTrue( + issubclass(synthesizer.MultiTableConfig, api.MechanismConfig) + ) + self.assertTrue( + issubclass(synthesizer.MultiTableMechanism, api.CalibratedMechanism) + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/relational/transformations_test.py b/tests/relational/transformations_test.py new file mode 100644 index 0000000..f2f5cda --- /dev/null +++ b/tests/relational/transformations_test.py @@ -0,0 +1,982 @@ +# 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 dpsynth.relational.transformations.""" + +from absl.testing import absltest +from dpsynth.relational import domain as rel_domain +from dpsynth.relational import transformations +import mbi +import numpy as np +import pandas as pd + + +class TransformationsTest(absltest.TestCase): + + def test_build_exploration_domain_empty_token(self): + parent_domain = mbi.Domain.fromdict({'income': 4, 'region': 3}) + child_domain = mbi.Domain.fromdict({'age': 10, 'gender': 2}) + domain = transformations._build_exploration_domain( + parent_domain=parent_domain, + child_domain=child_domain, + max_group_size=3, + num_permutation_slots=2, + strategy='empty_token', + ) + expected_attrs = ( + 'income', + 'region', + 'group_size', + 'slot_1.age', + 'slot_1.gender', + 'slot_2.age', + 'slot_2.gender', + ) + self.assertEqual(domain.attributes, expected_attrs) + self.assertEqual(domain.shape, (4, 3, 4, 11, 3, 11, 3)) + + def test_build_exploration_domain_size_sliced(self): + parent_domain = mbi.Domain.fromdict({'income': 4}) + child_domain = mbi.Domain.fromdict({'age': 10, 'gender': 2}) + domain = transformations._build_exploration_domain( + parent_domain=parent_domain, + child_domain=child_domain, + max_group_size=2, + num_permutation_slots=2, + strategy='size_sliced', + ) + self.assertEqual( + domain.attributes, + ( + 'income', + 'group_size', + 'slot_1.age', + 'slot_1.gender', + 'slot_2.age', + 'slot_2.gender', + ), + ) + self.assertEqual(domain.shape, (4, 3, 10, 2, 10, 2)) + + def test_build_exploration_domain_single_group_size_and_3slots(self): + parent_domain = mbi.Domain.fromdict({'p1': 2}) + child_domain = mbi.Domain.fromdict({'c1': 5}) + domain = transformations._build_exploration_domain( + parent_domain=parent_domain, + child_domain=child_domain, + max_group_size=1, + num_permutation_slots=3, + strategy='empty_token', + ) + self.assertEqual( + domain.attributes, + ('p1', 'group_size', 'slot_1.c1', 'slot_2.c1', 'slot_3.c1'), + ) + self.assertEqual(domain.shape, (2, 2, 6, 6, 6)) + + def test_get_slot_permutation_patterns_k0(self): + patterns_empty, w_empty = transformations._get_slot_permutation_patterns( + k=0, num_permutation_slots=2, strategy='empty_token' + ) + self.assertEqual(patterns_empty, [(-1, -1)]) + self.assertEqual(w_empty, 1.0) + + patterns_sliced, w_sliced = transformations._get_slot_permutation_patterns( + k=0, num_permutation_slots=2, strategy='size_sliced' + ) + self.assertEqual(patterns_sliced, [(0, 0)]) + self.assertEqual(w_sliced, 1.0) + + def test_get_slot_permutation_patterns_k1_o2(self): + patterns_empty, w_empty = transformations._get_slot_permutation_patterns( + k=1, num_permutation_slots=2, strategy='empty_token' + ) + self.assertCountEqual(patterns_empty, [(0, -1), (-1, 0)]) + self.assertEqual(w_empty, 0.5) + + patterns_sliced, w_sliced = transformations._get_slot_permutation_patterns( + k=1, num_permutation_slots=2, strategy='size_sliced' + ) + self.assertEqual(patterns_sliced, [(0, 0)]) + self.assertEqual(w_sliced, 1.0) + + def test_get_slot_permutation_patterns_k2_o2(self): + patterns_empty, w_empty = transformations._get_slot_permutation_patterns( + k=2, num_permutation_slots=2, strategy='empty_token' + ) + self.assertCountEqual(patterns_empty, [(0, 1), (1, 0)]) + self.assertEqual(w_empty, 0.5) + + patterns_sliced, w_sliced = transformations._get_slot_permutation_patterns( + k=2, num_permutation_slots=2, strategy='size_sliced' + ) + self.assertCountEqual(patterns_sliced, [(0, 1), (1, 0)]) + self.assertEqual(w_sliced, 0.5) + + def test_get_slot_permutation_patterns_k3_o2(self): + patterns, w = transformations._get_slot_permutation_patterns( + k=3, num_permutation_slots=2, strategy='empty_token' + ) + expected = [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)] + self.assertCountEqual(patterns, expected) + self.assertAlmostEqual(w, 1.0 / 6.0) + + def test_get_slot_permutation_patterns_k1_and_k2_o3(self): + # o = 3, k = 1: 3 permutations + patterns_k1, w_k1 = transformations._get_slot_permutation_patterns( + k=1, num_permutation_slots=3, strategy='empty_token' + ) + self.assertCountEqual(patterns_k1, [(0, -1, -1), (-1, 0, -1), (-1, -1, 0)]) + self.assertAlmostEqual(w_k1, 1.0 / 3.0) + + # o = 3, k = 2: 6 permutations + patterns_k2, w_k2 = transformations._get_slot_permutation_patterns( + k=2, num_permutation_slots=3, strategy='empty_token' + ) + self.assertLen(patterns_k2, 6) + self.assertAlmostEqual(w_k2, 1.0 / 6.0) + + def test_get_slot_permutation_patterns_k10_o3(self): + # P(10, 3) = 10 * 9 * 8 = 720 permutations + patterns_empty, w_empty = transformations._get_slot_permutation_patterns( + k=10, num_permutation_slots=3, strategy='empty_token' + ) + self.assertLen(patterns_empty, 720) + self.assertLen(set(patterns_empty), 720) + self.assertAlmostEqual(w_empty, 1.0 / 720.0) + for p in patterns_empty: + self.assertLen(p, 3) + self.assertTrue(all(0 <= idx < 10 for idx in p)) + self.assertLen(set(p), 3) + + patterns_sliced, w_sliced = transformations._get_slot_permutation_patterns( + k=10, num_permutation_slots=3, strategy='size_sliced' + ) + self.assertEqual(patterns_empty, patterns_sliced) + self.assertEqual(w_empty, w_sliced) + + def test_get_slot_permutation_patterns_weight_invariant(self): + for k in range(6): + for o in range(1, 5): + for strategy in ['empty_token', 'size_sliced']: + patterns, weight = transformations._get_slot_permutation_patterns( + k=k, num_permutation_slots=o, strategy=strategy + ) + self.assertAlmostEqual(len(patterns) * weight, 1.0) + + def test_build_permuted_exploration_dataset_running_example(self): + # 4 households: k=0, 1, 2, 3 children + parent_dom = mbi.Domain.fromdict({'income': 4}) + parent_data = {'income': np.array([0, 1, 2, 3], dtype=np.int64)} + parent_ds = mbi.Dataset(parent_data, parent_dom) + parent_pks = ['h0', 'h1', 'h2', 'h3'] + + child_dom = mbi.Domain.fromdict({'age': 10}) + # h1: c0(3); h2: c1(5), c2(8); h3: c3(2), c4(6), c5(9) + child_data = {'age': np.array([3, 5, 8, 2, 6, 9], dtype=np.int64)} + child_ds = mbi.Dataset(child_data, child_dom) + child_fks = ['h1', 'h2', 'h2', 'h3', 'h3', 'h3'] + + ds = transformations.build_permuted_exploration_dataset( + parent_dataset=parent_ds, + child_dataset=child_ds, + parent_primary_keys=parent_pks, + child_foreign_keys=child_fks, + max_group_size=3, + num_permutation_slots=2, + strategy='empty_token', + ) + + # 1. Domain verification: group_size cardinality=4, age cardinality=11 + self.assertEqual( + ds.domain.attributes, + ('income', 'group_size', 'slot_1.age', 'slot_2.age'), + ) + self.assertEqual(ds.domain.shape, (4, 4, 11, 11)) + + # 2. Row count and weight mass verification + # Rows: h0->1, h1->2, h2->2, h3->6 = 11 total rows + self.assertEqual(ds.records, 11) + self.assertAlmostEqual(float(np.sum(ds.weights)), 4.0) + + # 3. Exact slot exchangeability (Hermitian marginal symmetry P(S1) == P(S2)) + s1_hist = np.bincount( + ds.data['slot_1.age'], weights=ds.weights, minlength=11 + ) + s2_hist = np.bincount( + ds.data['slot_2.age'], weights=ds.weights, minlength=11 + ) + np.testing.assert_allclose(s1_hist, s2_hist) + + # 4. Total token mass: h0 has 1(both empty), h1 has 1(one empty) + empty_mass_s1 = float(np.sum(ds.weights[ds.data['slot_1.age'] == 10])) + empty_mass_s2 = float(np.sum(ds.weights[ds.data['slot_2.age'] == 10])) + self.assertAlmostEqual(empty_mass_s1, 1.5) + self.assertAlmostEqual(empty_mass_s2, 1.5) + + def test_build_permuted_exploration_dataset_size_sliced(self): + parent_dom = mbi.Domain.fromdict({'income': 4}) + parent_data = {'income': np.array([0, 1, 2, 3], dtype=np.int64)} + parent_ds = mbi.Dataset(parent_data, parent_dom) + parent_pks = ['h0', 'h1', 'h2', 'h3'] + + child_dom = mbi.Domain.fromdict({'age': 10}) + child_data = {'age': np.array([3, 5, 8, 2, 6, 9], dtype=np.int64)} + child_ds = mbi.Dataset(child_data, child_dom) + child_fks = ['h1', 'h2', 'h2', 'h3', 'h3', 'h3'] + + ds = transformations.build_permuted_exploration_dataset( + parent_dataset=parent_ds, + child_dataset=child_ds, + parent_primary_keys=parent_pks, + child_foreign_keys=child_fks, + max_group_size=3, + num_permutation_slots=2, + strategy='size_sliced', + ) + + # Domain shape has unextended child shape (10) + self.assertEqual(ds.domain.shape, (4, 4, 10, 10)) + # Rows: h0->1, h1->1 (clone tiled), h2->2, h3->6 = 10 total rows + self.assertEqual(ds.records, 10) + self.assertAlmostEqual(float(np.sum(ds.weights)), 4.0) + + # h1 (row index 1) has both slots clone tiled with child 0 (age 3) + h1_mask = ds.data['income'] == 1 + self.assertEqual(int(ds.data['slot_1.age'][h1_mask][0]), 3) + self.assertEqual(int(ds.data['slot_2.age'][h1_mask][0]), 3) + + def test_build_permuted_exploration_dataset_order1(self): + parent_dom = mbi.Domain.fromdict({'income': 2}) + parent_data = {'income': np.array([0, 1], dtype=np.int64)} + parent_ds = mbi.Dataset(parent_data, parent_dom) + + child_dom = mbi.Domain.fromdict({'age': 5}) + # h0 has 0 children; h1 has 2 children (ages 1, 4) + child_data = {'age': np.array([1, 4], dtype=np.int64)} + child_ds = mbi.Dataset(child_data, child_dom) + + ds = transformations.build_permuted_exploration_dataset( + parent_dataset=parent_ds, + child_dataset=child_ds, + parent_primary_keys=['h0', 'h1'], + child_foreign_keys=['h1', 'h1'], + max_group_size=2, + num_permutation_slots=1, + strategy='empty_token', + ) + + self.assertEqual( + ds.domain.attributes, ('income', 'group_size', 'slot_1.age') + ) + self.assertEqual(ds.domain.shape, (2, 3, 6)) + # Rows: h0->1 (=5), h1->2 (age 1, age 4 with w=0.5) = 3 rows + self.assertEqual(ds.records, 3) + self.assertAlmostEqual(float(np.sum(ds.weights)), 2.0) + + def test_build_permuted_exploration_dataset_order3(self): + parent_dom = mbi.Domain.fromdict({'income': 2}) + parent_data = {'income': np.array([0, 1], dtype=np.int64)} + parent_ds = mbi.Dataset(parent_data, parent_dom) + + child_dom = mbi.Domain.fromdict({'age': 5}) + child_data = {'age': np.array([1, 2], dtype=np.int64)} + child_ds = mbi.Dataset(child_data, child_dom) + + ds = transformations.build_permuted_exploration_dataset( + parent_dataset=parent_ds, + child_dataset=child_ds, + parent_primary_keys=['h0', 'h1'], + child_foreign_keys=['h1', 'h1'], + max_group_size=2, + num_permutation_slots=3, + strategy='empty_token', + ) + + self.assertEqual( + ds.domain.attributes, + ('income', 'group_size', 'slot_1.age', 'slot_2.age', 'slot_3.age'), + ) + self.assertEqual(ds.domain.shape, (2, 3, 6, 6, 6)) + # Rows: h0->1 (), h1 (k=2, o=3)->6 patterns with w=1/6 = 7 rows + self.assertEqual(ds.records, 7) + self.assertAlmostEqual(float(np.sum(ds.weights)), 2.0) + + def test_build_permuted_exploration_dataset_multi_column_child(self): + parent_dom = mbi.Domain.fromdict({'income': 3}) + parent_ds = mbi.Dataset( + {'income': np.array([0, 1], dtype=np.int64)}, parent_dom + ) + child_dom = mbi.Domain.fromdict({'age': 10, 'gender': 2}) + child_data = { + 'age': np.array([3, 7], dtype=np.int64), + 'gender': np.array([0, 1], dtype=np.int64), + } + child_ds = mbi.Dataset(child_data, child_dom) + + ds = transformations.build_permuted_exploration_dataset( + parent_dataset=parent_ds, + child_dataset=child_ds, + parent_primary_keys=['h0', 'h1'], + child_foreign_keys=['h1', 'h1'], + max_group_size=2, + num_permutation_slots=2, + strategy='empty_token', + ) + + self.assertEqual( + ds.domain.attributes, + ( + 'income', + 'group_size', + 'slot_1.age', + 'slot_1.gender', + 'slot_2.age', + 'slot_2.gender', + ), + ) + self.assertEqual(ds.domain.shape, (3, 3, 11, 3, 11, 3)) + # For h0 (childless), empty tokens are age=10, gender=2 + h0_mask = ds.data['income'] == 0 + self.assertEqual(int(ds.data['slot_1.age'][h0_mask][0]), 10) + self.assertEqual(int(ds.data['slot_1.gender'][h0_mask][0]), 2) + self.assertEqual(int(ds.data['slot_2.age'][h0_mask][0]), 10) + self.assertEqual(int(ds.data['slot_2.gender'][h0_mask][0]), 2) + + def test_build_permuted_exploration_dataset_edge_cases_and_validation(self): + parent_dom = mbi.Domain.fromdict({'p': 2}) + child_dom = mbi.Domain.fromdict({'c': 3}) + + # Empty parent dataset (Np = 0) + empty_parent = mbi.Dataset({'p': np.empty(0, dtype=np.int64)}, parent_dom) + empty_child = mbi.Dataset({'c': np.empty(0, dtype=np.int64)}, child_dom) + ds_empty = transformations.build_permuted_exploration_dataset( + parent_dataset=empty_parent, + child_dataset=empty_child, + parent_primary_keys=[], + child_foreign_keys=[], + max_group_size=2, + num_permutation_slots=2, + strategy='empty_token', + ) + self.assertEqual(ds_empty.records, 0) + self.assertAlmostEqual(float(np.sum(ds_empty.weights)), 0.0) + + # Orphaned foreign keys (dropped cleanly) + p_ds = mbi.Dataset({'p': np.array([0], dtype=np.int64)}, parent_dom) + c_ds = mbi.Dataset({'c': np.array([1, 2], dtype=np.int64)}, child_dom) + ds_orphans = transformations.build_permuted_exploration_dataset( + parent_dataset=p_ds, + child_dataset=c_ds, + parent_primary_keys=['h0'], + child_foreign_keys=['orphan_1', 'orphan_2'], + max_group_size=2, + num_permutation_slots=2, + strategy='empty_token', + ) + # Parent h0 has 0 valid children -> 1 row with + self.assertEqual(ds_orphans.records, 1) + self.assertEqual(int(ds_orphans.data['slot_1.c'][0]), 3) + + # Validation errors + with self.assertRaises(ValueError): + transformations.build_permuted_exploration_dataset( + p_ds, c_ds, ['h0'], ['h0', 'h0'], max_group_size=0 + ) + with self.assertRaises(ValueError): + transformations.build_permuted_exploration_dataset( + p_ds, + c_ds, + ['h0'], + ['h0', 'h0'], + max_group_size=2, + num_permutation_slots=0, + ) + with self.assertRaises(ValueError): + transformations.build_permuted_exploration_dataset( + p_ds, + c_ds, + ['h0'], + ['h0', 'h0'], + max_group_size=2, + strategy='invalid', + ) + with self.assertRaises(ValueError): + transformations.build_permuted_exploration_dataset( + p_ds, c_ds, ['h0', 'extra'], ['h0', 'h0'], max_group_size=2 + ) + with self.assertRaises(ValueError): + transformations.build_permuted_exploration_dataset( + p_ds, c_ds, ['h0'], ['h0'], max_group_size=2 + ) + + def test_transformations_import_and_callable(self): + self.assertTrue(callable(transformations.compute_hierarchical_weights)) + self.assertTrue( + callable(transformations.build_permuted_exploration_dataset) + ) + self.assertTrue( + callable(transformations.create_slot_linear_chain_constraints) + ) + self.assertTrue(callable(transformations.symmetrize_to_wide_domain)) + self.assertTrue(callable(transformations.quantile_copula_coupling)) + self.assertTrue(callable(transformations.unstack_wide_family_records)) + + def test_compute_row_root_mappings_single_table(self): + households = pd.DataFrame({ + 'household_id': ['h1', 'h2', 'h3'], + 'income': [50000.0, 75000.0, 100000.0], + }) + hierarchy = [(0, 'households', None)] + mapping = transformations._compute_row_root_mappings( + {'households': households}, hierarchy + ) + self.assertIsInstance(mapping['households'], pd.Series) + self.assertEqual( + mapping['households'].tolist(), + [('households', 0), ('households', 1), ('households', 2)], + ) + + def test_compute_row_root_mappings_2tier(self): + households = pd.DataFrame({'hid': ['h1', 'h2']}) + persons = pd.DataFrame({ + 'pid': ['p1', 'p2', 'p3'], + 'hid': ['h1', 'h1', 'h2'], + }) + fk = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=3, + ) + hierarchy = [(0, 'households', None), (1, 'persons', fk)] + mapping = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, hierarchy + ) + self.assertIsInstance(mapping['persons'], pd.Series) + self.assertEqual( + mapping['persons'].tolist(), + [('households', 0), ('households', 0), ('households', 1)], + ) + + def test_compute_row_root_mappings_truncation_and_cascading(self): + households = pd.DataFrame({'hid': ['h1']}) + # h1 has 3 persons, but max_children_per_parent is 2 -> exactly 2 chosen + persons = pd.DataFrame({ + 'pid': ['p1', 'p2', 'p3'], + 'hid': ['h1', 'h1', 'h1'], + }) + # activities for each person + activities = pd.DataFrame({ + 'aid': ['a1', 'a2', 'a3'], + 'pid': ['p1', 'p2', 'p3'], + }) + fk1 = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=2, + ) + fk2 = rel_domain.ForeignKeyRelation( + parent_table='persons', + parent_primary_key='pid', + child_table='activities', + child_foreign_key='pid', + max_children_per_parent=2, + ) + hierarchy = [ + (0, 'households', None), + (1, 'persons', fk1), + (2, 'activities', fk2), + ] + + rng = np.random.default_rng(42) + mapping = transformations._compute_row_root_mappings( + { + 'households': households, + 'persons': persons, + 'activities': activities, + }, + hierarchy, + rng=rng, + ) + # Exactly 2 persons are active (non-None), 1 is truncated (None) + active_persons = mapping['persons'].dropna().tolist() + self.assertLen(active_persons, 2) + self.assertEqual(mapping['persons'].isna().sum(), 1) + + # Subchildren of active persons active, subchild of truncated person is None + active_activities = mapping['activities'].dropna().tolist() + self.assertLen(active_activities, 2) + self.assertEqual(mapping['activities'].isna().sum(), 1) + + def test_compute_row_root_mappings_branching_tree(self): + households = pd.DataFrame({'hid': ['h1', 'h2']}) + persons = pd.DataFrame({'pid': ['p1', 'p2'], 'hid': ['h1', 'h2']}) + vehicles = pd.DataFrame( + {'vid': ['v1', 'v2', 'v3'], 'hid': ['h1', 'h1', 'h2']} + ) + + fk_p = rel_domain.ForeignKeyRelation( + 'households', 'hid', 'persons', 'hid', 2 + ) + fk_v = rel_domain.ForeignKeyRelation( + 'households', 'hid', 'vehicles', 'hid', 5 + ) + hierarchy = [ + (0, 'households', None), + (1, 'persons', fk_p), + (1, 'vehicles', fk_v), + ] + mapping = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons, 'vehicles': vehicles}, + hierarchy, + ) + self.assertEqual( + mapping['persons'].tolist(), + [('households', 0), ('households', 1)], + ) + self.assertEqual( + mapping['vehicles'].tolist(), + [('households', 0), ('households', 0), ('households', 1)], + ) + + def test_compute_row_root_mappings_multi_tree_forest(self): + households = pd.DataFrame({'hid': ['h1']}) + persons = pd.DataFrame({'pid': ['p1'], 'hid': ['h1']}) + companies = pd.DataFrame({'cid': ['c1']}) + departments = pd.DataFrame({'did': ['d1'], 'cid': ['c1']}) + + fk_h = rel_domain.ForeignKeyRelation( + 'households', 'hid', 'persons', 'hid', 2 + ) + fk_c = rel_domain.ForeignKeyRelation( + 'companies', 'cid', 'departments', 'cid', 5 + ) + hierarchy = [ + (0, 'households', None), + (0, 'companies', None), + (1, 'persons', fk_h), + (1, 'departments', fk_c), + ] + mapping = transformations._compute_row_root_mappings( + { + 'households': households, + 'persons': persons, + 'companies': companies, + 'departments': departments, + }, + hierarchy, + ) + self.assertEqual(mapping['households'].tolist(), [('households', 0)]) + self.assertEqual(mapping['companies'].tolist(), [('companies', 0)]) + self.assertEqual(mapping['persons'].tolist(), [('households', 0)]) + self.assertEqual(mapping['departments'].tolist(), [('companies', 0)]) + + def test_compute_row_root_mappings_custom_index_alignment(self): + # Non-standard indices (strings, custom obj) can't break positional mapping + households = pd.DataFrame( + {'hid': ['h1', 'h2']}, index=['custom_a', 'custom_b'] + ) + persons = pd.DataFrame( + {'pid': ['p1', 'p2', 'p3'], 'hid': ['h1', 'h2', 'h1']}, + index=[100, 200, 300], + ) + fk = rel_domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 5) + mapping = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, + [(0, 'households', None), (1, 'persons', fk)], + ) + # Positions are strictly 0 and 1 in households DataFrame + self.assertEqual( + mapping['persons'].tolist(), + [('households', 0), ('households', 1), ('households', 0)], + ) + + def test_compute_row_root_mappings_orphans_and_validation(self): + households = pd.DataFrame({'hid': ['h1']}) + persons = pd.DataFrame({ + 'pid': ['p1', 'p2'], + 'hid': ['h1', 'orphan_h'], + }) + fk = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=2, + ) + mapping = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, + [(0, 'households', None), (1, 'persons', fk)], + ) + self.assertEqual(mapping['persons'].tolist(), [('households', 0), None]) + + # Missing parent primary key + bad_fk1 = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='missing_id', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=2, + ) + with self.assertRaisesRegex(ValueError, 'Parent primary key'): + transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, + [(0, 'households', None), (1, 'persons', bad_fk1)], + ) + + # Missing child foreign key + bad_fk2 = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='missing_hid', + max_children_per_parent=2, + ) + with self.assertRaisesRegex(ValueError, 'Child foreign key'): + transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, + [(0, 'households', None), (1, 'persons', bad_fk2)], + ) + + def test_compute_row_root_mappings_nan_and_corrupt_data(self): + households = pd.DataFrame({'hid': ['h1', np.nan, 'h2', 'h1']}) + persons = pd.DataFrame({ + 'pid': ['p1', 'p2', 'p3', 'p4', 'p5'], + 'hid': ['h1', np.nan, 'h2', None, ['unhashable_list']], + }) + fk = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=2, + ) + hierarchy = [(0, 'households', None), (1, 'persons', fk)] + mapping = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, + hierarchy, + ) + self.assertEqual( + mapping['persons'].tolist(), + [('households', 0), None, ('households', 2), None, None], + ) + + def test_compute_row_root_mappings_empty_tables(self): + empty_h = pd.DataFrame({'hid': []}) + persons = pd.DataFrame({'pid': ['p1'], 'hid': ['h1']}) + fk = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=2, + ) + mapping = transformations._compute_row_root_mappings( + {'households': empty_h, 'persons': persons}, + [(0, 'households', None), (1, 'persons', fk)], + ) + self.assertEqual(mapping['households'].tolist(), []) + self.assertEqual(mapping['persons'].tolist(), [None]) + + empty_p = pd.DataFrame({'pid': [], 'hid': []}) + mapping2 = transformations._compute_row_root_mappings( + {'households': empty_h, 'persons': empty_p}, + [(0, 'households', None), (1, 'persons', fk)], + ) + self.assertEqual(mapping2['households'].tolist(), []) + self.assertEqual(mapping2['persons'].tolist(), []) + + def test_compute_row_root_mappings_random_subsampling_reproducibility(self): + # A single household with 10 persons, capacity s = 3 + households = pd.DataFrame({'hid': ['h1']}) + persons = pd.DataFrame({ + 'pid': [f'p{i}' for i in range(10)], + 'hid': ['h1'] * 10, + }) + fk = rel_domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 3) + hierarchy = [(0, 'households', None), (1, 'persons', fk)] + + # Same seed must yield identical active row selections + rng1 = np.random.default_rng(123) + mapping1 = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, hierarchy, rng=rng1 + ) + rng2 = np.random.default_rng(123) + mapping2 = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, hierarchy, rng=rng2 + ) + self.assertEqual(mapping1['persons'].tolist(), mapping2['persons'].tolist()) + self.assertEqual(mapping1['persons'].dropna().count(), 3) + self.assertEqual(mapping1['persons'].isna().sum(), 7) + + # Different seeds must produce valid selections of size exactly 3 + rng3 = np.random.default_rng(999) + mapping3 = transformations._compute_row_root_mappings( + {'households': households, 'persons': persons}, hierarchy, rng=rng3 + ) + self.assertEqual(mapping3['persons'].dropna().count(), 3) + self.assertEqual(mapping3['persons'].isna().sum(), 7) + + def test_compute_hierarchical_weights_single_table(self): + households = pd.DataFrame({ + 'household_id': ['h1', 'h2', 'h3'], + 'income': [50000.0, 75000.0, 100000.0], + }) + hierarchy = [(0, 'households', None)] + weights = transformations.compute_hierarchical_weights( + {'households': households}, hierarchy=hierarchy + ) + self.assertIn('households', weights) + self.assertEqual(weights['households'].shape, (3,)) + np.testing.assert_allclose(weights['households'], np.array([1.0, 1.0, 1.0])) + self.assertAlmostEqual(weights['households'].sum(), 3.0) + + def test_compute_hierarchical_weights_2tier(self): + # H1 has 2 persons (P1, P2) -> w = 0.5 each + # H2 has 1 person (P3) -> w = 1.0 + households = pd.DataFrame({'hid': ['h1', 'h2']}) + persons = pd.DataFrame({ + 'pid': ['p1', 'p2', 'p3'], + 'hid': ['h1', 'h1', 'h2'], + }) + fk = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=3, + ) + hierarchy = [(0, 'households', None), (1, 'persons', fk)] + weights = transformations.compute_hierarchical_weights( + {'households': households, 'persons': persons}, hierarchy=hierarchy + ) + np.testing.assert_allclose(weights['households'], np.array([1.0, 1.0])) + np.testing.assert_allclose(weights['persons'], np.array([0.5, 0.5, 1.0])) + # Total sum of weights in every table matches number of households (2.0) + self.assertAlmostEqual(weights['households'].sum(), 2.0) + self.assertAlmostEqual(weights['persons'].sum(), 2.0) + + def test_compute_hierarchical_weights_3tier_with_truncation(self): + # H1 has 3 persons, but s1 = 2 + # -> 2 active (w = 0.5 each), 1 truncated (w = 0.0) + # H1's active persons have 2 activities each (4 total) -> w = 0.25 each + # H1's truncated person has 2 activities -> both cascade to w = 0.0 + households = pd.DataFrame({'hid': ['h1']}) + persons = pd.DataFrame({ + 'pid': ['p1', 'p2', 'p3'], + 'hid': ['h1', 'h1', 'h1'], + }) + activities = pd.DataFrame({ + 'aid': ['a1', 'a2', 'a3', 'a4', 'a5', 'a6'], + 'pid': ['p1', 'p1', 'p2', 'p2', 'p3', 'p3'], + }) + fk1 = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=2, + ) + fk2 = rel_domain.ForeignKeyRelation( + parent_table='persons', + parent_primary_key='pid', + child_table='activities', + child_foreign_key='pid', + max_children_per_parent=2, + ) + hierarchy = [ + (0, 'households', None), + (1, 'persons', fk1), + (2, 'activities', fk2), + ] + + rng = np.random.default_rng(42) + weights = transformations.compute_hierarchical_weights( + { + 'households': households, + 'persons': persons, + 'activities': activities, + }, + hierarchy=hierarchy, + rng=rng, + ) + + # Household sum = 1.0 + self.assertAlmostEqual(weights['households'].sum(), 1.0) + # Person sum = 1.0 (2 active persons with 0.5, 1 truncated with 0.0) + self.assertAlmostEqual(weights['persons'].sum(), 1.0) + self.assertEqual((weights['persons'] == 0.0).sum(), 1) + self.assertEqual((weights['persons'] == 0.5).sum(), 2) + + # Activity sum = 1.0 (4 active activities with 0.25, 2 truncated with 0.0) + self.assertAlmostEqual(weights['activities'].sum(), 1.0) + self.assertEqual((weights['activities'] == 0.0).sum(), 2) + self.assertEqual((weights['activities'] == 0.25).sum(), 4) + + def test_compute_hierarchical_weights_empty_table(self): + empty_h = pd.DataFrame({'hid': []}) + empty_p = pd.DataFrame({'pid': [], 'hid': []}) + fk = rel_domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=2, + ) + hierarchy = [(0, 'households', None), (1, 'persons', fk)] + weights = transformations.compute_hierarchical_weights( + {'households': empty_h, 'persons': empty_p}, hierarchy=hierarchy + ) + self.assertEqual(weights['households'].shape, (0,)) + self.assertEqual(weights['persons'].shape, (0,)) + + def test_compute_hierarchical_weights_multi_tree_forest(self): + households = pd.DataFrame({'hid': ['h1', 'h2']}) + persons = pd.DataFrame({ + 'pid': ['p1', 'p2', 'p3', 'p4'], + 'hid': ['h1', 'h1', 'h2', 'h2'], + }) + companies = pd.DataFrame({'cid': ['c1']}) + departments = pd.DataFrame({ + 'did': ['d1', 'd2'], + 'cid': ['c1', 'c1'], + }) + + fk_h = rel_domain.ForeignKeyRelation( + 'households', 'hid', 'persons', 'hid', 5 + ) + fk_c = rel_domain.ForeignKeyRelation( + 'companies', 'cid', 'departments', 'cid', 5 + ) + hierarchy = [ + (0, 'households', None), + (0, 'companies', None), + (1, 'persons', fk_h), + (1, 'departments', fk_c), + ] + weights = transformations.compute_hierarchical_weights( + { + 'households': households, + 'persons': persons, + 'companies': companies, + 'departments': departments, + }, + hierarchy=hierarchy, + ) + self.assertAlmostEqual(weights['households'].sum(), 2.0) + self.assertAlmostEqual(weights['persons'].sum(), 2.0) + self.assertAlmostEqual(weights['companies'].sum(), 1.0) + self.assertAlmostEqual(weights['departments'].sum(), 1.0) + + def test_dp_adversarial_data_dependent_robustness(self): + """Stress tests DP safety: mechanism must never crash on adversarial data.""" + # Adversarial dataset with IEEE specials, mixed types, unhashable objects, + # duplicate keys, and missing references. + households = pd.DataFrame({ + 'hid': [ + 'h1', + np.nan, + None, + float('inf'), + float('-inf'), + 'h1', # duplicate + 12345, # integer in string column + True, # boolean + ], + 'val': [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + }) + persons = pd.DataFrame({ + 'pid': [f'p{i}' for i in range(12)], + 'hid': [ + 'h1', # matches first h1 + 'orphan_key', # orphan + np.nan, # NaN + None, # None + pd.NA, # pd.NA + float('inf'), # matches inf + float('-inf'), # matches -inf + 12345, # matches int + True, # matches boolean + ['unhashable_list'], # unhashable list + {'unhashable': 'dict'}, # unhashable dict + 'h1', # another match to h1 + ], + }) + activities = pd.DataFrame({ + 'aid': [f'a{i}' for i in range(5)], + 'pid': ['p0', 'p1', 'p9', 'missing_person', 'p11'], + }) + + fk1 = rel_domain.ForeignKeyRelation( + 'households', 'hid', 'persons', 'hid', 1 + ) + fk2 = rel_domain.ForeignKeyRelation( + 'persons', 'pid', 'activities', 'pid', 2 + ) + hierarchy = [ + (0, 'households', None), + (1, 'persons', fk1), + (2, 'activities', fk2), + ] + + rng = np.random.default_rng(100) + # Must execute cleanly (no exceptions) on this corrupted dataset + weights = transformations.compute_hierarchical_weights( + { + 'households': households, + 'persons': persons, + 'activities': activities, + }, + hierarchy=hierarchy, + rng=rng, + ) + + # Output shapes must be strictly aligned with input DataFrame row counts + self.assertEqual(weights['households'].shape, (len(households),)) + self.assertEqual(weights['persons'].shape, (len(persons),)) + self.assertEqual(weights['activities'].shape, (len(activities),)) + + # All weights must be finite non-negative floats + self.assertTrue(np.all(np.isfinite(weights['households']))) + self.assertTrue(np.all(weights['households'] >= 0.0)) + self.assertTrue(np.all(np.isfinite(weights['persons']))) + self.assertTrue(np.all(weights['persons'] >= 0.0)) + self.assertTrue(np.all(np.isfinite(weights['activities']))) + self.assertTrue(np.all(weights['activities'] >= 0.0)) + + # Sensitivity invariant: sum of weights for any single household <= 1.0 + mapping = transformations._compute_row_root_mappings( + { + 'households': households, + 'persons': persons, + 'activities': activities, + }, + hierarchy=hierarchy, + rng=rng, + ) + for table_name in ['households', 'persons', 'activities']: + t_roots = mapping[table_name] + t_weights = weights[table_name] + for root in t_roots.dropna().unique(): + root_mask = (t_roots == root).values + root_weight_sum = t_weights[root_mask].sum() + self.assertAlmostEqual(root_weight_sum, 1.0, places=5) + + +if __name__ == '__main__': + absltest.main()