Add optional fixed bin_edges to NumericalAttribute.

dpsynth/domain.py

@@ -188,6 +188,9 @@ class NumericalAttribute:
       interval (or returns the finite endpoint if the other is infinite).
       'interval' keeps the interval (string) in the output unchanged.
     description: An optional semantic description of the attribute.
+    bin_edges: Optional sequence of inner bin edges in [min_value, max_value).
+      When specified, fixed discretization using these edges is used instead of
+      computing DP quantiles.
   """
 
   min_value: float
@@ -197,6 +200,7 @@ class NumericalAttribute:
   dtype: str = 'float'
   interval_handling: str = 'midpoint'
   description: str | None = None
+  bin_edges: Sequence[float] | None = None
 
   def __post_init__(self):
     # Coerce to float to preserve the former attrs ``converter=float`` behavior;
@@ -230,6 +234,18 @@ def __post_init__(self):
             f' interval_handling={self.interval_handling!r}, got'
             f' sentinel={self.sentinel!r}.'
         )
+    if self.bin_edges is not None:
+      edges = [float(e) for e in self.bin_edges]
+      object.__setattr__(self, 'bin_edges', edges)
+      if not all(math.isfinite(e) for e in edges):
+        raise ValueError(f'bin_edges must contain finite numbers, got {edges}.')
+      lo, hi = self.min_value, self.max_value
+      if edges and (edges[0] < lo or edges[-1] >= hi):
+        raise ValueError(f'bin_edges must be in [{lo}, {hi}), got {edges}.')
+      if any(a >= b for a, b in zip(edges, edges[1:])):
+        raise ValueError(f'bin_edges must be strictly increasing: {edges}.')
+      if self.dtype == 'int' and not all(e.is_integer() for e in edges):
+        raise ValueError(f'bin_edges must be integers: {edges}.')
 
   @property
   def resolved_sentinel(self) -> float | int | str:

dpsynth/local_mode/initialization.py

@@ -159,6 +159,10 @@ def configure(
     budget_weights = rho_ratio ** np.arange(levels)[::-1]
     rho_levels = quantile_rho * budget_weights / budget_weights.sum()
     eps = np.sqrt(8.0 * rho_levels)
+    if attribute.bin_edges is not None:
+      # Skip DP quantiles and allocate the full budget to bin counts.
+      eps = np.array([])
+      count_rho = budget
     sigma = math.sqrt(0.5 / count_rho)
     return NumericalInitializer(
         config=self,
@@ -236,17 +240,23 @@ def from_summary(
       ood_count: float = 0.0,
   ) -> NumericalMeasurement:
     """Returns a NumericalMeasurement from pre-aggregated histogram counts."""
-    jitter_strategy = 'refine' if self.attribute.dtype == 'int' else 'symmetric'
-    indices = primitives.quantiles_from_histogram(
-        rng,
-        counts,
-        epsilon_levels=np.asarray(self.epsilon_levels),
-        jitter_strategy=jitter_strategy,
-        max_records_per_user=self.max_records_per_user,
-    )
     lower, upper, _ = self.grid_spec
     delta = (upper - lower) / max(1, np.asarray(counts).size - 1)
-    raw_edges = [lower + i * delta for i in indices]
+    if self.attribute.bin_edges is not None:
+      # Skip DP quantiles and aggregate the grid histogram at bin_edges.
+      raw_edges = self.attribute.bin_edges
+    else:
+      jitter_strategy = (
+          'refine' if self.attribute.dtype == 'int' else 'symmetric'
+      )
+      indices = primitives.quantiles_from_histogram(
+          rng,
+          counts,
+          epsilon_levels=np.asarray(self.epsilon_levels),
+          jitter_strategy=jitter_strategy,
+          max_records_per_user=self.max_records_per_user,
+      )
+      raw_edges = [lower + i * delta for i in indices]
 
     cm = edges_to_column_measurement(
         raw_edges=raw_edges,

tests/domain_test.py

@@ -15,11 +15,12 @@
 import math
 
 from absl.testing import absltest
+from absl.testing import parameterized
 from dpsynth import domain
 import numpy as np
 
 
-class TestDomain(absltest.TestCase):
+class TestDomain(parameterized.TestCase):
 
   def test_valid_attribute(self):
     attribute = domain.CategoricalAttribute(
@@ -80,6 +81,43 @@ def test_invalid_interval_handling(self):
     with self.assertRaises(ValueError):
       domain.NumericalAttribute(0, 10, interval_handling='bad')
 
+  def test_numerical_attribute_bin_edges_float(self):
+    attr = domain.NumericalAttribute(0, 10, bin_edges=(2, 5.5, 8))
+    self.assertEqual(attr.bin_edges, [2.0, 5.5, 8.0])
+
+  def test_numerical_attribute_bin_edges_int(self):
+    int_attr = domain.NumericalAttribute(
+        0, 10, dtype='int', bin_edges=[2, 5, 8]
+    )
+    self.assertEqual(int_attr.bin_edges, [2.0, 5.0, 8.0])
+
+  def test_numerical_attribute_bin_edges_yaml_roundtrip(self):
+    original = {
+        'float_num': domain.NumericalAttribute(0, 10, bin_edges=(2, 5.5, 8)),
+        'int_num': domain.NumericalAttribute(
+            0, 10, dtype='int', bin_edges=[2, 5, 8]
+        ),
+    }
+    temp_file = self.create_tempfile('bin_edges.yaml', mode='w+')
+    domain.to_yaml_file(original, temp_file.full_path)
+    loaded = domain.from_yaml_file(temp_file.full_path)
+    self.assertEqual(loaded, original)
+
+  @parameterized.named_parameters(
+      dict(testcase_name='below_min', bin_edges=[-1.0, 5.0]),
+      dict(testcase_name='above_max', bin_edges=[2.0, 10.0]),
+      dict(testcase_name='duplicate', bin_edges=[5.0, 5.0]),
+      dict(testcase_name='unsorted', bin_edges=[6.0, 4.0]),
+      dict(testcase_name='non_int_for_int', bin_edges=[2.5, 5.0], dtype='int'),
+      dict(testcase_name='nan', bin_edges=[math.nan, 5.0]),
+      dict(testcase_name='inf', bin_edges=[2.0, math.inf]),
+  )
+  def test_numerical_attribute_bin_edges_invalid(
+      self, bin_edges, dtype='float'
+  ):
+    with self.assertRaises(ValueError):
+      domain.NumericalAttribute(0, 10, dtype=dtype, bin_edges=bin_edges)
+
   def test_standardize_categorical(self):
     attribute = domain.CategoricalAttribute(
         possible_values=['a', 'b', 'c'], out_of_domain_index=1

tests/local_mode/initialization_test.py

@@ -237,6 +237,37 @@ def test_integer_heterogeneous_data_buckets(self):
         f' edges={result.bin_edges}',
     )
 
+  def test_numerical_initializer_fixed_bin_edges_dp_event(self):
+    attr = domain.NumericalAttribute(
+        min_value=0, max_value=10, bin_edges=[2.0, 5.0, 8.0]
+    )
+    initializer = initialization.NumericalInitializerConfig(num_partitions=4)
+    configured = initializer.configure(attr, budget=0.5)
+    expected = initialization.NumericalInitializer(
+        config=initializer,
+        attribute=attr,
+        epsilon_levels=(),
+        sigma=1.0,
+    )
+    self.assertEqual(configured, expected)
+    self.assertEqual(
+        configured.dp_event,
+        dp_accounting.ComposedDpEvent([dp_accounting.GaussianDpEvent(1.0)]),
+    )
+
+  def test_numerical_initializer_fixed_bin_edges_call(self):
+    attr = domain.NumericalAttribute(
+        min_value=0, max_value=10, bin_edges=[3.0, 7.0]
+    )
+    rng = np.random.default_rng(0)
+    initializer = initialization.NumericalInitializerConfig(num_partitions=8)
+    data = np.array([0.0, 1.0, 2.0, 3.0, 5.0, 7.0, 8.0, 9.0, 10.0])
+    measurement = initializer.configure(attr, budget=np.inf)(rng, data)
+
+    np.testing.assert_array_equal(measurement.bin_edges, [3.0, 7.0])
+    self.assertEqual(measurement.categorical_attribute.size, 3)
+    np.testing.assert_array_equal(measurement.noisy_counts, [4.0, 2.0, 3.0])
+
 
 class CategoricalInitializerTest(absltest.TestCase):