Integrated the random allocation directory

python/dp_accounting/dp_accounting/__init__.py

@@ -32,6 +32,7 @@
 from dp_accounting.dp_event import SampledWithReplacementDpEvent
 from dp_accounting.dp_event import SelfComposedDpEvent
 from dp_accounting.dp_event import SingleEpochTreeAggregationDpEvent
+from dp_accounting.dp_event import RandomAllocationDpEvent
 from dp_accounting.dp_event import TruncatedSubsampledGaussianDpEvent
 from dp_accounting.dp_event import UnsupportedDpEvent
 from dp_accounting.dp_event import ZCDpEvent

python/dp_accounting/dp_accounting/dp_event.py

@@ -494,3 +494,23 @@ class TruncatedSubsampledGaussianDpEvent(DpEvent):
   sampling_probability: float
   truncated_batch_size: int
   noise_multiplier: float
+
+
+@attr.s(frozen=True, slots=True, auto_attribs=True)
+class RandomAllocationDpEvent(DpEvent):
+  """Represents the random-allocation mechanism.
+
+  In this event, each element in the dataset is independently allocated to k
+  uniformly sampled calls to the mechanism out of t total calls.
+
+  Currently only GaussianDpEvent is supported as the inner event.
+
+  Attributes:
+    event: The DpEvent for each call to the mechanism (e.g., GaussianDpEvent).
+    k: Number of calls each element participates in.
+    t: Total number of calls to the mechanism.
+  """
+
+  event: DpEvent
+  k: int
+  t: int

python/dp_accounting/dp_accounting/dp_event_test.py

@@ -123,6 +123,14 @@ class DpEventTest(parameterized.TestCase):
               1.0,
           ),
       ),
+      (
+          'random_allocation',
+          dp_event.RandomAllocationDpEvent(
+              event=dp_event.GaussianDpEvent(2.0),
+              k=10,
+              t=100,
+          ),
+      ),
   )
   def test_to_from_named_tuple(self, event):
     named_tuple = event.to_named_tuple()
@@ -134,6 +142,12 @@ def test_to_from_named_tuple(self, event):
     assert_not_contains_named_tuples(reconstructed)
     self.assertEqual(event, reconstructed)
 
+  def test_random_allocation_event_has_no_discretization_override(self):
+    event = dp_event.RandomAllocationDpEvent(
+        event=dp_event.GaussianDpEvent(2.0), k=10, t=100
+    )
+    self.assertFalse(hasattr(event, 'loss_discretization'))
+
 
 if __name__ == '__main__':
   absltest.main()

python/dp_accounting/dp_accounting/pld/pld_privacy_accountant.py

@@ -20,6 +20,7 @@
 from dp_accounting import privacy_accountant
 from dp_accounting.pld import common
 from dp_accounting.pld import privacy_loss_distribution
+from dp_accounting.pld import random_allocation as _random_allocation
 
 NeighborRel = privacy_accountant.NeighboringRelation
 CompositionErrorDetails = (
@@ -263,6 +264,31 @@ def _maybe_compose(self, event: dp_event.DpEvent, count: int,
           )
           self._pld = self._pld.compose(truncated_subsampled_gaussian_pld)
       return None
+    elif isinstance(event, dp_event.RandomAllocationDpEvent):
+      ra_event: dp_event.RandomAllocationDpEvent = event
+      if not isinstance(ra_event.event, dp_event.GaussianDpEvent):
+        return CompositionErrorDetails(
+            invalid_event=event,
+            error_message=(
+                'Subevent of `RandomAllocationDpEvent` must be '
+                f'`GaussianDpEvent`. Found {ra_event.event}.'
+            ),
+        )
+      if do_compose:
+        if ra_event.event.noise_multiplier == 0:
+          self._contains_non_dp_event = True
+        else:
+          params = _random_allocation.PrivacyParams(
+              sigma=ra_event.event.noise_multiplier,
+              num_steps=ra_event.t,
+              num_selected=ra_event.k,
+          )
+          config = _random_allocation.AllocationSchemeConfig(
+              loss_discretization=self._value_discretization_interval,
+          )
+          alloc_pld = _random_allocation.gaussian_allocation_pld(params, config)
+          self._pld = self._pld.compose(alloc_pld.self_compose(count))
+      return None
     else:
       # Unsupported event (including `UnsupportedDpEvent`).
       return CompositionErrorDetails(

python/dp_accounting/dp_accounting/pld/privacy_loss_distribution_basic_example.py

@@ -16,7 +16,10 @@
 
 from absl import app
 
+from dp_accounting import GaussianDpEvent
+from dp_accounting import RandomAllocationDpEvent
 from dp_accounting.pld import privacy_loss_distribution
+from dp_accounting.pld.pld_privacy_accountant import PLDAccountant
 
 
 def main(argv):
@@ -67,6 +70,26 @@ def main(argv):
       f'{standard_deviation} is ({epsilon}, {delta})-DP.'
   )
 
+  # The PLDAccountant also supports the random allocation mechanism. Each
+  # element is independently assigned to k out of t calls to the mechanism.
+  noise_multiplier = 2.0
+  k = 10
+  t = 100
+  accountant = PLDAccountant()
+  accountant.compose(
+      RandomAllocationDpEvent(
+          event=GaussianDpEvent(noise_multiplier=noise_multiplier),
+          k=k,
+          t=t,
+      )
+  )
+  target_delta = 1e-6
+  epsilon = accountant.get_epsilon(target_delta)
+  print(
+      f'Random allocation with Gaussian noise (noise_multiplier={noise_multiplier},'
+      f' k={k}, t={t}) satisfies ({epsilon:.4f}, {target_delta})-DP.'
+  )
+
 
 if __name__ == '__main__':
   app.run(main)

python/dp_accounting/dp_accounting/pld/random_allocation/__init__.py

@@ -0,0 +1,29 @@
+"""Public entry points for random-allocation privacy accounting."""
+
+from .random_allocation_distributions import PLDRealization
+from .random_allocation_api import (
+    gaussian_allocation_pld,
+    general_allocation_pld,
+)
+from .random_allocation_types import (
+    DEFAULT_LOSS_DISCRETIZATION,
+    DEFAULT_TAIL_TRUNCATION,
+    AllocationSchemeConfig,
+    BoundType,
+    Direction,
+    PrivacyParams,
+    SpacingType,
+)
+
+__all__ = [
+    "PLDRealization",
+    "AllocationSchemeConfig",
+    "BoundType",
+    "DEFAULT_LOSS_DISCRETIZATION",
+    "DEFAULT_TAIL_TRUNCATION",
+    "Direction",
+    "PrivacyParams",
+    "SpacingType",
+    "gaussian_allocation_pld",
+    "general_allocation_pld",
+]

python/dp_accounting/dp_accounting/pld/random_allocation/random_allocation_api.py

@@ -0,0 +1,137 @@
+"""Public API surface for random-allocation PLD construction."""
+
+from __future__ import annotations
+
+from functools import partial
+
+from dp_accounting.pld import privacy_loss_distribution
+
+from .random_allocation_distributions import PLDRealization
+from .random_allocation_core import (
+    allocation_full_pld,
+    gaussian_allocation_pld_core,
+    geometric_allocation_pld_base_add,
+    geometric_allocation_pld_base_remove,
+    realization_add_base_distribution,
+    realization_remove_base_distributions,
+)
+from .random_allocation_types import (
+    AllocationSchemeConfig,
+    BoundType,
+    Direction,
+    PrivacyParams,
+)
+from .random_allocation_utils import (
+    validate_allocation_params,
+    validate_allocation_scheme_config,
+    validate_bound_type,
+    validate_privacy_params,
+)
+
+
+def gaussian_allocation_pld(
+    params: PrivacyParams,
+    config: AllocationSchemeConfig,
+    bound_type: BoundType = BoundType.DOMINATES,
+) -> privacy_loss_distribution.PrivacyLossDistribution:
+    """Compute upper / lower PLD for random-allocation with the Gaussian mechanism.
+
+    Args:
+        params: Privacy parameters describing noise scale, number of steps,
+            and optional delta/epsilon query target.
+        config: Discretization and convolution configuration.
+        bound_type: Whether to compute a dominating or dominated discretized bound.
+
+    Returns:
+        A ``dp_accounting`` ``PrivacyLossDistribution`` for both privacy directions.
+
+    """
+    # Input validation
+    validate_privacy_params(params)
+    validate_allocation_scheme_config(config)
+    validate_bound_type(bound_type)
+
+    compute_base_pld_remove = partial(
+        gaussian_allocation_pld_core,
+        direction=Direction.REMOVE,
+        sigma=params.sigma,
+    )
+    compute_base_pld_add = partial(
+        gaussian_allocation_pld_core,
+        direction=Direction.ADD,
+        sigma=params.sigma,
+    )
+    return allocation_full_pld(
+        compute_base_pld_remove=compute_base_pld_remove,
+        compute_base_pld_add=compute_base_pld_add,
+        num_steps=params.num_steps,
+        num_selected=params.num_selected,
+        num_epochs=params.num_epochs,
+        loss_discretization=config.loss_discretization,
+        tail_truncation=config.tail_truncation,
+        bound_type=bound_type,
+    )
+
+
+def general_allocation_pld(
+    num_steps: int,
+    num_selected: int,
+    num_epochs: int,
+    remove_realization: PLDRealization,
+    add_realization: PLDRealization,
+    config: AllocationSchemeConfig,
+    bound_type: BoundType = BoundType.DOMINATES,
+) -> privacy_loss_distribution.PrivacyLossDistribution:
+    """Build a random-allocation PLD from explicit PLD realizations.
+
+    Args:
+        num_steps: Total number of random-allocation steps.
+        num_selected: Number of selections per epoch.
+        num_epochs: Number of epochs.
+        remove_realization: Explicit remove-direction PLD realization.
+        add_realization: Explicit add-direction PLD realization.
+        config: Discretization and convolution configuration.
+        bound_type: Whether to compute a dominating or dominated discretized bound.
+
+    Returns:
+        A ``dp_accounting`` ``PrivacyLossDistribution`` for the composed realization.
+
+    Notes:
+        The delivery package always uses the geometric convolution path.
+
+    """
+    # Input validation
+    validate_allocation_params(num_steps, num_selected, num_epochs)
+    if not isinstance(remove_realization, PLDRealization):
+        raise TypeError(
+            f"remove_realization must be PLDRealization, got {type(remove_realization)}"
+        )
+    if not isinstance(add_realization, PLDRealization):
+        raise TypeError(f"add_realization must be PLDRealization, got {type(add_realization)}")
+    validate_allocation_scheme_config(config)
+    validate_bound_type(bound_type)
+
+    compute_base_pld_remove = partial(
+        geometric_allocation_pld_base_remove,
+        base_distributions_creation=partial(
+            realization_remove_base_distributions,
+            realization=remove_realization,
+        ),
+    )
+    compute_base_pld_add = partial(
+        geometric_allocation_pld_base_add,
+        base_distributions_creation=partial(
+            realization_add_base_distribution,
+            realization=add_realization,
+        ),
+    )
+    return allocation_full_pld(
+        compute_base_pld_remove=compute_base_pld_remove,
+        compute_base_pld_add=compute_base_pld_add,
+        num_steps=num_steps,
+        num_selected=num_selected,
+        num_epochs=num_epochs,
+        loss_discretization=config.loss_discretization,
+        tail_truncation=config.tail_truncation,
+        bound_type=bound_type,
+    )

python/dp_accounting/dp_accounting/pld/random_allocation/random_allocation_api_test.py

@@ -0,0 +1,68 @@
+"""Delivery API tests: only PLD builders are public."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from dp_accounting.pld import privacy_loss_distribution
+
+from .random_allocation_api import gaussian_allocation_pld, general_allocation_pld
+from .random_allocation_distributions import PLDRealization
+from .random_allocation_types import AllocationSchemeConfig, BoundType, PrivacyParams
+
+
+def _simple_realization() -> PLDRealization:
+    return PLDRealization(
+        x_min=0.0,
+        step=1.0,
+        prob_arr=np.array([0.7, 0.3], dtype=np.float64),
+        p_max=0.0,
+    )
+
+
+def test_gaussian_allocation_pld_returns_dp_accounting_pld():
+    params = PrivacyParams(sigma=2.0, num_steps=5, num_selected=1, num_epochs=1)
+    config = AllocationSchemeConfig(loss_discretization=0.05, tail_truncation=1e-6)
+
+    pld = gaussian_allocation_pld(params=params, config=config, bound_type=BoundType.DOMINATES)
+
+    assert isinstance(pld, privacy_loss_distribution.PrivacyLossDistribution)
+
+
+def test_general_allocation_pld_returns_dp_accounting_pld():
+    config = AllocationSchemeConfig(loss_discretization=0.05, tail_truncation=1e-6)
+
+    pld = general_allocation_pld(
+        num_steps=5,
+        num_selected=1,
+        num_epochs=1,
+        remove_realization=_simple_realization(),
+        add_realization=_simple_realization(),
+        config=config,
+        bound_type=BoundType.DOMINATES,
+    )
+
+    assert isinstance(pld, privacy_loss_distribution.PrivacyLossDistribution)
+
+
+def test_non_pld_exports_are_not_exposed():
+    import importlib
+
+    package = importlib.import_module(__package__)
+
+    assert not hasattr(package, "gaussian_allocation_epsilon_range")
+    assert not hasattr(package, "gaussian_distribution")
+    assert not hasattr(package, "subsample_pld")
+
+
+def test_general_allocation_rejects_non_realization_inputs():
+    config = AllocationSchemeConfig()
+    with pytest.raises(TypeError, match="remove_realization must be PLDRealization"):
+        general_allocation_pld(
+            num_steps=2,
+            num_selected=1,
+            num_epochs=1,
+            remove_realization=object(),
+            add_realization=_simple_realization(),
+            config=config,
+        )