GeometricDistribution.java

  1. /*
  2.  * Licensed to the Apache Software Foundation (ASF) under one or more
  3.  * contributor license agreements.  See the NOTICE file distributed with
  4.  * this work for additional information regarding copyright ownership.
  5.  * The ASF licenses this file to You under the Apache License, Version 2.0
  6.  * (the "License"); you may not use this file except in compliance with
  7.  * the License.  You may obtain a copy of the License at
  8.  *
  9.  *      http://www.apache.org/licenses/LICENSE-2.0
  10.  *
  11.  * Unless required by applicable law or agreed to in writing, software
  12.  * distributed under the License is distributed on an "AS IS" BASIS,
  13.  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14.  * See the License for the specific language governing permissions and
  15.  * limitations under the License.
  16.  */
  17. package org.apache.commons.statistics.distribution;

  19. import java.util.function.IntToDoubleFunction;
  20. import org.apache.commons.rng.UniformRandomProvider;
  21. import org.apache.commons.rng.sampling.distribution.GeometricSampler;

  23. /**
  24.  * Implementation of the geometric distribution.
  25.  *
  26.  * <p>The probability mass function of \( X \) is:
  27.  *
  28.  * <p>\[ f(k; p) = (1-p)^k \, p \]
  29.  *
  30.  * <p>for \( p \in (0, 1] \) the probability of success and
  31.  * \( k \in \{0, 1, 2, \dots\} \) the number of failures.
  32.  *
  33.  * <p>This parameterization is used to model the number of failures until
  34.  * the first success.
  35.  *
  36.  * @see <a href="https://en.wikipedia.org/wiki/Geometric_distribution">Geometric distribution (Wikipedia)</a>
  37.  * @see <a href="https://mathworld.wolfram.com/GeometricDistribution.html">Geometric distribution (MathWorld)</a>
  38.  */
  39. public final class GeometricDistribution extends AbstractDiscreteDistribution {
  40.     /** 1/2. */
  41.     private static final double HALF = 0.5;

  43.     /** The probability of success. */
  44.     private final double probabilityOfSuccess;
  45.     /** {@code log(p)} where p is the probability of success. */
  46.     private final double logProbabilityOfSuccess;
  47.     /** {@code log(1 - p)} where p is the probability of success. */
  48.     private final double log1mProbabilityOfSuccess;
  49.     /** Value of survival probability for x=0.
  50.      * Used in the survival functions. Equal to (1 - probability of success). */
  51.     private final double sf0;
  52.     /** Implementation of PMF(x). Assumes that {@code x > 0}. */
  53.     private final IntToDoubleFunction pmf;

  55.     /**
  56.      * @param p Probability of success.
  57.      */
  58.     private GeometricDistribution(double p) {
  59.         probabilityOfSuccess = p;
  60.         logProbabilityOfSuccess = Math.log(p);
  61.         log1mProbabilityOfSuccess = Math.log1p(-p);
  62.         sf0 = 1 - p;

  64.         // Choose the PMF implementation.
  65.         // When p >= 0.5 then 1 - p is exact and using the power function
  66.         // is consistently more accurate than the use of the exponential function.
  67.         // When p -> 0 then the exponential function avoids large error propagation
  68.         // of the power function used with an inexact 1 - p.
  69.         // Also compute the survival probability for use when x=0.
  70.         if (p >= HALF) {
  71.             pmf = x -> Math.pow(sf0, x) * probabilityOfSuccess;
  72.         } else {
  73.             pmf = x -> Math.exp(log1mProbabilityOfSuccess * x) * probabilityOfSuccess;
  74.         }
  75.     }

  77.     /**
  78.      * Creates a geometric distribution.
  79.      *
  80.      * @param p Probability of success.
  81.      * @return the geometric distribution
  82.      * @throws IllegalArgumentException if {@code p <= 0} or {@code p > 1}.
  83.      */
  84.     public static GeometricDistribution of(double p) {
  85.         if (p <= 0 || p > 1) {
  86.             throw new DistributionException(DistributionException.INVALID_NON_ZERO_PROBABILITY, p);
  87.         }
  88.         return new GeometricDistribution(p);
  89.     }

  91.     /**
  92.      * Gets the probability of success parameter of this distribution.
  93.      *
  94.      * @return the probability of success.
  95.      */
  96.     public double getProbabilityOfSuccess() {
  97.         return probabilityOfSuccess;
  98.     }

  100.     /** {@inheritDoc} */
  101.     @Override
  102.     public double probability(int x) {
  103.         if (x <= 0) {
  104.             // Special case of x=0 exploiting cancellation.
  105.             return x == 0 ? probabilityOfSuccess : 0;
  106.         }
  107.         return pmf.applyAsDouble(x);
  108.     }

  110.     /** {@inheritDoc} */
  111.     @Override
  112.     public double logProbability(int x) {
  113.         if (x <= 0) {
  114.             // Special case of x=0 exploiting cancellation.
  115.             return x == 0 ? logProbabilityOfSuccess : Double.NEGATIVE_INFINITY;
  116.         }
  117.         return x * log1mProbabilityOfSuccess + logProbabilityOfSuccess;
  118.     }

  120.     /** {@inheritDoc} */
  121.     @Override
  122.     public double cumulativeProbability(int x) {
  123.         if (x <= 0) {
  124.             // Note: CDF(x=0) = PDF(x=0) = probabilityOfSuccess
  125.             return x == 0 ? probabilityOfSuccess : 0;
  126.         }
  127.         // Note: Double addition avoids overflow. This may compute a value less than 1.0
  128.         // for the max integer value when p is very small.
  129.         return -Math.expm1(log1mProbabilityOfSuccess * (x + 1.0));
  130.     }

  132.     /** {@inheritDoc} */
  133.     @Override
  134.     public double survivalProbability(int x) {
  135.         if (x <= 0) {
  136.             // Note: SF(x=0) = 1 - PDF(x=0) = 1 - probabilityOfSuccess
  137.             // Use a pre-computed value to avoid cancellation when probabilityOfSuccess -> 0
  138.             return x == 0 ? sf0 : 1;
  139.         }
  140.         // Note: Double addition avoids overflow. This may compute a value greater than 0.0
  141.         // for the max integer value when p is very small.
  142.         return Math.exp(log1mProbabilityOfSuccess * (x + 1.0));
  143.     }

  145.     /** {@inheritDoc} */
  146.     @Override
  147.     public int inverseCumulativeProbability(double p) {
  148.         ArgumentUtils.checkProbability(p);
  149.         if (p == 1) {
  150.             return getSupportUpperBound();
  151.         }
  152.         if (p <= probabilityOfSuccess) {
  153.             return 0;
  154.         }
  155.         // p > probabilityOfSuccess
  156.         // => log(1-p) < log(1-probabilityOfSuccess);
  157.         // Both terms are negative as probabilityOfSuccess > 0.
  158.         // This should be lower bounded to (2 - 1) = 1
  159.         int x = (int) (Math.ceil(Math.log1p(-p) / log1mProbabilityOfSuccess) - 1);

  161.         // Correct rounding errors.
  162.         // This ensures x == icdf(cdf(x))

  164.         if (cumulativeProbability(x - 1) >= p) {
  165.             // No checks for x=0.
  166.             // If x=0; cdf(-1) = 0 and the condition is false as p>0 at this point.
  167.             x--;
  168.         } else if (cumulativeProbability(x) < p && x < Integer.MAX_VALUE) {
  169.             // The supported upper bound is max_value here as probabilityOfSuccess != 1
  170.             x++;
  171.         }

  173.         return x;
  174.     }

  176.     /** {@inheritDoc} */
  177.     @Override
  178.     public int inverseSurvivalProbability(double p) {
  179.         ArgumentUtils.checkProbability(p);
  180.         if (p == 0) {
  181.             return getSupportUpperBound();
  182.         }
  183.         if (p >= sf0) {
  184.             return 0;
  185.         }

  187.         // p < 1 - probabilityOfSuccess
  188.         // Inversion as for icdf using log(p) in place of log1p(-p)
  189.         int x = (int) (Math.ceil(Math.log(p) / log1mProbabilityOfSuccess) - 1);

  191.         // Correct rounding errors.
  192.         // This ensures x == isf(sf(x))

  194.         if (survivalProbability(x - 1) <= p) {
  195.             // No checks for x=0
  196.             // If x=0; sf(-1) = 1 and the condition is false as p<1 at this point.
  197.             x--;
  198.         } else if (survivalProbability(x) > p && x < Integer.MAX_VALUE) {
  199.             // The supported upper bound is max_value here as probabilityOfSuccess != 1
  200.             x++;
  201.         }

  203.         return x;
  204.     }

  206.     /**
  207.      * {@inheritDoc}
  208.      *
  209.      * <p>For probability parameter \( p \), the mean is:
  210.      *
  211.      * <p>\[ \frac{1 - p}{p} \]
  212.      */
  213.     @Override
  214.     public double getMean() {
  215.         return (1 - probabilityOfSuccess) / probabilityOfSuccess;
  216.     }

  218.     /**
  219.      * {@inheritDoc}
  220.      *
  221.      * <p>For probability parameter \( p \), the variance is:
  222.      *
  223.      * <p>\[ \frac{1 - p}{p^2} \]
  224.      */
  225.     @Override
  226.     public double getVariance() {
  227.         return (1 - probabilityOfSuccess) / (probabilityOfSuccess * probabilityOfSuccess);
  228.     }

  230.     /**
  231.      * {@inheritDoc}
  232.      *
  233.      * <p>The lower bound of the support is always 0.
  234.      *
  235.      * @return 0.
  236.      */
  237.     @Override
  238.     public int getSupportLowerBound() {
  239.         return 0;
  240.     }

  242.     /**
  243.      * {@inheritDoc}
  244.      *
  245.      * <p>The upper bound of the support is positive infinity except for the
  246.      * probability parameter {@code p = 1.0}.
  247.      *
  248.      * @return {@link Integer#MAX_VALUE} or 0.
  249.      */
  250.     @Override
  251.     public int getSupportUpperBound() {
  252.         return probabilityOfSuccess < 1 ? Integer.MAX_VALUE : 0;
  253.     }

  255.     /** {@inheritDoc} */
  256.     @Override
  257.     public Sampler createSampler(UniformRandomProvider rng) {
  258.         return GeometricSampler.of(rng, probabilityOfSuccess)::sample;
  259.     }
  260. }
Created with JaCoCo 0.8.12.202403310830