Merge pull request #39 from marimo-team/haleshot/09_random_variables

probability/09_random_variables.py

@@ -0,0 +1,552 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "marimo",
+#     "matplotlib==3.10.0",
+#     "numpy==2.2.3",
+#     "scipy==1.15.2",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.11.10"
+app = marimo.App(width="medium", app_title="Random Variables")
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+@app.cell
+def _():
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from scipy import stats
+    return np, plt, stats
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        # Random Variables
+
+        _This notebook is a computational companion to ["Probability for Computer Scientists"](https://chrispiech.github.io/probabilityForComputerScientists/en/part2/rvs/), by Stanford professor Chris Piech._
+
+        Random variables are functions that map outcomes from a probability space to numbers. This mathematical abstraction allows us to:
+
+        - Work with numerical outcomes in probability
+        - Calculate expected values and variances
+        - Model real-world phenomena quantitatively
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Types of Random Variables
+
+        ### Discrete Random Variables
+        - Take on countable values (finite or infinite)
+        - Described by a probability mass function (PMF)
+        - Example: Number of heads in 3 coin flips
+
+        ### Continuous Random Variables
+        - Take on uncountable values in an interval
+        - Described by a probability density function (PDF)
+        - Example: Height of a randomly selected person
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Properties of Random Variables
+
+        Each random variable has several key properties:
+
+        | Property | Description | Example |
+        |----------|-------------|---------|
+        | Meaning | Semantic description | Number of successes in n trials |
+        | Symbol | Notation used | $X$, $Y$, $Z$ |
+        | Support/Range | Possible values | $\{0,1,2,...,n\}$ for binomial |
+        | Distribution | PMF or PDF | $p_X(x)$ or $f_X(x)$ |
+        | Expectation | Weighted average | $E[X]$ |
+        | Variance | Measure of spread | $\text{Var}(X)$ |
+        | Standard Deviation | Square root of variance | $\sigma_X$ |
+        | Mode | Most likely value | argmax$_x$ $p_X(x)$ |
+
+        Additional properties include:
+
+        - [Entropy](https://en.wikipedia.org/wiki/Entropy_(information_theory)) (measure of uncertainty)
+        - [Median](https://en.wikipedia.org/wiki/Median) (middle value)
+        - [Skewness](https://en.wikipedia.org/wiki/Skewness) (asymmetry measure)
+        - [Kurtosis](https://en.wikipedia.org/wiki/Kurtosis) (tail heaviness measure)
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Probability Mass Functions (PMF)
+
+        For discrete random variables, the PMF $p_X(x)$ gives the probability that $X$ equals $x$:
+
+        $p_X(x) = P(X = x)$
+
+        Properties of a PMF:
+
+        1. $p_X(x) \geq 0$ for all $x$
+        2. $\sum_x p_X(x) = 1$
+
+        Let's implement a PMF for rolling a fair die:
+        """
+    )
+    return
+
+
+@app.cell
+def _(np, plt):
+    def die_pmf(x):
+        if x in [1, 2, 3, 4, 5, 6]:
+            return 1 / 6
+        return 0
+
+    # Plot the PMF
+    _x = np.arange(1, 7)
+    probabilities = [die_pmf(i) for i in _x]
+
+    plt.figure(figsize=(8, 2))
+    plt.bar(_x, probabilities)
+    plt.title("PMF of Rolling a Fair Die")
+    plt.xlabel("Outcome")
+    plt.ylabel("Probability")
+    plt.grid(True, alpha=0.3)
+    plt.gca()
+    return die_pmf, probabilities
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Probability Density Functions (PDF)
+
+        For continuous random variables, we use a PDF $f_X(x)$. The probability of $X$ falling in an interval $[a,b]$ is:
+
+        $P(a \leq X \leq b) = \int_a^b f_X(x)dx$
+
+        Properties of a PDF:
+
+        1. $f_X(x) \geq 0$ for all $x$
+        2. $\int_{-\infty}^{\infty} f_X(x)dx = 1$
+
+        Let's look at the normal distribution, a common continuous random variable:
+        """
+    )
+    return
+
+
+@app.cell
+def _(np, plt, stats):
+    # Generate points for plotting
+    _x = np.linspace(-4, 4, 100)
+    _pdf = stats.norm.pdf(_x, loc=0, scale=1)
+
+    plt.figure(figsize=(8, 4))
+    plt.plot(_x, _pdf, "b-", label="PDF")
+    plt.fill_between(_x, _pdf, where=(_x >= -1) & (_x <= 1), alpha=0.3)
+    plt.title("Standard Normal Distribution")
+    plt.xlabel("x")
+    plt.ylabel("Density")
+    plt.grid(True, alpha=0.3)
+    plt.legend()
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Expected Value
+
+        The expected value $E[X]$ is the long-run average of a random variable.
+
+        For discrete random variables:
+        $E[X] = \sum_x x \cdot p_X(x)$
+
+        For continuous random variables:
+        $E[X] = \int_{-\infty}^{\infty} x \cdot f_X(x)dx$
+
+        Properties:
+
+        1. $E[aX + b] = aE[X] + b$
+        2. $E[X + Y] = E[X] + E[Y]$
+        """
+    )
+    return
+
+
+@app.cell
+def _(np):
+    def expected_value_discrete(x_values, probabilities):
+        return sum(x * p for x, p in zip(x_values, probabilities))
+
+    # Example: Expected value of a fair die roll
+    die_values = np.arange(1, 7)
+    die_probs = np.ones(6) / 6
+
+    E_X = expected_value_discrete(die_values, die_probs)
+    return E_X, die_probs, die_values, expected_value_discrete
+
+
+@app.cell
+def _(E_X):
+    E_X
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Variance
+
+        The variance $\text{Var}(X)$ measures the spread of a random variable around its mean:
+
+        $\text{Var}(X) = E[(X - E[X])^2]$
+
+        This can be computed as:
+        $\text{Var}(X) = E[X^2] - (E[X])^2$
+
+        Properties:
+
+        1. $\text{Var}(aX) = a^2Var(X)$
+        2. $\text{Var}(X + b) = Var(X)$
+        """
+    )
+    return
+
+
+@app.cell
+def _(E_X, die_probs, die_values, np):
+    def variance_discrete(x_values, probabilities, expected_value):
+        squared_diff = [(x - expected_value) ** 2 for x in x_values]
+        return sum(d * p for d, p in zip(squared_diff, probabilities))
+
+    # Example: Variance of a fair die roll
+    var_X = variance_discrete(die_values, die_probs, E_X)
+    std_X = np.sqrt(var_X)
+    return std_X, var_X, variance_discrete
+
+
+@app.cell(hide_code=True)
+def _(mo, std_X, var_X):
+    mo.md(
+        f"""
+        ### Examples of Variance Calculation
+
+        For our fair die example:
+
+        - Variance: {var_X:.2f}
+        - Standard Deviation: {std_X:.2f}
+
+        This means that on average, a roll deviates from the mean (3.5) by about {std_X:.2f} units.
+
+        Let's look another example for a fair coin:
+        """
+    )
+    return
+
+
+@app.cell
+def _(variance_discrete):
+    # Fair coin (X = 0 or 1)
+    coin_values = [0, 1]
+    coin_probs = [0.5, 0.5]
+    coin_mean = sum(x * p for x, p in zip(coin_values, coin_probs))
+    coin_var = variance_discrete(coin_values, coin_probs, coin_mean)
+    return coin_mean, coin_probs, coin_values, coin_var
+
+
+@app.cell
+def _(np, stats, variance_discrete):
+    # Standard normal (discretized for example)
+    normal_values = np.linspace(-3, 3, 100)
+    normal_probs = stats.norm.pdf(normal_values)
+    normal_probs = normal_probs / sum(normal_probs)  # normalize
+    normal_mean = 0
+    normal_var = variance_discrete(normal_values, normal_probs, normal_mean)
+    return normal_mean, normal_probs, normal_values, normal_var
+
+
+@app.cell
+def _(np, variance_discrete):
+    # Uniform on [0,1] (discretized for example)
+    uniform_values = np.linspace(0, 1, 100)
+    uniform_probs = np.ones_like(uniform_values) / len(uniform_values)
+    uniform_mean = 0.5
+    uniform_var = variance_discrete(uniform_values, uniform_probs, uniform_mean)
+    return uniform_mean, uniform_probs, uniform_values, uniform_var
+
+
+@app.cell(hide_code=True)
+def _(coin_var, mo, normal_var, uniform_var):
+    mo.md(
+        rf"""
+        Let's look at some calculated variances:
+
+        - Fair coin (X = 0 or 1): $\text{{Var}}(X) = {coin_var:.4f}$
+        - Standard normal distribution (discretized): $\text{{Var(X)}} ≈ {normal_var:.4f}$
+        - Uniform distribution on $[0,1]$ (discretized): $\text{{Var(X)}} ≈ {uniform_var:.4f}$
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Common Distributions
+
+        1. Bernoulli Distribution
+            - Models a single success/failure experiment
+            - $P(X = 1) = p$, $P(X = 0) = 1-p$
+            - $E[X] = p$, $\text{Var}(X) = p(1-p)$
+
+        2. Binomial Distribution
+
+            - Models number of successes in $n$ independent trials
+            - $P(X = k) = \binom{n}{k}p^k(1-p)^{n-k}$
+            - $E[X] = np$, $\text{Var}(X) = np(1-p)$
+
+        3. Normal Distribution
+
+            - Bell-shaped curve defined by mean $\mu$ and variance $\sigma^2$
+            - PDF: $f_X(x) = \frac{1}{\sigma\sqrt{2\pi}}e^{-\frac{(x-\mu)^2}{2\sigma^2}}$
+            - $E[X] = \mu$, $\text{Var}(X) = \sigma^2$
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ### Example: Comparing Discrete and Continuous Distributions
+
+        This example shows the relationship between a Binomial distribution (discrete) and its Normal approximation (continuous).
+        The parameters control both distributions:
+
+        - **Number of Trials**: Controls the range of possible values and the shape's width
+        - **Success Probability**: Affects the distribution's center and skewness
+        """
+    )
+    return
+
+
+@app.cell
+def _(mo, n_trials, p_success):
+    mo.hstack([n_trials, p_success], justify="space-around")
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    # Distribution parameters
+    n_trials = mo.ui.slider(1, 20, value=10, label="Number of Trials")
+    p_success = mo.ui.slider(0, 1, value=0.5, step=0.05, label="Success Probability")
+    return n_trials, p_success
+
+
+@app.cell(hide_code=True)
+def _(n_trials, np, p_success, plt, stats):
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 3))
+
+    # Discrete: Binomial PMF
+    k = np.arange(0, n_trials.value + 1)
+    pmf = stats.binom.pmf(k, n_trials.value, p_success.value)
+    ax1.bar(k, pmf, alpha=0.8, color="#1f77b4", label="PMF")
+    ax1.set_title(f"Binomial PMF (n={n_trials.value}, p={p_success.value})")
+    ax1.set_xlabel("Number of Successes")
+    ax1.set_ylabel("Probability")
+    ax1.grid(True, alpha=0.3)
+
+    # Continuous: Normal PDF approx.
+    mu = n_trials.value * p_success.value
+    sigma = np.sqrt(n_trials.value * p_success.value * (1 - p_success.value))
+    x = np.linspace(max(0, mu - 4 * sigma), min(n_trials.value, mu + 4 * sigma), 100)
+    pdf = stats.norm.pdf(x, mu, sigma)
+
+    ax2.plot(x, pdf, "r-", linewidth=2, label="PDF")
+    ax2.fill_between(x, pdf, alpha=0.3, color="red")
+    ax2.set_title(f"Normal PDF (μ={mu:.1f}, σ={sigma:.1f})")
+    ax2.set_xlabel("Continuous Approximation")
+    ax2.set_ylabel("Density")
+    ax2.grid(True, alpha=0.3)
+
+    # Set consistent x-axis limits for better comparison
+    ax1.set_xlim(-0.5, n_trials.value + 0.5)
+    ax2.set_xlim(-0.5, n_trials.value + 0.5)
+
+    plt.tight_layout()
+    plt.gca()
+    return ax1, ax2, fig, k, mu, pdf, pmf, sigma, x
+
+
+@app.cell(hide_code=True)
+def _(mo, n_trials, np, p_success):
+    mo.md(f"""
+    **Current Distribution Properties:**
+
+    - Mean (μ) = {n_trials.value * p_success.value:.2f}
+    - Standard Deviation (σ) = {np.sqrt(n_trials.value * p_success.value * (1 - p_success.value)):.2f}
+
+    Notice how the Normal distribution (right) approximates the Binomial distribution (left) better when:
+
+    1. The number of trials is larger
+    2. The success probability is closer to 0.5
+    """)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Practice Problems
+
+        ### Problem 1: Discrete Random Variable
+        Let $X$ be the sum when rolling two fair dice. Find:
+
+        1. The support of $X$
+        2. The PMF $p_X(x)$
+        3. $E[X]$ and $\text{Var}(X)$
+
+        <details>
+        <summary>Solution</summary>
+        Let's solve this step by step:
+        ```python
+        def two_dice_pmf(x):
+            outcomes = [(i,j) for i in range(1,7) for j in range(1,7)]
+            favorable = [pair for pair in outcomes if sum(pair) == x]
+            return len(favorable)/36
+
+        # Support: {2,3,...,12}
+        # E[X] = 7
+        # Var(X) = 5.83
+        ```
+        </details>
+
+        ### Problem 2: Continuous Random Variable
+        For a uniform random variable on $[0,1]$, verify that:
+
+        1. The PDF integrates to 1
+        2. $E[X] = 1/2$
+        3. $\text{Var}(X) = 1/12$
+
+        Try solving this yourself first, then check the solution below.
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    # DIY
+    return
+
+
+@app.cell(hide_code=True)
+def _(mktext, mo):
+    mo.accordion({"Solution": mktext}, lazy=True)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mktext = mo.md(
+        r"""
+        Let's solve each part:
+
+        1. **PDF integrates to 1**:
+           $\int_0^1 1 \, dx = [x]_0^1 = 1 - 0 = 1$
+
+        2. **Expected Value**:
+           $E[X] = \int_0^1 x \cdot 1 \, dx = [\frac{x^2}{2}]_0^1 = \frac{1}{2} - 0 = \frac{1}{2}$
+
+        3. **Variance**:
+           $\text{Var}(X) = E[X^2] - (E[X])^2$
+
+           First calculate $E[X^2]$:
+           $E[X^2] = \int_0^1 x^2 \cdot 1 \, dx = [\frac{x^3}{3}]_0^1 = \frac{1}{3}$
+
+           Then:
+           $\text{Var}(X) = \frac{1}{3} - (\frac{1}{2})^2 = \frac{1}{3} - \frac{1}{4} = \frac{1}{12}$
+        """
+    )
+    return (mktext,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## 🤔 Test Your Understanding
+
+        Pick which of these statements about random variables you think are correct:
+
+        <details>
+        <summary>The probability density function can be greater than 1</summary>
+        ✅ Correct! Unlike PMFs, PDFs can exceed 1 as long as the total area equals 1.
+        </details>
+
+        <details>
+        <summary>The expected value of a random variable must equal one of its possible values</summary>
+        ❌ Incorrect! For example, the expected value of a fair die is 3.5, which is not a possible outcome.
+        </details>
+
+        <details>
+        <summary>Adding a constant to a random variable changes its variance</summary>
+        ❌ Incorrect! Adding a constant shifts the distribution but doesn't affect its spread.
+        </details>
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        """
+        ## Summary
+
+        You've learned:
+
+        - The difference between discrete and continuous random variables
+        - How PMFs and PDFs describe probability distributions
+        - Methods for calculating expected values and variances
+        - Properties of common probability distributions
+
+        In the next lesson, we'll explore Probability Mass Functions in detail, focusing on their properties and applications.
+        """
+    )
+    return
+
+
+if __name__ == "__main__":
+    app.run()