Python中为柱状图拟合几何分布曲线及绘制概率质量函数的方法

Got it, let's walk through how to fit your geometric distribution and overlay its probability mass function (PMF) onto your existing bar plot. I'll break this down into clear, actionable steps with code:

Step 1: Prep Your Data and Import Required Libraries #

First, make sure you have the necessary libraries, and extract your frequency data into structured arrays:

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import geometric

# Extract your existing data from the freq list
x_vals = np.array([val[0] for val in freq])  # Number of times a word appears
counts = np.array([val[1] for val in freq])  # Number of words with that frequency
total_words = counts.sum()  # Total number of unique words (to scale PMF to match bar heights)

Step 2: Fit the Geometric Distribution #

Scipy's geometric distribution has two common parameterizations: one where x represents the number of failures before the first success (starting at 0), and another where x is the number of trials until the first success (starting at 1). Since your data is word occurrence counts (starting at 1), we'll fit using the latter:

# Create a flattened dataset (repeat each x value by its count for fitting)
flattened_data = np.repeat(x_vals, counts)

# Fit the geometric distribution, fixing the location parameter to 1 (since x starts at 1)
p_fit, loc_fit = geometric.fit(flattened_data, floc=1)

The p_fit variable is the estimated success probability for the geometric distribution.

Step 3: Calculate the Fitted PMF (Scaled to Match Your Bar Plot) #

The PMF gives probabilities, but your bar plot shows raw counts. Multiply the PMF values by the total number of unique words to get scaled counts that align with your y-axis:

# Calculate the fitted PMF values for each x, then scale to counts
fitted_counts = geometric.pmf(x_vals, p_fit, loc=loc_fit) * total_words

Step 4: Plot the Bar Chart and Overlay the Fitted PMF #

Update your existing plotting code to add the fitted curve:

# Original bar plot
plt.bar(x_vals, counts, log=True, alpha=0.75, label='Observed Frequency')

# Overlay the fitted PMF (as a line with markers for clarity)
plt.plot(x_vals, fitted_counts, 'r-', marker='o', label=f'Fitted Geometric PMF (p={p_fit:.4f})')

# Keep your existing formatting
plt.xticks(x_vals, rotation=70)
plt.xlabel('Times a word appears in the collection', labelpad=1)
plt.ylabel('Number of words appearing x times')
plt.legend()
plt.tight_layout()  # Prevents x-axis labels from getting cut off
plt.show()

Quick Notes to Avoid Pitfalls #

• Log Scale Check: Since you're using log=True for the bar plot, the fitted geometric PMF should appear as a straight line (because geometric distributions have exponential decay, which becomes linear on a log scale). This is a good sanity check for your fit.
• Parameterization Confirmation: If your x_vals start at 0 instead of 1, just remove the floc=1 argument from geometric.fit()—Scipy will default to the failure-count parameterization.
• Goodness of Fit: If you want to validate how well the geometric distribution fits, you can use a chi-squared test or visualize the residuals between observed counts and fitted counts.

内容的提问来源于stack exchange，提问作者gd13