The geometric distribution models the number of trials needed to get the first success in a sequence of independent, binary trials, each with the same probability of success.
This distribution is used when we’re interested in knowing how many attempts it will take for an event to occur, such as how many times we need to roll a die to get a six or the number of website visits before a user makes a purchase.
In this tutorial, we’ll learn the basics of the geometric distribution, explore its probability mass function, and learn how to use it in Python with the Scipy stats module.
▶️ Link to the Google Colab notebook
Understanding the Geometric Distribution
The geometric distribution describes the probability of achieving the first success on the k-th trial and is characterized by a single parameter, p, which is the probability of success on each trial. The PMF of the geometric distribution is:
where:
- k is the trial on which you get the first success and
- p is the success probability for each trial
The geometric distribution has the following mean and variance:
- Mean = 1/p
- Variance = (1-p)/ p 2
The geometric distribution is memoryless, meaning the probability of success on the next trial is independent of prior failures.
Working with the Geometric Distribution in Python
You can use scipy.stats.geom to work with the geometric distribution in Python.
Importing Required Libraries
First, we import all the libraries we’ll need:
import numpy as np import matplotlib.pyplot as plt from scipy.stats import geom
Visualizing the Probability Mass Function (PMF)
Let’s visualize the PMF of the geometric distribution for a fixed success probability to understand how the probability of success changes with each trial.
# Set probability of success
p = 0.3
# Define the number of trials for the x-axis
trials = np.arange(1, 21)
# Calculate the PMF for each trial
pmf_values = geom.pmf(trials, p)
# Plot the PMF
plt.stem(trials, pmf_values, use_line_collection=True)
plt.title('Geometric Distribution PMF (p=0.3)')
plt.xlabel('Number of Trials (k)')
plt.ylabel('P(X = k)')
plt.show()
Here’s the output plot showing the probability of getting the first success on each successive trial. Note that it decreases as the number of trials increases.
Generating Geometric-Distributed Random Numbers
Generating random samples from a geometric distribution can give us a sense of the variability in trial counts needed to get the first success. This is useful for simulations or estimating real-world probabilities.
# Set the probability of success
p = 0.3
# Generate random samples
samples = geom.rvs(p, size=1000)
# Plot histogram of samples
plt.hist(samples, bins=range(1, max(samples) + 1), density=True, alpha=0.6, color='blue', align='left')
plt.title(f'Histogram of Geometric-distributed samples (p={p})')
plt.xlabel('Number of Trials to First Success')
plt.ylabel('Frequency')
plt.show()
# Sample mean and variance
sample_mean = np.mean(samples)
sample_variance = np.var(samples)
print(f'Sample Mean: {sample_mean:.3f}, Sample Variance: {sample_variance:.3f}')
This example generates 1,000 samples from a Geometric distribution with a 30% probability of success (histogram as shown).
The sample mean and variance:
Sample Mean: 3.297, Sample Variance: 7.195
Example: Modeling Customer Retention
Suppose we want to model how many visits it takes for a customer to make a purchase on a website. Given that each visit has a 10% probability of conversion, we can model the number of visits required to get the first purchase using the geometric distribution.
# Probability of purchase on any given visit
purchase_probability = 0.1
# Simulate customer visits until the first purchase
num_customers = 1000
visits_until_purchase = geom.rvs(purchase_probability, size=num_customers)
# Plot the distribution of visits until purchase
plt.hist(visits_until_purchase, bins=range(1, max(visits_until_purchase) + 1), density=True, color='purple', alpha=0.6, align='left')
plt.title('Distribution of Visits Until Purchase (p=0.1)')
plt.xlabel('Number of Visits')
plt.ylabel('Frequency')
plt.show()
# Calculate mean visits until purchase
avg_visits = np.mean(visits_until_purchase)
print(f'Average visits until purchase: {avg_visits}')
In this example, we simulate 1,000 customers and plot the distribution of visits until they make a purchase.
And get the average visit count as well:
Average visits until purchase: 9.776
Conclusion
In this tutorial, we explored the Geometric distribution, a discrete probability distribution often used to model the number of trials required for the first success.
We looked at the probability mass function and generated samples and also applied the distribution in a practical example. If you’d like more such probability tutorials, let us know.
Cover image: Created by Author | Segmind SSD-1B Model
