Simplify additive mechanism creation (#464)

analysis/tests/parameter_tuning_test.py

@@ -245,7 +245,7 @@ def test_find_candidate_parameters_constant_relative_step_strategy_number_of_can
         self.assertEqual(list(range(1, 51)),
                          candidates.max_contributions_per_partition)
 
-    def test_tune_count_new(self):
+    def test_tune_count(self):
         # Arrange.
         input = [(i % 10, f"pk{i/10}") for i in range(10)]
         public_partitions = [f"pk{i}" for i in range(10)]
@@ -282,7 +282,7 @@ def test_tune_count_new(self):
         self.assertEqual(utility_reports[0].metric_errors[0].metric,
                          pipeline_dp.Metrics.COUNT)
 
-    def test_tune_privacy_id_count_new(self):
+    def test_tune_privacy_id_count(self):
         # Arrange.
         input = [(i % 10, f"pk{i/10}") for i in range(10)]
         public_partitions = [f"pk{i}" for i in range(10)]

pipeline_dp/combiners.py

@@ -216,8 +216,7 @@ class AdditiveMechanismMixin(MechanismContainerMixin):
 
     def create_mechanism(self) -> dp_computations.AdditiveMechanism:
         return dp_computations.create_additive_mechanism(
-            dp_computations.to_additive_mechanism_spec(self.mechanism_spec()),
-            self.sensitivities())
+            self.mechanism_spec(), self.sensitivities())
 
     @abc.abstractmethod
     def sensitivities(self) -> dp_computations.Sensitivities:
@@ -432,13 +431,10 @@ def explain_computation(self) -> ExplainComputationReport:
     def create_mechanism(self) -> dp_computations.MeanMechanism:
         range_middle = dp_computations.compute_middle(self._min_value,
                                                       self._max_value)
-
-        count_spec = dp_computations.to_additive_mechanism_spec(
-            self._count_spec)
-        sum_spec = dp_computations.to_additive_mechanism_spec(self._sum_spec)
-        return dp_computations.create_mean_mechanism(range_middle, count_spec,
+        return dp_computations.create_mean_mechanism(range_middle,
+                                                     self._count_spec,
                                                      self._count_sensitivities,
-                                                     sum_spec,
+                                                     self._sum_spec,
                                                      self._sum_sensitivities)
 
     def mechanism_spec(

pipeline_dp/dp_computations.py

@@ -570,44 +570,32 @@ def check_is_positive(num: Any, name: str) -> bool:
                     raise ValueError(f"L2={self.l2} != sqrt(L0)*Linf={l2}")
 
 
-@dataclass
-class AdditiveMechanismSpec:
-    """Contains the budget and noise_kind."""
-    epsilon: float
-    delta: float
-    noise_kind: pipeline_dp.NoiseKind
-
-
-def to_additive_mechanism_spec(
-        spec: pipeline_dp.budget_accounting.MechanismSpec
-) -> AdditiveMechanismSpec:
-    noise_kind = spec.mechanism_type.to_noise_kind()
-    return AdditiveMechanismSpec(spec.eps, spec.delta, noise_kind)
-
-
 def create_additive_mechanism(
-        spec: AdditiveMechanismSpec,
+        mechanism_spec: pipeline_dp.budget_accounting.MechanismSpec,
         sensitivities: Sensitivities) -> AdditiveMechanism:
     """Creates AdditiveMechanism from a mechanism spec and sensitivities."""
-    if spec.noise_kind == pipeline_dp.NoiseKind.LAPLACE:
+    noise_kind = mechanism_spec.mechanism_type.to_noise_kind()
+    if noise_kind == pipeline_dp.NoiseKind.LAPLACE:
         if sensitivities.l1 is None:
             raise ValueError("L1 or (L0 and Linf) sensitivities must be set for"
                              " Laplace mechanism.")
-        return LaplaceMechanism(spec.epsilon, sensitivities.l1)
+        return LaplaceMechanism(mechanism_spec.eps, sensitivities.l1)
 
-    if spec.noise_kind == pipeline_dp.NoiseKind.GAUSSIAN:
+    if noise_kind == pipeline_dp.NoiseKind.GAUSSIAN:
         if sensitivities.l2 is None:
             raise ValueError("L2 or (L0 and Linf) sensitivities must be set for"
                              " Gaussian mechanism.")
-        return GaussianMechanism(spec.epsilon, spec.delta, sensitivities.l2)
+        return GaussianMechanism(mechanism_spec.eps, mechanism_spec.delta,
+                                 sensitivities.l2)
 
-    assert False, f"{spec.noise_kind} not supported."
+    assert False, f"{noise_kind} not supported."
 
 
 def create_mean_mechanism(
-        range_middle: float, count_spec: AdditiveMechanismSpec,
+        range_middle: float,
+        count_spec: pipeline_dp.budget_accounting.MechanismSpec,
         count_sensitivities: Sensitivities,
-        normalized_sum_spec: AdditiveMechanismSpec,
+        normalized_sum_spec: pipeline_dp.budget_accounting.MechanismSpec,
         normalized_sum_sensitivities: Sensitivities) -> MeanMechanism:
     """Creates MeanMechanism from a mechanism specs and sensitivities."""
     count_mechanism = create_additive_mechanism(count_spec, count_sensitivities)

tests/combiners_test.py

@@ -533,7 +533,7 @@ def test_compute_metrics_with_noise(self):
         self.assertGreater(np.var(noisy_counts), 1)  # check that noise is added
 
         noisy_sums = [noisy_value['sum'] for noisy_value in noisy_values]
-        self.assertAlmostEqual(sum_, np.mean(noisy_sums), delta=10)
+        self.assertAlmostEqual(sum_, np.mean(noisy_sums), delta=20)
         self.assertGreater(np.var(noisy_sums), 1)  # check that noise is added
 
         noisy_means = [noisy_value['mean'] for noisy_value in noisy_values]

tests/dp_computations_test.py

@@ -24,6 +24,8 @@
 import pipeline_dp
 import pipeline_dp.dp_computations as dp_computations
 from pipeline_dp.aggregate_params import NoiseKind
+from pipeline_dp import aggregate_params
+from pipeline_dp import budget_accounting
 
 N_ITERATIONS = 200000
 DUMMY_MIN_VALUE = 2.0
@@ -607,8 +609,9 @@ def test_sensitivities_post_init_l1_l2_computation(self):
     def test_create_laplace_mechanism(self, epsilon, l0_sensitivity,
                                       linf_sensitivity, l1_sensitivity,
                                       expected_noise_parameter):
-        spec = dp_computations.AdditiveMechanismSpec(
-            epsilon, delta=0, noise_kind=pipeline_dp.NoiseKind.LAPLACE)
+        spec = budget_accounting.MechanismSpec(
+            aggregate_params.MechanismType.LAPLACE)
+        spec.set_eps_delta(epsilon, delta=0)
         sensitivities = dp_computations.Sensitivities(l0=l0_sensitivity,
                                                       linf=linf_sensitivity,
                                                       l1=l1_sensitivity)
@@ -632,8 +635,9 @@ def test_create_laplace_mechanism(self, epsilon, l0_sensitivity,
     )
     def test_create_gaussian_mechanism(self, epsilon, delta, l2_sensitivity,
                                        expected_noise_parameter):
-        spec = dp_computations.AdditiveMechanismSpec(
-            epsilon, delta=delta, noise_kind=pipeline_dp.NoiseKind.GAUSSIAN)
+        spec = budget_accounting.MechanismSpec(
+            aggregate_params.MechanismType.GAUSSIAN)
+        spec.set_eps_delta(epsilon, delta)
         sensitivities = dp_computations.Sensitivities(l2=l2_sensitivity)
 
         mechanism = dp_computations.create_additive_mechanism(
@@ -715,11 +719,13 @@ def test_compute_sensitivities_for_sum_max_contributions(self):
 class MeanMechanismTests(parameterized.TestCase):
 
     def create_mean_mechanism(self) -> dp_computations.MeanMechanism:
-        count_spec = dp_computations.AdditiveMechanismSpec(
-            1.0, delta=1e-5, noise_kind=pipeline_dp.NoiseKind.GAUSSIAN)
+        count_spec = budget_accounting.MechanismSpec(
+            aggregate_params.MechanismType.GAUSSIAN)
+        count_spec.set_eps_delta(eps=1.0, delta=1e-5)
         count_sensitivities = dp_computations.Sensitivities(l0=4, linf=10)
-        sum_spec = dp_computations.AdditiveMechanismSpec(
-            2.0, delta=1e-7, noise_kind=pipeline_dp.NoiseKind.LAPLACE)
+        sum_spec = budget_accounting.MechanismSpec(
+            aggregate_params.MechanismType.LAPLACE)
+        sum_spec.set_eps_delta(eps=2.0, delta=1e-7)
         sum_sensitivities = dp_computations.Sensitivities(l0=3, linf=5)
         return dp_computations.create_mean_mechanism(5, count_spec,
                                                      count_sensitivities,