Updates to _find_optimal_rho reflecting new dp_accounting features (Poisson sampling of arbitrary mechanisms and smarter calibration).

dpsynth/api.py

@@ -186,11 +186,17 @@ def _find_optimal_rho(
     Raises:
       UnsupportedEventError: If no accountant supports the DpEvent.
     """
+    # Calibration fails for epsilon=0, we assume the make_event_fn correctly
+    # gives a non-DP event given rho = 0.
     rho = float('inf')
+    # Rho is roughly quadratic in epsilon, so we use epsilon^2 as a guess.
+    init_guess = target_epsilon**2
+    pld_error = None
     try:
-      # This is a heuristic to avoid excessively fine discretization in PLD
-      # accounting, which can cause OOM at extremely small target epsilons.
-      value_discretization_interval = max(1e-4, 1e-4 / (target_epsilon + 1e-5))
+      # Scale value_discretization_interval with target_epsilon to avoid
+      # the discretization dominating the epsilon at smaller budgets, which
+      # causes calibration to fail.
+      value_discretization_interval = min(1e-4, 1e-1 * target_epsilon)
       accountant_fn = functools.partial(
           dp_accounting.pld.PLDAccountant,
           value_discretization_interval=value_discretization_interval,
@@ -200,21 +206,36 @@ def _find_optimal_rho(
           make_event_from_param=make_event_fn,
           target_epsilon=target_epsilon,
           target_delta=target_delta,
+          bracket_interval=dp_accounting.LowerEndpointAndGuess(0.0, init_guess),  # pyrefly: ignore[bad-argument-count]
+      )
+    except (dp_accounting.UnsupportedEventError, NotImplementedError) as e:
+      # Okay if one of the accountants fails.
+      pld_error = e
+    rdp_error = None
+    try:
+      # Rho is roughly quadratic in epsilon, so we use epsilon^2 as a guess.
+      rho2 = dp_accounting.calibrate_dp_mechanism(
+          make_fresh_accountant=dp_accounting.rdp.RdpAccountant,
+          make_event_from_param=make_event_fn,
+          target_epsilon=target_epsilon,
+          target_delta=target_delta,
+          bracket_interval=dp_accounting.LowerEndpointAndGuess(0.0, init_guess),  # pyrefly: ignore[bad-argument-count]
+      )
+      rho = min(rho, rho2)
+    except (dp_accounting.UnsupportedEventError, NotImplementedError) as e:
+      # Okay if one of the accountants fails.
+      rdp_error = e
+
+    if rho == float('inf'):
+      raise dp_accounting.UnsupportedEventError(
+          'No accountant supports the mechanism:\n'
+          f'  PLDAccountant error: {pld_error}\n'
+          f'  RdpAccountant error: {rdp_error}'
       )
-    except (dp_accounting.UnsupportedEventError, NotImplementedError):
-      # If PLD accounting is not supported, fall back to RDP accounting.
-      pass
-
-    rho2 = dp_accounting.calibrate_dp_mechanism(
-        make_fresh_accountant=dp_accounting.rdp.RdpAccountant,
-        make_event_from_param=make_event_fn,
-        target_epsilon=target_epsilon,
-        target_delta=target_delta,
-    )
 
     # RDP can also be better than PLD in some cases due to looseness in the
     # handling of certain DpEvents like the ExponentialMechanismDpEvent.
-    return min(rho, rho2)
+    return rho
 
   def calibrate(
       self,

tests/data_generation_v3_test.py

@@ -171,7 +171,11 @@ def test_end_to_end_mixed_domain(self):
     self.assertIsInstance(synthetic_df, pd.DataFrame)
     self.assertListEqual(synthetic_df.columns.tolist(), ['A', 'B'])
 
-  def test_end_to_end_with_epsilon_delta(self):
+  @parameterized.named_parameters(
+      ('large_epsilon', 100, 0.1),
+      ('small_epsilon', 1e-5, 1e-6),
+  )
+  def test_end_to_end_with_epsilon_delta(self, epsilon: float, delta: float):
     domains = {
         'A': domain.CategoricalAttribute(
             possible_values=['a', 'b', 'c'], out_of_domain_index=0
@@ -182,7 +186,9 @@ def test_end_to_end_with_epsilon_delta(self):
     }
     df = pd.DataFrame({'A': ['a', 'b', 'c'], 'B': ['x', 'y', 'z']})
     rng = np.random.default_rng(0)
-    calibrated = TabularConfig().calibrate(domains, epsilon=100, delta=0.1)
+    calibrated = TabularConfig().calibrate(
+        domains, epsilon=epsilon, delta=delta
+    )
     result = calibrated(rng, df)
     synthetic_df = result.synthetic_data
     self.assertIsInstance(synthetic_df, pd.DataFrame)
@@ -447,13 +453,13 @@ def test_poisson_calibrate_with_mixed_domains(self):
         'C': domain.OpenSetCategoricalAttribute(),
     }
     config = TabularConfig()
-    with self.assertRaises(dp_accounting.UnsupportedEventError):
-      _ = config.calibrate(
-          domains,
-          epsilon=1.0,
-          delta=1e-3,
-          poisson_sampling_prob=0.1,
-      )
+    mechanism = config.calibrate(
+        domains,
+        epsilon=1.0,
+        delta=1e-6,
+        poisson_sampling_prob=0.1,
+    )
+    self.assertIsNotNone(mechanism)
 
   def test_configure_infinite_zcdp_rho(self):
     domains = {

tests/discrete_mechanisms/discrete_mechanisms_test.py

@@ -117,18 +117,10 @@ def test_compression_with_initial_measurements(self, config):
 class CalibrationTest(parameterized.TestCase):
   """Tests that calibration works across mechanisms."""
 
-  @parameterized.named_parameters(*_MECHANISMS.items())
-  def test_zero_epsilon_calibration(self, mechanism):
-    rng = np.random.default_rng(0)
-    data = _make_skewed_dataset(rng)
+  @parameterized.named_parameters(_MECHANISMS.items())
+  def test_low_epsilon_calibration(self, mechanism):
     if isinstance(mechanism, independent.IndependentConfig):
       return
-    result = mechanism.calibrate(epsilon=0.0, delta=0.01)(rng, data)
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
-
-  @parameterized.named_parameters(*_MECHANISMS.items())
-  def test_low_epsilon_calibration(self, mechanism):
-    self.skipTest('Low epsilon calibration is currently really slow, skipping.')
     rng = np.random.default_rng(0)
     data = _make_skewed_dataset(rng)
     result = mechanism.calibrate(epsilon=1e-3, delta=1e-5)(rng, data)