privatize swift query selection

dpsynth/pipeline_transformations/swift.py

@@ -89,15 +89,22 @@ def compile_workload_fn(dom):
       name='Swift Select Queries',
       weight=parameters.select_budget_frac,
   )
+  measure_spec = budget_accountant.request_budget(
+      pipeline_dp.budget_accounting.MechanismType.GAUSSIAN,
+      name='Swift Measure Marginals',
+      weight=1 - parameters.select_budget_frac,
+  )
+  errors = backend.map(
+      errors,
+      lambda errors_dict: _add_noise_to_errors(errors_dict, mechanism_spec),
+      'Add noise to Swift errors',
+  )
 
   def select_queries_fn(
       errors_dict, candidates_dict, domain_obj
   ) -> tuple[dict[mbi.Clique, float], nx.Graph]:
     """Selects queries using SWIFT algorithm."""
-    # `mechanism_spec` corresponds to the Gaussian mechanism that should be used
-    # to specify the total (epsilon, delta)-budget for the whole pipeline.
-    # Convert it to GDP budget.
-    gdp_budget = 1.0 / mechanism_spec.noise_standard_deviation**2
+    gdp_budget = 1.0 / measure_spec.noise_standard_deviation**2
     return swift.select_queries(
         errors_dict,
         candidates_dict,
@@ -182,6 +189,23 @@ def fit_model_fn(measurements_list, jtree_obj, domain_obj):
   )
 
 
+def _add_noise_to_errors(
+    errors_dict: dict[mbi.Clique, float],
+    mechanism_spec: pipeline_dp.budget_accounting.MechanismSpec,
+) -> dict[mbi.Clique, float]:
+  """Adds DP noise to SWIFT selection errors as one vector query."""
+  cliques = list(errors_dict)
+  errors = np.array([errors_dict[clique] for clique in cliques])
+  sensitivities = pipeline_dp.dp_computations.Sensitivities(
+      l2=np.sqrt(len(cliques))
+  )
+  mechanism = pipeline_dp.dp_computations.create_additive_mechanism(
+      mechanism_spec, sensitivities
+  )
+  noised_errors = mechanism.add_noise(errors)
+  return dict(zip(cliques, noised_errors))
+
+
 def _add_noise_fn(
     clique_marginal: tuple[mbi.Clique, np.ndarray],
     selected_dict: dict[mbi.Clique, float],

tests/pipeline_transformations/swift_test.py

@@ -35,6 +35,24 @@ def from_tuple(self, record, proto_object=None):
 
 class SwiftTest(absltest.TestCase):
 
+  def test_add_noise_to_errors(self):
+    mechanism_spec = pipeline_dp.budget_accounting.MechanismSpec(
+        mechanism_type=pipeline_dp.budget_accounting.MechanismType.GAUSSIAN,
+        name="Swift Select Queries",
+    )
+    mechanism_spec.set_noise_standard_deviation(1.0)
+    errors = {(0, 1): 2.0, (1, 2): 4.0}
+
+    samples = np.array([
+        [swift._add_noise_to_errors(errors, mechanism_spec)[cl] for cl in errors]
+        for _ in range(1000)
+    ])
+    noise = samples - np.array(list(errors.values()))
+
+    self.assertGreater(np.std(noise), 0.0)
+    self.assertAlmostEqual(np.mean(noise), 0.0, delta=0.25)
+    self.assertAlmostEqual(np.std(noise), np.sqrt(len(errors)), delta=0.25)
+
   def test_fit_model(self):
     backend = pipeline_dp.LocalBackend()
     data = [(0, 1), (0, 1), (1, 0), (1, 1)]