Introduces a ComputeMarginals Beam PTransform that takes raw rows + ColumnMeasurements from stage 1 and computes a workload of marginals as an mbi.CliqueVector.

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:

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

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

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