text: # Copyright 2021 The Private Cardinality Estimation Framework Authors
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Gamma-Poisson reach curve model for underreported counts."""
  
  from absl import logging
  import numpy as np
  import scipy.stats
  from typing import List
  from wfa_planning_evaluation_framework.models.reach_point import ReachPoint
  from wfa_planning_evaluation_framework.models.reach_curve import ReachCurve
  
  # The following constant defines the highest frequency that will be
  # used for parameter estimation.
  MAXIMUM_COMPUTATIONAL_FREQUENCY = 1000
  
  
  class GammaPoissonModel(ReachCurve):
      """Gamma-Poisson reach curve model for underreported counts.
  
      The Gamma-Poisson distribution is a discrete distribution defined on the
      non-negative integers.  The probably of observing a count of k is given
      by a Poisson distribution with parameter lambda, where lambda is drawn
      from a Gamma distribution with shape alpha and rate beta.  Explicitly,
         Pr(X = k) = lambda^k e^{-lambda} / k!, where lambda ~ Gamma(alpha, beta),
      k = 0, 1, 2, ...
  
      The Gamma-Poisson distribution is often used for modeling overdispersed
      data, which is common in the advertising industry.  In this setting, the
      random variable X represents the total number of impressions that could
      potentially be shown to a randomly chosen user.  This quantity is called
      the available inventory for X.
  
      In the underreported Gamma-Poisson model, we are given a set of
      observed impression counts, one value for each viewer that is
      reached.  Let these values be k_1,k_2,...,k_R.  Thus, R is the
      number of people that were reached, and the number of impressions
      that shown is I = k_1+k_2+...+k_R.  The unknown parameters that
      need to be estimated are:
  
        * the parameters alpha and beta of the associated Gamma distribution,
        * the total size I_max of the impression inventory, and
        * the total number N of viewers that could potentially be reached.
  
      This class implements code for estimating these parameters and for
      performing forecasts with those estimates.
  
      Let f(n | alpha, beta) be the probability that the size of the available
      inventory of a user is n.  Then, the probability that a randomly chosen
      user will have inventory n and be reached k times is given by
  
         g(k, n | alpha, beta, I, I_max) =
              C(n, k) (I / I_max)^k (1 - I/I_max)^{n-k} f(n | alpha, beta),
  
      where C(n, k) is the binomial coefficient n!/[k!(n-k)!].
  
      Therefore, the probability that a randomly chosen user will be reached
      exactly k times is
  
         g(k | alpha, beta, I, I_max) =
              sum_{n=k}^{infinity} g(k, n | alpha, beta, I, I_max).
  
      Thus, if a total of N viewers can potentially be reached, then the expected
      number of people who will be reached k times is
  
         hbar[i] = N * g(k | alpha, beta, I, I_max).
  
      Note that the value I is determined from the input and is not a
      model parameter that needs to be fit.  Thus, the apparent number
      of parameters that need to be fit is four.  However, we reduce
      the number of parameters to two as follows.  First, we note that
      the expected value of a Gamma-Poisson distribution with parameters
      alpha and beta is alpha/beta.  Thus, the total number of
      impressions I_max is related to the total number of potential
      viewers N via the relation
  
        E[I_max] = N * alpha/beta.
  
      We therefore set N = I_max * beta / alpha.  Next, we note that
      the size of the total impression inventory I_max is essentially
      a free parameter.  In fact, along the line I_max * beta = constant,
      the objective function is nearly constant.  To see this, observe
      that if the available inventory for each user is doubled, then
      the probability that any given user will receive k impressions
      remains essentially constant.  Thus, we set I / I_max to a
      constant value.  This leaves two remaining parameters to optimize,
      alpha and beta.
  
      To estimate these parameters we use a histogram matching approach.
      That is to say, let h[i] be the number of people that were reached
      i times, and let hbar[i] be the estimated number of people that
      would be reached i times, given alpha, beta and I_max, N.  The
      objective function that we compute is as follows:
  
        chi2(h, hbar) = \sum_i (h[i] - hbar[i])^2 / hbar[i]
  
      The perceptive reader might recognize this as the chi-squared
      statistic.  There are a couple advantages to using the chi-squared
      objective function.  First, in practice, the parameter estimates
      that are found often give a nearly exact match to the histograms
      that are observed.  Second, a low value of the objective function
      gives some confidence that a statistically good fit has been
      found.
      """
  
      def __init__(
          self, data: List[ReachPoint], max_reach=None, extrapolation_multiplier=1.0
      ):
          """Constructs a Gamma-Poisson model of underreported count data.
  
          Args:
            data:  A list of consisting of a single ReachPoint to which the model is
              to be fit.
            max_reach:  Optional.  If specified, the maximum possible reach that can
              be achieved.
            extrapolation_multiplier:  Int.  If specified, then a penalty term is
              introduced that penalizes models where the expected number of impressions
              in the inventory is less than extrapolation_multiplier times the
              observed number of impressions.  Thus, if the extrapolation_multiplier
              is 2, then optimizer prefers models where the number of impressions in
              the inventory is at least twice the number of impressions that were
              observed in the campaign.
          """
          super().__init__(data, max_reach)
          if len(data) != 1:
              raise ValueError("Exactly one ReachPoint must be specified")
          if data[0].impressions[0] < 0.001:
              raise ValueError("Attempt to create model with 0 impressions")
          if data[0].impressions[0] < data[0].reach(1):
              raise ValueError(
                  "Cannot have a model with fewer impressions than reached people"
              )
          self._reach_point = data[0]
          self._max_reach = max_reach
          self._extrapolation_multiplier = extrapolation_multiplier
          self._extrapolation_value = extrapolation_multiplier * data[0].impressions[0]
          if data[0].spends:
              self._cpi = data[0].spends[0] / data[0].impressions[0]
          else:
              self._cpi = None
  
      def _logpmf(self, n, alpha, beta):
          """Log of the PMF of the shifted Gamma-Poisson distribution.
  
          This implementation makes use of the equivalence between the
          Gamma-Poisson distribution with parameters (alpha, beta) and
          the negative binomial distribution with parameters (p, r) =
          (1 / (1 + beta), alpha).
  
          Args:
            n:  Value(s) at which the distribution is to be evaluated.
              This can be a scalar or a numpy array.
            alpha: Alpha parameter of the Gamma-Poisson distribution.
            beta: Beta parameter of the Gamma-Poisson distribution.
          Returns:
            f(n | alpha, beta), where f is the Gamma-Poisson distribution.
          """
          return scipy.stats.nbinom.logpmf(n - 1, alpha, 1.0 / (1.0 + beta))
  
      def _knreach(self, k, n, I, Imax, alpha, beta):
          """Probability that a random user has n impressions of which k are shown.
  
          Computes the function:
  
              C(n, k) (I / I_max)^k (1 - I/I_max)^{n-k} f(n | alpha, beta).
  
          where C(n, k) is the binomial coefficient n!/[k!(n-k)!].
  
          Args:
            k:  Scalar number of impressions that were seen by the user.
            n:  Available inventory for the user.
              This can be either a scalar or a numpy array.
            I:  Total number of impressions shown to all users.
            Imax: Total size of impression inventory.
            alpha: Parameter of Gamma-Poisson distribution.
            beta: Parameter of Gamma-Poisson distribution.
          Returns:
            Probability that a randomly chosen user will have an inventory
            of n impressions, of which k are shown.
          """
          kprob = scipy.stats.binom.logpmf(k, n, I / Imax)
          return np.exp(kprob + self._logpmf(n, alpha, beta))
  
      def _kreach(self, k, I, Imax, alpha, beta):
          """Probability that a random user receives k impressions.
  
          Computes the function:
              sum_{n=k}^{infinity} g(k, n | alpha, beta, I, I_max).
          where g(k, n | alpha, beta, I, I_max) is the probability that
          a randomly chosen user will have an available inventory of n
          impressions, of which k are shown.
  
          Args:
            k:  np.array specifying number of impressions that were seen by the user.
            I:  Total number of impressions shown to all users.
            Imax: Total size of impression inventory.
            alpha: Parameter of Gamma-Poisson distribution.
            beta: Parameter of Gamma-Poisson distribution.
          Returns:
            For each k, probability that a randomly chosen user will have an inventory
            of n impressions, of which k are shown.
          """
          return np.array(
              [
                  np.sum(
                      self._knreach(
                          kv,
                          np.arange(kv, MAXIMUM_COMPUTATIONAL_FREQUENCY + 1),
                          I,
                          Imax,
                          alpha,
                          beta,
                      )
                  )
                  for kv in k
              ]
          )
  
      def _expected_impressions(self, N, alpha, beta):
          """Estimates the expected size of impression inventory for N users.
  
          Makes use of the identity that the expected value of a negative binomial
          distribution with parameters p and r is pr/(1-p).
  
          Args:
            N:      total number of people that can potentially be reached.
            alpha:  shape parameter of Gamma distribution.
            beta:   rate parameter of Gamma distribution.
          Returns:
            Expected size of impression inventory for a population of N users.
          """
          # N * (1.0 + alpha * beta)
          return N * (1 + scipy.stats.nbinom.mean(alpha, 1.0 / (1.0 + beta)))
  
      def _expected_histogram(self, I, Imax, N, alpha, beta, max_freq=10):
          """Computes an estimated histogram.
  
          Args:
            I:      number of impressions that were purchased.
            Imax:   total size of impression inventory.
            N:      total number of people that can potentially be reached.
            alpha:  shape parameter of Gamma distribution.
            beta:   rate parameter of Gamma distribution.
            max_freq: maximum frequency to report in the output histogram.
          Returns:
            An np.array h where h[i] is the expected number of users who will see i+1
            impressions.
          """
          return N * self._kreach(np.arange(1, max_freq + 1), I, Imax, alpha, beta)
  
      def _impression_count(self, h):
          """Number of impressions recorded in the histogram h."""
          return np.sum([h[i] * (i + 1) for i in range(len(h))])
  
      def _histogram_from_reach_point(self, point):
          """Returns a frequency histogram from a reach point.
  
          Args:
            point: A ReachPoint
          Returns:
            A list h[], where h[i] is the number of people reached exactly i
            times.  Note that this histogram will have length one less than the
            list of k+ reach values in ReachPoint, because the last k-plus
            reach value in the ReachPoint represents people reached at various
            different frequencies.
          """
          h = [point.frequency(i) for i in range(1, point.max_frequency)]
          return h
  
      def _histogram_with_kplus_bucket(self, point):
          """Computes a histogram whose last bucket is k+ reach.
  
          Args:
            point:  A ReachPoint for which the histogram is to be
              computed.
          Returns:
            A list h[] where h[i] is the number of people reached
            exactly i+1 times, except for h[-1], which records the
            number of people reached len(h) or more times.
          """
          # A partial histogram of people who were reached exactly k times.
          h = self._histogram_from_reach_point(point)
  
          # The number of people who were reach len(h) or more times.  These
          # people are not counted in the histogram h.
          kplus = len(h) + 1
          kplus_reach = point.reach(kplus)
  
          # Append kplus_reach to h.  So, the last value of h is treated
          # specially.
          h.append(kplus_reach)
          h = np.array(h)
          return h
  
      def _chi2_distance_with_one_plus_reach(
          self, h, I, N, alpha, beta, one_plus_reach_weight=1.0
      ):
          """Returns distance between actual and expected histograms.
  
          Computes the metric
  
              chi2(h, hbar) = \sum_i (h[i] - hbar[i])^2 / hbar[i] +
                 w * (\sum_i (h[i] - hbar[i]))^2
  
          where
              hbar is the histogram that would be expected given parameters
                 I, N, alpha, beta
              w is the one_plus_reach_weight that is used to determine how
                 much to weight an exact match of 1+ reach.
  
          The first part of this objective formula is the chi-squared statistic
          used for determining whether two sets of categorical samples are drawn
          from the same underlying distribution.  Therefore, if the value returned
          by this function is relatively low (say < 20), then there is reason to
          believe that the Gamma-Poisson distribution adequately fits the data.
  
          Args:
            h: Actual frequency histogram.  h[i] is the number of people
               reached exactly i+1 times, except for h[-1] which is the
               number of people reached len(h) or more times.
            I: Total number of impressions purchased.
            N: Total number of unique users.
            alpha:  Parameter of the Gamma-Poisson distribution.
            beta:  Second parameter of the Gamma-Poisson distribution.
            one_plus_reach_weight:  Weight to be given to the squared
              difference in 1+ reach values for the actual and estimated
              histograms.
          Returns:
            chi^2 distance between the actual and expected histograms, plus
            a term for weighting the difference in 1+ reach.
          """
          if alpha <= 0 or beta <= 0 or N <= 0:
              return np.sum(np.array(h) ** 2)
  
          # Estimate total number of potential impressions
          Imax = self._expected_impressions(N, alpha, beta)
          if Imax <= I:
              return np.sum(np.array(h) ** 2)
  
          hbar = list(self._expected_histogram(I, Imax, N, alpha, beta, len(h) - 1))
  
          # Compute expected number of people in the k+ bucket.
          kplus_prob = 1.0 - np.sum(
              self._kreach(np.arange(0, len(h)), I, Imax, alpha, beta)
          )
          hbar.append(kplus_prob * N)
          hbar = np.array(hbar)
  
          actual_oneplus = np.sum(h)
          predicted_oneplus = np.sum(hbar)
          oneplus_error = (
              one_plus_reach_weight * (actual_oneplus - predicted_oneplus) ** 2
          )
  
          extrapolation_penalty = 0
          if Imax < self._extrapolation_value:
              extrapolation_penalty = (self._extrapolation_value - Imax) ** 2
  
          obj = (
              np.sum((hbar - h) ** 2 / (hbar + 1e-6))
              + oneplus_error
              + extrapolation_penalty
          )
  
          if np.isnan(obj):
              logging.vlog(2, f"alpha {alpha} beta {beta} N {N} Imax {Imax}")
              logging.vlog(2, f"h    {h}")
              logging.vlog(2, f"hbar {hbar}")
              raise RuntimeError("Invalid value of objective function")
  
          return obj
  
      def _fit_histogram_chi2_distance(
          self, point, one_plus_reach_weight=1.0, nstarting_points=10
      ):
          """Chi-squared fit to histogram h.
  
          Computes parameters alpha, beta, Imax, N such that the histogram hbar of
          of expected frequencies minimizes the metric
  
            chi2(h, hbar) = \sum_i (h[i] - hbar[i])^2 / hbar[i] +
               w * (\sum(h[i] - hbar[i]))^2
  
          where w is the one_plus_reach_weight.
  
          Experience shows that finding the optimal value depends on starting from
          a point where the estimated audience size is not too far from the true
          audience size.  So, this code retries the optimization for various different
          values of the estimated audience size.  The number of such attempts is given
          by nstarting_points.
  
          Args:
            point:  A ReachPoint whose histogram is to be approximated.
            one_plus_reach_weight:  Weight to be given to the squared
              difference in 1+ reach values for the actual and estimated
              histograms.
            nstarting_points: The number of different starting points that
              should be tried when performing optimizations.
          Returns:
            A tuple (Imax, N, alpha, beta) representing the parameters of the
            best fit that was found.
          """
          # There must be at least this many people in the final model
          Nmin = point.reach(1)
  
          # This is the number of impressions that were actually served.
          I0 = point.impressions[0]
  
          # The histogram of the data.
          h = self._histogram_with_kplus_bucket(point)
  
          # Best score found so far of the objective function.
          best_score = 1e99
  
          # Calculate the set of audience sizes that are used as initial
          # starting points.
          Nvalues = Nmin / (np.arange(1, nstarting_points + 1) / (nstarting_points + 1))
  
          # If the objective function for an optimization attempt is smaller than
          # this value, then we have found a model that fits the data.  There is no
          # need to continue searching for a better fit.
          early_stopping_value = scipy.stats.chi2.ppf(0.98, len(h))
  
          for N0 in Nvalues:
              # Choose a reasonable starting point for optimization.
              Imax0, alpha0, beta0 = self._fit_histogram_fixed_N(
                  point, N0, one_plus_reach_weight
              )
  
              def gamma_obj(params):
                  """Objective function for optimization."""
                  alpha, beta, N = params
                  return self._chi2_distance_with_one_plus_reach(
                      h, I0, N, alpha, beta, one_plus_reach_weight
                  )
  
              result = scipy.optimize.minimize(
                  gamma_obj,
                  (alpha0, beta0, N0),
                  method="L-BFGS-B",
                  bounds=((1e-6, 100.0), (1e-6, 100.0), (Nmin, 100.0 * Nmin)),
              )
  
              if result.fun < best_score:
                  best_alpha, best_beta, best_N = result.x
                  best_score = result.fun
  
              if best_score < early_stopping_value:
                  break
  
          best_Imax = self._expected_impressions(best_N, best_alpha, best_beta)
          return best_Imax, best_N, best_alpha, best_beta
  
      def _find_starting_point_for_fixed_N(self, h, I0, N, one_plus_reach_weight):
          """Searches for a good initial starting point for optimization.
  
          The standard optimizers unfortunately often fail to find the
          global minimum if they are started from suboptimal starting points.
          The surface represented by the objective function has large flat
          regions, and it often contains a valley with a false minimum that
          is adjacent to the basin containing the true minimum.  However,
          it seems to be the case that the basin containing the true minimum
          is often located relatively close to the origin.  This function performs
          a simple downhill search for a local optimum starting from a point
          close to the origin.  The value returned by this function can then
          be used as the starting point for a more thorough optimization
          attempt.
  
          Args:
            h:  Histogram with kplus reach bucket for the reach point
                that is to be fit.
            I0: The number of impressions that were shown to users.
            N:  The total audience size.
            one_plus_reach_weight:  Weight to be given to the squared
              difference in 1+ reach values for the actual and estimated
              histograms.
          Returns:
            alpha, beta: The values of the parameters that gave the best
            score of the objective function when performing a simple
            downhill search.
          """
  
          # Search for local optimum on a grid whose initial spacing is
          # given by grid_width
          grid_width = 0.05
  
          # Try to start from a point that is near the origin but also
          # in the feasible region
          alpha = beta = 0.1
          if I0 > N:
              alpha = beta = 1.1 * np.sqrt(I0 / N - 1) + grid_width
  
          score = 1e99
          while grid_width > 0.01:
              improvement_found = False
              local_best_score = score
              for delta in [(-1, -1), (-1, 1), (1, -1), (1, 1)]:
                  test_alpha = alpha + delta[0] * grid_width
                  test_beta = beta + delta[1] * grid_width
                  test_score = self._chi2_distance_with_one_plus_reach(
                      h, I0, N, test_alpha, test_beta, one_plus_reach_weight
                  )
                  if test_score < local_best_score:
                      improvement_found = True
                      local_best_alpha = test_alpha
                      local_best_beta = test_beta
                      local_best_score = test_score
              if improvement_found:
                  alpha, beta, score = local_best_alpha, local_best_beta, local_best_score
              else:
                  grid_width *= 0.5
  
          return alpha, beta
  
      def _fit_histogram_fixed_N(self, point, N, one_plus_reach_weight=1.0):
          """Chi-squared fit to histogram h.
  
          Computes parameters alpha, beta, Imax, N such that the histogram hbar of
          of expected frequencies minimizes the metric
  
            chi2(h, hbar) = \sum_i (h[i] - hbar[i])^2 / hbar[i] +
               w * (\sum(h[i] - hbar[i]))^2
  
          where w is the one_plus_reach_weight.
  
          Args:
            point:  A ReachPoint whose histogram is to be approximated.
            N:  int, total reachable audience.
            one_plus_reach_weight:  Weight to be given to the squared
              difference in 1+ reach values for the actual and estimated
              histograms.
          Returns:
            A tuple (Imax, alpha, beta) representing the parameters of the
            best fit that was found.
          """
          # This is the number of impressions that were actually served.
          I0 = point.impressions[0]
  
          # The histogram of the data.
          h = self._histogram_with_kplus_bucket(point)
  
          # Look for a good starting point for optimization
          alpha0, beta0 = self._find_starting_point_for_fixed_N(
              h, I0, N, one_plus_reach_weight
          )
  
          def gamma_obj(params):
              """Objective function for optimization."""
              alpha, beta = params
              score = self._chi2_distance_with_one_plus_reach(
                  h, I0, N, alpha, beta, one_plus_reach_weight
              )
              return score
  
          result = scipy.optimize.minimize(
              gamma_obj,
              (alpha0, beta0),
              method="L-BFGS-B",
              bounds=((1e-6, 100.0), (1e-6, 100.0)),
          )
  
          best_alpha, best_beta = result.x
          best_Imax = self._expected_impressions(N, best_alpha, best_beta)
          return best_Imax, best_alpha, best_beta
  
      def _fit(self) -> None:
          """Fits a model to the data that was provided in the constructor."""
  
          if self._max_reach is None:
              Imax, N, alpha, beta = self._fit_histogram_chi2_distance(self._reach_point)
              self._max_impressions = Imax
              self._max_reach = N
              self._alpha = alpha
              self._beta = beta
          else:
              Imax, alpha, beta = self._fit_histogram_fixed_N(
                  self._reach_point, self._max_reach
              )
              self._max_impressions = Imax
              self._alpha = alpha
              self._beta = beta
  
      def by_impressions(
          self, impressions: List[int], max_frequency: int = 1
      ) -> ReachPoint:
          """Returns the estimated reach as a function of impressions.
  
          Args:
            impressions: list of ints of length 1, specifying the hypothetical number
              of impressions that are shown.
            max_frequency: int, specifies the number of frequencies for which reach
              will be reported.
          Returns:
            A ReachPoint specifying the estimated reach for this number of impressions.
          """
          if len(impressions) != 1:
              raise ValueError("Impressions vector must have a length of 1.")
          hist = self._expected_histogram(
              min(impressions[0], self._max_impressions - 1),
              self._max_impressions,
              self.max_reach,
              self._alpha,
              self._beta,
              MAXIMUM_COMPUTATIONAL_FREQUENCY,
          )
          kplus_reach = self._kplus_reaches_from_frequencies(hist)
          if self._cpi:
              return ReachPoint(
                  impressions, kplus_reach[:max_frequency], [impressions[0] * self._cpi]
              )
          else:
              return ReachPoint(impressions, kplus_reach[:max_frequency])
  
      def by_spend(self, spends: List[int], max_frequency: int = 1) -> ReachPoint:
          """Returns the estimated reach as a function of spend assuming constant CPM
  
          Args:
            spend: list of floats of length 1, specifying the hypothetical spend.
            max_frequency: int, specifies the number of frequencies for which reach
              will be reported.
          Returns:
            A ReachPoint specifying the estimated reach for this number of impressions.
          """
          if len(spends) != 1:
              raise ValueError("Spend vector must have a length of 1.")
          return self.by_impressions(
              [self.impressions_for_spend(spends[0])], max_frequency
          )
  
      def impressions_for_spend(self, spend: float) -> int:
          if not self._cpi:
              raise ValueError("Impression cost is not known for this ReachPoint.")
          return spend / self._cpi
  