Merge branch 'marimo-team:main' into feat/issue#18/polars-data-wrangling

README.md

@@ -26,6 +26,7 @@ notebooks for educators, students, and practitioners.
 
 - 🎲 Probability
 - 📏 Linear algebra
+- ⚖️ Optimization
 - ❄️ Polars
 - 🔥 Pytorch
 

optimization/01_least_squares.py

@@ -0,0 +1,123 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "cvxpy==1.6.0",
+#     "marimo",
+#     "numpy==2.2.2",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.11.0"
+app = marimo.App()
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        # Least squares
+
+        In a least-squares problem, we have measurements $A \in \mathcal{R}^{m \times
+        n}$ (i.e., $m$ rows and $n$ columns) and $b \in \mathcal{R}^m$. We seek a vector
+        $x \in \mathcal{R}^{n}$ such that $Ax$ is close to $b$. The matrices $A$ and $b$ are problem data or constants, and $x$ is the variable we are solving for.
+
+        Closeness is defined as the sum of the squared differences:
+
+        \[ \sum_{i=1}^m (a_i^Tx - b_i)^2, \]
+
+        also known as the $\ell_2$-norm squared, $\|Ax - b\|_2^2$.
+
+        For example, we might have a dataset of $m$ users, each represented by $n$ features. Each row $a_i^T$ of $A$ is the feature vector for user $i$, while the corresponding entry $b_i$ of $b$ is the measurement we want to predict from $a_i^T$, such as ad spending. The prediction for user $i$ is given by $a_i^Tx$.
+
+        We find the optimal value of $x$ by solving the optimization problem
+
+        \[
+            \begin{array}{ll}
+            \text{minimize}   & \|Ax - b\|_2^2.
+            \end{array}
+        \]
+
+        Let $x^\star$ denote the optimal $x$. The quantity $r = Ax^\star - b$ is known as the residual. If $\|r\|_2 = 0$, we have a perfect fit.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Example
+
+        In this example, we use the Python library [CVXPY](https://github.com/cvxpy/cvxpy) to construct and solve a least-squares problems.
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    import cvxpy as cp
+    import numpy as np
+    return cp, np
+
+
+@app.cell
+def _():
+    m = 20
+    n = 15
+    return m, n
+
+
+@app.cell
+def _(m, n, np):
+    np.random.seed(0)
+    A = np.random.randn(m, n)
+    b = np.random.randn(m)
+    return A, b
+
+
+@app.cell
+def _(A, b, cp, n):
+    x = cp.Variable(n)
+    objective = cp.sum_squares(A @ x - b)
+    problem = cp.Problem(cp.Minimize(objective))
+    optimal_value = problem.solve()
+    return objective, optimal_value, problem, x
+
+
+@app.cell
+def _(A, b, cp, mo, optimal_value, x):
+    mo.md(
+        f"""
+        - The optimal value is **{optimal_value:.04f}**.
+        - The optimal value of $x$ is {mo.as_html(list(x.value))}
+        - The norm of the residual is **{cp.norm(A @ x - b, p=2).value:0.4f}**
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Further reading
+
+        For a primer on least squares, with many real-world examples, check out the free book
+        [Vectors, Matrices, and Least Squares](https://web.stanford.edu/~boyd/vmls/), which is used for undergraduate linear algebra education at Stanford.
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+if __name__ == "__main__":
+    app.run()

optimization/02_linear_program.py

@@ -0,0 +1,269 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = [
+#     "cvxpy==1.6.0",
+#     "marimo",
+#     "matplotlib==3.10.0",
+#     "numpy==2.2.2",
+#     "wigglystuff==0.1.9",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.11.0"
+app = marimo.App()
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        # Linear program
+
+        A linear program is an optimization problem with a linear objective and affine
+        inequality constraints. A common standard form is the following:
+
+        \[  
+            \begin{array}{ll}
+            \text{minimize}   & c^Tx \\
+            \text{subject to} & Ax \leq b.
+            \end{array}
+        \]
+
+        Here $A \in \mathcal{R}^{m \times n}$, $b \in \mathcal{R}^m$, and $c \in \mathcal{R}^n$ are problem data and $x \in \mathcal{R}^{n}$ is the optimization variable. The inequality constraint $Ax \leq b$ is elementwise.
+
+        For example, we might have $n$ different products, each constructed out of $m$ components. Each entry $A_{ij}$ is the amount of component $i$ required to build one unit of product $j$. Each entry $b_i$ is the total amount of component $i$ available. We lose $c_j$ for each unit of product $j$ ($c_j < 0$ indicates profit). Our goal then is to choose how many units of each product $j$ to make, $x_j$, in order to minimize loss without exceeding our budget for any component.
+
+        In addition to a solution $x^\star$, we obtain a dual solution $\lambda^\star$. A positive entry $\lambda^\star_i$ indicates that the constraint $a_i^Tx \leq b_i$ holds with equality for $x^\star$ and suggests that changing $b_i$ would change the optimal value.
+
+        **Why linear programming?** Linear programming is a way to achieve an optimal outcome, such as maximum utility or lowest cost, subject to a linear objective function and affine constraints. Developed in the 20th century, linear programming is widely used today to solve problems in resource allocation, scheduling, transportation, and more. The discovery of polynomial-time algorithms to solve linear programs was of tremendous worldwide importance and entered the public discourse, even making the front page of the New York Times.
+
+        In the late 20th and early 21st century, researchers generalized linear programming to a much wider class of problems called convex optimization problems. Nearly all convex optimization problems can be solved efficiently and reliably, and even more difficult problems are readily solved by a sequence of convex optimization problems. Today, convex optimization is used to fit machine learning models, land rockets in real-time at SpaceX, plan trajectories for self-driving cars at Waymo, execute many billions of dollars of financial trades a day, and much more.
+
+        This marimo learn course uses CVXPY, a modeling language for convex optimization problems developed originally at Stanford, to construct and solve convex programs.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        f"""
+        {mo.image("https://www.debugmind.com/wp-content/uploads/2020/01/spacex-1.jpg")}
+        _SpaceX solves convex optimization problems onboard to land its rockets, using CVXGEN, a code generator for quadratic programming developed at Stephen Boyd’s Stanford lab. Photo by SpaceX, licensed CC BY-NC 2.0._
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Example
+
+        Here we use CVXPY to construct and solve a linear program.
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    import cvxpy as cp
+    import numpy as np
+    return cp, np
+
+
+@app.cell(hide_code=True)
+def _(np):
+    A = np.array(
+        [
+            [0.76103773, 0.12167502],
+            [0.44386323, 0.33367433],
+            [1.49407907, -0.20515826],
+            [0.3130677, -0.85409574],
+            [-2.55298982, 0.6536186],
+            [0.8644362, -0.74216502],
+            [2.26975462, -1.45436567],
+            [0.04575852, -0.18718385],
+            [1.53277921, 1.46935877],
+            [0.15494743, 0.37816252],
+        ]
+    )
+
+    b = np.array(
+        [
+            2.05062369,
+            0.94934659,
+            0.89559424,
+            1.04389978,
+            2.45035643,
+            -0.95479445,
+            -0.83801349,
+            -0.26562529,
+            2.35763652,
+            0.98286942,
+        ]
+    )
+    return A, b
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""We've randomly generated problem data $A$ and $B$. The vector for $c$ is shown below. Try playing with the value of $c$ by dragging the components, and see how the level curves change in the visualization below.""")
+    return
+
+
+@app.cell
+def _(mo, np):
+    from wigglystuff import Matrix
+
+    c_widget = mo.ui.anywidget(Matrix(matrix=np.array([[0.1, -0.2]]), step=0.01))
+    c_widget
+    return Matrix, c_widget
+
+
+@app.cell
+def _(c_widget, np):
+    c = np.array(c_widget.value["matrix"][0])
+    return (c,)
+
+
+@app.cell
+def _(A, b, c, cp):
+    x = cp.Variable(A.shape[1])
+    prob = cp.Problem(cp.Minimize(c.T @ x), [A @ x <= b])
+    _ = prob.solve()
+    x_star = x.value
+    return prob, x, x_star
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""Below, we plot the feasible region of the problem — the intersection of the inequalities — and the level curves of the objective function. The optimal value $x^\star$ is the point farthest in the feasible region in the direction $-c$.""")
+    return
+
+
+@app.cell(hide_code=True)
+def _(A, b, c, make_plot, x_star):
+    make_plot(A, b, c, x_star)
+    return
+
+
+@app.cell(hide_code=True)
+def _(np):
+    import matplotlib.pyplot as plt
+
+
+    def make_plot(A, b, c, x_star):
+        # Define a grid over a region that covers the feasible set.
+        # You might need to adjust these limits.
+        x_vals = np.linspace(-1, 1, 400)
+        y_vals = np.linspace(1, 3, 400)
+        X, Y = np.meshgrid(x_vals, y_vals)
+
+        # Flatten the grid points into an (N,2) array.
+        points = np.vstack([X.ravel(), Y.ravel()]).T
+
+        # For each point, check if it satisfies all the constraints: A @ x <= b.
+        # A dot product: shape of A @ x.T will be (m, N). We add a little tolerance.
+        feasible = np.all(np.dot(A, points.T) <= (b[:, None] + 1e-8), axis=0)
+        feasible = feasible.reshape(X.shape)
+
+        # Create the figure with a white background.
+        fig = plt.figure(figsize=(8, 6), facecolor="white")
+        ax = fig.add_subplot(111)
+        ax.set_facecolor("white")
+
+        # Plot the feasible region.
+        # Since "feasible" is a boolean array (False=0, True=1), we set contour levels so that
+        # the region with value 1 (feasible) gets filled.
+        ax.contourf(
+            X,
+            Y,
+            feasible.astype(float),
+            levels=[-0.5, 0.5, 1.5],
+            colors=["white", "gray"],
+            alpha=0.5,
+        )
+
+        # Plot level curves of the objective function c^T x.
+        # Compute c^T x over the grid:
+        Z = c[0] * X + c[1] * Y
+        # Choose several levels for the iso-cost lines.
+        levels = np.linspace(np.min(Z), np.max(Z), 20)
+        contours = ax.contour(
+            X, Y, Z, levels=levels, colors="gray", linestyles="--", linewidths=1
+        )
+
+        # Draw the vector -c as an arrow starting at x_star.
+        norm_c = np.linalg.norm(c)
+        if norm_c > 0:
+            head_width = norm_c * 0.1
+            head_length = norm_c * 0.1
+            # The arrow starts at x_star and points in the -c direction.
+            ax.arrow(
+                x_star[0],
+                x_star[1],
+                -c[0],
+                -c[1],
+                head_width=head_width,
+                head_length=head_length,
+                fc="black",
+                ec="black",
+                length_includes_head=True,
+            )
+            # Label the arrow near its tip.
+            ax.text(
+                x_star[0] - c[0] * 1.05,
+                x_star[1] - c[1] * 1.05,
+                r"$-c$",
+                color="black",
+                fontsize=12,
+            )
+
+        # Optionally, mark and label the point x_star.
+        ax.plot(x_star[0], x_star[1], "ko", markersize=5)
+        ax.text(
+            x_star[0],
+            x_star[1],
+            r"$\mathbf{x}^\star$",
+            color="black",
+            fontsize=12,
+            verticalalignment="bottom",
+            horizontalalignment="right",
+        )
+        # Label the axes and set title.
+        ax.set_xlabel("$x_1$")
+        ax.set_ylabel("$x_2$")
+        ax.set_title("Feasible Region and Level Curves of $c^Tx$")
+        ax.set_xlim(np.min(x_vals), np.max(x_vals))
+        ax.set_ylim(np.min(y_vals), np.max(y_vals))
+        return ax
+    return make_plot, plt
+
+
+@app.cell(hide_code=True)
+def _(mo, prob, x):
+    mo.md(
+        f"""
+        The optimal value is {prob.value:.04f}.
+        
+        A solution $x$ is {mo.as_html(list(x.value))}
+        A dual solution is is {mo.as_html(list(prob.constraints[0].dual_value))}
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+if __name__ == "__main__":
+    app.run()

optimization/03_minimum_fuel_optimal_control.py

@@ -0,0 +1,172 @@
+import marimo
+
+__generated_with = "0.11.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        # Minimal fuel optimal control
+
+        This notebook includes an application of linear programming to controlling a
+        physical system, adapted from [Convex
+        Optimization](https://web.stanford.edu/~boyd/cvxbook/) by Boyd and Vandenberghe.
+
+        We consider a linear dynamical system with state $x(t) \in \mathbf{R}^n$, for $t = 0, \ldots, T$. At each time step $t = 0, \ldots, T - 1$, an actuator or input signal $u(t)$ is applied, affecting the state. The dynamics
+        of the system is given by the linear recurrence
+
+        \[
+            x(t + 1) = Ax(t) + bu(t), \quad t = 0, \ldots, T - 1,
+        \]
+
+        where $A \in \mathbf{R}^{n \times n}$ and $b \in \mathbf{R}^n$ are given and encode how the system evolves. The initial state $x(0)$ is also given.
+
+        The _minimum fuel optimal control problem_ is to choose the inputs $u(0), \ldots, u(T - 1)$ so as to achieve
+        a given desired state $x_\text{des} = x(T)$ while minimizing the total fuel consumed
+
+        \[
+        F = \sum_{t=0}^{T - 1} f(u(t)).
+        \]
+
+        The function $f : \mathbf{R} \to \mathbf{R}$ tells us how much fuel is consumed as a function of the input, and is given by
+
+        \[
+            f(a) = \begin{cases}
+            |a| & |a| \leq 1 \\
+            2|a| - 1 & |a| > 1.
+            \end{cases}
+        \]
+
+        This means the fuel use is proportional to the magnitude of the signal between $-1$ and $1$, but for larger signals the marginal fuel efficiency is half.
+
+        **This notebook.** In this notebook we use CVXPY to formulate the minimum fuel optimal control problem as a linear program. The notebook lets you play with the initial and target states, letting you see how they affect the planned trajectory of inputs $u$.
+
+        First, we create the **problem data**.
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    import numpy as np
+    return (np,)
+
+
+@app.cell
+def _():
+    n, T = 3, 30
+    return T, n
+
+
+@app.cell
+def _(np):
+    A = np.array([[-1, 0.4, 0.8], [1, 0, 0], [0, 1, 0]])
+    b = np.array([[1, 0, 0.3]]).T
+    return A, b
+
+
+@app.cell(hide_code=True)
+def _(mo, n, np):
+    import wigglystuff
+
+    x0_widget = mo.ui.anywidget(wigglystuff.Matrix(np.zeros((1, n))))
+    xdes_widget = mo.ui.anywidget(wigglystuff.Matrix(np.array([[7, 2, -6]])))
+
+    _a = mo.md(
+        rf"""
+
+        Choose a value for $x_0$ ...
+        
+        {x0_widget}
+        """
+    )
+
+    _b = mo.md(
+        rf"""
+        ... and for $x_\text{{des}}$
+
+        {xdes_widget}
+        """
+    )
+
+    mo.hstack([_a, _b], justify="space-around")
+    return wigglystuff, x0_widget, xdes_widget
+
+
+@app.cell
+def _(x0_widget, xdes_widget):
+    x0 = x0_widget.matrix
+    xdes = xdes_widget.matrix
+    return x0, xdes
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""**Next, we specify the problem as a linear program using CVXPY.** This problem is linear because the objective and constraints are affine. (In fact, the objective is piecewise affine, but CVXPY rewrites it to be affine for you.)""")
+    return
+
+
+@app.cell
+def _():
+    import cvxpy as cp
+    return (cp,)
+
+
+@app.cell
+def _(A, T, b, cp, mo, n, x0, xdes):
+    X, u = cp.Variable(shape=(n, T + 1)), cp.Variable(shape=(1, T))
+
+    objective = cp.sum(cp.maximum(cp.abs(u), 2 * cp.abs(u) - 1))
+    constraints = [
+        X[:, 1:] == A @ X[:, :-1] + b @ u,
+        X[:, 0] == x0,
+        X[:, -1] == xdes,
+    ]
+
+    fuel_used = cp.Problem(cp.Minimize(objective), constraints).solve()
+    mo.md(f"Achieved a fuel usage of {fuel_used:.02f}. 🚀")
+    return X, constraints, fuel_used, objective, u
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        """
+        Finally, we plot the chosen inputs over time.
+
+        **🌊 Try it!** Change the initial and desired states; how do fuel usage and controls change? Can you explain what you see? You can also try experimenting with the value of $T$.
+        """
+    )
+    return
+
+
+@app.cell
+def _(plot_solution, u):
+    plot_solution(u)
+    return
+
+
+@app.cell
+def _(T, cp, np):
+    def plot_solution(u: cp.Variable):
+        import matplotlib.pyplot as plt
+
+        plt.step(np.arange(T), u.T.value)
+        plt.axis("tight")
+        plt.xlabel("$t$")
+        plt.ylabel("$u$")
+        return plt.gca()
+    return (plot_solution,)
+
+
+if __name__ == "__main__":
+    app.run()

optimization/04_quadratic_program.py

@@ -0,0 +1,265 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = [
+#     "cvxpy==1.6.0",
+#     "marimo",
+#     "matplotlib==3.10.0",
+#     "numpy==2.2.2",
+#     "wigglystuff==0.1.9",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.11.0"
+app = marimo.App()
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        # Quadratic program
+
+        A quadratic program is an optimization problem with a quadratic objective and
+        affine equality and inequality constraints. A common standard form is the
+        following:
+
+        \[
+            \begin{array}{ll}
+            \text{minimize}   & (1/2)x^TPx + q^Tx\\
+            \text{subject to} & Gx \leq h \\
+                              & Ax = b.
+            \end{array}
+        \]
+
+        Here $P \in \mathcal{S}^{n}_+$, $q \in \mathcal{R}^n$, $G \in \mathcal{R}^{m \times n}$, $h \in \mathcal{R}^m$, $A \in \mathcal{R}^{p \times n}$, and $b \in \mathcal{R}^p$ are problem data and $x \in \mathcal{R}^{n}$ is the optimization variable. The inequality constraint $Gx \leq h$ is elementwise.
+
+        **Why quadratic programming?** Quadratic programs are convex optimization problems that generalize both least-squares and linear programming.They can be solved efficiently and reliably, even in real-time.
+
+        **An example from finance.** A simple example of a quadratic program arises in finance. Suppose we have $n$ different stocks, an estimate $r \in \mathcal{R}^n$ of the expected return on each stock, and an estimate $\Sigma \in \mathcal{S}^{n}_+$ of the covariance of the returns. Then we solve the optimization problem
+
+        \[
+            \begin{array}{ll}
+            \text{minimize}   & (1/2)x^T\Sigma x - r^Tx\\
+            \text{subject to} & x \geq 0 \\
+                              & \mathbf{1}^Tx = 1,
+            \end{array}
+        \]
+
+        to find a nonnegative portfolio allocation $x \in \mathcal{R}^n_+$ that optimally balances expected return and variance of return.
+
+        When we solve a quadratic program, in addition to a solution $x^\star$, we obtain a dual solution $\lambda^\star$ corresponding to the inequality constraints. A positive entry $\lambda^\star_i$ indicates that the constraint $g_i^Tx \leq h_i$ holds with equality for $x^\star$ and suggests that changing $h_i$ would change the optimal value.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Example
+
+        In this example, we use CVXPY to construct and solve a quadratic program.
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    import cvxpy as cp
+    import numpy as np
+    return cp, np
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md("""First we generate synthetic data. In this problem, we don't include equality constraints, only inequality.""")
+    return
+
+
+@app.cell
+def _(np):
+    m = 4
+    n = 2
+
+    np.random.seed(1)
+    q = np.random.randn(n)
+    G = np.random.randn(m, n)
+    h = G @ np.random.randn(n)
+    return G, h, m, n, q
+
+
+@app.cell(hide_code=True)
+def _(mo, np):
+    import wigglystuff
+
+    P_widget = mo.ui.anywidget(
+        wigglystuff.Matrix(np.array([[4.0, -1.4], [-1.4, 4]]), step=0.1)
+    )
+
+    mo.md(
+        f"""
+        The quadratic form $P$ is equal to the symmetrized version of this
+        matrix:
+
+        {P_widget.center()}
+        """
+    )
+    return P_widget, wigglystuff
+
+
+@app.cell
+def _(P_widget, np):
+    P = 0.5 * (np.array(P_widget.matrix) + np.array(P_widget.matrix).T)
+    return (P,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""Next, we specify the problem. Notice that we use the `quad_form` function from CVXPY to create the quadratic form $x^TPx$.""")
+    return
+
+
+@app.cell
+def _(G, P, cp, h, n, q):
+    x = cp.Variable(n)
+
+    problem = cp.Problem(
+        cp.Minimize((1 / 2) * cp.quad_form(x, P) + q.T @ x),
+        [G @ x <= h],
+    )
+    _ = problem.solve()
+    return problem, x
+
+
+@app.cell(hide_code=True)
+def _(mo, problem, x):
+    mo.md(
+        f"""
+        The optimal value is {problem.value:.04f}.
+
+        A solution $x$ is {mo.as_html(list(x.value))}
+        A dual solution is is {mo.as_html(list(problem.constraints[0].dual_value))}
+        """
+    )
+    return
+
+
+@app.cell
+def _(G, P, h, plot_contours, q, x):
+    plot_contours(P, G, h, q, x.value)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        In this plot, the gray shaded region is the feasible region (points satisfying the inequality), and the ellipses are level curves of the quadratic form.
+
+        **🌊 Try it!** Try changing the entries of $P$ above with your mouse. How do the
+        level curves and the optimal value of $x$ change? Can you explain what you see?
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(P, mo):
+    mo.md(
+        rf"""
+        The above contour lines were generated with
+        
+        \[
+        P= \begin{{bmatrix}}
+        {P[0, 0]:.01f} & {P[0, 1]:.01f} \\
+        {P[1, 0]:.01f} & {P[1, 1]:.01f} \\
+        \end{{bmatrix}}
+        \]
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(np):
+    def plot_contours(P, G, h, q, x_star):
+        import matplotlib.pyplot as plt
+
+        # Create a grid of x and y values.
+        x = np.linspace(-5, 5, 400)
+        y = np.linspace(-5, 5, 400)
+        X, Y = np.meshgrid(x, y)
+
+        # Compute the quadratic form Q(x, y) = a*x^2 + 2*b*x*y + c*y^2.
+        # Here, a = P[0,0], b = P[0,1] (and P[1,0]), c = P[1,1]
+        Z = (
+            0.5 * (P[0, 0] * X**2 + 2 * P[0, 1] * X * Y + P[1, 1] * Y**2)
+            + q[0] * X
+            + q[1] * Y
+        )
+
+        # --- Evaluate the constraints on the grid ---
+        # We stack X and Y to get a list of (x,y) points.
+        points = np.vstack([X.ravel(), Y.ravel()]).T
+
+        # Start with all points feasible
+        feasible = np.ones(points.shape[0], dtype=bool)
+
+        # Apply the inequality constraints Gx <= h.
+        # Each row of G and corresponding h defines a condition.
+        for i in range(G.shape[0]):
+            # For a given point x, the condition is: G[i,0]*x + G[i,1]*y <= h[i]
+            feasible &= points.dot(G[i]) <= h[i] + 1e-8  # small tolerance
+        # Reshape the boolean mask back to grid shape.
+        feasible_grid = feasible.reshape(X.shape)
+
+        # --- Plot the feasible region and contour lines---
+        plt.figure(figsize=(8, 6))
+
+        # Use contourf to fill the region where feasible_grid is True.
+        # We define two levels, so that points that are True (feasible) get one
+        # color.
+        plt.contourf(
+            X,
+            Y,
+            feasible_grid,
+            levels=[-0.5, 0.5, 1.5],
+            colors=["white", "gray"],
+            alpha=0.5,
+        )
+
+        contours = plt.contour(X, Y, Z, levels=10, cmap="viridis")
+        plt.clabel(contours, inline=True, fontsize=8)
+        plt.title("Feasible region and level curves")
+        plt.xlabel("$x_1$")
+        plt.ylabel("$y_2$")
+        # plt.colorbar(contours, label='Q(x, y)')
+
+        ax = plt.gca()
+        # Optionally, mark and label the point x_star.
+        ax.plot(x_star[0], x_star[1], "ko", markersize=5)
+        ax.text(
+            x_star[0],
+            x_star[1],
+            r"$\mathbf{x}^\star$",
+            color="black",
+            fontsize=12,
+            verticalalignment="bottom",
+            horizontalalignment="right",
+        )
+        return plt.gca()
+    return (plot_contours,)
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+if __name__ == "__main__":
+    app.run()

optimization/README.md

@@ -0,0 +1,31 @@
+# Learn optimization
+
+🚧 _This collection is a work in progress. Check back later for new notebooks!_
+
+This collection of marimo notebooks teaches you the basics of mathematical
+optimization.
+
+After working through these notebooks, you'll understand how to create
+and solve optimization problems using the Python library
+[CVXPY](https://github.com/cvxpy/cvxpy), as well as how to apply what you've
+learned to real-world problems such as portfolio allocation in finance,
+resource allocation, and more.
+
+![SpaceX](https://www.debugmind.com/wp-content/uploads/2020/01/spacex-1.jpg)
+
+_SpaceX solves convex optimization problems onboard to land its rockets, using CVXGEN, a code generator for quadratic programming developed at Stephen Boyd’s Stanford lab. Photo by SpaceX, licensed CC BY-NC 2.0._
+
+**Running notebooks.** To run a notebook locally, use
+
+```bash
+uvx marimo edit <URL>
+```
+
+For example, run the least-squares tutorial with
+
+```bash
+uvx marimo edit https://github.com/marimo-team/learn/blob/main/optimization/01_least_squares.py
+```
+
+You can also open notebooks in our online playground by appending `marimo.app/`
+to a notebook's URL: [marimo.app/github.com/marimo-team/learn/blob/main/optimization/01_least_squares.py](https://marimo.app/https://github.com/marimo-team/learn/blob/main/optimization/01_least_squares.py).

probability/01_sets.py

@@ -0,0 +1,297 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "marimo",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.11.0"
+app = marimo.App()
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        # Sets
+
+        Probability is the study of "events", assigning numerical values to how likely
+        events are to occur. For example, probability lets us quantify how likely it is for it to rain or shine on a given day.
+
+
+        Typically we reason about _sets_ of events. In mathematics,
+        a set is a collection of elements, with no element included more than once.
+        Elements can be any kind of object.
+
+        For example:
+
+        - ☀️ Weather events: $\{\text{Rain}, \text{Overcast}, \text{Clear}\}$
+        - 🎲 Die rolls: $\{1, 2, 3, 4, 5, 6\}$
+        - 🪙 Pairs of coin flips = $\{ \text{(Heads, Heads)}, \text{(Heads, Tails)}, \text{(Tails, Tails)} \text{(Tails, Heads)}\}$
+
+        Sets are the building blocks of probability, and will arise frequently in our study.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""## Set operations""")
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""In Python, sets are made with the `set` function:""")
+    return
+
+
+@app.cell
+def _():
+    A = set([2, 3, 5, 7])
+    A
+    return (A,)
+
+
+@app.cell
+def _():
+    B = set([0, 1, 2, 3, 5, 8])
+    B
+    return (B,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        Below we explain common operations on sets.
+
+        _**Try it!** Try modifying the definitions of `A` and `B` above, and see how the results change below._
+
+        The **union** $A \cup B$ of sets $A$ and $B$ is the set of elements in $A$, $B$, or both.
+        """
+    )
+    return
+
+
+@app.cell
+def _(A, B):
+    A | B
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""The **intersection** $A \cap B$ is the set of elements in both $A$ and $B$""")
+    return
+
+
+@app.cell
+def _(A, B):
+    A & B
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""The **difference** $A \setminus B$ is the set of elements in $A$ that are not in $B$.""")
+    return
+
+
+@app.cell
+def _(A, B):
+    A - B
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        """
+        ### 🎬 An interactive example
+
+        Here's a simple example that classifies TV shows into sets by genre, and uses these sets to recommend shows to a user based on their preferences.
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    viewer_type = mo.ui.radio(
+        options={
+            "I like action and drama!": "New Viewer",
+            "I only like action shows": "Action Fan",
+            "I only like dramas": "Drama Fan",
+        },
+        value="I like action and drama!",
+        label="Which genre do you prefer?",
+    )
+    return (viewer_type,)
+
+
+@app.cell(hide_code=True)
+def _(viewer_type):
+    viewer_type
+    return
+
+
+@app.cell
+def _():
+    action_shows = {"Stranger Things", "The Witcher", "Money Heist"}
+    drama_shows = {"The Crown", "Money Heist", "Bridgerton"}
+    return action_shows, drama_shows
+
+
+@app.cell
+def _(action_shows, drama_shows):
+    recommendations = {
+        "New Viewer": action_shows | drama_shows,  # Union for new viewers
+        "Action Fan": action_shows - drama_shows,  # Unique action shows
+        "Drama Fan": drama_shows - action_shows,  # Unique drama shows
+    }
+    return (recommendations,)
+
+
+@app.cell(hide_code=True)
+def _(mo, recommendations, viewer_type):
+    result = recommendations[viewer_type.value]
+
+    explanation = {
+        "New Viewer": "You get everything to explore!",
+        "Action Fan": "Pure action, no drama!",
+        "Drama Fan": "Drama-focused selections!",
+    }
+
+    mo.md(f"""
+    **🎬 Recommended shows.** Based on your preference for **{viewer_type.value}**,
+    we recommend:
+
+    {", ".join(result)}
+
+    **Why these shows?** 
+    {explanation[viewer_type.value]}
+    """)
+    return explanation, result
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md("""
+    ### Exercise
+
+    Given these sets:
+
+    - A = {🎮, 📱, 💻}
+
+    - B = {📱, 💻, 🖨️}
+
+    - C = {💻, 🖨️, ⌨️}
+
+    Can you:
+
+    1. Find all elements that are in A or B
+
+    2. Find elements common to all three sets
+
+    3. Find elements in A that aren't in C
+
+    <details>
+
+    <summary>Check your answers!</summary>
+
+    1. A ∪ B = {🎮, 📱, 💻, 🖨️}<br>
+    2. A ∩ B ∩ C = {💻}<br>
+    3. A - C = {🎮, 📱}
+
+    </details>
+    """)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## 🧮 Set properties
+
+        Here are some important properties of the set operations:
+
+        1. **Commutative**: $A \cup B = B \cup A$
+        2. **Associative**: $(A \cup B) \cup C = A \cup (B \cup C)$
+        3. **Distributive**: $A \cup (B \cap C) = (A \cup B) \cap (A \cup C)$
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(
+        r"""
+        ## Set builder notation
+
+        To compactly describe the elements in a set, we can use **set builder notation**, which specifies conditions that must be true for elements to be in the set.
+
+        For example, here is how to specify the set of positive numbers less than 10:
+
+        \[
+        \{x \mid 0 < x < 10 \}
+        \]
+
+        The predicate to the right of the vertical bar $\mid$ specifies conditions that must be true for an element to be in the set; the expression to the left of $\mid$ specifies the value being included.
+
+        In Python, set builder notation is called a "set comprehension."
+        """
+    )
+    return
+
+
+@app.cell
+def _():
+    def predicate(x):
+        return x > 0 and x < 10
+    return (predicate,)
+
+
+@app.cell
+def _(predicate):
+    set(x for x in range(100) if predicate(x))
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md("""**Try it!** Try modifying the `predicate` function above and see how the set changes.""")
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md("""
+    ## Summary
+
+    You've learned:
+
+    - Basic set operations
+    - Set properties
+    - Real-world applications
+
+    In the next lesson, we'll define probability from the ground up, using sets.
+
+    Remember: In probability, every event is a set, and every set can be an event!
+    """)
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+if __name__ == "__main__":
+    app.run()

probability/README.md

@@ -0,0 +1,22 @@
+# Learn probability
+
+🚧 _This collection is a work in progress. Check back later for new noteboks._
+
+This collection of marimo notebooks teaches the fundamentals of probability,
+with an emphasis on computation with Python.
+
+
+**Running notebooks.** To run a notebook locally, use
+
+```bash
+uvx marimo edit <URL>
+```
+
+For example, run the numbers tutorial with
+
+```bash
+uvx marimo edit https://github.com/marimo-team/learn/blob/main/probability/01_sets.py
+```
+
+You can also open notebooks in our online playground by appending `marimo.app/`
+to a notebook's URL: [marimo.app/https://github.com/marimo-team/learn/blob/main/probability/01_sets.py](https://marimo.app/https://github.com/marimo-team/learn/blob/main/probability/01_sets.py).