tests

OpenMined · Sep 11, 2024 · cf0f5c2 · cf0f5c2
1 parent 1675b5f
commit cf0f5c2
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/analysis/data_structures.py b/analysis/data_structures.py
@@ -120,7 +120,7 @@ def get_aggregate_params(
             if isinstance(min_sum, Sequence):
                 min_max_sum = list(zip(min_sum, max_sum))
             else:
-                min_max_sum = [min_sum, max_sum]
+                min_max_sum = [[min_sum, max_sum]]
         return params, min_max_sum
 
 

diff --git a/analysis/parameter_tuning.py b/analysis/parameter_tuning.py
@@ -244,7 +244,12 @@ def _add_dp_strategy_to_multi_parameter_configuration(
         noise_kind: Optional[pipeline_dp.NoiseKind],
         strategy_selector: dp_strategy_selector.DPStrategySelector) -> None:
     params = [
-        configuration.get_aggregate_params(blueprint_params, i)
+        # get_aggregate_params returns a tuple (AggregateParams,
+        # min_max_sum_per_partitions)
+        # for multi-columns. DP Strategy (i.e. noise_kind, partition_selection)
+        # is independent from min_max_sum_per_partitions, so it's fine just to
+        # AggregateParams with [0]
+        configuration.get_aggregate_params(blueprint_params, i)[0]
         for i in range(configuration.size)
     ]
     # Initialize fields corresponding to DP strategy configuration

diff --git a/analysis/utility_analysis_engine.py b/analysis/utility_analysis_engine.py
@@ -101,6 +101,7 @@ def _create_compound_combiner(
     ) -> combiners.CompoundCombiner:
         # Create Utility analysis combiners.
         internal_combiners = [per_partition_combiners.RawStatisticsCombiner()]
+        n_sum_aggregations = 0
         for params, min_max_sum_per_partition in data_structures.get_aggregate_params(
                 self._options):
             # Each parameter configuration has own BudgetAccountant which allows
@@ -124,15 +125,18 @@ def _create_compound_combiner(
                         budget_accountant.request_budget(
                             pipeline_dp.MechanismType.GENERIC), params))
             if pipeline_dp.Metrics.SUM in aggregate_params.metrics:
+                n_sum_aggregations = len(min_max_sum_per_partition)
                 for i_column, (min_sum,
                                max_sum) in enumerate(min_max_sum_per_partition):
+                    if len(min_max_sum_per_partition) == 1:
+                        i_column = None  # 1 column, no need to set index
                     sum_params = copy.deepcopy(params)
                     sum_params.min_sum_per_partition = min_sum
                     sum_params.max_sum_per_partition = max_sum
                     internal_combiners.append(
                         per_partition_combiners.SumCombiner(
                             budget_accountant.request_budget(mechanism_type),
-                            params,
+                            sum_params,
                             i_column=i_column))
             if pipeline_dp.Metrics.COUNT in aggregate_params.metrics:
                 internal_combiners.append(
@@ -147,7 +151,7 @@ def _create_compound_combiner(
             budget_accountant.compute_budgets()
 
         return per_partition_combiners.CompoundCombiner(
-            internal_combiners, return_named_tuple=False)
+            internal_combiners, n_sum_aggregations=n_sum_aggregations)
 
     def _select_private_partitions_internal(
             self, col, max_partitions_contributed: int,