From 0443923c6f48db64c0a15af06bb4f0c82d1dd8cf Mon Sep 17 00:00:00 2001
From: Antonio Sartori <antoniosartori@chromium.org>
Date: Fri, 13 Jan 2023 14:45:07 +0000
Subject: [PATCH] [Identifiability Study] Fix construction of mesa distribution

The parameter of the geometric distribution should be always comprised
between 0 and 1.

Bug: 1405579,957519
Change-Id: I6db291d8b9f9cbfc166c591281e554b05ee9dc49
---

diff --git a/chrome/browser/privacy_budget/identifiability_study_state.cc b/chrome/browser/privacy_budget/identifiability_study_state.cc
index dc6f4d5..0453861 100644
--- a/chrome/browser/privacy_budget/identifiability_study_state.cc
+++ b/chrome/browser/privacy_budget/identifiability_study_state.cc
@@ -65,7 +65,8 @@
               // However, `MesaDistribution` needs a `pivot_point` parameter
               // bigger than 0.
               : 1,
-          kMesaDistributionRatio),
+          kMesaDistributionRatio,
+          kMesaDistributionGeometricDistributionParam),
       reid_estimator_(
           PrivacyBudgetReidScoreEstimator(&settings_, pref_service)) {
   InitializeGlobalStudySettings();
diff --git a/chrome/browser/privacy_budget/identifiability_study_state.h b/chrome/browser/privacy_budget/identifiability_study_state.h
index 8996e8c..1fa3ff2 100644
--- a/chrome/browser/privacy_budget/identifiability_study_state.h
+++ b/chrome/browser/privacy_budget/identifiability_study_state.h
@@ -135,6 +135,10 @@
   // of random numbers for selecting identifiable surface for measurement.
   static constexpr double kMesaDistributionRatio = 0.9;
 
+  // The parameter of the geometric distribution used for the tail of the Mesa
+  // distribution.
+  static constexpr double kMesaDistributionGeometricDistributionParam = 0.5;
+
  private:
   friend class test_utils::InspectableIdentifiabilityStudyState;
 
diff --git a/chrome/browser/privacy_budget/mesa_distribution.h b/chrome/browser/privacy_budget/mesa_distribution.h
index 01427a0..635d358 100644
--- a/chrome/browser/privacy_budget/mesa_distribution.h
+++ b/chrome/browser/privacy_budget/mesa_distribution.h
@@ -15,14 +15,13 @@
 //
 // The PDF is:
 //
-//            ⎧    0                            ... if x < 0
+//            ⎧    0                               ... if x < 0
 //            ⎪
-//            ⎪    λ                            ... if 0 <= x < T
-//     P(x) = ⎨
-//            ⎪    λ (1 - γ)⁽ˣ⁻ᵀ⁾               ... otherwise
+//     P(x) = ⎨    λ                               ... if 0 <= x < T
 //            ⎪
-//            ⎩
-// where ...
+//            ⎩    (1 - τ) * γ * (1 - γ)^{X - T}   ... otherwise
+//
+// where
 //
 //   T = Value at which the PDF switches from a uniform to a geometric
 //       distribution. Referred to in code as the `pivot_point`.
@@ -31,16 +30,14 @@
 //       then 90% of the probability space is in the linear region. The ratio is
 //       referred to in code as `dist_ratio`.
 //
+//   γ = Parameter of the geometric distribution.
+//
 //        τ
 //   λ = ───
 //        T
 //
-//         λ          τ
-//   γ = ───── = ───────────
-//       1 - τ   T * (1 - τ)
-//
 // In otherwords, the PDF is uniform up to T with a probability of λ, and then
-// switches to a geometric distribution with parameter λ that extends to
+// switches to a geometric distribution with parameter γ that extends to
 // infinity.
 //
 // It looks like this in the form of a graph which should make a little bit more
@@ -73,11 +70,12 @@
           std::enable_if_t<std::is_integral<ResultType>::value, int> = 0>
 class MesaDistribution {
  public:
-  MesaDistribution(ResultType pivot_point, double dist_ratio)
+  MesaDistribution(ResultType pivot_point,
+                   double dist_ratio,
+                   double geometric_distribution_param)
       : pivot_point_(pivot_point),
         uniform_distribution_(0, std::ceil(pivot_point / dist_ratio)),
-        geometric_distribution_(dist_ratio /
-                                (pivot_point * (1.0l - dist_ratio))) {
+        geometric_distribution_(geometric_distribution_param) {
     DCHECK_GT(pivot_point, static_cast<ResultType>(0));
     DCHECK_GT(dist_ratio, 0);
     DCHECK_LT(dist_ratio, 1);
diff --git a/chrome/browser/privacy_budget/mesa_distribution_unittest.cc b/chrome/browser/privacy_budget/mesa_distribution_unittest.cc
index 7e16815..039e148c 100644
--- a/chrome/browser/privacy_budget/mesa_distribution_unittest.cc
+++ b/chrome/browser/privacy_budget/mesa_distribution_unittest.cc
@@ -17,10 +17,12 @@
 constexpr auto kSeed = 3;
 constexpr auto kPivotPoint = 300;
 constexpr auto kDistRatio = 0.9l;
+constexpr auto kGeometricDistributionParam = 0.5l;
 }  // namespace
 
 TEST(MesaDistributionTest, Get) {
-  MesaDistribution<int> mesa(kPivotPoint, kDistRatio);
+  MesaDistribution<int> mesa(kPivotPoint, kDistRatio,
+                             kGeometricDistributionParam);
   std::mt19937 g(kSeed);
   auto v1 = mesa.Get(g);
   g.seed(kSeed);
@@ -69,7 +71,7 @@
   std::array<double, kMaxOffset + 1> occurrences = {0.0};
 
   // The distribution under test:
-  MesaDistribution<int> mesa(kPivotPoint, kDistRatio);
+  MesaDistribution<int> mesa(kPivotPoint, kDistRatio, kGamma);
 
   std::mt19937 random_bit_generator(kSeed);