From b047ff35f7e7f6ef78f2273722e590311424838e Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Wed, 7 Oct 2026 01:42:54 -0700 Subject: [PATCH] Add optional fixed `bin_edges` to `NumericalAttribute`. When `bin_edges` is specified on `NumericalAttribute`, `NumericalInitializer` skips DP quantile estimation and allocates the full budget to the Gaussian bin count measurement over the fixed discretization. PiperOrigin-RevId: 994938341 --- dpsynth/domain.py | 16 ++++++++++ dpsynth/local_mode/initialization.py | 28 +++++++++++------ tests/domain_test.py | 40 ++++++++++++++++++++++++- tests/local_mode/initialization_test.py | 31 +++++++++++++++++++ 4 files changed, 105 insertions(+), 10 deletions(-) diff --git a/dpsynth/domain.py b/dpsynth/domain.py index 6abe66f2..6445e13d 100644 --- a/dpsynth/domain.py +++ b/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: diff --git a/dpsynth/local_mode/initialization.py b/dpsynth/local_mode/initialization.py index f50b5f89..9799f183 100644 --- a/dpsynth/local_mode/initialization.py +++ b/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, diff --git a/tests/domain_test.py b/tests/domain_test.py index a5e740bd..89cc6274 100644 --- a/tests/domain_test.py +++ b/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 diff --git a/tests/local_mode/initialization_test.py b/tests/local_mode/initialization_test.py index 3bb5c48d..68a3cacd 100644 --- a/tests/local_mode/initialization_test.py +++ b/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):