Support NumericalAttribute intervals in dpsynth.constraints.Constraint:

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,15 +142,15 @@ 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
       attribute values. Requires exactly two attributes.
   """
 
   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)

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()

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()