Merge pull request #75 from marimo-team/haleshot/14_binomial

probability/14_binomial_distribution.py

@@ -0,0 +1,545 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "marimo",
+#     "matplotlib==3.10.0",
+#     "numpy==2.2.4",
+#     "scipy==1.15.2",
+#     "altair==5.2.0",
+#     "wigglystuff==0.1.10",
+#     "pandas==2.2.3",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.11.24"
+app = marimo.App(width="medium", app_title="Binomial Distribution")
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        # Binomial Distribution
+
+        _This notebook is a computational companion to ["Probability for Computer Scientists"](https://chrispiech.github.io/probabilityForComputerScientists/en/part2/binomial/), by Stanford professor Chris Piech._
+
+        In this section, we will discuss the binomial distribution. To start, imagine the following example:
+
+        Consider $n$ independent trials of an experiment where each trial is a "success" with probability $p$. Let $X$ be the number of successes in $n$ trials.
+
+        This situation is truly common in the natural world, and as such, there has been a lot of research into such phenomena. Random variables like $X$ are called **binomial random variables**. If you can identify that a process fits this description, you can inherit many already proved properties such as the PMF formula, expectation, and variance!
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Binomial Random Variable Definition
+
+        $X \sim \text{Bin}(n, p)$ represents a binomial random variable where:
+
+        - $X$ is our random variable (number of successes)
+        - $\text{Bin}$ indicates it follows a binomial distribution
+        - $n$ is the number of trials
+        - $p$ is the probability of success in each trial
+
+        ```
+        X ~ Bin(n, p)
+         ↑    ↑  ↑
+         |    |  +-- Probability of
+         |    |      success on each
+         |    |      trial
+         |    +-- Number of trials
+         |
+        Our random variable
+          is distributed
+          as a Binomial
+        ```
+
+        Here are a few examples of binomial random variables:
+
+        - Number of heads in $n$ coin flips
+        - Number of 1's in randomly generated length $n$ bit string
+        - Number of disk drives crashed in 1000 computer cluster, assuming disks crash independently
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Properties of Binomial Distribution
+
+        | Property | Formula |
+        |----------|---------|
+        | Notation | $X \sim \text{Bin}(n, p)$ |
+        | Description | Number of "successes" in $n$ identical, independent experiments each with probability of success $p$ |
+        | Parameters | $n \in \{0, 1, \dots\}$, the number of experiments<br>$p \in [0, 1]$, the probability that a single experiment gives a "success" |
+        | Support | $x \in \{0, 1, \dots, n\}$ |
+        | PMF equation | $P(X=x) = {n \choose x}p^x(1-p)^{n-x}$ |
+        | Expectation | $E[X] = n \cdot p$ |
+        | Variance | $\text{Var}(X) = n \cdot p \cdot (1-p)$ |
+
+        Let's explore how the binomial distribution changes with different parameters.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(TangleSlider, mo):
+    # Interactive elements using TangleSlider
+    n_slider = mo.ui.anywidget(TangleSlider(
+        amount=10, 
+        min_value=1, 
+        max_value=30, 
+        step=1,
+        digits=0,
+        suffix=" trials"
+    ))
+
+    p_slider = mo.ui.anywidget(TangleSlider(
+        amount=0.5, 
+        min_value=0.01, 
+        max_value=0.99, 
+        step=0.01,
+        digits=2,
+        suffix=" probability"
+    ))
+
+    # Grid layout for the interactive controls
+    controls = mo.vstack([
+        mo.md("### Adjust Parameters to See How Binomial Distribution Changes"),
+        mo.hstack([
+            mo.md("**Number of trials (n):** "),
+            n_slider
+        ], justify="start"),
+        mo.hstack([
+            mo.md("**Probability of success (p):** "),
+            p_slider
+        ], justify="start"),
+    ])
+    return controls, n_slider, p_slider
+
+
+@app.cell(hide_code=True)
+def _(controls):
+    controls
+    return
+
+
+@app.cell(hide_code=True)
+def _(n_slider, np, p_slider, plt, stats):
+    # Parameters from sliders
+    _n = int(n_slider.amount)
+    _p = p_slider.amount
+
+    # Calculate PMF
+    _x = np.arange(0, _n + 1)
+    _pmf = stats.binom.pmf(_x, _n, _p)
+
+    # Relevant stats
+    _mean = _n * _p
+    _variance = _n * _p * (1 - _p)
+    _std_dev = np.sqrt(_variance)
+
+    _fig, _ax = plt.subplots(figsize=(10, 6))
+
+    # Plot PMF as bars
+    _ax.bar(_x, _pmf, color='royalblue', alpha=0.7, label=f'PMF: P(X=k)')
+
+    # Add a line
+    _ax.plot(_x, _pmf, 'ro-', alpha=0.6, label='PMF line')
+
+    # Add vertical lines
+    _ax.axvline(x=_mean, color='green', linestyle='--', linewidth=2, 
+               label=f'Mean: {_mean:.2f}')
+
+    # Shade the stdev region
+    _ax.axvspan(_mean - _std_dev, _mean + _std_dev, alpha=0.2, color='green',
+               label=f'±1 Std Dev: {_std_dev:.2f}')
+
+    # Add labels and title
+    _ax.set_xlabel('Number of Successes (k)')
+    _ax.set_ylabel('Probability: P(X=k)')
+    _ax.set_title(f'Binomial Distribution with n={_n}, p={_p:.2f}')
+
+    # Annotations
+    _ax.annotate(f'E[X] = {_mean:.2f}', 
+                xy=(_mean, stats.binom.pmf(int(_mean), _n, _p)), 
+                xytext=(_mean + 1, max(_pmf) * 0.8),
+                arrowprops=dict(facecolor='black', shrink=0.05, width=1))
+
+    _ax.annotate(f'Var(X) = {_variance:.2f}', 
+                xy=(_mean, stats.binom.pmf(int(_mean), _n, _p) / 2), 
+                xytext=(_mean + 1, max(_pmf) * 0.6),
+                arrowprops=dict(facecolor='black', shrink=0.05, width=1))
+
+    # Grid and legend
+    _ax.grid(alpha=0.3)
+    _ax.legend()
+
+    plt.tight_layout()
+    plt.gca()
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Relationship to Bernoulli Random Variables
+
+        One way to think of the binomial is as the sum of $n$ Bernoulli variables. Say that $Y_i$ is an indicator Bernoulli random variable which is 1 if experiment $i$ is a success. Then if $X$ is the total number of successes in $n$ experiments, $X \sim \text{Bin}(n, p)$:
+
+        $$X = \sum_{i=1}^n Y_i$$
+
+        Recall that the outcome of $Y_i$ will be 1 or 0, so one way to think of $X$ is as the sum of those 1s and 0s.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Binomial Probability Mass Function (PMF)
+
+        The most important property to know about a binomial is its [Probability Mass Function](https://marimo.app/https://github.com/marimo-team/learn/blob/main/probability/10_probability_mass_function.py):
+
+        $$P(X=k) = {n \choose k}p^k(1-p)^{n-k}$$
+
+        ```
+        P(X = k) = (n) p^k(1-p)^(n-k)
+         ↑           (k)
+         |            ↑
+         |            +-- Binomial coefficient:
+         |                number of ways to choose
+         |                k successes from n trials
+         |
+        Probability that our
+        variable takes on the
+        value k
+        ```
+
+        Recall, we derived this formula in Part 1. There is a complete example on the probability of $k$ heads in $n$ coin flips, where each flip is heads with probability $p$.
+
+        To briefly review, if you think of each experiment as being distinct, then there are ${n \choose k}$ ways of permuting $k$ successes from $n$ experiments. For any of the mutually exclusive permutations, the probability of that permutation is $p^k \cdot (1-p)^{n-k}$.
+
+        The name binomial comes from the term ${n \choose k}$ which is formally called the binomial coefficient.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Expectation of Binomial
+
+        There is an easy way to calculate the expectation of a binomial and a hard way. The easy way is to leverage the fact that a binomial is the sum of Bernoulli indicator random variables $X = \sum_{i=1}^{n} Y_i$ where $Y_i$ is an indicator of whether the $i$-th experiment was a success: $Y_i \sim \text{Bernoulli}(p)$. 
+
+        Since the [expectation of the sum](http://marimo.app/https://github.com/marimo-team/learn/blob/main/probability/11_expectation.py) of random variables is the sum of expectations, we can add the expectation, $E[Y_i] = p$, of each of the Bernoulli's:
+
+        \begin{align}
+        E[X] &= E\Big[\sum_{i=1}^{n} Y_i\Big] && \text{Since }X = \sum_{i=1}^{n} Y_i \\
+        &= \sum_{i=1}^{n}E[ Y_i] && \text{Expectation of sum} \\
+        &= \sum_{i=1}^{n}p && \text{Expectation of Bernoulli} \\
+        &= n \cdot p && \text{Sum $n$ times}
+        \end{align}
+
+        The hard way is to use the definition of expectation:
+
+        \begin{align}
+        E[X] &= \sum_{i=0}^n i \cdot P(X = i) && \text{Def of expectation} \\
+        &= \sum_{i=0}^n i \cdot {n \choose i} p^i(1-p)^{n-i} && \text{Sub in PMF} \\
+        & \cdots && \text{Many steps later} \\
+        &= n \cdot p
+        \end{align}
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Binomial Distribution in Python
+
+        As you might expect, you can use binomial distributions in code. The standardized library for binomials is `scipy.stats.binom`.
+
+        One of the most helpful methods that this package provides is a way to calculate the PMF. For example, say $n=5$, $p=0.6$ and you want to find $P(X=2)$, you could use the following code:
+        """
+    )
+    return
+
+
+@app.cell
+def _(stats):
+    # define variables for x, n, and p
+    _n = 5  # Integer value for n
+    _p = 0.6
+    _x = 2
+
+    # use scipy to compute the pmf
+    p_x = stats.binom.pmf(_x, _n, _p)
+
+    # use the probability for future work
+    print(f'P(X = {_x}) = {p_x:.4f}')
+    return (p_x,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""Another particularly helpful function is the ability to generate a random sample from a binomial. For example, say $X$ represents the number of requests to a website. We can draw 100 samples from this distribution using the following code:""")
+    return
+
+
+@app.cell
+def _(n, p, stats):
+    n_int = int(n)
+
+    # samples from the binomial distribution
+    samples = stats.binom.rvs(n_int, p, size=100)
+
+    # Print the samples
+    print(samples)
+    return n_int, samples
+
+
+@app.cell(hide_code=True)
+def _(n_int, np, p, plt, samples, stats):
+    # Plot histogram of samples
+    plt.figure(figsize=(10, 5))
+    plt.hist(samples, bins=np.arange(-0.5, n_int+1.5, 1), alpha=0.7, color='royalblue', 
+             edgecolor='black', density=True)
+
+    # Overlay the PMF
+    x_values = np.arange(0, n_int+1)
+    pmf_values = stats.binom.pmf(x_values, n_int, p)
+    plt.plot(x_values, pmf_values, 'ro-', ms=8, label='Theoretical PMF')
+
+    # Add labels and title
+    plt.xlabel('Number of Successes')
+    plt.ylabel('Relative Frequency / Probability')
+    plt.title(f'Histogram of 100 Samples from Bin({n_int}, {p})')
+    plt.legend()
+    plt.grid(alpha=0.3)
+
+    # Annotate
+    plt.annotate('Sample mean: %.2f' % np.mean(samples), 
+                xy=(0.7, 0.9), xycoords='axes fraction',
+                bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.3))
+    plt.annotate('Theoretical mean: %.2f' % (n_int*p), 
+                xy=(0.7, 0.8), xycoords='axes fraction',
+                bbox=dict(boxstyle='round,pad=0.5', fc='lightgreen', alpha=0.3))
+
+    plt.tight_layout()
+    plt.gca()
+    return pmf_values, x_values
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        You might be wondering what a random sample is! A random sample is a randomly chosen assignment for our random variable. Above we have 100 such assignments. The probability that value $k$ is chosen is given by the PMF: $P(X=k)$. 
+
+        There are also functions for getting the mean, the variance, and more. You can read the [scipy.stats.binom documentation](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.binom.html), especially the list of methods.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Interactive Exploration of Binomial vs. Negative Binomial
+
+        The standard binomial distribution is a special case of a broader family of distributions. One related distribution is the negative binomial, which can model count data with overdispersion (where the variance is larger than the mean).
+
+        Below, you can explore how the negative binomial distribution compares to a Poisson distribution (which can be seen as a limiting case of the binomial as $n$ gets large and $p$ gets small, with $np$ held constant).
+
+        Adjust the sliders to see how the parameters affect the distribution:
+
+        *Note: The interactive visualization in this section was inspired by work from [liquidcarbon on GitHub](https://github.com/liquidcarbon).*
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(alpha_slider, chart, equation, mo, mu_slider):
+    mo.vstack(
+        [
+            mo.md(f"## Negative Binomial Distribution (Poisson + Overdispersion)\n{equation}"),
+            mo.hstack([mu_slider, alpha_slider], justify="start"),
+            chart,
+        ], justify='space-around'
+    ).center()
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## 🤔 Test Your Understanding
+        Pick which of these statements about binomial distributions you think are correct:
+
+        /// details | The variance of a binomial distribution is always equal to its mean
+        ❌ Incorrect! The variance is $np(1-p)$ while the mean is $np$. They're only equal when $p=1$ (which is a degenerate case).
+        ///
+
+        /// details | If $X \sim \text{Bin}(n, p)$ and $Y \sim \text{Bin}(n, 1-p)$, then $X$ and $Y$ have the same variance
+        ✅ Correct! $\text{Var}(X) = np(1-p)$ and $\text{Var}(Y) = n(1-p)p$, which are the same.
+        ///
+
+        /// details | As the number of trials increases, the binomial distribution approaches a normal distribution
+        ✅ Correct! For large $n$, the binomial distribution can be approximated by a normal distribution with the same mean and variance.
+        ///
+
+        /// details | The PMF of a binomial distribution is symmetric when $p = 0.5$
+        ✅ Correct! When $p = 0.5$, the PMF is symmetric around $n/2$.
+        ///
+
+        /// details | The sum of two independent binomial random variables with the same $p$ is also a binomial random variable
+        ✅ Correct! If $X \sim \text{Bin}(n_1, p)$ and $Y \sim \text{Bin}(n_2, p)$ are independent, then $X + Y \sim \text{Bin}(n_1 + n_2, p)$.
+        ///
+
+        /// details | The maximum value of the PMF for $\text{Bin}(n,p)$ always occurs at $k = np$
+        ❌ Incorrect! The mode (maximum value of PMF) is either $\lfloor (n+1)p \rfloor$ or $\lceil (n+1)p-1 \rceil$ depending on whether $(n+1)p$ is an integer.
+        ///
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Summary
+
+        So we've explored the binomial distribution, and honestly, it's one of the most practical probability distributions you'll encounter. Think about it — anytime you're counting successes in a fixed number of trials (like those coin flips we discussed), this is your go-to distribution.
+
+        I find it fascinating how the expectation is simply $np$. Such a clean, intuitive formula! And remember that neat visualization we saw earlier? When we adjusted the parameters, you could actually see how the distribution shape changes—becoming more symmetric as $n$ increases.
+
+        The key things to take away:
+
+        - The binomial distribution models the number of successes in $n$ independent trials, each with probability $p$ of success
+
+        - Its PMF is given by the formula $P(X=k) = {n \choose k}p^k(1-p)^{n-k}$, which lets us calculate exactly how likely any specific number of successes is
+
+        - The expected value is $E[X] = np$ and the variance is $Var(X) = np(1-p)$
+
+        - It's related to other distributions: it's essentially a sum of Bernoulli random variables, and connects to both the negative binomial and Poisson distributions
+
+        - In Python, the `scipy.stats.binom` module makes working with binomial distributions straightforward—you can generate random samples and calculate probabilities with just a few lines of code
+
+        You'll see the binomial distribution pop up everywhere—from computer science to quality control, epidemiology, and data science. Any time you have scenarios with binary outcomes over multiple trials, this distribution has you covered.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""Appendix code (helper functions, variables, etc.):""")
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+@app.cell(hide_code=True)
+def _():
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import scipy.stats as stats
+    import pandas as pd
+    import altair as alt
+    from wigglystuff import TangleSlider
+    return TangleSlider, alt, np, pd, plt, stats
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    alpha_slider = mo.ui.slider(
+        value=0.1,
+        steps=[0, 0.01, 0.02, 0.03, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1],
+        label="α (overdispersion)",
+        show_value=True,
+    )
+    mu_slider = mo.ui.slider(
+        value=100, start=1, stop=100, step=1, label="μ (mean)", show_value=True
+    )
+    return alpha_slider, mu_slider
+
+
+@app.cell(hide_code=True)
+def _():
+    equation = """
+    $$
+    P(X = k) = \\frac{\\Gamma(k + \\frac{1}{\\alpha})}{\\Gamma(k + 1) \\Gamma(\\frac{1}{\\alpha})} \\left( \\frac{1}{\\mu \\alpha + 1} \\right)^{\\frac{1}{\\alpha}} \\left( \\frac{\\mu \\alpha}{\\mu \\alpha + 1} \\right)^k
+    $$
+
+    $$
+    \\sigma^2 = \\mu + \\alpha \\mu^2
+    $$
+    """
+    return (equation,)
+
+
+@app.cell(hide_code=True)
+def _(alpha_slider, alt, mu_slider, np, pd, stats):
+    mu = mu_slider.value
+    alpha = alpha_slider.value
+    n = 1000 - mu if alpha == 0 else 1 / alpha
+    p = n / (mu + n)
+    x = np.arange(0, mu * 3 + 1, 1)
+    df = pd.DataFrame(
+        {
+            "x": x,
+            "y": stats.nbinom.pmf(x, n, p),
+            "y_poi": stats.nbinom.pmf(x, 1000 - mu, 1 - mu / 1000),
+        }
+    )
+    r1k = stats.nbinom.rvs(n, p, size=1000)
+    df["in 95% CI"] = df["x"].between(*np.percentile(r1k, q=[2.5, 97.5]))
+    base = alt.Chart(df)
+
+    chart_poi = base.mark_bar(
+        fillOpacity=0.25, width=100 / mu, fill="magenta"
+    ).encode(
+        x=alt.X("x").scale(domain=(-0.4, x.max() + 0.4), nice=False),
+        y=alt.Y("y_poi").scale(domain=(0, df.y_poi.max() * 1.1)).title(None),
+    )
+    chart_nb = base.mark_bar(fillOpacity=0.75, width=100 / mu).encode(
+        x="x",
+        y="y",
+        fill=alt.Fill("in 95% CI")
+        .scale(domain=[False, True], range=["#aaa", "#7c7"])
+        .legend(orient="bottom-right"),
+    )
+
+    chart = (chart_poi + chart_nb).configure_view(continuousWidth=450)
+    return alpha, base, chart, chart_nb, chart_poi, df, mu, n, p, r1k, x
+
+
+if __name__ == "__main__":
+    app.run()