Add group_size parameter to calibrate and export dpsynth.with_group_size.

dpsynth/__init__.py

@@ -26,6 +26,7 @@
 from dpsynth import reporting
 from dpsynth import serialize
 from dpsynth._calibration import calibrate
+from dpsynth._calibration import with_group_size
 from dpsynth._checkpoint import checkpoint
 from dpsynth.data_generation_v3 import TabularConfig
 from dpsynth.data_generation_v3 import TabularMechanism
@@ -80,4 +81,5 @@
     'reporting',
     'serialize',
     'to_yaml',
+    'with_group_size',
 ]

dpsynth/_calibration.py

@@ -74,6 +74,7 @@ def calibrate(
     delta: float,
     delta_split: float = 0.5,
     poisson_sampling_prob: float = 1.0,
+    group_size: int = 1,
     accountant_fn: Callable[[], dp_accounting.PrivacyAccountant] | None = None,
 ) -> Any:
   """Calibrates a mechanism config to a target (epsilon, delta)-DP guarantee.
@@ -93,28 +94,42 @@ def calibrate(
       partition selection). Defaults to 0.5.
     poisson_sampling_prob: If specified, calibrate the mechanism assuming the
       input data is subsampled with the given probability. The actual sampling
-      is **NOT** handled internally by the calibrated mechanism.
+      is **NOT** handled internally by the calibrated mechanism. When combined
+      with ``group_size > 1``, group-size scaling is applied before wrapping in
+      ``PoissonSampledDpEvent``, which assumes group-level (not record-level)
+      Poisson subsampling.
+    group_size: Assumed upper bound on the number of records a single group or
+      individual contributes. Lifts the candidate mechanism's ``dp_event`` via
+      ``with_group_size`` before applying ``poisson_sampling_prob``. Mechanisms
+      that already contain internal ``PoissonSampledDpEvent``s (such as DP-SGD
+      training) cannot be scaled with ``group_size > 1`` and will raise
+      ``UnsupportedEventError``. Soundness relies on the caller enforcing this
+      bound.
     accountant_fn: Optional zero-argument callable returning a fresh
       ``PrivacyAccountant``. If specified, calibrate using this accountant.
 
   Returns:
     A calibrated, runnable mechanism.
 
   Raises:
-    ValueError: If epsilon is not positive or delta_split is not in (0, 1).
+    ValueError: If epsilon is not positive, delta_split is not in (0, 1), or
+      group_size < 1.
     UnsupportedEventError: If no accountant supports the mechanism.
   """
   if epsilon <= 0:
     raise ValueError(f'Target epsilon must be positive, got {epsilon}.')
   if not 0 < delta_split < 1:
     raise ValueError(f'delta_split must be in (0, 1), got {delta_split}.')
+  if group_size < 1:
+    raise ValueError(f'group_size < 1: {group_size}.')
 
   def make_event_fn(rho: float) -> dp_accounting.DpEvent:
     base = config.configure(
         domain,
         budget=rho,
         delta=delta * delta_split,
     ).dp_event
+    base = with_group_size(base, group_size)
     sampled = dp_accounting.PoissonSampledDpEvent(poisson_sampling_prob, base)
     return base if poisson_sampling_prob == 1.0 else sampled
 

dpsynth/api.py

@@ -49,14 +49,27 @@ class CalibratedMechanism(abc.ABC):
 
   Subclasses must implement:
 
-  - ``dp_event``: return the exact ``DpEvent`` characterizing the mechanism.
+  - ``dp_event``: return the exact record-level ``DpEvent`` characterizing the
+    mechanism.
   - ``__call__``: run the mechanism on data.
   """
 
   @property
   @abc.abstractmethod
   def dp_event(self) -> dp_accounting.DpEvent:
-    """The DpEvent characterizing the privacy cost of this mechanism."""
+    """The record-level DpEvent characterizing the privacy cost of this mechanism.
+
+    This property always reports the base record-level (``group_size=1``),
+    un-subsampled (``poisson_sampling_prob=1.0``) privacy guarantee of the
+    mechanism itself, even if the mechanism was produced via
+    ``dpsynth.calibrate(..., group_size=k, poisson_sampling_prob=q)``.
+
+    To compute group-level or subsampled privacy guarantees from this event
+    (e.g. when composing mechanisms manually in an external accountant), apply
+    ``dpsynth.with_group_size(mech.dp_event, group_size)`` and/or wrap the
+    result in ``dp_accounting.PoissonSampledDpEvent(poisson_sampling_prob,
+    ...)``.
+    """
 
   @abc.abstractmethod
   def __call__(self, *args: Any, **kwargs: Any) -> Any:
@@ -167,6 +180,7 @@ def calibrate(
       delta: float,
       delta_split: float = 0.5,
       poisson_sampling_prob: float = 1.0,
+      group_size: int = 1,
       accountant_fn: (
           Callable[[], dp_accounting.PrivacyAccountant] | None
       ) = None,
@@ -185,6 +199,7 @@ def calibrate(
         delta=delta,
         delta_split=delta_split,
         poisson_sampling_prob=poisson_sampling_prob,
+        group_size=group_size,
         accountant_fn=accountant_fn,
     )
 

tests/data_generation_v3_test.py

@@ -550,5 +550,72 @@ def test_mbi_callbacks_logging_configured(self):
     self.assertTrue(any('test | message' in output for output in logs.output))
 
 
+class GroupSizeTest(parameterized.TestCase):
+  """Tests the ``group_size`` parameter of ``dpsynth.calibrate`` end to end."""
+
+  def _categorical_domains(self):
+    return {
+        'A': domain.CategoricalAttribute(['a', 'b', 'c']),
+        'B': domain.CategoricalAttribute(['x', 'y', 'z']),
+    }
+
+  def test_calibrate_scales_noise_with_group_size(self):
+    k = 5
+    config = TabularConfig()
+    domains = self._categorical_domains()
+    cal_base = dpsynth.calibrate(
+        config, domains, epsilon=1.0, delta=1e-5, group_size=1
+    )
+    cal_scaled = dpsynth.calibrate(
+        config, domains, epsilon=1.0, delta=1e-5, group_size=k
+    )
+    self.assertAlmostEqual(
+        cal_scaled.initializers['A'].sigma,
+        k * cal_base.initializers['A'].sigma,
+        places=4,
+    )
+    lifted = dpsynth.with_group_size(cal_scaled.initializers['A'].dp_event, k)
+    self.assertAlmostEqual(
+        lifted.noise_multiplier,
+        cal_base.initializers['A'].dp_event.noise_multiplier,
+        places=4,
+    )
+
+  def test_end_to_end_with_group_size(self):
+    df = pd.DataFrame({'A': ['a', 'b', 'c'], 'B': [1.0, 5.0, 10.0]})
+    domains = {
+        'A': domain.CategoricalAttribute(['a', 'b', 'c']),
+        'B': domain.NumericalAttribute(0.0, 10.0),
+    }
+    config = TabularConfig()
+    calibrated = dpsynth.calibrate(
+        config,
+        domains,
+        epsilon=10.0,
+        delta=1e-5,
+        group_size=3,
+    )
+    synthetic_df = calibrated(np.random.default_rng(0), df).synthetic_data
+    self.assertListEqual(synthetic_df.columns.tolist(), ['A', 'B'])
+
+  def test_open_set_with_group_size_greater_than_one_raises(self):
+    domains = {'A': domain.OpenSetCategoricalAttribute()}
+    config = TabularConfig()
+    with self.assertRaises(dp_accounting.UnsupportedEventError):
+      dpsynth.calibrate(config, domains, epsilon=1.0, delta=1e-5, group_size=3)
+
+  @parameterized.named_parameters(('zero', 0), ('negative', -3))
+  def test_invalid_group_size_raises(self, k):
+    config = TabularConfig()
+    with self.assertRaises(ValueError):
+      _ = dpsynth.calibrate(
+          config,
+          self._categorical_domains(),
+          epsilon=1.0,
+          delta=1e-5,
+          group_size=k,
+      )
+
+
 if __name__ == '__main__':
   absltest.main()