diff --git a/dpsynth/domain.py b/dpsynth/domain.py index 10d8f18..843bc89 100644 --- a/dpsynth/domain.py +++ b/dpsynth/domain.py @@ -105,6 +105,11 @@ def standardize(self, value: Any) -> CategoricalValue: return value return self.possible_values[self.out_of_domain_index] + @functools.cached_property + def lookup(self) -> dict[str, int]: + """Returns a mapping from stringified values to their indices.""" + return {str(v): i for i, v in enumerate(self.possible_values)} + @attr.define(frozen=True) class OpenSetCategoricalAttribute: diff --git a/dpsynth/local_mode/beam_initializers.py b/dpsynth/local_mode/beam_initializers.py index 6fed006..ba348f7 100644 --- a/dpsynth/local_mode/beam_initializers.py +++ b/dpsynth/local_mode/beam_initializers.py @@ -32,6 +32,7 @@ import apache_beam as beam from dpsynth.local_mode import initialization +import mbi import numpy as np # A single row of tabular data: column name -> raw value. @@ -63,10 +64,10 @@ def __init__(self, initializers: dict[str, Initializer]): self._specs.append((column, 'numerical', meta)) elif isinstance(init, initialization.CategoricalInitializer): - lookup = { - str(v): i for i, v in enumerate(init.attribute.possible_values) + meta = { + 'lookup': init.attribute.lookup, + 'default': init.attribute.out_of_domain_index, } - meta = {'lookup': lookup, 'default': init.attribute.out_of_domain_index} self._specs.append((column, 'categorical', meta)) elif isinstance(init, initialization.OpenSetCategoricalInitializer): self._specs.append((column, 'openset', {})) @@ -240,3 +241,135 @@ def expand( rng=self._rng, ) ) + + +class _EncodeAndProject(beam.DoFn): + """Integer-encodes each row and emits (clique_index, linear_index) pairs.""" + + def __init__( + self, + column_measurements: dict[str, initialization.ColumnMeasurement], + domains: dict[str, Any], + workload: list[mbi.Clique], + ): + super().__init__() + self._cms = column_measurements + self._domains = domains + self._clique_meta: list[tuple[int, tuple[str, ...], tuple[int, ...]]] = [] + for idx, clique in enumerate(workload): + shape = tuple( + column_measurements[c].categorical_attribute.size for c in clique + ) + self._clique_meta.append((idx, clique, shape)) + + def _encode_value(self, col: str, raw_value: Any) -> int: + """Encodes a single raw value to an integer index.""" + cm = self._cms[col] + if cm.bin_edges is not None: + attr = self._domains[col] + value = attr.standardize(raw_value) + if math.isnan(value): + return 0 # OOD bucket (clip_to_range=False). + offset = 0 if attr.clip_to_range else 1 + return int(np.searchsorted(cm.bin_edges, value, side='left')) + offset + else: + cat = cm.categorical_attribute + return cat.lookup.get(str(raw_value), cat.out_of_domain_index) + + def process(self, row: Row): + encoded = {col: self._encode_value(col, row.get(col)) for col in self._cms} + for clique_idx, clique_cols, shape in self._clique_meta: + multi_index = tuple(encoded[c] for c in clique_cols) + linear = int(np.ravel_multi_index(multi_index, shape)) + yield clique_idx, linear + + +def _unpack_marginal_count(element): + """Restructures ((clique_idx, linear_idx), count) for GroupByKey.""" + (clique_idx, linear_idx), count = element + return clique_idx, (linear_idx, count) + + +def _assemble_dense_marginal(element, clique_meta, mbi_domain): + """Converts sparse counts to an mbi.Factor for one clique.""" + clique_idx, sparse_pairs = element + _, clique_cols, shape = clique_meta[clique_idx] + total_size = math.prod(shape) + dense = np.zeros(total_size, dtype=np.float64) + for linear_idx, count in sparse_pairs: + dense[linear_idx] = count + return mbi.Factor(mbi_domain.project(clique_cols), dense.reshape(shape)) + + +def _build_mbi_domain(column_measurements): + """Builds an mbi.Domain from ColumnMeasurement results.""" + attrs = tuple(column_measurements.keys()) + shape = tuple( + r.categorical_attribute.size for r in column_measurements.values() + ) + labels = tuple( + tuple(r.categorical_attribute.possible_values) + for r in column_measurements.values() + ) + return mbi.Domain(attributes=attrs, shape=shape, labels=labels) + + +class ComputeMarginals(beam.PTransform): + """Computes a workload of marginals over integer-encoded rows. + + Takes raw rows plus the ``ColumnMeasurement`` results from stage 1, + integer-encodes each row, and computes the contingency table for each + clique in the workload. The output is a singleton ``PCollection`` + containing one ``mbi.CliqueVector``. + + Attributes: + column_measurements: Per-column results from stage 1 initialization. + domains: Original attribute domain specs (needed for numerical encoding). + workload: List of cliques (tuples of column names) to measure. + """ + + def __init__( + self, + column_measurements: dict[str, initialization.ColumnMeasurement], + domains: dict[str, Any], + workload: list[mbi.Clique], + ): + super().__init__() + self._column_measurements = column_measurements + self._domains = domains + self._workload = workload + self._mbi_domain = _build_mbi_domain(column_measurements) + self._clique_meta = [] + for idx, clique in enumerate(workload): + shape = self._mbi_domain.project(clique).shape + self._clique_meta.append((idx, clique, shape)) + + def expand(self, rows: beam.PCollection[Row]): + mbi_domain = self._mbi_domain + + def _to_clique_vector(factors): + cliques = tuple(f.domain.attributes for f in factors) + arrays = {cl: f for cl, f in zip(cliques, factors)} + return mbi.CliqueVector(mbi_domain, cliques, arrays) + + return ( + rows + | 'EncodeProject' + >> beam.ParDo( + _EncodeAndProject( + self._column_measurements, self._domains, self._workload + ) + ) + | 'CountPerElement' >> beam.combiners.Count.PerElement() + | 'Unpack' >> beam.Map(_unpack_marginal_count) + | 'GroupByClique' >> beam.GroupByKey() + | 'ToLists' >> beam.MapTuple(_materialize_pairs) + | 'ToFactor' + >> beam.Map( + _assemble_dense_marginal, + clique_meta=self._clique_meta, + mbi_domain=mbi_domain, + ) + | 'ToList' >> beam.combiners.ToList() + | 'BuildCliqueVector' >> beam.Map(_to_clique_vector) + ) diff --git a/dpsynth/local_mode/vectorized_transformations.py b/dpsynth/local_mode/vectorized_transformations.py index fcc969a..d9e7d77 100644 --- a/dpsynth/local_mode/vectorized_transformations.py +++ b/dpsynth/local_mode/vectorized_transformations.py @@ -50,7 +50,7 @@ def discrete_encode( Returns: A 1-D integer array of indices into ``attribute_domain.possible_values``. """ - lookup = {str(v): i for i, v in enumerate(attribute_domain.possible_values)} + lookup = attribute_domain.lookup default = attribute_domain.out_of_domain_index # Loop over unique values only (typically ≪ len(data) for categoricals), # then remap via pure numpy fancy indexing. Normalizing to str avoids diff --git a/tests/local_mode/beam_initializers_test.py b/tests/local_mode/beam_initializers_test.py index ae19f04..f09ef91 100644 --- a/tests/local_mode/beam_initializers_test.py +++ b/tests/local_mode/beam_initializers_test.py @@ -20,6 +20,7 @@ from dpsynth import domain from dpsynth.local_mode import beam_initializers from dpsynth.local_mode import initialization +import mbi import numpy as np _TEST_RESULTS = [] @@ -33,7 +34,10 @@ class NumericalHistogramTest(absltest.TestCase): def _run(self, rows, attr, grid_size=101): init = initialization.NumericalInitializer( - name='x', num_partitions=4, attribute=attr, grid_size=grid_size + name='x', + num_partitions=4, + attribute=attr, + grid_size=grid_size, ).calibrate(zcdp_rho=np.inf) _TEST_RESULTS.clear() with beam.Pipeline() as p: @@ -80,10 +84,12 @@ class CategoricalCountsTest(absltest.TestCase): def test_basic_counts(self): attr = domain.CategoricalAttribute( - possible_values=['unk', 'a', 'b', 'c'], out_of_domain_index=0 + possible_values=['unk', 'a', 'b', 'c'], + out_of_domain_index=0, ) init = initialization.CategoricalInitializer( - name='col', attribute=attr + name='col', + attribute=attr, ).calibrate(zcdp_rho=np.inf) rows = [ {'col': 'a'}, @@ -187,5 +193,73 @@ def test_end_to_end_mixed(self): self.assertIsInstance(cm, initialization.ColumnMeasurement) +class ComputeMarginalsTest(absltest.TestCase): + + def test_marginals_match_manual_counts(self): + cat_attr = domain.CategoricalAttribute(possible_values=['a', 'b', 'c']) + num_attr = domain.NumericalAttribute(min_value=0, max_value=10) + cat_init = initialization.CategoricalInitializer( + name='color', + attribute=cat_attr, + ).calibrate(zcdp_rho=np.inf) + num_init = initialization.NumericalInitializer( + name='size', + num_partitions=4, + attribute=num_attr, + grid_size=11, + ).calibrate(zcdp_rho=np.inf) + domains = {'color': cat_attr, 'size': num_attr} + rows = [ + {'color': 'a', 'size': 0}, + {'color': 'a', 'size': 5}, + {'color': 'b', 'size': 5}, + {'color': 'b', 'size': 10}, + {'color': 'c', 'size': 0}, + {'color': 'c', 'size': 0}, + ] + + # Stage 1: get ColumnMeasurements. + inits = {'color': cat_init, 'size': num_init} + rng = np.random.default_rng(42) + _TEST_RESULTS.clear() + with beam.Pipeline() as p: + stats = ( + p + | 'Create1' >> beam.Create(rows) + | beam_initializers.ComputeSufficientStats(inits) + ) + _ = stats | 'ToDict1' >> beam.combiners.ToDict() | beam.Map(_store) + cms = beam_initializers.run_from_summary(_TEST_RESULTS[0], inits, rng) + + # Stage 2: compute marginals. + workload = [('color',), ('size',), ('color', 'size')] + _TEST_RESULTS.clear() + with beam.Pipeline() as p: + result = ( + p + | 'Create2' >> beam.Create(rows) + | beam_initializers.ComputeMarginals(cms, domains, workload) + ) + _ = result | beam.Map(_store) + + cv = _TEST_RESULTS[0] + self.assertIsInstance(cv, mbi.CliqueVector) + self.assertLen(cv.cliques, 3) + + # 1-way: color [a=2, b=2, c=2]. + np.testing.assert_array_equal( + cv.arrays[('color',)].datavector(), + [2, 2, 2], + ) + # 1-way: size total equals number of rows. + self.assertEqual(cv.arrays[('size',)].datavector().sum(), 6) + # 2-way: shape matches product of column sizes, total equals rows. + joint = cv.arrays[('color', 'size')] + expected_size = cms['color'].categorical_attribute.size + expected_size *= cms['size'].categorical_attribute.size + self.assertEqual(joint.domain.size(), expected_size) + self.assertEqual(joint.datavector().sum(), 6) + + if __name__ == '__main__': absltest.main()