Notebook 1a: Vector Fields and ODEs¶

Spin-off from Notebook 1 — Classical Vector Fields

---

Overview¶

Vector fields aren't just pretty pictures — they define ordinary differential equations (ODEs). Following a vector field means solving an ODE!

This notebook bridges:

• GRL: Gradient ascent on $Q^+$ is an ODE
• Flow Matching: Sampling = solving the flow ODE
• Diffusion Models: Probability flow ODEs

📚 Cross-reference: For a deep dive into flow matching and diffusion models, see the genai-lab project:

• docs/flow_matching/01_flow_matching_foundations.md — Flow matching theory
• docs/DDPM/01_ddpm_foundations.md — Diffusion model foundations

Both projects share the same mathematical foundation: vector fields define dynamics.

Learning Objectives¶

1. ODEs as vector fields — $\dot{x} = F(x)$ means "follow the arrows"
2. Numerical solvers — Euler, RK4, and why accuracy matters
3. Phase portraits — Visualizing solution families
4. Fixed points — Attractors, repellers, and stability
5. Applications — Gradient flow, flow matching, GRL

Prerequisites¶

• Notebook 1 (Classical Vector Fields)
• Basic calculus (derivatives)

Time¶

~25-30 minutes

In [21]:

# Setup
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import seaborn as sns

sns.set_theme(style='whitegrid', context='notebook')
plt.rcParams['figure.figsize'] = (12, 8)
%matplotlib inline
print("Libraries loaded.")

# Setup import numpy as np import matplotlib.pyplot as plt from matplotlib import cm import seaborn as sns sns.set_theme(style='whitegrid', context='notebook') plt.rcParams['figure.figsize'] = (12, 8) %matplotlib inline print("Libraries loaded.")

Libraries loaded.

---

Part 1: ODEs as Vector Fields¶

The Key Insight¶

An ordinary differential equation (ODE):

$$\frac{dx}{dt} = F(x)$$

says: "The velocity at position $x$ is given by $F(x)$"

This IS a vector field! The field $F$ tells you which direction to move at each point.

Solving an ODE = Following the Field¶

Given initial condition $x(0) = x_0$, the solution $x(t)$ is the trajectory that:

1. Starts at $x_0$
2. Always moves in the direction given by $F(x)$

Why This Matters¶

Application	ODE	Vector Field
GRL	$\theta_{t+1} = \theta_t + \eta \nabla_\theta Q^+$	Gradient of $Q^+$
Flow Matching	$\frac{dx}{dt} = v_\theta(x, t)$	Learned velocity field
Gradient Descent	$\frac{dx}{dt} = -\nabla f(x)$	Negative gradient

In [22]:

# Example 1.1: Simple Linear ODE
# dx/dt = -x  (exponential decay)

def F_decay(x):
    """Vector field for dx/dt = -x"""
    return -x

# Analytical solution: x(t) = x_0 * exp(-t)
def analytical_decay(x0, t):
    return x0 * np.exp(-t)

# Visualize in 1D
x = np.linspace(-3, 3, 20)
t_vals = np.linspace(0, 3, 100)

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Left: Vector field (1D arrows)
ax1 = axes[0]
ax1.quiver(x, np.zeros_like(x), F_decay(x), np.zeros_like(x), 
           np.abs(F_decay(x)), cmap='coolwarm', scale=20)
ax1.axhline(0, color='k', lw=0.5)
ax1.axvline(0, color='k', lw=2, label='Fixed point (x=0)')
ax1.set_xlabel('$x$'); ax1.set_ylabel('')
ax1.set_title(r'Vector Field: $F(x) = -x$')
ax1.set_ylim(-0.5, 0.5); ax1.legend()

# Right: Solutions (trajectories)
ax2 = axes[1]
for x0 in [-2, -1, 1, 2]:
    ax2.plot(t_vals, analytical_decay(x0, t_vals), lw=2, label=f'$x_0 = {x0}$')
ax2.axhline(0, color='k', linestyle='--', alpha=0.5)
ax2.set_xlabel('$t$'); ax2.set_ylabel('$x(t)$')
ax2.set_title(r'Solutions: $x(t) = x_0 e^{-t}$')
ax2.legend(); ax2.grid(True, alpha=0.3)

plt.tight_layout(); plt.show()
print("All trajectories converge to x=0 (the fixed point/attractor).")

# Example 1.1: Simple Linear ODE # dx/dt = -x (exponential decay) def F_decay(x): """Vector field for dx/dt = -x""" return -x # Analytical solution: x(t) = x_0 * exp(-t) def analytical_decay(x0, t): return x0 * np.exp(-t) # Visualize in 1D x = np.linspace(-3, 3, 20) t_vals = np.linspace(0, 3, 100) fig, axes = plt.subplots(1, 2, figsize=(14, 5)) # Left: Vector field (1D arrows) ax1 = axes[0] ax1.quiver(x, np.zeros_like(x), F_decay(x), np.zeros_like(x), np.abs(F_decay(x)), cmap='coolwarm', scale=20) ax1.axhline(0, color='k', lw=0.5) ax1.axvline(0, color='k', lw=2, label='Fixed point (x=0)') ax1.set_xlabel('$x$'); ax1.set_ylabel('') ax1.set_title(r'Vector Field: $F(x) = -x$') ax1.set_ylim(-0.5, 0.5); ax1.legend() # Right: Solutions (trajectories) ax2 = axes[1] for x0 in [-2, -1, 1, 2]: ax2.plot(t_vals, analytical_decay(x0, t_vals), lw=2, label=f'$x_0 = {x0}$') ax2.axhline(0, color='k', linestyle='--', alpha=0.5) ax2.set_xlabel('$t$'); ax2.set_ylabel('$x(t)$') ax2.set_title(r'Solutions: $x(t) = x_0 e^{-t}$') ax2.legend(); ax2.grid(True, alpha=0.3) plt.tight_layout(); plt.show() print("All trajectories converge to x=0 (the fixed point/attractor).")

All trajectories converge to x=0 (the fixed point/attractor).

In [23]:

# Example 1.2: 2D ODE System
# dx/dt = -y
# dy/dt = x
# This is circular motion!

def F_rotation(state):
    """Vector field for circular motion"""
    x, y = state
    return np.array([-y, x])

# Create grid
x = np.linspace(-2, 2, 15)
y = np.linspace(-2, 2, 15)
X, Y = np.meshgrid(x, y)
U, V = -Y, X

# Analytical solution: x(t) = r*cos(t + φ), y(t) = r*sin(t + φ)
def solve_rotation(x0, y0, t_max=2*np.pi, n_steps=100):
    t = np.linspace(0, t_max, n_steps)
    r = np.sqrt(x0**2 + y0**2)
    phi = np.arctan2(y0, x0)
    return r * np.cos(t + phi), r * np.sin(t + phi)

fig, ax = plt.subplots(figsize=(10, 10))

# Vector field
ax.quiver(X, Y, U, V, np.sqrt(U**2 + V**2), cmap='viridis', alpha=0.6, scale=25)

# Trajectories
colors = plt.cm.tab10(np.linspace(0, 1, 4))
for i, (x0, y0) in enumerate([(1.5, 0), (0, 1), (0.5, 0.5), (1, 1)]):
    traj_x, traj_y = solve_rotation(x0, y0)
    ax.plot(traj_x, traj_y, '-', color=colors[i], lw=2)
    ax.plot(x0, y0, 'o', color=colors[i], markersize=10, mec='k')

ax.plot(0, 0, 'k*', markersize=15, label='Fixed point')
ax.set_xlabel('$x$'); ax.set_ylabel('$y$')
ax.set_title(r'2D ODE: $\dot{x} = -y, \, \dot{y} = x$ (Circular Motion)')
ax.set_aspect('equal'); ax.legend(); ax.grid(True, alpha=0.3)
plt.show()

print("Trajectories are circles — the ODE describes rotation!")
print("The fixed point at origin is a 'center' (neutral stability).")

# Example 1.2: 2D ODE System # dx/dt = -y # dy/dt = x # This is circular motion! def F_rotation(state): """Vector field for circular motion""" x, y = state return np.array([-y, x]) # Create grid x = np.linspace(-2, 2, 15) y = np.linspace(-2, 2, 15) X, Y = np.meshgrid(x, y) U, V = -Y, X # Analytical solution: x(t) = r*cos(t + φ), y(t) = r*sin(t + φ) def solve_rotation(x0, y0, t_max=2*np.pi, n_steps=100): t = np.linspace(0, t_max, n_steps) r = np.sqrt(x0**2 + y0**2) phi = np.arctan2(y0, x0) return r * np.cos(t + phi), r * np.sin(t + phi) fig, ax = plt.subplots(figsize=(10, 10)) # Vector field ax.quiver(X, Y, U, V, np.sqrt(U**2 + V**2), cmap='viridis', alpha=0.6, scale=25) # Trajectories colors = plt.cm.tab10(np.linspace(0, 1, 4)) for i, (x0, y0) in enumerate([(1.5, 0), (0, 1), (0.5, 0.5), (1, 1)]): traj_x, traj_y = solve_rotation(x0, y0) ax.plot(traj_x, traj_y, '-', color=colors[i], lw=2) ax.plot(x0, y0, 'o', color=colors[i], markersize=10, mec='k') ax.plot(0, 0, 'k*', markersize=15, label='Fixed point') ax.set_xlabel('$x$'); ax.set_ylabel('$y$') ax.set_title(r'2D ODE: $\dot{x} = -y, \, \dot{y} = x$ (Circular Motion)') ax.set_aspect('equal'); ax.legend(); ax.grid(True, alpha=0.3) plt.show() print("Trajectories are circles — the ODE describes rotation!") print("The fixed point at origin is a 'center' (neutral stability).")

Trajectories are circles — the ODE describes rotation!
The fixed point at origin is a 'center' (neutral stability).

---

Part 2: Numerical ODE Solvers¶

Most ODEs can't be solved analytically. We need numerical methods.

Euler Method (Simplest)¶

$$x_{n+1} = x_n + \Delta t \cdot F(x_n)$$

"Take a step in the direction of the field."

Problem: Error accumulates! Step size $\Delta t$ must be small.

Runge-Kutta 4 (RK4) — The Workhorse¶

$$\begin{align} k_1 &= F(x_n) \\ k_2 &= F(x_n + \frac{\Delta t}{2} k_1) \\ k_3 &= F(x_n + \frac{\Delta t}{2} k_2) \\ k_4 &= F(x_n + \Delta t \cdot k_3) \\ x_{n+1} &= x_n + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4) \end{align}$$

"Sample the field at multiple points, then average."

Much more accurate for the same step size!

In [24]:

# Implement Euler and RK4 solvers

def euler_step(F, x, dt):
    """Single Euler step: x_new = x + dt * F(x)"""
    return x + dt * F(x)

def rk4_step(F, x, dt):
    """Single RK4 step"""
    k1 = F(x)
    k2 = F(x + 0.5 * dt * k1)
    k3 = F(x + 0.5 * dt * k2)
    k4 = F(x + dt * k3)
    return x + (dt / 6) * (k1 + 2*k2 + 2*k3 + k4)

def solve_ode(F, x0, t_span, dt, method='euler'):
    """Solve ODE from t_span[0] to t_span[1]"""
    t = np.arange(t_span[0], t_span[1] + dt, dt)
    x = np.zeros((len(t), len(x0)))
    x[0] = x0
    
    step_fn = euler_step if method == 'euler' else rk4_step
    
    for i in range(1, len(t)):
        x[i] = step_fn(F, x[i-1], dt)
    
    return t, x

print("Solvers defined: euler_step, rk4_step, solve_ode")

# Implement Euler and RK4 solvers def euler_step(F, x, dt): """Single Euler step: x_new = x + dt * F(x)""" return x + dt * F(x) def rk4_step(F, x, dt): """Single RK4 step""" k1 = F(x) k2 = F(x + 0.5 * dt * k1) k3 = F(x + 0.5 * dt * k2) k4 = F(x + dt * k3) return x + (dt / 6) * (k1 + 2*k2 + 2*k3 + k4) def solve_ode(F, x0, t_span, dt, method='euler'): """Solve ODE from t_span[0] to t_span[1]""" t = np.arange(t_span[0], t_span[1] + dt, dt) x = np.zeros((len(t), len(x0))) x[0] = x0 step_fn = euler_step if method == 'euler' else rk4_step for i in range(1, len(t)): x[i] = step_fn(F, x[i-1], dt) return t, x print("Solvers defined: euler_step, rk4_step, solve_ode")

Solvers defined: euler_step, rk4_step, solve_ode

In [25]:

# Example 2.1: Compare Euler vs RK4 on circular motion
# True solution should be a perfect circle!

x0 = np.array([1.0, 0.0])
t_span = (0, 4 * np.pi)  # Two full rotations

fig, axes = plt.subplots(1, 3, figsize=(15, 5))

# Different step sizes
step_sizes = [0.5, 0.2, 0.05]

for ax, dt in zip(axes, step_sizes):
    # Euler
    t_e, x_e = solve_ode(F_rotation, x0, t_span, dt, method='euler')
    # RK4
    t_r, x_r = solve_ode(F_rotation, x0, t_span, dt, method='rk4')
    # True circle
    theta = np.linspace(0, 2*np.pi, 100)
    
    ax.plot(np.cos(theta), np.sin(theta), 'k--', lw=1, label='True circle')
    ax.plot(x_e[:, 0], x_e[:, 1], 'r-', lw=1.5, alpha=0.7, label='Euler')
    ax.plot(x_r[:, 0], x_r[:, 1], 'b-', lw=1.5, alpha=0.7, label='RK4')
    ax.plot(x0[0], x0[1], 'go', markersize=10, label='Start')
    
    ax.set_title(f'$\Delta t = {dt}$\n({int(t_span[1]/dt)} steps)')
    ax.set_xlabel('$x$'); ax.set_ylabel('$y$')
    ax.set_xlim(-2, 2); ax.set_ylim(-2, 2)
    ax.set_aspect('equal'); ax.legend(loc='upper right', fontsize=8)
    ax.grid(True, alpha=0.3)

plt.suptitle('Euler vs RK4: Solving Circular Motion ODE', y=1.02)
plt.tight_layout(); plt.show()

print("Euler spirals outward (energy drift). RK4 stays on the circle!")
print("This is why flow matching uses RK4 or adaptive solvers.")

# Example 2.1: Compare Euler vs RK4 on circular motion # True solution should be a perfect circle! x0 = np.array([1.0, 0.0]) t_span = (0, 4 * np.pi) # Two full rotations fig, axes = plt.subplots(1, 3, figsize=(15, 5)) # Different step sizes step_sizes = [0.5, 0.2, 0.05] for ax, dt in zip(axes, step_sizes): # Euler t_e, x_e = solve_ode(F_rotation, x0, t_span, dt, method='euler') # RK4 t_r, x_r = solve_ode(F_rotation, x0, t_span, dt, method='rk4') # True circle theta = np.linspace(0, 2*np.pi, 100) ax.plot(np.cos(theta), np.sin(theta), 'k--', lw=1, label='True circle') ax.plot(x_e[:, 0], x_e[:, 1], 'r-', lw=1.5, alpha=0.7, label='Euler') ax.plot(x_r[:, 0], x_r[:, 1], 'b-', lw=1.5, alpha=0.7, label='RK4') ax.plot(x0[0], x0[1], 'go', markersize=10, label='Start') ax.set_title(f'$\Delta t = {dt}$\n({int(t_span[1]/dt)} steps)') ax.set_xlabel('$x$'); ax.set_ylabel('$y$') ax.set_xlim(-2, 2); ax.set_ylim(-2, 2) ax.set_aspect('equal'); ax.legend(loc='upper right', fontsize=8) ax.grid(True, alpha=0.3) plt.suptitle('Euler vs RK4: Solving Circular Motion ODE', y=1.02) plt.tight_layout(); plt.show() print("Euler spirals outward (energy drift). RK4 stays on the circle!") print("This is why flow matching uses RK4 or adaptive solvers.")

<>:25: SyntaxWarning: invalid escape sequence '\D'
<>:25: SyntaxWarning: invalid escape sequence '\D'
/var/folders/jt/h4k6wdyx36nbnjk_rmkwdk280000gn/T/ipykernel_55695/524608160.py:25: SyntaxWarning: invalid escape sequence '\D'
  ax.set_title(f'$\Delta t = {dt}$\n({int(t_span[1]/dt)} steps)')

Euler spirals outward (energy drift). RK4 stays on the circle!
This is why flow matching uses RK4 or adaptive solvers.

In [26]:

# Example 2.2: Error Analysis

def compute_error(method, dt, t_final=2*np.pi):
    """Compute final position error for circular motion"""
    x0 = np.array([1.0, 0.0])
    t, x = solve_ode(F_rotation, x0, (0, t_final), dt, method=method)
    # True final position (one full rotation = back to start)
    true_final = x0
    return np.linalg.norm(x[-1] - true_final)

dt_values = np.logspace(-2, 0, 20)  # 0.01 to 1.0
euler_errors = [compute_error('euler', dt) for dt in dt_values]
rk4_errors = [compute_error('rk4', dt) for dt in dt_values]

fig, ax = plt.subplots(figsize=(10, 6))
ax.loglog(dt_values, euler_errors, 'ro-', lw=2, label='Euler (1st order)')
ax.loglog(dt_values, rk4_errors, 'bs-', lw=2, label='RK4 (4th order)')

# Reference lines
ax.loglog(dt_values, 0.5 * dt_values, 'r--', alpha=0.5, label=r'$O(\Delta t)$')
ax.loglog(dt_values, 0.1 * dt_values**4, 'b--', alpha=0.5, label=r'$O(\Delta t^4)$')

ax.set_xlabel(r'Step size $\Delta t$')
ax.set_ylabel('Error (distance from true solution)')
ax.set_title('Error vs Step Size: Euler vs RK4')
ax.legend(); ax.grid(True, alpha=0.3, which='both')
plt.show()

print("Euler: Error ~ Δt (1st order). RK4: Error ~ Δt⁴ (4th order).")
print("RK4 with Δt=0.1 beats Euler with Δt=0.01!")

# Example 2.2: Error Analysis def compute_error(method, dt, t_final=2*np.pi): """Compute final position error for circular motion""" x0 = np.array([1.0, 0.0]) t, x = solve_ode(F_rotation, x0, (0, t_final), dt, method=method) # True final position (one full rotation = back to start) true_final = x0 return np.linalg.norm(x[-1] - true_final) dt_values = np.logspace(-2, 0, 20) # 0.01 to 1.0 euler_errors = [compute_error('euler', dt) for dt in dt_values] rk4_errors = [compute_error('rk4', dt) for dt in dt_values] fig, ax = plt.subplots(figsize=(10, 6)) ax.loglog(dt_values, euler_errors, 'ro-', lw=2, label='Euler (1st order)') ax.loglog(dt_values, rk4_errors, 'bs-', lw=2, label='RK4 (4th order)') # Reference lines ax.loglog(dt_values, 0.5 * dt_values, 'r--', alpha=0.5, label=r'$O(\Delta t)$') ax.loglog(dt_values, 0.1 * dt_values**4, 'b--', alpha=0.5, label=r'$O(\Delta t^4)$') ax.set_xlabel(r'Step size $\Delta t$') ax.set_ylabel('Error (distance from true solution)') ax.set_title('Error vs Step Size: Euler vs RK4') ax.legend(); ax.grid(True, alpha=0.3, which='both') plt.show() print("Euler: Error ~ Δt (1st order). RK4: Error ~ Δt⁴ (4th order).") print("RK4 with Δt=0.1 beats Euler with Δt=0.01!")

<>:20: SyntaxWarning: invalid escape sequence '\D'
<>:21: SyntaxWarning: invalid escape sequence '\D'
<>:23: SyntaxWarning: invalid escape sequence '\D'
<>:20: SyntaxWarning: invalid escape sequence '\D'
<>:21: SyntaxWarning: invalid escape sequence '\D'
<>:23: SyntaxWarning: invalid escape sequence '\D'
/var/folders/jt/h4k6wdyx36nbnjk_rmkwdk280000gn/T/ipykernel_55695/2141023215.py:20: SyntaxWarning: invalid escape sequence '\D'
  ax.loglog(dt_values, 0.5 * dt_values, 'r--', alpha=0.5, label='$O(\Delta t)$')
/var/folders/jt/h4k6wdyx36nbnjk_rmkwdk280000gn/T/ipykernel_55695/2141023215.py:21: SyntaxWarning: invalid escape sequence '\D'
  ax.loglog(dt_values, 0.1 * dt_values**4, 'b--', alpha=0.5, label='$O(\Delta t^4)$')
/var/folders/jt/h4k6wdyx36nbnjk_rmkwdk280000gn/T/ipykernel_55695/2141023215.py:23: SyntaxWarning: invalid escape sequence '\D'
  ax.set_xlabel('Step size $\Delta t$')

Euler: Error ~ Δt (1st order). RK4: Error ~ Δt⁴ (4th order).
RK4 with Δt=0.1 beats Euler with Δt=0.01!

---

Part 3: Phase Portraits and Fixed Points¶

Phase Portrait¶

A phase portrait shows the vector field plus representative trajectories. It reveals the qualitative behavior of all solutions.

Fixed Points¶

A fixed point (equilibrium) is where $F(x^*) = 0$. The system stays there forever.

Types of fixed points in 2D:

Type	Behavior	Eigenvalues
Stable node	Trajectories converge	Both negative real
Unstable node	Trajectories diverge	Both positive real
Saddle	Some converge, some diverge	Opposite signs
Center	Closed orbits	Pure imaginary
Spiral	Spiraling in/out	Complex with real part

In [27]:

# Example 3.1: Different Types of Fixed Points

def stable_node(state):
    """Stable node: eigenvalues -1, -2"""
    x, y = state
    return np.array([-x, -2*y])

def unstable_node(state):
    """Unstable node: eigenvalues +1, +2"""
    x, y = state
    return np.array([x, 2*y])

def saddle(state):
    """Saddle point: eigenvalues +1, -1"""
    x, y = state
    return np.array([x, -y])

def stable_spiral(state):
    """Stable spiral: eigenvalues -0.2 ± i"""
    x, y = state
    return np.array([-0.2*x - y, x - 0.2*y])

systems = [
    (stable_node, 'Stable Node', 'Attractor'),
    (unstable_node, 'Unstable Node', 'Repeller'),
    (saddle, 'Saddle Point', 'Mixed'),
    (stable_spiral, 'Stable Spiral', 'Attractor'),
]

fig, axes = plt.subplots(2, 2, figsize=(12, 12))

x = np.linspace(-2, 2, 15)
y = np.linspace(-2, 2, 15)
X, Y = np.meshgrid(x, y)

for ax, (F, title, ftype) in zip(axes.flat, systems):
    # Vector field
    U = np.zeros_like(X)
    V = np.zeros_like(Y)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            vec = F(np.array([X[i,j], Y[i,j]]))
            U[i,j], V[i,j] = vec
    
    ax.streamplot(X, Y, U, V, color='gray', density=1.5, linewidth=0.8)
    
    # Sample trajectories
    starts = [np.array([1.5, 1.5]), np.array([-1.5, 1.0]), 
              np.array([1.0, -1.5]), np.array([-1.0, -1.0])]
    colors = plt.cm.tab10(np.linspace(0, 0.4, len(starts)))
    
    for x0, c in zip(starts, colors):
        t, traj = solve_ode(F, x0, (0, 5), 0.05, method='rk4')
        ax.plot(traj[:, 0], traj[:, 1], '-', color=c, lw=2)
        ax.plot(x0[0], x0[1], 'o', color=c, markersize=8)
    
    ax.plot(0, 0, 'k*', markersize=15)
    ax.set_title(f'{title}\n({ftype})')
    ax.set_xlim(-2, 2); ax.set_ylim(-2, 2)
    ax.set_aspect('equal'); ax.grid(True, alpha=0.3)

plt.tight_layout(); plt.show()
print("Fixed point type determines long-term behavior of ALL trajectories.")

# Example 3.1: Different Types of Fixed Points def stable_node(state): """Stable node: eigenvalues -1, -2""" x, y = state return np.array([-x, -2*y]) def unstable_node(state): """Unstable node: eigenvalues +1, +2""" x, y = state return np.array([x, 2*y]) def saddle(state): """Saddle point: eigenvalues +1, -1""" x, y = state return np.array([x, -y]) def stable_spiral(state): """Stable spiral: eigenvalues -0.2 ± i""" x, y = state return np.array([-0.2*x - y, x - 0.2*y]) systems = [ (stable_node, 'Stable Node', 'Attractor'), (unstable_node, 'Unstable Node', 'Repeller'), (saddle, 'Saddle Point', 'Mixed'), (stable_spiral, 'Stable Spiral', 'Attractor'), ] fig, axes = plt.subplots(2, 2, figsize=(12, 12)) x = np.linspace(-2, 2, 15) y = np.linspace(-2, 2, 15) X, Y = np.meshgrid(x, y) for ax, (F, title, ftype) in zip(axes.flat, systems): # Vector field U = np.zeros_like(X) V = np.zeros_like(Y) for i in range(X.shape[0]): for j in range(X.shape[1]): vec = F(np.array([X[i,j], Y[i,j]])) U[i,j], V[i,j] = vec ax.streamplot(X, Y, U, V, color='gray', density=1.5, linewidth=0.8) # Sample trajectories starts = [np.array([1.5, 1.5]), np.array([-1.5, 1.0]), np.array([1.0, -1.5]), np.array([-1.0, -1.0])] colors = plt.cm.tab10(np.linspace(0, 0.4, len(starts))) for x0, c in zip(starts, colors): t, traj = solve_ode(F, x0, (0, 5), 0.05, method='rk4') ax.plot(traj[:, 0], traj[:, 1], '-', color=c, lw=2) ax.plot(x0[0], x0[1], 'o', color=c, markersize=8) ax.plot(0, 0, 'k*', markersize=15) ax.set_title(f'{title}\n({ftype})') ax.set_xlim(-2, 2); ax.set_ylim(-2, 2) ax.set_aspect('equal'); ax.grid(True, alpha=0.3) plt.tight_layout(); plt.show() print("Fixed point type determines long-term behavior of ALL trajectories.")

Fixed point type determines long-term behavior of ALL trajectories.

---

Part 4: Gradient Flow — Optimization as ODE¶

The Connection¶

Gradient descent is just following the ODE:

$$\frac{dx}{dt} = -\nabla f(x)$$

The vector field is the negative gradient of the loss function!

Properties¶

• Fixed points = critical points of $f$ (minima, maxima, saddles)
• Stable fixed points = local minima
• Trajectories always decrease $f$: $\frac{df}{dt} = \nabla f \cdot \dot{x} = -\|\nabla f\|^2 \leq 0$

GRL Connection¶

In GRL, we do gradient ascent on $Q^+$:

$$\frac{d\theta}{dt} = +\nabla_\theta Q^+(s, \theta)$$

This is the same math — just maximizing instead of minimizing!

In [28]:

# Example 4.1: Gradient Descent as ODE

def rosenbrock(state):
    """Rosenbrock function: f(x,y) = (1-x)² + 100(y-x²)²"""
    x, y = state
    return (1 - x)**2 + 100 * (y - x**2)**2

def rosenbrock_grad(state):
    """Gradient of Rosenbrock"""
    x, y = state
    dfdx = -2*(1-x) - 400*x*(y - x**2)
    dfdy = 200*(y - x**2)
    return np.array([dfdx, dfdy])

def neg_grad_rosenbrock(state):
    """Negative gradient (for descent)"""
    return -rosenbrock_grad(state)

# Create contour plot
x = np.linspace(-2, 2, 100)
y = np.linspace(-1, 3, 100)
X, Y = np.meshgrid(x, y)
Z = np.zeros_like(X)
for i in range(X.shape[0]):
    for j in range(X.shape[1]):
        Z[i,j] = rosenbrock(np.array([X[i,j], Y[i,j]]))

fig, ax = plt.subplots(figsize=(12, 8))

# Contours (log scale for better visualization)
levels = np.logspace(-1, 3, 20)
c = ax.contour(X, Y, Z, levels=levels, cmap='viridis', alpha=0.7)
ax.contourf(X, Y, Z, levels=levels, cmap='viridis', alpha=0.3)

# Gradient descent trajectories
starts = [np.array([-1.5, 2.5]), np.array([1.5, 2.5]), 
          np.array([-0.5, -0.5]), np.array([0.0, 2.0])]
colors = ['red', 'blue', 'orange', 'purple']

for x0, c in zip(starts, colors):
    # Use small step size for stability
    t, traj = solve_ode(neg_grad_rosenbrock, x0, (0, 50), 0.001, method='rk4')
    ax.plot(traj[:, 0], traj[:, 1], '-', color=c, lw=2, alpha=0.8)
    ax.plot(x0[0], x0[1], 'o', color=c, markersize=10, mec='k')

ax.plot(1, 1, 'w*', markersize=20, mec='k', mew=2, label='Global minimum (1,1)')
ax.set_xlabel('$x$'); ax.set_ylabel('$y$')
ax.set_title('Gradient Descent on Rosenbrock Function\n(Optimization as ODE)')
ax.legend(); ax.grid(True, alpha=0.3)
plt.colorbar(ax.contourf(X, Y, Z, levels=levels, cmap='viridis', alpha=0), 
             ax=ax, label='$f(x,y)$')
plt.show()

print("Gradient descent = following the negative gradient field.")
print("All trajectories flow toward the minimum at (1, 1).")

# Example 4.1: Gradient Descent as ODE def rosenbrock(state): """Rosenbrock function: f(x,y) = (1-x)² + 100(y-x²)²""" x, y = state return (1 - x)**2 + 100 * (y - x**2)**2 def rosenbrock_grad(state): """Gradient of Rosenbrock""" x, y = state dfdx = -2*(1-x) - 400*x*(y - x**2) dfdy = 200*(y - x**2) return np.array([dfdx, dfdy]) def neg_grad_rosenbrock(state): """Negative gradient (for descent)""" return -rosenbrock_grad(state) # Create contour plot x = np.linspace(-2, 2, 100) y = np.linspace(-1, 3, 100) X, Y = np.meshgrid(x, y) Z = np.zeros_like(X) for i in range(X.shape[0]): for j in range(X.shape[1]): Z[i,j] = rosenbrock(np.array([X[i,j], Y[i,j]])) fig, ax = plt.subplots(figsize=(12, 8)) # Contours (log scale for better visualization) levels = np.logspace(-1, 3, 20) c = ax.contour(X, Y, Z, levels=levels, cmap='viridis', alpha=0.7) ax.contourf(X, Y, Z, levels=levels, cmap='viridis', alpha=0.3) # Gradient descent trajectories starts = [np.array([-1.5, 2.5]), np.array([1.5, 2.5]), np.array([-0.5, -0.5]), np.array([0.0, 2.0])] colors = ['red', 'blue', 'orange', 'purple'] for x0, c in zip(starts, colors): # Use small step size for stability t, traj = solve_ode(neg_grad_rosenbrock, x0, (0, 50), 0.001, method='rk4') ax.plot(traj[:, 0], traj[:, 1], '-', color=c, lw=2, alpha=0.8) ax.plot(x0[0], x0[1], 'o', color=c, markersize=10, mec='k') ax.plot(1, 1, 'w*', markersize=20, mec='k', mew=2, label='Global minimum (1,1)') ax.set_xlabel('$x$'); ax.set_ylabel('$y$') ax.set_title('Gradient Descent on Rosenbrock Function\n(Optimization as ODE)') ax.legend(); ax.grid(True, alpha=0.3) plt.colorbar(ax.contourf(X, Y, Z, levels=levels, cmap='viridis', alpha=0), ax=ax, label='$f(x,y)$') plt.show() print("Gradient descent = following the negative gradient field.") print("All trajectories flow toward the minimum at (1, 1).")

Gradient descent = following the negative gradient field.
All trajectories flow toward the minimum at (1, 1).

---

Part 5: Flow Matching — Transporting Distributions¶

The Big Picture¶

In flow matching (generative modeling), we learn a velocity field that transports noise to data:

$$\frac{dx}{dt} = v_\theta(x, t), \quad t \in [0, 1]$$

• $t=0$: Data distribution
• $t=1$: Noise distribution
• Sampling: Start from noise, solve ODE backward to get data

Why Vector Fields?¶

The velocity field $v_\theta(x, t)$ IS a (time-dependent) vector field!

• Training: Learn $v_\theta$ to match the true transport
• Sampling: Solve the ODE using Euler/RK4
• Quality: Better ODE solver = better samples

In [29]:

# Example 5.1: Simulating Flow Matching (Toy Example)
# Transport a Gaussian to a mixture of Gaussians

np.random.seed(42)

# "Data" distribution: mixture of 4 Gaussians
def sample_data(n):
    centers = np.array([[2, 2], [-2, 2], [-2, -2], [2, -2]])
    idx = np.random.randint(0, 4, n)
    return centers[idx] + 0.3 * np.random.randn(n, 2)

# "Noise" distribution: standard Gaussian
def sample_noise(n):
    return np.random.randn(n, 2)

# Linear interpolation path (rectified flow)
def interpolate(x0, x1, t):
    """x_t = (1-t)*x0 + t*x1"""
    return (1 - t) * x0 + t * x1

# True velocity for linear interpolation
def true_velocity(x0, x1):
    """v = x1 - x0 (constant velocity)"""
    return x1 - x0

# Visualize the transport
n_samples = 200
x0 = sample_data(n_samples)  # Data
x1 = sample_noise(n_samples)  # Noise

fig, axes = plt.subplots(1, 5, figsize=(20, 4))
times = [0, 0.25, 0.5, 0.75, 1.0]

for ax, t in zip(axes, times):
    x_t = interpolate(x0, x1, t)
    ax.scatter(x_t[:, 0], x_t[:, 1], alpha=0.5, s=20)
    ax.set_xlim(-4, 4); ax.set_ylim(-4, 4)
    ax.set_title(f'$t = {t}$')
    ax.set_aspect('equal'); ax.grid(True, alpha=0.3)

axes[0].set_ylabel('$y$')
plt.suptitle('Flow Matching: Transporting Data → Noise', y=1.02)
plt.tight_layout(); plt.show()

print("t=0: Data (4 clusters). t=1: Noise (single Gaussian).")
print("Flow matching learns the velocity field that creates this transport.")

# Example 5.1: Simulating Flow Matching (Toy Example) # Transport a Gaussian to a mixture of Gaussians np.random.seed(42) # "Data" distribution: mixture of 4 Gaussians def sample_data(n): centers = np.array([[2, 2], [-2, 2], [-2, -2], [2, -2]]) idx = np.random.randint(0, 4, n) return centers[idx] + 0.3 * np.random.randn(n, 2) # "Noise" distribution: standard Gaussian def sample_noise(n): return np.random.randn(n, 2) # Linear interpolation path (rectified flow) def interpolate(x0, x1, t): """x_t = (1-t)*x0 + t*x1""" return (1 - t) * x0 + t * x1 # True velocity for linear interpolation def true_velocity(x0, x1): """v = x1 - x0 (constant velocity)""" return x1 - x0 # Visualize the transport n_samples = 200 x0 = sample_data(n_samples) # Data x1 = sample_noise(n_samples) # Noise fig, axes = plt.subplots(1, 5, figsize=(20, 4)) times = [0, 0.25, 0.5, 0.75, 1.0] for ax, t in zip(axes, times): x_t = interpolate(x0, x1, t) ax.scatter(x_t[:, 0], x_t[:, 1], alpha=0.5, s=20) ax.set_xlim(-4, 4); ax.set_ylim(-4, 4) ax.set_title(f'$t = {t}$') ax.set_aspect('equal'); ax.grid(True, alpha=0.3) axes[0].set_ylabel('$y$') plt.suptitle('Flow Matching: Transporting Data → Noise', y=1.02) plt.tight_layout(); plt.show() print("t=0: Data (4 clusters). t=1: Noise (single Gaussian).") print("Flow matching learns the velocity field that creates this transport.")

t=0: Data (4 clusters). t=1: Noise (single Gaussian).
Flow matching learns the velocity field that creates this transport.

In [30]:

# Example 5.2: Visualize Velocity Field at Different Times

# For visualization, we'll show the average velocity at each point
# (In practice, this is what the neural network learns)

def estimate_velocity_field(x0_samples, x1_samples, t, grid_x, grid_y):
    """Estimate velocity field by averaging over nearby samples"""
    X, Y = np.meshgrid(grid_x, grid_y)
    U = np.zeros_like(X)
    V = np.zeros_like(Y)
    
    # Interpolated positions at time t
    x_t = interpolate(x0_samples, x1_samples, t)
    velocities = true_velocity(x0_samples, x1_samples)
    
    # For each grid point, average velocities of nearby samples
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            point = np.array([X[i,j], Y[i,j]])
            dists = np.linalg.norm(x_t - point, axis=1)
            weights = np.exp(-dists**2 / 0.5)  # Gaussian weighting
            if weights.sum() > 1e-6:
                U[i,j] = np.average(velocities[:, 0], weights=weights)
                V[i,j] = np.average(velocities[:, 1], weights=weights)
    
    return X, Y, U, V

fig, axes = plt.subplots(1, 3, figsize=(15, 5))
grid = np.linspace(-4, 4, 15)

for ax, t in zip(axes, [0.1, 0.5, 0.9]):
    X, Y, U, V = estimate_velocity_field(x0, x1, t, grid, grid)
    x_t = interpolate(x0, x1, t)
    
    ax.scatter(x_t[:, 0], x_t[:, 1], alpha=0.3, s=10, c='blue')
    ax.quiver(X, Y, U, V, np.sqrt(U**2+V**2), cmap='Reds', scale=50, alpha=0.8)
    ax.set_xlim(-4, 4); ax.set_ylim(-4, 4)
    ax.set_title(f'Velocity Field at $t = {t}$')
    ax.set_aspect('equal'); ax.grid(True, alpha=0.3)

plt.suptitle('Flow Matching: Learned Velocity Field', y=1.02)
plt.tight_layout(); plt.show()

print("The velocity field points from data toward noise.")
print("To SAMPLE: start from noise, follow NEGATIVE velocity (reverse ODE).")

# Example 5.2: Visualize Velocity Field at Different Times # For visualization, we'll show the average velocity at each point # (In practice, this is what the neural network learns) def estimate_velocity_field(x0_samples, x1_samples, t, grid_x, grid_y): """Estimate velocity field by averaging over nearby samples""" X, Y = np.meshgrid(grid_x, grid_y) U = np.zeros_like(X) V = np.zeros_like(Y) # Interpolated positions at time t x_t = interpolate(x0_samples, x1_samples, t) velocities = true_velocity(x0_samples, x1_samples) # For each grid point, average velocities of nearby samples for i in range(X.shape[0]): for j in range(X.shape[1]): point = np.array([X[i,j], Y[i,j]]) dists = np.linalg.norm(x_t - point, axis=1) weights = np.exp(-dists**2 / 0.5) # Gaussian weighting if weights.sum() > 1e-6: U[i,j] = np.average(velocities[:, 0], weights=weights) V[i,j] = np.average(velocities[:, 1], weights=weights) return X, Y, U, V fig, axes = plt.subplots(1, 3, figsize=(15, 5)) grid = np.linspace(-4, 4, 15) for ax, t in zip(axes, [0.1, 0.5, 0.9]): X, Y, U, V = estimate_velocity_field(x0, x1, t, grid, grid) x_t = interpolate(x0, x1, t) ax.scatter(x_t[:, 0], x_t[:, 1], alpha=0.3, s=10, c='blue') ax.quiver(X, Y, U, V, np.sqrt(U**2+V**2), cmap='Reds', scale=50, alpha=0.8) ax.set_xlim(-4, 4); ax.set_ylim(-4, 4) ax.set_title(f'Velocity Field at $t = {t}$') ax.set_aspect('equal'); ax.grid(True, alpha=0.3) plt.suptitle('Flow Matching: Learned Velocity Field', y=1.02) plt.tight_layout(); plt.show() print("The velocity field points from data toward noise.") print("To SAMPLE: start from noise, follow NEGATIVE velocity (reverse ODE).")

The velocity field points from data toward noise.
To SAMPLE: start from noise, follow NEGATIVE velocity (reverse ODE).

---

Summary¶

Key Concepts¶

Concept	Meaning	Application
ODE	$\dot{x} = F(x)$	Follow the vector field
Euler method	$x_{n+1} = x_n + \Delta t \cdot F(x_n)$	Simple, 1st order
RK4	Weighted average of 4 samples	Accurate, 4th order
Fixed point	$F(x^*) = 0$	Equilibrium
Gradient flow	$\dot{x} = -\nabla f$	Optimization
Flow matching	$\dot{x} = v_\theta(x, t)$	Generative modeling

Key Equations¶

Euler step: $$x_{n+1} = x_n + \Delta t \cdot F(x_n)$$

RK4 step: $$x_{n+1} = x_n + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4)$$

Gradient descent: $$\frac{dx}{dt} = -\nabla f(x)$$

Flow matching: $$\frac{dx}{dt} = v_\theta(x, t)$$

Connections¶

• GRL: Policy improvement via gradient ascent on $Q^+$
• Flow Matching: Sample generation via ODE integration
• Diffusion Models: Probability flow ODE

Next: Notebook 2 — Functional Fields (functions as vectors in RKHS)