Add _as_zcdp and _parallel_compose_event to dpsynth._calibration

dpsynth/_calibration.py

@@ -16,13 +16,55 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 import functools
+import math
 from typing import Any
 
 import dp_accounting
 
 
+def _compose_deltas(deltas: Sequence[float]) -> float:
+  """Composes failure probabilities via 1 - prod_i (1 - delta_i)."""
+  # Use log1p/expm1 to avoid catastrophic cancellation for small delta_i.
+  return -math.expm1(sum(math.log1p(-d) for d in deltas))
+
+
+def _as_zcdp(event: dp_accounting.DpEvent) -> tuple[float, float]:
+  """Returns (rho, delta) for an event satisfying delta-approximate rho-zCDP."""
+  if isinstance(event, dp_accounting.NoOpDpEvent):
+    return 0.0, 0.0
+  if isinstance(event, dp_accounting.GaussianDpEvent):
+    return 0.5 / event.noise_multiplier**2, 0.0
+  if isinstance(event, dp_accounting.ExponentialMechanismDpEvent):
+    return event.epsilon**2 / 8.0, 0.0
+  if isinstance(event, dp_accounting.ZCDpEvent) and event.xi == 0:
+    return event.rho, 0.0
+  if isinstance(event, dp_accounting.dp_event.EpsilonDeltaDpEvent):
+    return 0.5 * event.epsilon**2, event.delta
+  if isinstance(event, dp_accounting.SelfComposedDpEvent):
+    rho, delta = _as_zcdp(event.event)
+    return event.count * rho, _compose_deltas([delta] * event.count)
+  if isinstance(event, dp_accounting.ComposedDpEvent):
+    pairs = [_as_zcdp(e) for e in event.events]
+    return sum(r for r, _ in pairs), _compose_deltas([d for _, d in pairs])
+  raise dp_accounting.UnsupportedEventError(f'Unsupported event: {event}.')
+
+
+def _parallel_compose_event(
+    events: Sequence[dp_accounting.DpEvent],
+) -> dp_accounting.DpEvent:
+  """Returns the worst-case (rho, delta) zCDP event across parallel events."""
+  if not events:
+    return dp_accounting.NoOpDpEvent()
+  pairs = [_as_zcdp(e) for e in events]
+  rho = max(r for r, _ in pairs)
+  delta = max(d for _, d in pairs)
+  base = dp_accounting.ZCDpEvent(rho)
+  failure = dp_accounting.dp_event.EpsilonDeltaDpEvent(0.0, delta)
+  return base if delta == 0 else dp_accounting.ComposedDpEvent([base, failure])
+
+
 def with_group_size(
     event: dp_accounting.DpEvent, group_size: int
 ) -> dp_accounting.DpEvent:

dpsynth/experimental/nested.py

@@ -31,6 +31,7 @@
 
 from absl import logging
 import dp_accounting
+from dpsynth import _calibration
 from dpsynth import api
 from dpsynth import data_generation_v3
 from dpsynth import discrete_mechanisms
@@ -117,27 +118,15 @@ class NestedTabularMechanism(api.CalibratedMechanism):
   shared_synth: data_generation_v3.TabularMechanism
   detail_synths: Mapping[str, data_generation_v3.TabularMechanism]
 
-  # Note: `detail_rho` must exactly match the zCDP guarantee of the individual
-  # mechanisms in `detail_synths`. This is because NestedTabularMechanism
-  # implements parallel privacy accounting across the detail_synths manually
-  # utilizing this detail_rho float. Consequently, directly constructing this
-  # mechanism manually (without going through the `configure(...)` API) is
-  # potentially dangerous/incorrect if detail_rho is not perfectly aligned with
-  # detail_synths.
-  detail_rho: float
-  detail_delta: float
-
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the composed DpEvent for the full mechanism."""
     # supports it, instead of falling back to a ZCDpEvent.
-    events = [self.shared_synth.dp_event]
-    if self.detail_rho is not None and self.detail_rho > 0:
-      base = dp_accounting.ZCDpEvent(self.detail_rho)
-      failure = dp_accounting.dp_event.EpsilonDeltaDpEvent(0, self.detail_delta)
-      composed = dp_accounting.ComposedDpEvent([base, failure])
-      events.append(base if self.detail_delta == 0 else composed)
-    return dp_accounting.ComposedDpEvent(events)
+    detail_events = [s.dp_event for s in self.detail_synths.values()]
+    detail_event = _calibration._parallel_compose_event(detail_events)  # pylint: disable=protected-access
+    return dp_accounting.ComposedDpEvent(
+        [self.shared_synth.dp_event, detail_event]
+    )
 
   def __call__(
       self,
@@ -296,8 +285,6 @@ def configure(  # pyrefly: ignore[bad-override]
         type_vocabulary=schema.type_vocabulary,
         shared_synth=shared_synth,
         detail_synths=detail_synths,
-        detail_rho=rho_detail,
-        detail_delta=delta_detail,
     )
 
 

tests/_calibration_test.py

@@ -78,6 +78,59 @@ def test_with_group_size(self):
     with self.assertRaises(ValueError):
       with_group_size(g, 0)
 
+  def test_as_zcdp(self):
+    as_zcdp = _calibration._as_zcdp
+
+    self.assertEqual(as_zcdp(dp_accounting.NoOpDpEvent()), (0.0, 0.0))
+    self.assertEqual(
+        as_zcdp(dp_accounting.GaussianDpEvent(noise_multiplier=0.5)),
+        (2.0, 0.0),
+    )
+    self.assertEqual(
+        as_zcdp(dp_accounting.ExponentialMechanismDpEvent(epsilon=4.0)),
+        (2.0, 0.0),
+    )
+    self.assertEqual(as_zcdp(dp_accounting.ZCDpEvent(rho=1.5)), (1.5, 0.0))
+    self.assertEqual(
+        as_zcdp(dp_accounting.dp_event.EpsilonDeltaDpEvent(2.0, 1e-5)),
+        (2.0, 1e-5),
+    )
+
+    composed = dp_accounting.ComposedDpEvent([
+        dp_accounting.SelfComposedDpEvent(
+            dp_accounting.ExponentialMechanismDpEvent(epsilon=2.0), 3
+        ),
+        dp_accounting.GaussianDpEvent(noise_multiplier=1.0),
+        dp_accounting.dp_event.EpsilonDeltaDpEvent(0.0, 1e-5),
+    ])
+    self.assertEqual(as_zcdp(composed), (2.0, 1e-5))
+
+    with self.assertRaises(dp_accounting.UnsupportedEventError):
+      as_zcdp(dp_accounting.LaplaceDpEvent(noise_multiplier=1.0))
+    with self.assertRaises(dp_accounting.UnsupportedEventError):
+      as_zcdp(dp_accounting.ZCDpEvent(rho=1.0, xi=0.1))
+
+  def test_parallel_compose_event(self):
+    parallel = _calibration._parallel_compose_event
+
+    self.assertEqual(parallel([]), dp_accounting.NoOpDpEvent())
+
+    e1 = dp_accounting.GaussianDpEvent(noise_multiplier=1.0)  # rho = 0.5
+    e2 = dp_accounting.ExponentialMechanismDpEvent(epsilon=4.0)  # rho = 2.0
+    self.assertEqual(parallel([e1, e2]), dp_accounting.ZCDpEvent(rho=2.0))
+
+    e3 = dp_accounting.ComposedDpEvent([
+        dp_accounting.GaussianDpEvent(noise_multiplier=1.0),
+        dp_accounting.dp_event.EpsilonDeltaDpEvent(0.0, 1e-4),
+    ])
+    self.assertEqual(
+        parallel([e2, e3]),
+        dp_accounting.ComposedDpEvent([
+            dp_accounting.ZCDpEvent(rho=2.0),
+            dp_accounting.dp_event.EpsilonDeltaDpEvent(0.0, 1e-4),
+        ]),
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

tests/experimental/nested_test.py

@@ -69,6 +69,38 @@ def test_dp_event_is_composed(self):
     # Detail level should be a ZCDpEvent (conservative parallel composition).
     self.assertLen(event.events, 2)
     self.assertIsInstance(event.events[1], dp_accounting.ZCDpEvent)
+    self.assertAlmostEqual(event.events[1].rho, 5.0)
+
+  def test_dp_event_with_open_set_detail(self):
+    shared_schema = domain.Schema({
+        'platform': domain.CategoricalAttribute(
+            possible_values=['web', 'mobile']
+        ),
+    })
+    per_type_schemas = {
+        'click': domain.Schema({
+            'url': domain.OpenSetCategoricalAttribute(),
+        }),
+        'purchase': domain.Schema({
+            'item': domain.OpenSetCategoricalAttribute(),
+            'coupon': domain.OpenSetCategoricalAttribute(),
+        }),
+    }
+    schema = nested.NestedSchema(
+        shared_schema=shared_schema,
+        per_type_schemas=per_type_schemas,
+    )
+    synth = nested.NestedTabularConfig()
+    calibrated = synth.configure(schema, budget=10.0, delta=1e-4)
+    event = calibrated.dp_event
+    self.assertIsInstance(event, dp_accounting.ComposedDpEvent)
+    detail_event = event.events[1]
+    self.assertIsInstance(detail_event, dp_accounting.ComposedDpEvent)
+    self.assertEqual(detail_event.events[0], dp_accounting.ZCDpEvent(5.0))
+    self.assertIsInstance(
+        detail_event.events[1], dp_accounting.dp_event.EpsilonDeltaDpEvent
+    )
+    self.assertGreater(detail_event.events[1].delta, 0.0)
 
   def test_end_to_end(self):
     schema = self._make_schema()