GeometricDistribution.java
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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.
*/
package org.apache.commons.statistics.distribution;
import java.util.function.IntToDoubleFunction;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.sampling.distribution.GeometricSampler;
/**
* Implementation of the geometric distribution.
*
* <p>The probability mass function of \( X \) is:
*
* <p>\[ f(k; p) = (1-p)^k \, p \]
*
* <p>for \( p \in (0, 1] \) the probability of success and
* \( k \in \{0, 1, 2, \dots\} \) the number of failures.
*
* <p>This parameterization is used to model the number of failures until
* the first success.
*
* @see <a href="https://en.wikipedia.org/wiki/Geometric_distribution">Geometric distribution (Wikipedia)</a>
* @see <a href="https://mathworld.wolfram.com/GeometricDistribution.html">Geometric distribution (MathWorld)</a>
*/
public final class GeometricDistribution extends AbstractDiscreteDistribution {
/** 1/2. */
private static final double HALF = 0.5;
/** The probability of success. */
private final double probabilityOfSuccess;
/** {@code log(p)} where p is the probability of success. */
private final double logProbabilityOfSuccess;
/** {@code log(1 - p)} where p is the probability of success. */
private final double log1mProbabilityOfSuccess;
/** Value of survival probability for x=0.
* Used in the survival functions. Equal to (1 - probability of success). */
private final double sf0;
/** Implementation of PMF(x). Assumes that {@code x > 0}. */
private final IntToDoubleFunction pmf;
/**
* @param p Probability of success.
*/
private GeometricDistribution(double p) {
probabilityOfSuccess = p;
logProbabilityOfSuccess = Math.log(p);
log1mProbabilityOfSuccess = Math.log1p(-p);
sf0 = 1 - p;
// Choose the PMF implementation.
// When p >= 0.5 then 1 - p is exact and using the power function
// is consistently more accurate than the use of the exponential function.
// When p -> 0 then the exponential function avoids large error propagation
// of the power function used with an inexact 1 - p.
// Also compute the survival probability for use when x=0.
if (p >= HALF) {
pmf = x -> Math.pow(sf0, x) * probabilityOfSuccess;
} else {
pmf = x -> Math.exp(log1mProbabilityOfSuccess * x) * probabilityOfSuccess;
}
}
/**
* Creates a geometric distribution.
*
* @param p Probability of success.
* @return the geometric distribution
* @throws IllegalArgumentException if {@code p <= 0} or {@code p > 1}.
*/
public static GeometricDistribution of(double p) {
if (p <= 0 || p > 1) {
throw new DistributionException(DistributionException.INVALID_NON_ZERO_PROBABILITY, p);
}
return new GeometricDistribution(p);
}
/**
* Gets the probability of success parameter of this distribution.
*
* @return the probability of success.
*/
public double getProbabilityOfSuccess() {
return probabilityOfSuccess;
}
/** {@inheritDoc} */
@Override
public double probability(int x) {
if (x <= 0) {
// Special case of x=0 exploiting cancellation.
return x == 0 ? probabilityOfSuccess : 0;
}
return pmf.applyAsDouble(x);
}
/** {@inheritDoc} */
@Override
public double logProbability(int x) {
if (x <= 0) {
// Special case of x=0 exploiting cancellation.
return x == 0 ? logProbabilityOfSuccess : Double.NEGATIVE_INFINITY;
}
return x * log1mProbabilityOfSuccess + logProbabilityOfSuccess;
}
/** {@inheritDoc} */
@Override
public double cumulativeProbability(int x) {
if (x <= 0) {
// Note: CDF(x=0) = PDF(x=0) = probabilityOfSuccess
return x == 0 ? probabilityOfSuccess : 0;
}
// Note: Double addition avoids overflow. This may compute a value less than 1.0
// for the max integer value when p is very small.
return -Math.expm1(log1mProbabilityOfSuccess * (x + 1.0));
}
/** {@inheritDoc} */
@Override
public double survivalProbability(int x) {
if (x <= 0) {
// Note: SF(x=0) = 1 - PDF(x=0) = 1 - probabilityOfSuccess
// Use a pre-computed value to avoid cancellation when probabilityOfSuccess -> 0
return x == 0 ? sf0 : 1;
}
// Note: Double addition avoids overflow. This may compute a value greater than 0.0
// for the max integer value when p is very small.
return Math.exp(log1mProbabilityOfSuccess * (x + 1.0));
}
/** {@inheritDoc} */
@Override
public int inverseCumulativeProbability(double p) {
ArgumentUtils.checkProbability(p);
if (p == 1) {
return getSupportUpperBound();
}
if (p <= probabilityOfSuccess) {
return 0;
}
// p > probabilityOfSuccess
// => log(1-p) < log(1-probabilityOfSuccess);
// Both terms are negative as probabilityOfSuccess > 0.
// This should be lower bounded to (2 - 1) = 1
int x = (int) (Math.ceil(Math.log1p(-p) / log1mProbabilityOfSuccess) - 1);
// Correct rounding errors.
// This ensures x == icdf(cdf(x))
if (cumulativeProbability(x - 1) >= p) {
// No checks for x=0.
// If x=0; cdf(-1) = 0 and the condition is false as p>0 at this point.
x--;
} else if (cumulativeProbability(x) < p && x < Integer.MAX_VALUE) {
// The supported upper bound is max_value here as probabilityOfSuccess != 1
x++;
}
return x;
}
/** {@inheritDoc} */
@Override
public int inverseSurvivalProbability(double p) {
ArgumentUtils.checkProbability(p);
if (p == 0) {
return getSupportUpperBound();
}
if (p >= sf0) {
return 0;
}
// p < 1 - probabilityOfSuccess
// Inversion as for icdf using log(p) in place of log1p(-p)
int x = (int) (Math.ceil(Math.log(p) / log1mProbabilityOfSuccess) - 1);
// Correct rounding errors.
// This ensures x == isf(sf(x))
if (survivalProbability(x - 1) <= p) {
// No checks for x=0
// If x=0; sf(-1) = 1 and the condition is false as p<1 at this point.
x--;
} else if (survivalProbability(x) > p && x < Integer.MAX_VALUE) {
// The supported upper bound is max_value here as probabilityOfSuccess != 1
x++;
}
return x;
}
/**
* {@inheritDoc}
*
* <p>For probability parameter \( p \), the mean is:
*
* <p>\[ \frac{1 - p}{p} \]
*/
@Override
public double getMean() {
return (1 - probabilityOfSuccess) / probabilityOfSuccess;
}
/**
* {@inheritDoc}
*
* <p>For probability parameter \( p \), the variance is:
*
* <p>\[ \frac{1 - p}{p^2} \]
*/
@Override
public double getVariance() {
return (1 - probabilityOfSuccess) / (probabilityOfSuccess * probabilityOfSuccess);
}
/**
* {@inheritDoc}
*
* <p>The lower bound of the support is always 0.
*
* @return 0.
*/
@Override
public int getSupportLowerBound() {
return 0;
}
/**
* {@inheritDoc}
*
* <p>The upper bound of the support is positive infinity except for the
* probability parameter {@code p = 1.0}.
*
* @return {@link Integer#MAX_VALUE} or 0.
*/
@Override
public int getSupportUpperBound() {
return probabilityOfSuccess < 1 ? Integer.MAX_VALUE : 0;
}
/** {@inheritDoc} */
@Override
public Sampler createSampler(UniformRandomProvider rng) {
return GeometricSampler.of(rng, probabilityOfSuccess)::sample;
}
}