From c9a6a8e5be1bf5d58639aad8671654e9af9a76ee Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Fri, 2 Oct 2026 10:52:01 -0700 Subject: [PATCH] Support NumericalAttribute intervals in dpsynth.constraints.Constraint: - Extend Constraint to accept (low, high) interval tuples and None (for clip_to_range=False) on NumericalAttribute columns with fixed bin_edges in possible_combinations and impossible_combinations. - Validate at configuration time that NumericalAttribute domains specify bin_edges and that interval endpoints match attr.bin_edges or [min_value, max_value]. - Discretize numerical intervals in Constraint.to_mbi() using attr.bin_edges before delegating to the categorical to_mbi() path. PiperOrigin-RevId: 992417631 --- dpsynth/constraints.py | 83 +++++++++++++++++++++++++----- tests/constraints_test.py | 88 ++++++++++++++++++++++++++++++++ tests/data_generation_v3_test.py | 52 +++++++++++++++++++ 3 files changed, 209 insertions(+), 14 deletions(-) diff --git a/dpsynth/constraints.py b/dpsynth/constraints.py index f6bf7b00..a4e9dce6 100644 --- a/dpsynth/constraints.py +++ b/dpsynth/constraints.py @@ -18,6 +18,7 @@ from collections.abc import Mapping, Sequence import dataclasses +import itertools from typing import Any from dpsynth import domain @@ -26,6 +27,50 @@ import numpy as np +def _validate_numerical(val, attr): + """Validates a numerical constraint value at configuration time.""" + # Entries must be None (out-of-domain bin when clip_to_range=False) or a + # (low, high) tuple whose endpoints match attr.bin_edges or the domain + # bounds. Validating this in __post_init__ catches invalid bounds before + # touching private data and guarantees exact alignment with discrete bins. + if val is None: + if attr.clip_to_range: + raise ValueError('None requires clip_to_range=False.') + return + if not (isinstance(val, tuple) and len(val) == 2): + raise ValueError(f'Expected (low, high) tuple or None, got {val!r}.') + low, high = float(val[0]), float(val[1]) + if low >= high: + raise ValueError(f'Interval ({low}, {high}) requires low < high.') + edges = {attr.min_value, *attr.bin_edges, attr.max_value} + if low not in edges or high not in edges: + raise ValueError(f'Interval ({low}, {high}) must use {sorted(edges)}.') + + +def _domain_size(attr) -> int: + """Returns the discrete MBI domain size for a categorical or numerical attribute.""" + if isinstance(attr, domain.NumericalAttribute): + assert attr.bin_edges is not None + return len(attr.bin_edges) + (1 if attr.clip_to_range else 2) + return attr.size + + +def _encode_value(val, attr) -> list[int]: + """Maps a categorical value or numerical interval/None to discrete bin indices.""" + # Categorical values map to a single index, None maps to out-of-domain bin 0, + # and a validated (low, high) interval selects all in-domain bins inside + # [low, high], shifted by 1 when index 0 is reserved for OOD. + if not isinstance(attr, domain.NumericalAttribute): + return [transformations.discrete_encoder(attr)(val)] + if val is None: + return [0] + low, high = val + full_edges = np.r_[attr.min_value, attr.bin_edges, attr.max_value] + mask = (full_edges[:-1] >= low) & (full_edges[1:] <= high) + offset = 0 if attr.clip_to_range else 1 + return (np.flatnonzero(mask) + offset).tolist() + + def _validate(c: Constraint) -> None: """Validate a Constraint's fields.""" if len(c.attribute_names) != len(c.attribute_domains): @@ -44,12 +89,19 @@ def _validate(c: Constraint) -> None: "Specify exactly one of 'possible_combinations'," " 'impossible_combinations', or 'functional_dependency'." ) - if c.functional_dependency is not None and len(c.attribute_names) != 2: - raise ValueError( - 'functional_dependency requires exactly 2 attributes (fine, coarse),' - f' got {len(c.attribute_names)}.' - ) + if c.functional_dependency is not None: + if len(c.attribute_names) != 2: + raise ValueError( + 'functional_dependency requires exactly 2 attributes (fine, coarse),' + f' got {len(c.attribute_names)}.' + ) + types = [type(d) for d in c.attribute_domains] + if domain.NumericalAttribute in types: + raise ValueError('functional_dependency requires CategoricalAttribute.') if c.functional_dependency is None: + for attr in c.attribute_domains: + if isinstance(attr, domain.NumericalAttribute) and attr.bin_edges is None: + raise ValueError('NumericalAttribute in Constraint requires bin_edges.') combos = c.possible_combinations or c.impossible_combinations n_attrs = len(c.attribute_names) for combo in combos: # pyrefly: ignore[not-iterable] @@ -58,6 +110,9 @@ def _validate(c: Constraint) -> None: 'Each combination must have length equal to the number of' f' attributes ({n_attrs}), got {len(combo)}.' ) + for val, attr in zip(combo, c.attribute_domains): + if isinstance(attr, domain.NumericalAttribute): + _validate_numerical(val, attr) @dataclasses.dataclass(frozen=True) @@ -87,7 +142,7 @@ class Constraint: Attributes: attribute_names: Names of the constrained attributes. - attribute_domains: Categorical domain for each attribute. + attribute_domains: Categorical or numerical domain for each attribute. possible_combinations: Allowed value combinations. impossible_combinations: Forbidden value combinations. functional_dependency: Dict mapping fine attribute values to coarse @@ -95,7 +150,7 @@ class Constraint: """ attribute_names: tuple[str, ...] - attribute_domains: tuple[domain.CategoricalAttribute, ...] + attribute_domains: Sequence[Any] possible_combinations: Sequence[tuple[Any, ...]] | None = None impossible_combinations: Sequence[tuple[Any, ...]] | None = None functional_dependency: Mapping[Any, Any] | None = None @@ -105,14 +160,11 @@ def __post_init__(self): def to_mbi(self) -> mbi.Constraint: """Convert to an mbi.Constraint.""" - shape = tuple(d.size for d in self.attribute_domains) + shape = tuple(_domain_size(d) for d in self.attribute_domains) mbi_domain = mbi.Domain(self.attribute_names, shape) - encoders = [ - transformations.discrete_encoder(d) for d in self.attribute_domains - ] if self.functional_dependency is not None: - _, coarse_enc = encoders + coarse_enc = transformations.discrete_encoder(self.attribute_domains[1]) fine_values = self.attribute_domains[0].possible_values coarse_indices = [ coarse_enc(self.functional_dependency[v]) for v in fine_values @@ -122,8 +174,11 @@ def to_mbi(self) -> mbi.Constraint: ) combos = self.possible_combinations or self.impossible_combinations - encoded = [[enc(v) for enc, v in zip(encoders, c)] for c in combos] # pyrefly: ignore[not-iterable] - indices = np.array(encoded, dtype=np.int32) + encoded = [] + for c in combos: # pyrefly: ignore[not-iterable] + bins = [_encode_value(v, d) for v, d in zip(c, self.attribute_domains)] + encoded.extend(itertools.product(*bins)) + indices = np.array(encoded, dtype=np.int32).reshape(-1, len(shape)) if self.possible_combinations is not None: return mbi.Constraint(domain=mbi_domain, valid=indices) return mbi.Constraint(domain=mbi_domain, invalid=indices) diff --git a/tests/constraints_test.py b/tests/constraints_test.py index 1f71c8a0..ab5e21bc 100644 --- a/tests/constraints_test.py +++ b/tests/constraints_test.py @@ -138,6 +138,94 @@ def test_functional_dependency(self): self.assertEqual(vals[0, 1], -np.inf) self.assertEqual(vals[2, 0], -np.inf) + def test_numerical_categorical_possible_combinations(self): + age = domain.NumericalAttribute( + min_value=0, max_value=100, dtype='int', bin_edges=[10, 17, 50] + ) + status = domain.CategoricalAttribute(['minor', 'adult']) + c = constraints.Constraint( + attribute_names=('age', 'status'), + attribute_domains=(age, status), + possible_combinations=[((0, 17), 'minor'), ((17, 100), 'adult')], + ) + mbi_c = c.to_mbi() + vals = np.asarray(mbi_c.potential.values) + # Bins 0 ([0, 10]) and 1 ((10, 17]) allow 'minor' (0) and forbid 'adult' (1) + self.assertEqual(vals[0, 0], 0.0) + self.assertEqual(vals[0, 1], -np.inf) + self.assertEqual(vals[1, 0], 0.0) + self.assertEqual(vals[1, 1], -np.inf) + # Bins 2 ((17, 50]) and 3 ((50, 100]) allow 'adult' (1) + self.assertEqual(vals[2, 1], 0.0) + self.assertEqual(vals[2, 0], -np.inf) + self.assertEqual(vals[3, 1], 0.0) + self.assertEqual(vals[3, 0], -np.inf) + + def test_numerical_none_possible_combinations(self): + status = domain.CategoricalAttribute(['missing', 'valid']) + score = domain.NumericalAttribute( + min_value=0.0, max_value=100.0, clip_to_range=False, bin_edges=[50.0] + ) + c = constraints.Constraint( + attribute_names=('status', 'score'), + attribute_domains=(status, score), + possible_combinations=[ + ('missing', None), + ('valid', (0.0, 100.0)), + ], + ) + mbi_c = c.to_mbi() + vals = np.asarray(mbi_c.potential.values) + # shape is (2, 3): score has OOD at 0, [0, 50] at 1, (50, 100] at 2. + expected = np.full((2, 3), -np.inf) + expected[0, 0] = 0.0 + expected[1, 1] = 0.0 + expected[1, 2] = 0.0 + np.testing.assert_array_equal(vals, expected) + + def test_numerical_validation_errors(self): + num_no_edges = domain.NumericalAttribute(min_value=0.0, max_value=10.0) + num_clipped = domain.NumericalAttribute( + min_value=0.0, max_value=10.0, clip_to_range=True, bin_edges=[5.0] + ) + cat = domain.CategoricalAttribute(['a', 'b']) + with self.assertRaisesRegex(ValueError, 'requires bin_edges'): + constraints.Constraint( + attribute_names=('num', 'cat'), + attribute_domains=(num_no_edges, cat), + possible_combinations=[((0.0, 10.0), 'a')], + ) + with self.assertRaisesRegex(ValueError, 'clip_to_range=False'): + constraints.Constraint( + attribute_names=('num', 'cat'), + attribute_domains=(num_clipped, cat), + possible_combinations=[(None, 'a')], + ) + with self.assertRaisesRegex(ValueError, 'must use'): + constraints.Constraint( + attribute_names=('num', 'cat'), + attribute_domains=(num_clipped, cat), + possible_combinations=[((-1.0, 5.0), 'a')], + ) + with self.assertRaisesRegex(ValueError, 'must use'): + constraints.Constraint( + attribute_names=('num', 'cat'), + attribute_domains=(num_clipped, cat), + possible_combinations=[((0.0, 3.0), 'a')], + ) + with self.assertRaisesRegex(ValueError, 'requires low < high'): + constraints.Constraint( + attribute_names=('num', 'cat'), + attribute_domains=(num_clipped, cat), + possible_combinations=[((5.0, 5.0), 'a')], + ) + with self.assertRaisesRegex(ValueError, 'CategoricalAttribute'): + constraints.Constraint( + attribute_names=('num', 'cat'), + attribute_domains=(num_clipped, cat), + functional_dependency={(0.0, 5.0): 'a', (5.0, 10.0): 'b'}, + ) + if __name__ == '__main__': absltest.main() diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py index 6b9f2671..db258917 100644 --- a/tests/data_generation_v3_test.py +++ b/tests/data_generation_v3_test.py @@ -606,6 +606,58 @@ def test_mbi_callbacks_logging_configured(self): mbi.callbacks.log('test', 'message', sep=' | ') self.assertTrue(any('test | message' in output for output in logs.output)) + def test_numerical_categorical_constraint_end_to_end(self): + domains = { + 'age': domain.NumericalAttribute( + min_value=0, max_value=100, dtype='int', bin_edges=[10, 17, 50] + ), + 'group': domain.CategoricalAttribute(['child', 'adult']), + } + c = constraints.Constraint( + attribute_names=('age', 'group'), + attribute_domains=(domains['age'], domains['group']), + possible_combinations=[((0, 17), 'child'), ((17, 100), 'adult')], + ) + schema = domain.Schema(domains, constraints=[c]) + df = pd.DataFrame({ + 'age': [5, 12, 17, 18, 35, 60] * 20, + 'group': ['child', 'child', 'child', 'adult', 'adult', 'adult'] * 20, + }) + rng = np.random.default_rng(0) + calibrated = TabularConfig(numerical_bins=4).configure(schema, budget=10.0) + syn = calibrated(rng, df).synthetic_data + self.assertTrue((syn.loc[syn['group'] == 'child', 'age'] <= 17).all()) + self.assertTrue((syn.loc[syn['group'] == 'adult', 'age'] >= 18).all()) + + def test_numerical_none_constraint_end_to_end(self): + domains = { + 'status': domain.CategoricalAttribute(['missing', 'present']), + 'score': domain.NumericalAttribute( + min_value=0.0, + max_value=100.0, + clip_to_range=False, + bin_edges=[25.0, 50.0, 75.0], + ), + } + c = constraints.Constraint( + attribute_names=('status', 'score'), + attribute_domains=(domains['status'], domains['score']), + possible_combinations=[ + ('missing', None), + ('present', (0.0, 100.0)), + ], + ) + schema = domain.Schema(domains, constraints=[c]) + df = pd.DataFrame({ + 'status': ['missing', 'present'] * 50, + 'score': [np.nan, 42.0] * 50, + }) + rng = np.random.default_rng(0) + calibrated = TabularConfig(numerical_bins=4).configure(schema, budget=10.0) + syn = calibrated(rng, df).synthetic_data + self.assertTrue(syn.loc[syn['status'] == 'missing', 'score'].isna().all()) + self.assertTrue(syn.loc[syn['status'] == 'present', 'score'].notna().all()) + if __name__ == '__main__': absltest.main()