diff --git a/dpsynth/local_mode/_quantiles.py b/dpsynth/local_mode/_quantiles.py index a76c668..d754a14 100644 --- a/dpsynth/local_mode/_quantiles.py +++ b/dpsynth/local_mode/_quantiles.py @@ -15,17 +15,32 @@ """DP quantiles from dense histograms via recursive median bisection. This module computes differentially private quantile edges from a dense -histogram of counts, using the discrete exponential mechanism. The primary -use case is a two-pass pipeline: a first pass computes a dense histogram -over a fine-grained grid, then ``quantiles_from_histogram`` finds DP -quantiles from that histogram without touching individual records. - -Public API: - - ``quantiles_from_histogram``: DP quantiles via recursive median splits. +histogram of counts, using the discrete exponential mechanism. It works purely +in index space -- ``quantiles_from_histogram`` returns cell indices into the +histogram, and the caller maps those indices to domain values. The primary use +case is a two-pass pipeline: a first pass computes a dense histogram over a +fine-grained grid, then ``quantiles_from_histogram`` finds DP quantile indices +from that histogram without touching individual records. + +Tie handling via jitter +----------------------- +Recursive median bisection needs each record assigned to one side of every +split independently. A "spike" of records tied on one grid cell breaks this: a +whole-cell split sends all that mass to one side, biasing the quantiles and +collapsing sub-ranges (dropping edges). We fix this by breaking ties directly +in the histogram domain rather than over the raw data values -- each cell's +count is redistributed to nearby cells as ``Multinomial(count, kernel)`` (one +draw per non-empty cell), which is distributionally identical to independently +perturbing each record and so needs no extra privacy budget. The ``refine`` +strategy uses a strictly-positive kernel over refined sub-cells (value- +preserving); the ``symmetric`` strategy uses a symmetric kernel over neighboring +grid cells. """ from __future__ import annotations +from typing import Literal + import numpy as np import scipy.special @@ -64,51 +79,67 @@ def _median_from_histogram( return int(rng.choice(total_points, p=probs)) +def jitter_factor(num_partitions): + """Returns a data-independent jitter resolution m from num_partitions.""" + # m >= num_partitions keeps each jittered cell below one partition's mass; + # the 4x absorbs multinomial fluctuation. + return max(1, 4 * num_partitions) + + def quantiles_from_histogram( rng: np.random.Generator, counts: np.ndarray, - lower: float, - upper: float, epsilon_levels: np.ndarray, - grid_size: int = 10_000_000, -) -> list[float]: - """Computes DP quantiles from a dense histogram via recursive median splits. + jitter_strategy: Literal['symmetric', 'refine'], +) -> list[int]: + """DP quantile edge indices into ``counts`` via jittered median bisection. - Uses the discrete exponential mechanism to recursively find medians, - splitting the histogram at each level to produce ``num_buckets - 1`` - quantile edges. The number of buckets is ``2 ** len(epsilon_levels)``. + Operates purely in index space: it returns cell indices into ``counts`` and + leaves the mapping from index to domain value to the caller. Args: rng: A numpy random number generator. - counts: Dense 1D histogram of shape ``(grid_size,)``. - lower: Lower bound of the data domain. - upper: Upper bound of the data domain (exclusive). + counts: Dense 1D histogram counts. epsilon_levels: Per-level exponential mechanism epsilons, ordered from the deepest (finest) level to the shallowest (coarsest). - grid_size: Number of uniformly spaced grid points spanning ``[lower, - upper]``. + jitter_strategy: Specifies the pre-processing jitter strategy, - + 'symmetric': jitter mass to +/- m//2 neighbors on the same grid. - + 'refine': jitter mass to m equivalent sub-cells. Returns: - A sorted list of ``2 ** len(epsilon_levels) - 1`` quantile edge values. + A sorted list of ``2 ** len(epsilon_levels) - 1`` cell indices. """ - levels = len(epsilon_levels) - if levels == 0: - return [] - - # Uniform grid: counts[i] corresponds to value lower + i * delta. - delta = (upper - lower) / (grid_size - 1) + counts = np.asarray(counts) + m = jitter_factor(2 ** len(epsilon_levels)) + + if jitter_strategy == 'refine': + stride, offsets = m, np.arange(m) + else: + half = m // 2 + stride, offsets = 1, np.arange(-half, half + 1) + + # Scatter each cell's mass over its jittered targets: same law as perturbing + # each record, so it breaks ties without spending extra privacy budget. + num_cells = counts.size * stride + nz = np.flatnonzero(counts) + probas = np.full(offsets.size, 1.0 / offsets.size) + split = rng.multinomial(counts[nz].astype(np.int64), probas) + targets = np.clip(nz[:, None] * stride + offsets, 0, num_cells - 1) + jittered = np.bincount( + targets.flatten(), weights=split.flatten(), minlength=num_cells + ) def _rec(lo_idx, hi_idx, depth): - if depth == 0 or lo_idx >= hi_idx: + if depth == 0: return [] - sub_counts = counts[lo_idx:hi_idx] - median_local = _median_from_histogram( - rng, sub_counts, epsilon_levels[depth - 1] + median_idx = lo_idx + _median_from_histogram( + rng, jittered[lo_idx:hi_idx], epsilon_levels[depth - 1] ) - median_global_idx = lo_idx + median_local - median_value = lower + median_global_idx * delta - left = _rec(lo_idx, median_global_idx, depth - 1) - right = _rec(median_global_idx, hi_idx, depth - 1) - return left + [median_value] + right - - return _rec(0, len(counts), levels) + left = _rec(lo_idx, median_idx, depth - 1) + right = _rec(median_idx, hi_idx, depth - 1) + return left + [median_idx] + right + + result = _rec(0, jittered.size, len(epsilon_levels)) + if jitter_strategy == 'refine': + result = [idx // m for idx in result] + return result diff --git a/dpsynth/local_mode/initialization.py b/dpsynth/local_mode/initialization.py index ea49c57..0cba735 100644 --- a/dpsynth/local_mode/initialization.py +++ b/dpsynth/local_mode/initialization.py @@ -22,6 +22,7 @@ import dp_accounting from dpsynth import domain +from dpsynth.local_mode import _quantiles from dpsynth.local_mode import primitives from dpsynth.local_mode import vectorized_transformations as vtx import mbi @@ -95,8 +96,11 @@ def _grid_spec(self) -> tuple[float, float, int]: """Returns (lower, upper, grid_size) for the quantile candidate grid.""" attr = self.attribute if attr.dtype == 'int': + # Reserve budget for the m-fold refinement so the refined grid fits. + m = _quantiles.jitter_factor(self.num_partitions) + budget = max(2, self.max_grid_size // m) int_range = int(attr.max_value - attr.min_value + 1) - step = max(1, math.ceil(int_range / self.max_grid_size)) + step = max(1, math.ceil(int_range / budget)) gs = math.ceil(int_range / step) return (attr.min_value, attr.min_value + (gs - 1) * step, gs) return (attr.min_value, attr.exclusive_max_value, self.max_grid_size) @@ -110,12 +114,14 @@ def configure( self, *, zcdp_rho: float, delta: float = 0.0, epsilon_ratio: float = 2.0 ) -> NumericalInitializer: """Returns a copy calibrated to the given zCDP budget.""" - lower, upper, gs = self._grid_spec + lower, upper, _ = self._grid_spec mechanism = primitives.DPQuantiles( num_partitions=self.num_partitions, lower=lower, upper=upper, - grid_size=gs, + jitter_strategy=( + 'refine' if self.attribute.dtype == 'int' else 'symmetric' + ), ).configure(zcdp_rho=zcdp_rho, epsilon_ratio=epsilon_ratio) return dataclasses.replace(self, mechanism=mechanism) diff --git a/dpsynth/local_mode/primitives.py b/dpsynth/local_mode/primitives.py index a5fd635..de44fde 100644 --- a/dpsynth/local_mode/primitives.py +++ b/dpsynth/local_mode/primitives.py @@ -22,6 +22,7 @@ import dataclasses import math +from typing import Literal import dp_accounting from dpsynth import api @@ -274,13 +275,14 @@ class DPQuantiles(DPMechanism): num_partitions: Number of quantile partitions (must be a power of 2). lower: Lower bound of the data domain. upper: Upper bound of the data domain (exclusive). - grid_size: Number of uniformly spaced grid points. + jitter_strategy: Tie-breaking jitter passed to ``quantiles_from_histogram``: + ``'refine'`` for integer attributes, ``'symmetric'`` for continuous ones. """ num_partitions: int lower: float upper: float - grid_size: int = 10_000_000 + jitter_strategy: Literal['symmetric', 'refine'] = 'symmetric' _epsilon_levels: tuple[float, ...] | None = dataclasses.field( default=None, repr=False ) @@ -336,14 +338,16 @@ def __call__( """Returns quantile edges from a dense histogram of counts.""" if self._epsilon_levels is None: raise ValueError(_UNCALIBRATED_MSG.format(param='_epsilon_levels')) - return _quantiles.quantiles_from_histogram( + indices = _quantiles.quantiles_from_histogram( rng, counts, - self.lower, - self.upper, epsilon_levels=np.asarray(self._epsilon_levels), - grid_size=self.grid_size, + jitter_strategy=self.jitter_strategy, ) + # Map cell indices back to domain values; delta is the grid step, which + # equals the integer step for integer attributes so edges stay integer. + delta = (self.upper - self.lower) / max(1, np.asarray(counts).size - 1) + return [self.lower + i * delta for i in indices] @dataclasses.dataclass diff --git a/tests/local_mode/_quantiles_test.py b/tests/local_mode/_quantiles_test.py index 13b08ca..e5dd7ad 100644 --- a/tests/local_mode/_quantiles_test.py +++ b/tests/local_mode/_quantiles_test.py @@ -23,29 +23,113 @@ class QuantilesFromHistogramTest(parameterized.TestCase): def test_no_levels_returns_empty(self): rng = np.random.default_rng(0) counts = np.array([10]) + for jitter_strategy in ('symmetric', 'refine'): + edges = _quantiles.quantiles_from_histogram( + rng, counts, np.array([]), jitter_strategy + ) + self.assertEmpty(edges) + + @parameterized.product( + levels=(1, 2, 3, 4), + jitter_strategy=('symmetric', 'refine'), + ) + def test_edge_count_matches_levels(self, levels, jitter_strategy): + rng = np.random.default_rng(0) + grid_size = 10001 + counts = rng.integers(0, 20, size=grid_size) edges = _quantiles.quantiles_from_histogram( rng, counts, - 0.0, - 10.0, - epsilon_levels=np.array([]), + epsilon_levels=np.ones(levels), + jitter_strategy=jitter_strategy, ) - self.assertEmpty(edges) + self.assertLen(edges, 2**levels - 1) @parameterized.parameters(1, 2, 3, 4) - def test_edge_count_matches_levels(self, levels): + def test_edge_count_matches_levels_with_spike(self, levels): + # A large tied spike must not collapse split ranges and drop edges. With + # whole-cell splits this dropped edges; jitter breaks up the spike so the + # recursion always emits the full 2**levels - 1 edges. + counts = np.zeros(101, dtype=np.int64) + counts[:40] = 1 + counts[40] = 1000 + counts[41:80] = 1 + edges = _quantiles.quantiles_from_histogram( + np.random.default_rng(0), + counts, + epsilon_levels=np.array([np.inf] * levels), + jitter_strategy='refine', + ) + self.assertLen(edges, 2**levels - 1) + + def test_integer_edges_are_integer_indices(self): + # Integer attributes must return integer cell indices into ``counts``. + counts = np.zeros(101, dtype=np.int64) + counts[40] = 5000 + counts[:40] = 50 + counts[41:] = 50 + edges = _quantiles.quantiles_from_histogram( + np.random.default_rng(0), + counts, + epsilon_levels=np.array([np.inf] * 3), + jitter_strategy='refine', + ) + for edge in edges: + self.assertEqual(edge, int(edge)) + self.assertBetween(edge, 0, counts.size - 1) + + def test_exact_budget_matches_numpy_smooth(self): rng = np.random.default_rng(0) - grid_size = 10001 - counts = rng.integers(0, 20, size=grid_size) + data = rng.integers(0, 100, size=50000) + counts = np.bincount(data, minlength=101) edges = _quantiles.quantiles_from_histogram( rng, counts, - 0.0, - 10.0, - epsilon_levels=np.ones(levels), - grid_size=grid_size, + epsilon_levels=np.array([np.inf] * 3), + jitter_strategy='refine', ) - self.assertLen(edges, 2**levels - 1) + # Cell indices map 1:1 to values here (delta == 1), so compare directly. + expected = np.quantile(data, [0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875]) + np.testing.assert_allclose(edges, expected, atol=1.0) + + def test_exact_budget_matches_numpy_with_spike(self): + # Reproduces the failure mode from the 'hours-per-week' column of the adult + # census dataset: a dominant spike at 40 carrying ~45% of the mass, with + # lighter integer support on either side. The pre-jitter whole-cell split + # lumped all tied mass into one subtree, biasing the low quantiles (e.g. the + # 0.25 edge came out as 18 instead of 40) and dropping upper edges. Jitter + # breaks up the spike so the recursive medians match numpy. + below = np.arange(1, 40).repeat(230) + spike = np.full(13500, 40) + above = np.arange(41, 80).repeat(190) + data = np.concatenate([below, spike, above]) + counts = np.bincount(data, minlength=101) + edges = _quantiles.quantiles_from_histogram( + np.random.default_rng(0), + counts, + epsilon_levels=np.array([np.inf] * 3), + jitter_strategy='refine', + ) + # Cell indices map 1:1 to values here (delta == 1), so compare directly. + expected = np.quantile(data, [0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875]) + np.testing.assert_allclose(edges, expected, atol=1.0) + + def test_spike_owns_consecutive_edges(self): + # When a single value holds a majority of the mass, the quantiles on both + # sides of the median should collapse onto that value. This is the key + # correctness property the deterministic whole-cell split got wrong. + counts = np.zeros(101, dtype=np.int64) + counts[:40] = 20 + counts[40] = 20000 # ~96% of the mass. + counts[41:80] = 20 + edges = _quantiles.quantiles_from_histogram( + np.random.default_rng(0), + counts, + epsilon_levels=np.array([np.inf] * 3), + jitter_strategy='refine', + ) + # The interior quantiles (0.25 through 0.75) must all be cell 40. + self.assertEqual(edges[1:6], [40, 40, 40, 40, 40]) if __name__ == '__main__': diff --git a/tests/local_mode/initialization_test.py b/tests/local_mode/initialization_test.py index 5d070af..aa1d449 100644 --- a/tests/local_mode/initialization_test.py +++ b/tests/local_mode/initialization_test.py @@ -16,6 +16,7 @@ from absl.testing import parameterized import dp_accounting from dpsynth import domain +from dpsynth.local_mode import _quantiles from dpsynth.local_mode import initialization from dpsynth.local_mode import vectorized_transformations as vtx import numpy as np @@ -158,6 +159,19 @@ def test_max_grid_size_two_float(self): result = init(rng, np.linspace(0, 100, 100)) self.assertIsNotNone(result.categorical_attribute) + def test_int_grid_reserves_budget_for_jitter_refinement(self): + # A wide integer range with many partitions must not blow past + # max_grid_size once the m-fold jitter refinement is applied. + attr = domain.NumericalAttribute( + min_value=0, max_value=10_000_000, dtype='int' + ) + max_grid_size = 100_000 + init = initialization.NumericalInitializer( + name='x', num_partitions=64, attribute=attr, max_grid_size=max_grid_size + ) + m = _quantiles.jitter_factor(init.num_partitions) + self.assertLessEqual(init.grid_size * m, max_grid_size) + def test_numerical_initializer_measurement_with_estimated_total(self): attr = domain.NumericalAttribute(min_value=0, max_value=10) rng = np.random.default_rng(0) @@ -198,10 +212,12 @@ def test_integer_edges_at_max_value_absorbed_into_last_bin(self): initializer = initialization.NumericalInitializer( name='test', num_partitions=8, attribute=attr ) - # All data at max_value: all edges should land at or near max_value. - data = np.array([10] * 100) + # A spread of lower values carrying most of the mass, plus a moderate spike + # at max_value. The lower values form genuine interior bins, while the top + # quantile edges land at max_value and must be absorbed into the last bin. + data = np.concatenate([np.repeat(np.arange(0, 10), 20), np.full(100, 10)]) result = initializer.configure(zcdp_rho=100.0)( - rng, data, estimated_total=100.0 + rng, data, estimated_total=len(data) ) # No edge should equal max_value (they get absorbed). if len(result.bin_edges) > 0: @@ -210,8 +226,7 @@ def test_integer_edges_at_max_value_absorbed_into_last_bin(self): np.testing.assert_allclose( result.measurement.noisy_measurement.sum(), 1.0, atol=1e-10 ) - # The bin containing max_value=10 should get the most mass (either the - # last bin, or a bin that absorbed all degenerate edges). + # The last bin absorbs the max_value spike, so it should dominate. counts = result.measurement.noisy_measurement self.assertGreater(counts.max(), 1.0 / len(counts))