Update _compute_budgets_for_aggregation (OpenMined#567)

Author: miracvbasaran

pipeline_dp/aggregate_params.py

@@ -144,6 +144,14 @@ def is_partition_selection(self) -> bool:
             MechanismType.GAUSSIAN_THRESHOLDING
         ]
 
+    @property
+    def uses_delta(self) -> bool:
+        return self in [
+            MechanismType.GAUSSIAN, MechanismType.TRUNCATED_GEOMETRIC,
+            MechanismType.LAPLACE_THRESHOLDING,
+            MechanismType.GAUSSIAN_THRESHOLDING
+        ]
+
 
 def noise_to_thresholding(noise_kind: NoiseKind) -> MechanismType:
     if noise_kind == NoiseKind.LAPLACE:

pipeline_dp/budget_accounting.py

@@ -130,7 +130,8 @@ class MechanismSpecInternal:
     mechanism_spec: MechanismSpec
 
 
-Budget = collections.namedtuple("Budget", ["epsilon", "delta"])
+Budget = collections.namedtuple("Budget",
+                                ["epsilon", "delta", "mechanism_type"])
 
 
 class BudgetAccountant(abc.ABC):
@@ -197,10 +198,12 @@ def scope(self, weight: float):
         """
         return BudgetAccountantScope(self, weight)
 
-    def _compute_budget_for_aggregation(self, weight: float) -> Budget:
+    def _compute_budgets_for_aggregation(self, weight: float) -> Budget:
         """Computes budget per aggregation.
 
-        It splits the budget using the naive composition.
+        It splits the budget using the naive composition. If neither
+        'num_aggregations' nor 'aggregation_weights' is set in the constructor,
+        it returns an empty list.
 
         Warning: This function changes the 'self' internal state. It can be
         called only from the API function of DPEngine, like aggregate() or
@@ -210,18 +213,48 @@ def _compute_budget_for_aggregation(self, weight: float) -> Budget:
             weight: the budget weight of the aggregation.
 
         Returns:
-            the budget.
+            budgets for each mechanism used in the aggregation.
         """
+        budgets = []
         self._actual_aggregation_weights.append(weight)
         if self._expected_num_aggregations:
-            return Budget(self._total_epsilon / self._expected_num_aggregations,
-                          self._total_delta / self._expected_num_aggregations)
-        if self._expected_aggregation_weights:
+            epsilon = self._total_epsilon / self._expected_num_aggregations
+            delta = self._total_delta / self._expected_num_aggregations
+        elif self._expected_aggregation_weights:
             budget_ratio = weight / sum(self._expected_aggregation_weights)
-            return Budget(self._total_epsilon * budget_ratio,
-                          self._total_delta * budget_ratio)
-        # No expectations on aggregations, no way to compute budget.
-        return None
+            epsilon = self._total_epsilon * budget_ratio
+            delta = self._total_delta * budget_ratio
+        else:
+            # No expectations on aggregations, no way to compute budget.
+            return budgets
+
+        if not self._scopes_stack:
+            raise ValueError(
+                "No scope is active. Did you forget to use 'with'?")
+        current_scope = self._scopes_stack[-1]
+        current_scope._normalise_mechanism_weights()
+        eps_weights = []
+        del_weights = []
+        for mechanism in current_scope.mechanisms:
+            if mechanism.mechanism_spec.mechanism_type.uses_delta:
+                eps_weights.append(mechanism.weight)
+                del_weights.append(mechanism.weight)
+            else:
+                eps_weights.append(mechanism.weight)
+                del_weights.append(0)
+        sum_eps_weights = sum(eps_weights)
+        sum_del_weights = sum(del_weights)
+        if sum_del_weights == 0:
+            sum_del_weights = 1
+
+        for i, mechanism in enumerate(current_scope.mechanisms):
+            budgets.append(
+                Budget(
+                    epsilon=epsilon * eps_weights[i] / sum_eps_weights,
+                    delta=delta * del_weights[i] / sum_del_weights,
+                    mechanism_type=mechanism.mechanism_spec.mechanism_type,
+                ))
+        return budgets
 
     def _check_aggregation_restrictions(self):
         if self._expected_num_aggregations:

pipeline_dp/dp_engine.py

@@ -106,9 +106,9 @@ def aggregate(self,
                     self._current_report_generator)
             col = self._aggregate(col, params, data_extractors,
                                   public_partitions)
-            budget = self._budget_accountant._compute_budget_for_aggregation(
+            budgets = self._budget_accountant._compute_budgets_for_aggregation(
                 params.budget_weight)
-            return self._annotate(col, params=params, budget=budget)
+            return self._annotate(col, params=params, budget=budgets)
 
     def _aggregate(self, col, params: pipeline_dp.AggregateParams,
                    data_extractors: pipeline_dp.DataExtractors,
@@ -233,9 +233,9 @@ def select_partitions(self, col, params: pipeline_dp.SelectPartitionsParams,
         with self._budget_accountant.scope(weight=params.budget_weight):
             self._add_report_generator(params, "select_partitions")
             col = self._select_partitions(col, params, data_extractors)
-            budget = self._budget_accountant._compute_budget_for_aggregation(
+            budgets = self._budget_accountant._compute_budgets_for_aggregation(
                 params.budget_weight)
-            return self._annotate(col, params=params, budget=budget)
+            return self._annotate(col, params=params, budget=budgets)
 
     def _select_partitions(self, col,
                            params: pipeline_dp.SelectPartitionsParams,
@@ -597,6 +597,20 @@ def add_dp_noise(self,
             collection of (partition_key, value + noise).
             Output partition keys are the same as in the input collection.
         """
+        with self._budget_accountant.scope(weight=params.budget_weight):
+            self._add_report_generator(params,
+                                       "add_dp_noise",
+                                       is_public_partition=True)
+            if out_explain_computation_report is not None:
+                out_explain_computation_report._set_report_generator(
+                    self._current_report_generator)
+            anonymized_col = self._add_dp_noise(col, params)
+            budgets = self._budget_accountant._compute_budgets_for_aggregation(
+                params.budget_weight)
+            return self._annotate(anonymized_col, params=params, budget=budgets)
+
+    def _add_dp_noise(self, col,
+                      params: pipeline_dp.aggregate_params.AddDPNoiseParams):
         # Request budget and create Sensitivities object
         mechanism_type = params.noise_kind.convert_to_mechanism_type()
         mechanism_spec = self._budget_accountant.request_budget(mechanism_type)
@@ -606,15 +620,6 @@ def add_dp_noise(self,
             l1=params.l1_sensitivity,
             l2=params.l2_sensitivity)
 
-        # Initialize ReportGenerator.
-        self._report_generators.append(
-            report_generator.ReportGenerator(params,
-                                             "add_dp_noise",
-                                             is_public_partition=True))
-        if out_explain_computation_report is not None:
-            out_explain_computation_report._set_report_generator(
-                self._current_report_generator)
-
         # Add noise to values.
         def create_mechanism() -> dp_computations.AdditiveMechanism:
             return dp_computations.create_additive_mechanism(
@@ -635,11 +640,7 @@ def add_noise(value: Union[int, float]) -> DpNoiseAdditionResult:
             def add_noise(value: Union[int, float]) -> float:
                 return create_mechanism().add_noise(value)
 
-        anonymized_col = self._backend.map_values(col, add_noise, "Add noise")
-
-        budget = self._budget_accountant._compute_budget_for_aggregation(
-            params.budget_weight)
-        return self._annotate(anonymized_col, params=params, budget=budget)
+        return self._backend.map_values(col, add_noise, "Add noise")
 
     def _annotate(self, col, params: Union[pipeline_dp.AggregateParams,
                                            pipeline_dp.SelectPartitionsParams,

tests/budget_accounting_test.py

@@ -153,11 +153,14 @@ def test_num_aggregations(self, num_aggregations):
             total_delta=total_delta,
             num_aggregations=num_aggregations)
         for _ in range(num_aggregations):
-            budget = budget_accountant._compute_budget_for_aggregation(1)
-            expected_epsilon = total_epsilon / num_aggregations
-            expected_delta = total_delta / num_aggregations
-            self.assertAlmostEqual(expected_epsilon, budget.epsilon)
-            self.assertAlmostEqual(expected_delta, budget.delta)
+            with budget_accountant.scope(weight=1):
+                budget_accountant.request_budget(
+                    mechanism_type=MechanismType.GAUSSIAN)
+                budgets = budget_accountant._compute_budgets_for_aggregation(1)
+                expected_epsilon = total_epsilon / num_aggregations
+                expected_delta = total_delta / num_aggregations
+                self.assertAlmostEqual(expected_epsilon, budgets[0].epsilon)
+                self.assertAlmostEqual(expected_delta, budgets[0].delta)
 
         budget_accountant.compute_budgets()
 
@@ -169,11 +172,15 @@ def test_aggregation_weights(self):
                                                   total_delta=total_delta,
                                                   aggregation_weights=weights)
         for weight in weights:
-            budget = budget_accountant._compute_budget_for_aggregation(weight)
-            expected_epsilon = total_epsilon * weight / sum(weights)
-            expected_delta = total_delta * weight / sum(weights)
-            self.assertAlmostEqual(expected_epsilon, budget.epsilon)
-            self.assertAlmostEqual(expected_delta, budget.delta)
+            with budget_accountant.scope(weight=weight):
+                budget_accountant.request_budget(
+                    mechanism_type=MechanismType.GAUSSIAN)
+                budgets = budget_accountant._compute_budgets_for_aggregation(
+                    weight)
+                expected_epsilon = total_epsilon * weight / sum(weights)
+                expected_delta = total_delta * weight / sum(weights)
+                self.assertAlmostEqual(expected_epsilon, budgets[0].epsilon)
+                self.assertAlmostEqual(expected_delta, budgets[0].delta)
 
         budget_accountant.compute_budgets()
 
@@ -190,7 +197,8 @@ def test_not_enough_aggregations(self, use_num_aggregations):
             num_aggregations=num_aggregations,
             aggregation_weights=weights)
 
-        budget_accountant._compute_budget_for_aggregation(1)
+        with budget_accountant.scope(weight=1):
+            budget_accountant._compute_budgets_for_aggregation(1)
         with self.assertRaises(ValueError):
             # num_aggregations = 2, but only 1 aggregation_scope was created
             budget_accountant.compute_budgets()

tests/combiners_test.py

@@ -213,7 +213,9 @@ def test_create_compound_combiner_with_post_aggregation(self):
         # Act.
         compound_combiner = dp_combiners.create_compound_combiner(
             params, budget_accountant)
-        budget_accountant._compute_budget_for_aggregation(params.budget_weight)
+        with budget_accountant.scope(weight=params.budget_weight):
+            budget_accountant._compute_budgets_for_aggregation(
+                params.budget_weight)
         budget_accountant.compute_budgets()
 
         # Assert

tests/dp_engine_test.py

@@ -1304,7 +1304,12 @@ def test_annotate_call(self, mock_annotate_fn):
                                                   total_delta,
                                                   num_aggregations=3)
         dp_engine = self._create_dp_engine_default(budget_accountant)
-        aggregate_params, public_partitions = self._create_params_default()
+        aggregate_params = pipeline_dp.AggregateParams(
+            noise_kind=pipeline_dp.NoiseKind.GAUSSIAN,
+            metrics=[agg.Metrics.COUNT],
+            max_partitions_contributed=1,
+            max_contributions_per_partition=1)
+        public_partitions = ["pk0", "pk10", "pk11"]
         select_partition_params = pipeline_dp.SelectPartitionsParams(2)
         extractors = self._get_default_extractors()
         input = [1, 2, 3]
@@ -1320,9 +1325,9 @@ def test_annotate_call(self, mock_annotate_fn):
         # Assert
         self.assertEqual(3, mock_annotate_fn.call_count)
         for i_call in range(3):
-            budget = mock_annotate_fn.call_args_list[i_call][1]['budget']
-            self.assertEqual(total_epsilon / 3, budget.epsilon)
-            self.assertEqual(total_delta / 3, budget.delta)
+            budgets = mock_annotate_fn.call_args_list[i_call][1]['budget']
+            self.assertEqual(total_epsilon / 3, budgets[0].epsilon)
+            self.assertEqual(total_delta / 3, budgets[0].delta)
 
     def test_min_max_sum_per_partition(self):
         dp_engine, budget_accountant = self._create_dp_engine_default(