diff --git a/dpsynth/adapters/beam.py b/dpsynth/adapters/beam.py index 715a45c6..99b00761 100644 --- a/dpsynth/adapters/beam.py +++ b/dpsynth/adapters/beam.py @@ -42,7 +42,6 @@ from dpsynth import domain from dpsynth.discrete_mechanisms import common as dm_common from dpsynth.local_mode import initialization -from dpsynth.local_mode import primitives import mbi import numpy as np @@ -193,6 +192,8 @@ def run_from_summary( sparse_stats: dict[str, list[tuple[Any, int]]], initializers: dict[str, CalibratedInitializer], rng: np.random.Generator, + *, + num_rows: int | None = None, ) -> dict[str, initialization.ColumnMeasurement]: """Converts materialized sparse stats to ColumnMeasurements on the driver. @@ -204,6 +205,8 @@ def run_from_summary( produced by ``ComputeSufficientStats``. initializers: Calibrated initializers keyed by column name. rng: NumPy random generator for DP noise. + num_rows: Optional total row count, used to compute out-of-domain counts for + numerical columns with ``clip_to_range=False``. Returns: Per-column ``ColumnMeasurement`` results. @@ -213,7 +216,12 @@ def run_from_summary( sparse = sparse_stats[column] if isinstance(init, initialization.NumericalInitializer): counts = _sparse_to_dense_numerical(sparse, init.grid_spec[2]) - results[column] = init.from_summary(rng, counts) + ood_count = ( + float(num_rows - counts.sum()) + if (not init.attribute.clip_to_range and num_rows is not None) + else 0.0 + ) + results[column] = init.from_summary(rng, counts, ood_count=ood_count) elif isinstance(init, initialization.CategoricalInitializer): counts = _sparse_to_dense_categorical(sparse, init.attribute.size) results[column] = init.from_summary(rng, counts) @@ -382,7 +390,6 @@ def generate_from_marginals( rng: np.random.Generator, column_measurements: dict[str, initialization.ColumnMeasurement], marginals: mbi.CliqueVector, - total_measurement: mbi.LinearMeasurement, ) -> data_generation_v3.DataGenerationResult: """Runs the discrete mechanism and decoding from pre-computed marginals. @@ -391,7 +398,6 @@ def generate_from_marginals( rng: NumPy random generator for the discrete mechanism's DP noise. column_measurements: Per-column results from pass 1 initialization. marginals: The exact joint marginals computed by pass 2. - total_measurement: The DP-noised total-count measurement (clique ``()``). Returns: A DataGenerationResult containing the synthetic DataFrame. @@ -402,7 +408,7 @@ def generate_from_marginals( column_measurements, synth.schema ) - initial_measurements = [total_measurement, *codec.one_way_measurements()] + initial_measurements = codec.one_way_measurements() logging.info('[DPSynth/Beam]: Running discrete mechanism.') # pyrefly: ignore[missing-attribute,not-callable] mechanism_result = synth.base_mechanism( @@ -439,7 +445,6 @@ def _run_two_pass( ) -> data_generation_v3.DataGenerationResult: """Two-pass Beam pipeline that delegates to a local TabularConfig.""" - sigma = synth.total_count_sigma inits = cast(dict[str, CalibratedInitializer], synth.initializers) if pipeline_kwargs is None: pipeline_kwargs = {} @@ -465,15 +470,11 @@ def _run_two_pass( _ = count | 'WriteRowCount' >> beam.Map(_write, path=count_path) # We run this on the driver so we don't have to track worker-side RNGs. sparse_stats = _read(summary_path) - column_measurements = run_from_summary(sparse_stats, inits, rng) num_rows = int(_read(count_path)) - logging.info('[DPSynth/Beam]: Pass 1 complete.') - # pyrefly: ignore[missing-attribute] - total = primitives.add_gaussian_noise( - rng, float(num_rows), sigma, cast(int, synth.max_records_per_user) + column_measurements = run_from_summary( + sparse_stats, inits, rng, num_rows=num_rows ) - total = float(max(1.0, total)) - total_measurement = mbi.LinearMeasurement(np.array([total]), (), sigma) + logging.info('[DPSynth/Beam]: Pass 1 complete.') # Ask the configured discrete mechanism which marginals it needs. mbi_domain = data_generation_v3.TabularCodec.from_measurements( @@ -499,7 +500,7 @@ def _run_two_pass( # Run the discrete mechanism and decode on the driver. return generate_from_marginals( - synth, rng, column_measurements, clique_vector, total_measurement + synth, rng, column_measurements, clique_vector ) finally: # Only remove a temp dir we created; never a user-supplied temp_location. diff --git a/dpsynth/data_generation_v3.py b/dpsynth/data_generation_v3.py index 7aee93e1..21408bc5 100644 --- a/dpsynth/data_generation_v3.py +++ b/dpsynth/data_generation_v3.py @@ -29,7 +29,6 @@ from dpsynth import domain from dpsynth.discrete_mechanisms import common as dm_common from dpsynth.local_mode import initialization -from dpsynth.local_mode import primitives from dpsynth.local_mode import vectorized_transformations as vtx import mbi import numpy as np @@ -157,8 +156,6 @@ def one_way_measurements(self) -> list[mbi.LinearMeasurement]: cm = c.column_measurement if cm.noisy_counts is None: continue - elif isinstance(cm, initialization.NumericalMeasurement): - query = mbi.DatavectorQuery(use_for_total_estimation=False) elif isinstance(cm, initialization.OpenSetMeasurement): query = mbi.SlicedQuery(start=1) else: @@ -208,7 +205,6 @@ class TabularMechanism(api.CalibratedMechanism): schema: The dataset schema and constraints. base_mechanism: The calibrated discrete mechanism. initializers: Per-column calibrated initializers. - total_count_sigma: Sigma for the total-count mechanism. max_records_per_user: Assumed upper bound on the number of records a single user contributes. """ @@ -217,16 +213,12 @@ class TabularMechanism(api.CalibratedMechanism): schema: domain.Schema base_mechanism: discrete_mechanisms.CalibratedMechanism initializers: dict[str, api.CalibratedMechanism] - total_count_sigma: float = dataclasses.field(repr=False) max_records_per_user: int = 1 @property def dp_event(self) -> dp_accounting.DpEvent: """Returns the composed DpEvent for all sub-mechanisms.""" events = [init.dp_event for init in self.initializers.values()] - events.append( - dp_accounting.GaussianDpEvent(noise_multiplier=self.total_count_sigma) - ) events.append(self.base_mechanism.dp_event) events = [e for e in events if not isinstance(e, dp_accounting.NoOpDpEvent)] @@ -267,32 +259,13 @@ def __call__( ) # Phase 1: Per-column initialization. - # Measure total count first, then run per-column initializers. def _run_initializers(): - noisy_total = primitives.add_gaussian_noise( - rng, - len(data), - self.total_count_sigma, - self.max_records_per_user, - ) - total = max(1.0, noisy_total) - total_measurement = mbi.LinearMeasurement( - noisy_measurement=np.array([total]), - clique=(), - stddev=self.max_records_per_user * self.total_count_sigma, - ) + return { + col: init(rng, data[col].values) + for col, init in self.initializers.items() + } - results: dict[str, initialization.ColumnMeasurement] = {} - for col, init in self.initializers.items(): - if isinstance(init, initialization.NumericalInitializer): - results[col] = init( - rng, data[col].values, estimated_total=float(total) - ) - else: - results[col] = init(rng, data[col].values) - return total_measurement, results - - total_measurement, results = _checkpoint.get_or_compute( + results = _checkpoint.get_or_compute( 'column_measurements', _run_initializers ) @@ -306,7 +279,6 @@ def _run_initializers(): rng, results, discrete, - total_measurement, cross_attribute_constraints=cross_attribute_constraints, num_rows=num_rows, ) @@ -316,7 +288,6 @@ def from_summary( rng: np.random.Generator, column_measurements: Mapping[str, initialization.ColumnMeasurement], data: mbi.Dataset | mbi.CliqueVector, - total_measurement: mbi.LinearMeasurement, *, cross_attribute_constraints: Sequence[constraints.Constraint] = (), num_rows: int | None = None, @@ -328,11 +299,11 @@ def from_summary( ) mbi_constraints = tuple(c.to_mbi() for c in cross_attribute_constraints) - # Feed the noisy total (clique ()) and one-way column measurements as - # initial measurements so the mechanism does not re-measure them. + # Feed one-way column measurements as initial measurements so the mechanism + # does not re-measure them. codec = TabularCodec.from_measurements(column_measurements, self.schema) one_way_measurements = codec.one_way_measurements() - initial_measurements = [total_measurement, *one_way_measurements] + initial_measurements = list(one_way_measurements) mappings = {} if isinstance(data, mbi.Dataset): @@ -539,10 +510,8 @@ def configure( schema, self.numerical_bins, self.numerical_epsilon_ratio ) init_rho = self.init_budget_fraction * zcdp_rho - # +1 for the DPGaussianCount that always measures the total. - per_col_rho = init_rho / (len(inits) + 1) + per_col_rho = init_rho / len(inits) discrete_rho = (1 - self.init_budget_fraction) * zcdp_rho - total_count_sigma = (0.5 / per_col_rho) ** 0.5 calibrated_inits: dict[str, api.CalibratedMechanism] @@ -566,7 +535,6 @@ def configure( schema=schema, base_mechanism=calibrated_discrete, initializers=calibrated_inits, - total_count_sigma=total_count_sigma, max_records_per_user=max_records_per_user, ) diff --git a/dpsynth/local_mode/initialization.py b/dpsynth/local_mode/initialization.py index 7e47bcaa..abb7f16a 100644 --- a/dpsynth/local_mode/initialization.py +++ b/dpsynth/local_mode/initialization.py @@ -122,17 +122,26 @@ class NumericalInitializerConfig(api.MechanismConfig): deeper level of recursive bisection. Defaults to 1.0 (uniform budget split across levels). Setting to sqrt(2) approx 1.414 can provide minor accuracy gains on smooth continuous data. + quantile_budget_fraction: Fraction of the zCDP budget allocated to the + quantile tree; the remainder is allocated to the Gaussian mechanism for + measuring discretized bin counts. Defaults to 0.5. """ num_partitions: int max_grid_size: int = 10_000_000 epsilon_ratio: float = 1.0 + quantile_budget_fraction: float = 0.5 def __post_init__(self): if self.max_grid_size < 2: raise ValueError(f'max_grid_size must be >= 2, got {self.max_grid_size}.') if self.num_partitions >= self.max_grid_size: raise ValueError(f'{self.num_partitions=} >= {self.max_grid_size=}') + if not 0.0 < self.quantile_budget_fraction < 1.0: + raise ValueError( + 'quantile_budget_fraction must be in (0, 1), got' + f' {self.quantile_budget_fraction}.' + ) def configure( self, attribute=None, *, zcdp_rho, delta=0, max_records_per_user=1 @@ -144,14 +153,18 @@ def configure( if 2**levels != self.num_partitions: raise ValueError(f'{self.num_partitions=} must be a power of 2.') + quantile_rho = zcdp_rho * self.quantile_budget_fraction + count_rho = zcdp_rho * (1.0 - self.quantile_budget_fraction) rho_ratio = self.epsilon_ratio**2 budget_weights = rho_ratio ** np.arange(levels)[::-1] - rho_levels = zcdp_rho * budget_weights / budget_weights.sum() + rho_levels = quantile_rho * budget_weights / budget_weights.sum() eps = np.sqrt(8.0 * rho_levels) + sigma = math.sqrt(0.5 / count_rho) return NumericalInitializer( config=self, attribute=attribute, epsilon_levels=tuple(eps.tolist()), + sigma=sigma, max_records_per_user=max_records_per_user, ) @@ -163,6 +176,7 @@ class NumericalInitializer(api.CalibratedMechanism): config: NumericalInitializerConfig attribute: domain.NumericalAttribute epsilon_levels: tuple[float, ...] + sigma: float max_records_per_user: int = 1 def __post_init__(self): @@ -183,28 +197,29 @@ def grid_spec(self) -> tuple[float, float, int]: def grid_size(self) -> int: return self.grid_spec[2] - @property - def zcdp_rho(self) -> float: - return sum(e**2 / 8.0 for e in self.epsilon_levels) - @property def dp_event(self) -> dp_accounting.DpEvent: - """Returns the composed privacy event for the quantile computation.""" - return dp_accounting.ComposedDpEvent([ + """Returns the composed privacy event for quantiles and bin counts.""" + events: list[dp_accounting.DpEvent] = [ dp_accounting.ExponentialMechanismDpEvent(epsilon=float(eps)) for eps in self.epsilon_levels - ]) + ] + events.append(dp_accounting.GaussianDpEvent(noise_multiplier=self.sigma)) + return dp_accounting.ComposedDpEvent(events) def __call__( self, rng: np.random.Generator, data: np.ndarray, - *, - estimated_total: float | None = None, ) -> NumericalMeasurement: - """Returns a NumericalMeasurement with the discretization transform.""" + """Returns a NumericalMeasurement with the discretization and noisy counts.""" counts = self._grid_histogram(data) - return self.from_summary(rng, counts, estimated_total=estimated_total) + ood_count = ( + float(len(data) - counts.sum()) + if not self.attribute.clip_to_range + else 0.0 + ) + return self.from_summary(rng, counts, ood_count=ood_count) def _grid_histogram(self, data): """Returns the quantile candidate-grid histogram (length grid_size).""" @@ -218,7 +233,7 @@ def from_summary( rng: np.random.Generator, counts: np.ndarray, *, - estimated_total: float | None = None, + ood_count: float = 0.0, ) -> NumericalMeasurement: """Returns a NumericalMeasurement from pre-aggregated histogram counts.""" jitter_strategy = 'refine' if self.attribute.dtype == 'int' else 'symmetric' @@ -233,21 +248,29 @@ def from_summary( delta = (upper - lower) / max(1, np.asarray(counts).size - 1) raw_edges = [lower + i * delta for i in indices] - return edges_to_column_measurement( + cm = edges_to_column_measurement( raw_edges=raw_edges, attribute=self.attribute, - zcdp_rho=self.zcdp_rho, - estimated_total=estimated_total, - max_records_per_user=self.max_records_per_user, + ) + splits = np.round((cm.bin_edges - lower) / delta).astype(np.int64) + 1 + segments = np.split(np.asarray(counts, dtype=float), splits) + in_domain_counts = np.array([seg.sum() for seg in segments], dtype=float) + if self.attribute.clip_to_range: + bin_counts = in_domain_counts + else: + bin_counts = np.r_[ood_count, in_domain_counts] + noisy = primitives.add_gaussian_noise( + rng, bin_counts, self.sigma, self.max_records_per_user + ) + stddev = self.max_records_per_user * self.sigma + return dataclasses.replace( + cm, noisy_counts=np.asarray(noisy), stddev=stddev ) def edges_to_column_measurement( raw_edges, attribute, - zcdp_rho, - estimated_total=None, - max_records_per_user=1, ) -> NumericalMeasurement: """Converts raw quantile edges into a NumericalMeasurement. @@ -257,37 +280,16 @@ def edges_to_column_measurement( Args: raw_edges: Quantile edge values (unsorted duplicates are fine). attribute: The ``NumericalAttribute`` defining the data domain. - zcdp_rho: Total zCDP rho consumed by the quantile mechanism. - estimated_total: If provided, a heuristic one-way measurement is included. - max_records_per_user: Assumed upper bound on the number of records a single - user contributes. Returns: - A ``NumericalMeasurement`` with bin edges and optionally noisy counts. + A ``NumericalMeasurement`` with deduplicated bin edges. """ raw_edges = np.asarray(raw_edges, dtype=float) - bin_edges, edge_counts = np.unique(raw_edges, return_counts=True) - max_val = attribute.max_value - if len(bin_edges) > 0 and bin_edges[-1] >= max_val: - tail_count = edge_counts[-1] + bin_edges = np.unique(raw_edges) + if len(bin_edges) > 0 and bin_edges[-1] >= attribute.max_value: bin_edges = bin_edges[:-1] - edge_counts = edge_counts[:-1] - bin_weights = np.append(edge_counts, tail_count + 1) - else: - bin_weights = np.append(edge_counts, 1) cat_attr = vtx.categorical_attribute_from_edges(bin_edges, attribute) - - noisy_counts = None - stddev = np.nan - if estimated_total is not None: - if not attribute.clip_to_range: - bin_weights = np.r_[0, bin_weights] - noisy_counts = estimated_total * bin_weights / bin_weights.sum() - stddev = max_records_per_user / np.sqrt(zcdp_rho) - - return NumericalMeasurement( - cat_attr, bin_edges, noisy_counts=noisy_counts, stddev=stddev - ) + return NumericalMeasurement(cat_attr, bin_edges) @dataclasses.dataclass(frozen=True, kw_only=True) diff --git a/dpsynth/relational/synthesizer.py b/dpsynth/relational/synthesizer.py index 29c50c4f..d9418c73 100644 --- a/dpsynth/relational/synthesizer.py +++ b/dpsynth/relational/synthesizer.py @@ -166,7 +166,6 @@ def _run_single_col_initializer( rng: np.random.Generator, data: np.ndarray, weights: np.ndarray, - estimated_total: float | None = None, ) -> initialization.ColumnMeasurement: """Runs a single column initializer on weighted standalone table data. @@ -181,8 +180,6 @@ def _run_single_col_initializer( rng: NumPy random generator. data: 1D array of column values. weights: 1D float array of row sensitivity weights. - estimated_total: Optional root table estimated total count for - NumericalInitializer heuristic one-way measurements. Returns: A ColumnMeasurement containing the discovered categorical attribute, @@ -196,8 +193,10 @@ def _run_single_col_initializer( values = np.asarray(data, dtype=float) if attr.clip_to_range: values = np.where(np.isnan(values), attr.min_value, values) + ood_count = 0.0 else: in_domain = (values >= attr.min_value) & (values <= attr.max_value) + ood_count = float(weights[~in_domain].sum()) values, weights = values[in_domain], weights[in_domain] if attr.dtype == 'int': values = np.round(values) @@ -205,7 +204,7 @@ def _run_single_col_initializer( delta = (upper - lower) / (gs - 1) indices = initialization.encode_to_grid(values, lower, upper, delta) counts = np.bincount(indices, weights=weights, minlength=gs) - return init.from_summary(rng, counts, estimated_total=estimated_total) + return init.from_summary(rng, counts, ood_count=ood_count) if isinstance(init, initialization.CategoricalInitializer): encoded = vtx.discrete_encode(data, init.attribute) @@ -230,7 +229,6 @@ def _run_table_initializers( rng: np.random.Generator, tables: Mapping[str, pd.DataFrame], weights: Mapping[str, np.ndarray], - estimated_total: float | None = None, ) -> dict[str, dict[str, initialization.ColumnMeasurement]]: """Runs column initializers across all tables on weighted data. @@ -240,8 +238,6 @@ def _run_table_initializers( rng: NumPy random generator. tables: Mapping from table name to filtered active DataFrames. weights: Mapping from table name to 1D sensitivity weights. - estimated_total: Optional root table estimated total count for - NumericalInitializer heuristic one-way measurements. Returns: A nested mapping from table and column name to its ColumnMeasurement. @@ -257,7 +253,6 @@ def _run_table_initializers( rng=rng, data=table_df[col_name].to_numpy(), weights=table_w, - estimated_total=estimated_total, ) results[table_name] = table_results return results @@ -397,7 +392,6 @@ def _run_table_preprocessing( rng=rng, tables=filtered_tables, weights=weights, - estimated_total=noisy_root_total, ) codecs, datasets, mappings, one_ways = _encode_and_compress_tables( diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py index 2f4896a8..c488bafc 100644 --- a/tests/data_generation_v3_test.py +++ b/tests/data_generation_v3_test.py @@ -241,38 +241,6 @@ def test_calibrate_domain_positional_only(self): with self.assertRaises(TypeError): dpsynth.calibrate(config, domain=domains, epsilon=1.0, delta=1e-5) # pyrefly: ignore[unexpected-keyword] - def test_numerical_only_uses_dp_count(self): - """Numerical-only domains should allocate a DPGaussianCount for total.""" - domains = { - 'A': domain.NumericalAttribute(min_value=0, max_value=10), - 'B': domain.NumericalAttribute(min_value=-10, max_value=10), - } - df = pd.DataFrame({'A': [5, 5, 0], 'B': [5, -10, -5]}, dtype=float) - rng = np.random.default_rng(0) - calibrated = TabularConfig().configure(domains, zcdp_rho=100.0) - - # total_count_sigma should be set for numerical-only domains. - self.assertIsNotNone(calibrated.total_count_sigma) - synthetic_df = calibrated(rng, df).synthetic_data - self.assertListEqual(synthetic_df.columns.tolist(), ['A', 'B']) - - def test_mixed_domain_always_has_dp_count(self): - """Mixed domains also allocate a DPGaussianCount for total.""" - domains = { - 'A': domain.CategoricalAttribute( - possible_values=['a', 'b', 'c'], out_of_domain_index=0 - ), - 'B': domain.NumericalAttribute(min_value=0, max_value=10), - } - df = pd.DataFrame({'A': ['a', 'b', 'c'], 'B': [1.0, 5.0, 10.0]}) - rng = np.random.default_rng(0) - calibrated = TabularConfig().configure(domains, zcdp_rho=100.0) - - # DPGaussianCount is always allocated. - self.assertIsNotNone(calibrated.total_count_sigma) - synthetic_df = calibrated(rng, df).synthetic_data - self.assertListEqual(synthetic_df.columns.tolist(), ['A', 'B']) - @parameterized.product( sentinel=[np.nan, None], clip_to_range=[True, False], @@ -491,7 +459,6 @@ def test_configure_infinite_zcdp_rho(self): config = TabularConfig() mechanism = config.configure(domains, zcdp_rho=np.inf) self.assertIsNotNone(mechanism) - self.assertEqual(mechanism.total_count_sigma, 0.0) def test_configure_with_schema(self): domains = { diff --git a/tests/local_mode/initialization_test.py b/tests/local_mode/initialization_test.py index 240517bd..b96ebd81 100644 --- a/tests/local_mode/initialization_test.py +++ b/tests/local_mode/initialization_test.py @@ -29,9 +29,11 @@ def test_numerical_initializer_dp_event(self): initializer = initialization.NumericalInitializerConfig(num_partitions=4) event = initializer.configure(attr, zcdp_rho=1.0).dp_event self.assertIsInstance(event, dp_accounting.ComposedDpEvent) - self.assertLen(event.events, 2) - for e in event.events: + self.assertLen(event.events, 3) + for e in event.events[:2]: self.assertIsInstance(e, dp_accounting.ExponentialMechanismDpEvent) + self.assertIsInstance(event.events[2], dp_accounting.GaussianDpEvent) + self.assertAlmostEqual(event.events[2].noise_multiplier, 1.0) def test_numerical_initializer_call(self): attr = domain.NumericalAttribute(min_value=0, max_value=10) @@ -43,7 +45,8 @@ def test_numerical_initializer_call(self): self.assertIsInstance(measurement, initialization.NumericalMeasurement) self.assertEqual(measurement.categorical_attribute.size, 4) - self.assertIsNone(measurement.noisy_counts) + self.assertIsNotNone(measurement.noisy_counts) + self.assertEqual(measurement.stddev, 0.0) self.assertIsNotNone(measurement.bin_edges) encoded_data = vtx.discretize(data, measurement.bin_edges, attr) @@ -53,6 +56,26 @@ def test_numerical_initializer_call(self): self.assertEqual(counts.sum(), 9) self.assertLen(counts, 4) self.assertTrue(np.all(counts > 0)) + np.testing.assert_array_equal(measurement.noisy_counts, counts) + + def test_numerical_initializer_out_of_domain_values(self): + attr = domain.NumericalAttribute( + min_value=0, max_value=10, clip_to_range=False + ) + rng = np.random.default_rng(0) + initializer = initialization.NumericalInitializerConfig(num_partitions=4) + + data = np.array( + [-5.0, np.nan, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 15.0] + ) + measurement = initializer.configure(attr, zcdp_rho=np.inf)(rng, data) + + self.assertEqual(measurement.categorical_attribute.size, 5) + self.assertIsNotNone(measurement.noisy_counts) + self.assertLen(measurement.noisy_counts, 5) + # 3 OOD values (-5.0, nan, 15.0) map to bin 0. + self.assertEqual(measurement.noisy_counts[0], 3.0) + self.assertEqual(measurement.noisy_counts.sum(), 11.0) def test_numerical_initializer_deduplicates_bin_edges(self): """Concentrated data can make quantiles return duplicate edges.""" @@ -103,20 +126,6 @@ def test_numerical_initializer_integer_edges_are_floored(self): self.assertGreaterEqual(result.bin_edges[0], 0) self.assertLess(result.bin_edges[-1], 100) - def test_numerical_initializer_measurement_with_merged_bins(self): - """When integer edges collapse, merged bins get proportionally more mass.""" - attr = domain.NumericalAttribute(min_value=0, max_value=100, dtype='int') - rng = np.random.default_rng(0) - initializer = initialization.NumericalInitializerConfig(num_partitions=8) - # Concentrated data will cause edge collisions. - data = np.array([50] * 100 + [1, 99]) - result = initializer.configure(attr, zcdp_rho=1.0)( - rng, data, estimated_total=100.0 - ) - self.assertIsNotNone(result.noisy_counts) - # Counts should sum to the estimated total. - np.testing.assert_allclose(result.noisy_counts.sum(), 100.0, atol=1e-10) - def test_max_grid_size_below_two_raises(self): for bad in (0, 1): with self.assertRaises(ValueError): @@ -139,54 +148,29 @@ def test_int_grid_reserves_budget_for_jitter_refinement(self): init.configure(attr, zcdp_rho=1.0).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) - initializer = initialization.NumericalInitializerConfig(num_partitions=4) - data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]) - result = initializer.configure(attr, zcdp_rho=1.0)( - rng, data, estimated_total=100.0 - ) - - self.assertIsNotNone(result.noisy_counts) - # Measurement should be uniform counts: estimated_total / num_bins each. - num_bins = result.categorical_attribute.size - expected_count = 100.0 / num_bins - np.testing.assert_allclose( - result.noisy_counts, - np.full(num_bins, expected_count), - ) - # stddev should be 1/sqrt(rho) = 1.0 in count space. - self.assertAlmostEqual(result.stddev, 1.0) - - def test_numerical_initializer_no_measurement_without_estimated_total(self): - attr = domain.NumericalAttribute(min_value=0, max_value=10) - rng = np.random.default_rng(0) - initializer = initialization.NumericalInitializerConfig(num_partitions=4) - data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]) - result = initializer.configure(attr, zcdp_rho=1.0)(rng, data) - self.assertIsNone(result.noisy_counts) - def test_numerical_initializer_epsilon_ratio_default_and_custom(self): attr = domain.NumericalAttribute(min_value=0, max_value=10) - # Default is 1.0 (uniform) + # Default is 1.0 (uniform), and quantile_budget_fraction defaults to 0.5. init_default = initialization.NumericalInitializerConfig(num_partitions=4) self.assertEqual(init_default.epsilon_ratio, 1.0) - calibrated_default = init_default.configure(attr, zcdp_rho=2.0) - # For 4 partitions (2 levels), uniform budget splits zcdp_rho equally: - # rho/2 each. eps = sqrt(8 * rho_level) = sqrt(8 * 1.0) = sqrt(8). + self.assertEqual(init_default.quantile_budget_fraction, 0.5) + calibrated_default = init_default.configure(attr, zcdp_rho=4.0) + # quantile_rho = 2.0, count_rho = 2.0. + # For 4 partitions (2 levels), uniform budget splits quantile_rho equally: + # 1.0 each. eps = sqrt(8 * rho_level) = sqrt(8 * 1.0) = sqrt(8). np.testing.assert_allclose( calibrated_default.epsilon_levels, (np.sqrt(8.0), np.sqrt(8.0)), ) + self.assertAlmostEqual(calibrated_default.sigma, 0.5) # Custom ratio (e.g. sqrt(2)) init_custom = initialization.NumericalInitializerConfig( num_partitions=4, epsilon_ratio=np.sqrt(2) ) self.assertEqual(init_custom.epsilon_ratio, np.sqrt(2)) - calibrated_custom = init_custom.configure(attr, zcdp_rho=2.0) - # rho_ratio = 2.0. budget_weights = [2.0, 1.0]. + calibrated_custom = init_custom.configure(attr, zcdp_rho=4.0) + # quantile_rho = 2.0. rho_ratio = 2.0. budget_weights = [2.0, 1.0]. # Leaves get 2/3, root gets 1/3. # rho_levels = [2.0 * 2/3, 2.0 * 1/3] = [4/3, 2/3] # eps_levels = [sqrt(8 * 4/3), sqrt(8 * 2/3)] @@ -204,36 +188,12 @@ def test_integer_edges_at_max_value_absorbed_into_last_bin(self): # 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(attr, zcdp_rho=100.0)( - rng, data, estimated_total=len(data) - ) + result = initializer.configure(attr, zcdp_rho=100.0)(rng, data) # No edge should equal max_value (they get absorbed). if len(result.bin_edges) > 0: self.assertLess(result.bin_edges[-1], 10) - # Counts must still sum to the estimated total. - np.testing.assert_allclose(result.noisy_counts.sum(), len(data), atol=1e-10) - # The last bin absorbs the max_value spike, so it should dominate the - # uniform per-bin baseline (estimated_total / num_bins). - counts = result.noisy_counts - self.assertGreater(counts.max(), len(data) / len(counts)) - - def test_bin_weights_sum_to_num_partitions(self): - """bin_weights must always sum to num_partitions regardless of dedup.""" - attr = domain.NumericalAttribute(min_value=0, max_value=20, dtype='int') - for seed in range(10): - rng = np.random.default_rng(seed) - initializer = initialization.NumericalInitializerConfig(num_partitions=8) - data = np.array([5] * 50 + [15] * 50) - result = initializer.configure(attr, zcdp_rho=1.0)( - rng, data, estimated_total=100.0 - ) - # Counts sum to the estimated total. - np.testing.assert_allclose( - result.noisy_counts.sum(), - 100.0, - atol=1e-10, - err_msg=f'seed={seed}: counts do not sum to estimated_total', - ) + self.assertIsNotNone(result.noisy_counts) + self.assertLen(result.noisy_counts, result.categorical_attribute.size) def test_integer_jitter_prevents_spurious_splits(self): """Positive jitter should prevent edges from splitting across integers.""" @@ -255,9 +215,7 @@ def test_integer_heterogeneous_data_buckets(self): initializer = initialization.NumericalInitializerConfig(num_partitions=4) # Deliberately lumpy distribution: 45 points across 4 distinct values. data = np.array([0] * 10 + [3] * 10 + [5] * 17 + [6] * 8) - result = initializer.configure(attr, zcdp_rho=np.inf)( - rng, data, estimated_total=len(data) - ) + result = initializer.configure(attr, zcdp_rho=np.inf)(rng, data) # Edges should be integers strictly inside [min_value, max_value). for e in result.bin_edges: self.assertEqual(e, int(e), f'edge {e} is not an integer') @@ -280,213 +238,6 @@ def test_integer_heterogeneous_data_buckets(self): ) -class MeasurementApproximationTest(parameterized.TestCase): - """Property test: measurement counts ≈ true histogram. - - The heuristic uniform measurement should satisfy - L1(measurement, true_histogram) < max(3 / sqrt(rho), 2 - 1/K) - where the 3/sqrt(rho) term covers quantile noise and 2-1/K covers the - worst-case uniform-vs-delta misspecification error (K = num_bins). - """ - - @parameterized.named_parameters( - # --- High budget (rho=1000): misspecification-dominated --- - dict( - testcase_name='uniform_int', - attr=domain.NumericalAttribute( - min_value=0, max_value=50, dtype='int' - ), - data=np.arange(51), - num_partitions=4, - rho=1000.0, - ), - dict( - testcase_name='heterogeneous_int', - attr=domain.NumericalAttribute( - min_value=0, max_value=10, dtype='int' - ), - data=np.array([0] * 10 + [3] * 10 + [5] * 17 + [6] * 8 + [9] * 5), - num_partitions=4, - rho=1000.0, - ), - dict( - testcase_name='uniform_float', - attr=domain.NumericalAttribute( - min_value=0, max_value=100, dtype='float' - ), - data=np.linspace(0, 100, 200), - num_partitions=8, - rho=1000.0, - ), - dict( - testcase_name='boundary_heavy_int', - attr=domain.NumericalAttribute( - min_value=0, max_value=100, dtype='int' - ), - data=np.array([0] * 30 + [100] * 30 + [50] * 40), - num_partitions=4, - rho=1000.0, - ), - dict( - testcase_name='bimodal_int', - attr=domain.NumericalAttribute( - min_value=0, max_value=100, dtype='int' - ), - data=np.array([10] * 50 + [90] * 50), - num_partitions=8, - rho=1000.0, - ), - dict( - testcase_name='sparse_int', - attr=domain.NumericalAttribute( - min_value=0, max_value=1000, dtype='int' - ), - data=np.array([1] * 40 + [500] * 10 + [999] * 50), - num_partitions=4, - rho=1000.0, - ), - # --- Low budget (rho=0.5): noise-dominated --- - dict( - testcase_name='uniform_int_low_rho', - attr=domain.NumericalAttribute( - min_value=0, max_value=50, dtype='int' - ), - data=np.arange(51), - num_partitions=4, - rho=0.5, - ), - dict( - testcase_name='uniform_float_low_rho', - attr=domain.NumericalAttribute( - min_value=0, max_value=100, dtype='float' - ), - data=np.linspace(0, 100, 200), - num_partitions=8, - rho=0.5, - ), - dict( - testcase_name='heterogeneous_int_low_rho', - attr=domain.NumericalAttribute( - min_value=0, max_value=10, dtype='int' - ), - data=np.array([0] * 10 + [3] * 10 + [5] * 17 + [6] * 8 + [9] * 5), - num_partitions=4, - rho=0.5, - ), - ) - def test_measurement_approximates_true_histogram( - self, attr, data, num_partitions, rho - ): - rng = np.random.default_rng(0) - initializer = initialization.NumericalInitializerConfig( - num_partitions=num_partitions - ) - result = initializer.configure(attr, zcdp_rho=rho)( - rng, data, estimated_total=len(data) - ) - # -- Structural checks (must always hold) -- - num_bins = result.categorical_attribute.size - self.assertGreaterEqual(num_bins, 2) - noisy_counts = result.noisy_counts - self.assertIsNotNone(noisy_counts) - # Measurement counts must sum to the estimated total. - np.testing.assert_allclose(noisy_counts.sum(), len(data), atol=1e-10) - # All measurement counts should be positive. - self.assertTrue( - np.all(noisy_counts > 0), - f'non-positive measurement counts: {noisy_counts}', - ) - # -- Statistical approximation check -- - encoded = vtx.discretize(data, result.bin_edges, attr) - true_counts = np.bincount(encoded, minlength=num_bins).astype(float) - true_prob = true_counts / true_counts.sum() - meas = noisy_counts - meas_prob = meas / meas.sum() - l1_dist = np.abs(true_prob - meas_prob).sum() - # 3/sqrt(rho) covers quantile noise; 2-1/K covers uniform-vs-delta. - max_l1 = max(3.0 / np.sqrt(rho), 2.0 - 1.0 / num_bins) - self.assertLess( - l1_dist, - max_l1, - f'Measurement too far from true histogram (L1={l1_dist:.3f},' - f' bound={max_l1:.3f}, rho={rho}):\n' - f' true_prob = {true_prob}\n' - f' meas_prob = {meas_prob}', - ) - - def test_measurement_property_random_configs(self): - """Randomized property test: max(3/sqrt(rho), 2-1/K) over many configs.""" - master_rng = np.random.default_rng(20260619) - num_trials = 50 - - for trial in range(num_trials): - with self.subTest(trial=trial): - # -- Random configuration -- - rho = float(10 ** master_rng.uniform(-1, 3)) # 0.1 to 1000 - num_partitions = int(master_rng.choice([2, 4, 8, 16])) - min_val = int(master_rng.integers(0, 50)) - max_val = min_val + int(master_rng.integers(5, 200)) - is_int = bool(master_rng.random() < 0.5) - - # -- Random data as a mixture of 1-5 point masses -- - num_modes = int(master_rng.integers(1, 6)) - modes = master_rng.integers(min_val, max_val + 1, size=num_modes) - weights = master_rng.dirichlet(np.ones(num_modes)) - n = int(master_rng.integers(50, 300)) - counts = np.round(weights * n).astype(int) - counts[-1] = n - counts[:-1].sum() # ensure exact total - data = np.concatenate([np.full(c, m) for c, m in zip(counts, modes)]) - if not is_int: - data = data.astype(float) - - dtype = 'int' if is_int else 'float' - attr = domain.NumericalAttribute( - min_value=min_val, max_value=max_val, dtype=dtype - ) - - # -- Run initializer -- - rng = np.random.default_rng(trial) - initializer = initialization.NumericalInitializerConfig( - num_partitions=num_partitions - ) - result = initializer.configure(attr, zcdp_rho=rho)( - rng, data, estimated_total=len(data) - ) - - # -- Structural checks -- - noisy_counts = result.noisy_counts - self.assertIsNotNone(noisy_counts) - np.testing.assert_allclose( - noisy_counts.sum(), - len(data), - atol=1e-10, - err_msg=f'trial={trial}: counts do not sum to estimated_total', - ) - self.assertTrue( - np.all(noisy_counts > 0), - f'trial={trial}: non-positive probabilities {noisy_counts}', - ) - - # -- L1 property check -- - num_bins = result.categorical_attribute.size - encoded = vtx.discretize(data, result.bin_edges, attr) - true_counts = np.bincount(encoded, minlength=num_bins).astype(float) - true_prob = true_counts / true_counts.sum() - meas = noisy_counts - meas_prob = meas / meas.sum() - l1_dist = np.abs(true_prob - meas_prob).sum() - max_l1 = max(3.0 / np.sqrt(rho), 2.0 - 1.0 / num_bins) - self.assertLess( - l1_dist, - max_l1, - f'trial={trial} (rho={rho:.2f}, K={num_partitions},' - f' modes={modes}, is_int={is_int}):\n' - f' L1={l1_dist:.3f}, bound={max_l1:.3f}\n' - f' true_prob={true_prob}\n' - f' meas_prob={meas_prob}', - ) - - class CategoricalInitializerTest(absltest.TestCase): def test_dp_event(self): @@ -595,8 +346,8 @@ def test_calibrate_sets_dp_event(self): ).configure(attr, zcdp_rho=1.0) event = init.dp_event self.assertIsInstance(event, dp_accounting.ComposedDpEvent) - # 4 partitions = 2 levels. - self.assertLen(event.events, 2) + # 4 partitions = 2 levels + 1 Gaussian count event = 3 events. + self.assertLen(event.events, 3) def test_integer_attribute_snaps_edges(self): rng = np.random.default_rng(42) @@ -621,15 +372,19 @@ def test_call_and_from_summary_produce_same_structure(self): ).configure(attr, zcdp_rho=1.0) data = np.linspace(0.0, 100.0, 500) - cm_call = init(rng, data, estimated_total=500.0) + cm_call = init(rng, data) self.assertIsNotNone(cm_call.bin_edges) self.assertIsNotNone(cm_call.noisy_counts) + self.assertLen(cm_call.noisy_counts, cm_call.categorical_attribute.size) rng2 = np.random.default_rng(42) counts = np.random.randint(0, 100, size=max_grid_size) - cm_summary = init.from_summary(rng2, counts, estimated_total=500.0) + cm_summary = init.from_summary(rng2, counts) self.assertIsNotNone(cm_summary.bin_edges) self.assertIsNotNone(cm_summary.noisy_counts) + self.assertLen( + cm_summary.noisy_counts, cm_summary.categorical_attribute.size + ) class MaxRecordsPerUserTest(parameterized.TestCase): diff --git a/tests/relational/synthesizer_test.py b/tests/relational/synthesizer_test.py index 359148f0..c671aeb2 100644 --- a/tests/relational/synthesizer_test.py +++ b/tests/relational/synthesizer_test.py @@ -747,7 +747,7 @@ def test_run_single_col_initializer(self): data_num = np.array([10.0, 20.0, 30.0, 80.0, 90.0]) weights_num = np.array([0.5, 0.5, 1.0, 0.5, 0.5]) res_num = synthesizer._run_single_col_initializer( - num_init, rng, data_num, weights_num, estimated_total=100.0 + num_init, rng, data_num, weights_num ) self.assertIsInstance(res_num, initialization.NumericalMeasurement) self.assertIsNotNone(res_num.bin_edges) @@ -840,7 +840,6 @@ def test_run_table_initializers(self): rng=rng, tables=tables, weights=weights, - estimated_total=2.0, ) self.assertIn('Household', results)