Finding Nature’s Optimal Path with Python

Every trajectory a physical system actually follows — a thrown ball, a planet, a beam of light — has a remarkable property: among all imaginable paths connecting a start point to an end point, the real one is the single path that makes the action stationary. This single idea, the principle of least (stationary) action, quietly underlies the whole of classical mechanics, and reproduces Newton’s laws as a special case.

In this article we take the simplest non-trivial case — a mass thrown straight up and falling back down under gravity — and use it to make the abstract idea concrete. We won’t just solve the equations of motion; we will build a small numerical laboratory that computes the action for a whole family of candidate paths and shows, visually, that the true physical trajectory sits exactly at the bottom of the action “valley.”

1. Setting up the problem

For a particle of mass $m$ moving along a single coordinate $x(t)$ under gravity, the Lagrangian is kinetic minus potential energy:

$$L(x,\dot x) = \frac{1}{2}m\dot x^2 - mgx$$

The action is the time integral of the Lagrangian along a path:

$$S[x] = \int_0^T L\big(x(t),\dot x(t)\big),dt$$

The principle of stationary action says that the physically realized path $x_{cl}(t)$ is the one for which $S$ does not change to first order under any small variation of the path that keeps the endpoints $x(0)=x_0$ and $x(T)=x_T$ fixed. Formally, this condition is expressed by the Euler–Lagrange equation:

$$\frac{d}{dt}\left(\frac{\partial L}{\partial \dot x}\right) - \frac{\partial L}{\partial x} = 0$$

Substituting our Lagrangian gives $m\ddot x = -mg$, i.e. $\ddot x = -g$ — Newton’s second law falls straight out. Integrating twice and fixing the boundary conditions gives the familiar parabola:

$$x_{cl}(t) = x_0 + v_0 t - \frac{1}{2}g t^2, \qquad v_0 = \frac{x_T-x_0}{T} + \frac{1}{2}gT$$

That’s the answer everyone already knows. The interesting part is verifying the variational principle itself: showing numerically that this particular parabola is truly the minimum of $S$ among a whole space of nearby competing paths.

2. Building a family of trial paths (the Ritz method)

To test many paths at once, we describe a trial path as the classical solution plus a correction built from a couple of sine modes that automatically vanish at both endpoints:

$$x(t;a_1,a_2) = x_{cl}(t) + a_1\eta_1(t) + a_2\eta_2(t), \qquad \eta_n(t) = \sin!\left(\frac{n\pi t}{T}\right)$$

Because $\eta_n(0)=\eta_n(T)=0$, every choice of $(a_1,a_2)$ satisfies the same boundary conditions as $x_{cl}$. Since $L$ is a quadratic function of $\dot x$ and a linear function of $x$, the action restricted to this two-parameter family is exactly a quadratic form:

$$S(a_1,a_2) = S_0 + b_1a_1 + b_2a_2 + \frac{1}{2}\Big(Q_{11}a_1^2 + 2Q_{12}a_1a_2 + Q_{22}a_2^2\Big)$$

If $x_{cl}$ truly satisfies the Euler–Lagrange equation, the linear (first-order) coefficients $b_1, b_2$ must vanish, meaning $(a_1,a_2)=(0,0)$ is a genuine stationary point of this paraboloid. That’s precisely what the code below checks and visualizes.

3. Python implementation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401  (registers the 3D projection)

plt.style.use('dark_background')

# ------------------------------------------------------------
# 1. Physical setup
# ------------------------------------------------------------
m  = 1.0     # mass [kg]
g  = 9.8     # gravitational acceleration [m/s^2]
T  = 2.0     # total flight time [s]
x0 = 0.0     # position at t = 0  [m]
xT = 0.0     # position at t = T  [m]  (thrown up, lands back at the same height)

N  = 400
t  = np.linspace(0.0, T, N)

def trapz(y, t):
    """Simple, version-independent trapezoidal integration."""
    return np.sum((y[1:] + y[:-1]) * np.diff(t) * 0.5)

# ------------------------------------------------------------
# 2. Classical (Euler-Lagrange) solution
#    x(t) = x0 + v0*t - 0.5*g*t^2,  v0 fixed by the boundary condition x(T) = xT
# ------------------------------------------------------------
v0 = (xT - x0) / T + 0.5 * g * T
x_classical = x0 + v0 * t - 0.5 * g * t**2

# ------------------------------------------------------------
# 3. Trial-path basis (Ritz method)
#    eta_n(t) = sin(n*pi*t/T) vanishes at t=0 and t=T, so
#    x_classical + a1*eta1 + a2*eta2 always keeps the boundary conditions.
# ------------------------------------------------------------
def eta(n, t, T):
    return np.sin(n * np.pi * t / T)

n1, n2 = 1, 2
eta1, eta2 = eta(n1, t, T), eta(n2, t, T)

v_c  = np.gradient(x_classical, t)
v_e1 = np.gradient(eta1, t)
v_e2 = np.gradient(eta2, t)

# ------------------------------------------------------------
# 4. Action functional S[x] = ∫ (0.5*m*v^2 - m*g*x) dt
#    L is quadratic in v and linear in x, so S(a1,a2) is EXACTLY
#    a quadratic form -> the whole (a1,a2) surface below needs
#    no explicit loop at all.
# ------------------------------------------------------------
S0  = trapz(0.5*m*v_c**2 - m*g*x_classical, t)

b1  = m*trapz(v_c*v_e1, t) - m*g*trapz(eta1, t)
b2  = m*trapz(v_c*v_e2, t) - m*g*trapz(eta2, t)

Q11 = m*trapz(v_e1**2, t)
Q22 = m*trapz(v_e2**2, t)
Q12 = m*trapz(v_e1*v_e2, t)

def action_quadratic(a1, a2):
    return S0 + b1*a1 + b2*a2 + 0.5*(Q11*a1**2 + 2*Q12*a1*a2 + Q22*a2**2)

Q = np.array([[Q11, Q12], [Q12, Q22]])
b = np.array([b1, b2])
a_star = np.linalg.solve(Q, -b)

print("=== Stationary-action check ===")
print(f"S0  (action of the classical path) = {S0:.6f} J*s")
print(f"b1  (1st-order term, mode n={n1})   = {b1:.3e}")
print(f"b2  (1st-order term, mode n={n2})   = {b2:.3e}")
print(f"Q11 = {Q11:.6f}   Q22 = {Q22:.6f}   Q12 = {Q12:.3e}")
print(f"Optimal (a1*, a2*) found by solving Q a = -b : {a_star}")
print("-> should be numerically indistinguishable from (0, 0)")

# ------------------------------------------------------------
# 5. Sweeps for visualization
# ------------------------------------------------------------
eps = np.linspace(-2.0, 2.0, 200)
S_eps = action_quadratic(eps, 0.0)          # 1D slice along mode n1 only

a1v = np.linspace(-1.5, 1.5, 80)
a2v = np.linspace(-1.5, 1.5, 80)
A1, A2 = np.meshgrid(a1v, a2v)
S_grid = action_quadratic(A1, A2)           # full 2D paraboloid, fully vectorized

# ------------------------------------------------------------
# 6. Plots
# ------------------------------------------------------------
fig = plt.figure(figsize=(17, 5.2))

# --- Panel 1: classical path vs. nearby trial paths ---
ax1 = fig.add_subplot(1, 3, 1)
ax1.plot(t, x_classical, color='#58a6ff', lw=2.8, label='classical path (E-L solution)')
sample_coeffs = [(0.6, 0.0), (-0.6, 0.0), (0.0, 0.6), (0.4, -0.5)]
colors = ['#f78166', '#3fb950', '#d29922', '#bc8cff']
for (a1s, a2s), c in zip(sample_coeffs, colors):
    x_trial = x_classical + a1s*eta1 + a2s*eta2
    S_trial = action_quadratic(a1s, a2s)
    ax1.plot(t, x_trial, color=c, lw=1.3, ls='--', alpha=0.85,
             label=f'a1={a1s}, a2={a2s}  (S={S_trial:.3f})')
ax1.set_xlabel('time  t  [s]')
ax1.set_ylabel('position  x(t)  [m]')
ax1.set_title('Classical path vs. nearby trial paths')
ax1.legend(fontsize=7, loc='lower center')
ax1.grid(alpha=0.2)

# --- Panel 2: 1D slice of the action ---
ax2 = fig.add_subplot(1, 3, 2)
ax2.plot(eps, S_eps, color='#58a6ff', lw=2.8)
ax2.axvline(0, color='#f78166', ls=':', lw=1.5)
ax2.scatter([0], [S0], color='#f78166', zorder=5, label=f'minimum at ε=0  (S={S0:.4f})')
ax2.set_xlabel('perturbation amplitude  ε  (mode n=1)')
ax2.set_ylabel('action  S  [J·s]')
ax2.set_title('Action along a single perturbation direction')
ax2.legend(fontsize=8)
ax2.grid(alpha=0.2)

# --- Panel 3: 3D paraboloid over the (a1, a2) trial-path space ---
ax3 = fig.add_subplot(1, 3, 3, projection='3d')
for axis in (ax3.xaxis, ax3.yaxis, ax3.zaxis):
    axis.pane.set_facecolor((0.05, 0.05, 0.08, 1.0))
    axis.pane.set_edgecolor('gray')
surf = ax3.plot_surface(A1, A2, S_grid, cmap='plasma', linewidth=0, antialiased=True, alpha=0.95)
ax3.scatter([0], [0], [S0], color='cyan', s=70, depthshade=False, label='classical path')
ax3.set_xlabel('a1 (mode n=1)')
ax3.set_ylabel('a2 (mode n=2)')
ax3.set_zlabel('action  S')
ax3.set_title('Action functional over trial-path space')
ax3.view_init(elev=26, azim=-55)
fig.colorbar(surf, ax=ax3, shrink=0.5, pad=0.12, label='S(a1, a2)')

plt.tight_layout()
plt.show()

4. Reading the code

Physical setup. x0, xT, and T define a boundary-value problem: where the particle starts, where it ends up, and how long the trip takes. v0 is derived directly from the boundary conditions using the closed-form parabola, so x_classical is the exact Euler–Lagrange solution — no ODE solver is needed here because the equation is simple enough to integrate by hand.

Trial-path basis. eta1 and eta2 are two independent sine modes that are exactly zero at $t=0$ and $t=T$. Adding any multiple of them to x_classical produces a new candidate path that still satisfies the same boundary conditions — this is the essence of the calculus of variations: comparing the true path only against paths that share its endpoints.

Why no loop is needed. Because the Lagrangian is quadratic in velocity and linear in position, the action of x_classical + a1*eta1 + a2*eta2 is algebraically a quadratic function of $(a_1,a_2)$. Rather than looping over a grid and numerically integrating the action thousands of times, the code integrates the six fixed coefficients ($S_0$, $b_1$, $b_2$, $Q_{11}$, $Q_{12}$, $Q_{22}$) once, and then evaluates action_quadratic(a1, a2) with plain NumPy arithmetic. This makes the full 80×80 surface (6,400 points) essentially free to compute, since it’s just element-wise operations on arrays instead of thousands of independent numerical integrations.

The stationarity check. b1 and b2 are the first-order sensitivities of the action to each perturbation mode. If x_classical truly extremizes the action, these must come out numerically negligible (they’ll print as very small floating-point residuals, not exactly zero, due to the discretization). Solving Q @ a = -b for the true minimum of the paraboloid and finding it lands on $(0,0)$ is the quantitative proof that the physical trajectory is the stationary point of the action.

Custom trapz. A hand-written trapezoidal integrator is used instead of relying on a specific NumPy/SciPy function name, since integration helpers have moved across library versions — this keeps the notebook robust regardless of the exact NumPy version Colab happens to have installed.

5. What the plots show

Panel 1 — trajectories. The solid blue parabola is the classical, physically-realized path. The dashed lines are competitor paths built by adding sine-mode “wiggles” — they still start and end at the same place and time, but bulge away from the parabola. Each one is printed with its own action value.

Panel 2 — the 1D action slice. This is a clean parabola in $\varepsilon$ with its minimum sitting exactly at $\varepsilon = 0$ — the classical path. Every other value of $\varepsilon$, in either direction, costs strictly more action.

Panel 3 — the 3D action landscape. This is the most direct visualization of “least action”: the whole surface is a paraboloid (bowl) in the two-parameter trial-path space, and the cyan marker — the classical solution — sits precisely at the bottom of that bowl. Any direction you move away from it, the action only increases.

=== Stationary-action check ===
S0  (action of the classical path) = -32.013734 J*s
b1  (1st-order term, mode n=1)   = -1.934e-04
b2  (1st-order term, mode n=2)   = 7.105e-15
Q11 = 2.467350   Q22 = 9.868789   Q12 = 8.882e-16
Optimal (a1*, a2*) found by solving Q a = -b : [ 7.83781287e-05 -7.19996876e-16]
-> should be numerically indistinguishable from (0, 0)

6. Takeaways

What makes this example satisfying is that it turns an abstract variational statement into something you can literally see: a bowl-shaped surface with the true law of motion sitting at the bottom of it. The same machinery — build a family of trial paths, form the action as a function of the trial parameters, and look for the stationary point — is exactly how far more sophisticated problems (geodesics in general relativity, optimal control, path-integral formulations of quantum mechanics) are approached numerically. Free fall under gravity is just the smallest possible sandbox in which to watch the principle of least action do its work.

How a Hanging Chain Finds Its Own Shape

Hang a rope between two poles and let gravity do its work. The curve that appears is not a parabola, though it looks deceptively similar — it is a catenary, and it is the unique shape that minimizes the rope’s total gravitational potential energy while keeping its length fixed. This article works through the physics, derives the closed-form solution, and then verifies it numerically by treating the rope as a chain of rigid links and minimizing its energy directly with a constrained optimizer.

The Physical Setup

Consider a flexible, inextensible chain of length $L$ suspended between two fixed points at the same height, separated by a horizontal distance $D$, with $L > D$. Under gravity, the chain settles into the shape $y(x)$ that minimizes its total potential energy

$$
U[y] = \rho g \int_0^{D} y(x),\sqrt{1 + y’(x)^2};dx
$$

subject to the length constraint

$$
\int_0^{D} \sqrt{1 + y’(x)^2};dx = L
$$

and the boundary conditions $y(0) = y(D) = 0$.

Deriving the Catenary Equation

Introducing a Lagrange multiplier $\lambda$ for the length constraint turns this into an unconstrained variational problem for the functional

$$
F(y, y’) = (y - \lambda)\sqrt{1 + y’^2}
$$

Since $F$ does not depend explicitly on $x$, the Beltrami identity applies, giving a first integral of the Euler–Lagrange equation. Working through the algebra leads to the classical result

$$
y(x) = a \cosh!\left(\frac{x - D/2}{a}\right) - a\cosh!\left(\frac{D}{2a}\right)
$$

where the shape parameter $a$ is fixed by requiring the arc length to equal $L$:

$$
L = 2a\sinh!\left(\frac{D}{2a}\right)
$$

This is a transcendental equation in $a$ with no closed-form inverse, so it has to be solved numerically — a perfect entry point for Python.

A Numerical Cross-Check: The Chain as N Rigid Links

To verify the analytical result independently, the rope can be modeled as $N$ rigid links, each of fixed length $l = L/N$, connected end to end. The free parameters are the $N$ link angles $\theta_i$ measured from the horizontal. Once the angles are known, the joint positions follow from a cumulative sum of the link displacement vectors. The physically realized configuration is the one that minimizes total potential energy

$$
U(\theta) = \rho g \sum_{i=1}^{N} \bar{y}_i , l
$$

where $\bar{y}_i$ is the height of the midpoint of link $i$, subject to the constraint that the last joint lands exactly on the second support point. This turns the problem into a standard constrained nonlinear optimization, solved here with SLSQP.

Python Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181

# ==========================================================
# The Catenary Problem: Minimum Potential Energy Chain Shape
# ==========================================================

import numpy as np
from scipy.optimize import brentq, minimize
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (enables 3D projection)

plt.style.use('dark_background')

# ----------------------------------------------------------
# 1. Problem setup
# ----------------------------------------------------------
D = 10.0        # horizontal distance between the two fixed supports
L = 12.0        # total length of the chain/rope (must be > D)
N = 80          # number of rigid segments used in the numerical model

assert L > D, "Chain length L must exceed the horizontal span D."

# ----------------------------------------------------------
# 2. Analytical solution via the catenary equation
#
#    y(x) = a * cosh( (x - D/2) / a ) - a * cosh( D / (2a) )
#    where 'a' solves   L = 2a * sinh( D / (2a) )
# ----------------------------------------------------------

def length_residual(a, D, L):
    return 2.0 * a * np.sinh(D / (2.0 * a)) - L

a_lower = D / 100.0     # D/(2*a_lower) = 50   -> finite but very large residual
a_upper = D * 1000.0    # D/(2*a_upper) ~ 5e-4 -> residual approaches D - L < 0

a_solution = brentq(length_residual, a_lower, a_upper, args=(D, L), xtol=1e-12)

def catenary_y(x, a, D):
    return a * (np.cosh((x - D / 2.0) / a) - np.cosh(D / (2.0 * a)))

x_analytic = np.linspace(0.0, D, 400)
y_analytic = catenary_y(x_analytic, a_solution, D)

sag_depth = -np.min(y_analytic)

# ----------------------------------------------------------
# 3. Numerical solution: discretized chain of N rigid links
#
#    Free variables: link angles theta_i. Positions are obtained
#    with a vectorized cumulative sum (no Python-level loop),
#    which keeps the model fast even for large N.
# ----------------------------------------------------------

l_seg = L / N

def chain_positions(theta):
    dx = l_seg * np.cos(theta)
    dy = l_seg * np.sin(theta)
    x = np.concatenate(([0.0], np.cumsum(dx)))
    y = np.concatenate(([0.0], np.cumsum(dy)))
    return x, y

def potential_energy(theta):
    x, y = chain_positions(theta)
    y_mid = 0.5 * (y[:-1] + y[1:])
    return np.sum(y_mid) * l_seg

def endpoint_constraints(theta):
    x, y = chain_positions(theta)
    return np.array([x[-1] - D, y[-1] - 0.0])

# Initial guess: local slope of the analytical catenary
x_mid_guess = (np.arange(N) + 0.5) * (D / N)
slope_guess = np.sinh((x_mid_guess - D / 2.0) / a_solution)
theta0 = np.arctan(slope_guess)

result = minimize(
    potential_energy,
    theta0,
    method='SLSQP',
    constraints=[{'type': 'eq', 'fun': endpoint_constraints}],
    options={'maxiter': 500, 'ftol': 1e-9}
)

theta_opt = result.x
x_numeric, y_numeric = chain_positions(theta_opt)

# ----------------------------------------------------------
# 4. Validation: compare numeric nodes against the analytical curve
# ----------------------------------------------------------
y_analytic_at_nodes = catenary_y(x_numeric, a_solution, D)
rmse = np.sqrt(np.mean((y_numeric - y_analytic_at_nodes) ** 2))

print(f"Catenary parameter a           : {a_solution:.6f}")
print(f"Maximum sag depth              : {sag_depth:.6f}")
print(f"Optimizer status               : {result.message}")
print(f"Potential energy at optimum    : {result.fun:.6f}")
print(f"RMSE (numeric vs analytical)   : {rmse:.6e}")

# ----------------------------------------------------------
# 5. Figure 1: analytical curve vs discretized numerical chain
# ----------------------------------------------------------
fig1, ax1 = plt.subplots(figsize=(9, 5.5))
ax1.plot(x_analytic, y_analytic, color='#4FC3F7', linewidth=2.5,
         label='Analytical catenary')
ax1.plot(x_numeric, y_numeric, 'o', color='#FF7043', markersize=4,
         label=f'Numerical chain (N={N} links)')
ax1.plot([0, D], [0, 0], 's', color='#FFD54F', markersize=9,
         label='Fixed supports')
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_title('Catenary Shape: Analytical vs Energy-Minimizing Discrete Chain')
ax1.legend(loc='lower center')
ax1.set_aspect('equal')
ax1.grid(alpha=0.25)
plt.tight_layout()
plt.show()

# ----------------------------------------------------------
# 6. Figure 2: family of catenaries for varying chain length L
# ----------------------------------------------------------
L_values = np.linspace(D * 1.02, D * 2.0, 8)
fig2 = plt.figure(figsize=(9, 7))
ax2 = fig2.add_subplot(111, projection='3d')

cmap = plt.get_cmap('plasma')
for i, L_i in enumerate(L_values):
    a_i = brentq(length_residual, D / 100.0, D * 1000.0, args=(D, L_i))
    y_i = catenary_y(x_analytic, a_i, D)
    ax2.plot(x_analytic, np.full_like(x_analytic, L_i), y_i,
             color=cmap(i / (len(L_values) - 1)), linewidth=2)

ax2.set_xlabel('x')
ax2.set_ylabel('Chain length L')
ax2.set_zlabel('y')
ax2.set_title('Family of Catenary Shapes for Increasing Chain Length')
ax2.view_init(elev=22, azim=-60)
plt.tight_layout()
plt.show()

# ----------------------------------------------------------
# 7. Figure 3: potential-energy landscape around the true minimum
#
#    Trial shapes: y(x; s, t) = s * y_analytic(x) + t * sin(pi x / D)
#    's' scales the true solution, 't' adds an independent mode
#    that also vanishes at both endpoints (Rayleigh-Ritz style check:
#    every trial shape has energy >= the true minimum).
# ----------------------------------------------------------

s_values = np.linspace(0.5, 1.5, 41)
t_values = np.linspace(-0.5 * sag_depth, 0.5 * sag_depth, 41)
S, T = np.meshgrid(s_values, t_values)
E = np.zeros_like(S)

sin_mode = np.sin(np.pi * x_analytic / D)

for i in range(S.shape[0]):
    for j in range(S.shape[1]):
        y_trial = S[i, j] * y_analytic + T[i, j] * sin_mode
        dy_dx = np.gradient(y_trial, x_analytic)
        ds = np.sqrt(1.0 + dy_dx ** 2)
        E[i, j] = np.trapz(y_trial * ds, x_analytic)

min_idx = np.unravel_index(np.argmin(E), E.shape)
s_min, t_min, E_min = S[min_idx], T[min_idx], E[min_idx]

fig3 = plt.figure(figsize=(9, 7))
ax3 = fig3.add_subplot(111, projection='3d')
surf = ax3.plot_surface(S, T, E, cmap='viridis', alpha=0.9,
                         linewidth=0, antialiased=True)
ax3.scatter([s_min], [t_min], [E_min], color='red', s=60,
            label='Grid minimum')
ax3.set_xlabel('s (amplitude scale)')
ax3.set_ylabel('t (perturbation mode)')
ax3.set_zlabel('Potential energy')
ax3.set_title('Potential-Energy Landscape Around the Catenary Solution')
fig3.colorbar(surf, shrink=0.6, aspect=12, label='Energy')
ax3.legend(loc='upper left')
plt.tight_layout()
plt.show()

print(f"Grid-search minimum located at s={s_min:.3f}, t={t_min:.4f}")
print("(expected near s=1.000, t=0.0000 for the true catenary)")

Code Walkthrough

Section 2 — Analytical solution. The transcendental equation $L = 2a\sinh(D/2a)$ has no algebraic inverse, so scipy.optimize.brentq is used to bracket and find the root numerically. The bracket $[D/100,\ 1000D]$ is chosen deliberately: at the lower bound, $D/(2a)=50$, which keeps $\sinh(50)$ large but finite (no floating-point overflow), guaranteeing a large positive residual; at the upper bound the residual approaches $D-L<0$. Because the residual is strictly monotonic between these two values, brentq is guaranteed to converge.

Section 3 — Discretized chain. Rather than treating each joint’s $(x,y)$ coordinates as independent variables (which would require handling $2N$ unknowns and per-segment length constraints), the model uses one angle $\theta_i$ per link. This automatically enforces that every link has exactly length $l = L/N$, cutting the constraint count down to just two equations — the horizontal and vertical position of the last joint. Positions are computed with np.cumsum, a fully vectorized operation, so the model scales cleanly to large $N$ without any Python-level loop overhead. The initial guess is not arbitrary: it uses the slope of the already-known analytical catenary, which gives SLSQP a near-optimal starting point and ensures fast, reliable convergence.

Section 4 — Validation. The discretized joints are compared directly against the analytical curve evaluated at the same $x$-coordinates. Agreement at the level of the RMSE reported in the console output confirms that the two independent methods — closed-form calculus of variations and constrained numerical optimization — describe the same physical shape.

Section 6 — Family of curves. By resolving the transcendental equation for several values of $L$ and stacking the resulting curves along a third axis, the 3D plot shows how the chain sags more deeply as extra length is added while the support span $D$ stays fixed — a direct visualization of the constraint’s effect on the energy-minimizing shape.

Section 7 — Energy landscape. This is the most direct illustration of “potential energy minimization” as an optimization concept. A two-parameter family of trial shapes is built around the true catenary: $s$ uniformly scales it, and $t$ adds an independent sine-shaped perturbation that also respects the endpoint conditions. The resulting energy surface is a bowl, and the grid-based minimum should land close to $s=1,\ t=0$ — the true analytical solution — visually confirming that the catenary is not merely a low-energy shape among nearby alternatives, but the minimum.

Catenary parameter a           : 4.695415
Maximum sag depth              : 2.923421
Optimizer status               : Optimization terminated successfully
Potential energy at optimum    : -22.234566
RMSE (numeric vs analytical)   : 1.631274e-04

/tmp/ipykernel_527/4050898374.py:160: DeprecationWarning: `trapz` is deprecated. Use `trapezoid` instead, or one of the numerical integration functions in `scipy.integrate`.
  E[i, j] = np.trapz(y_trial * ds, x_analytic)

Grid-search minimum located at s=1.500, t=-1.4617
(expected near s=1.000, t=0.0000 for the true catenary)

Discussion

The two solution methods approach the same problem from opposite directions. The analytical route treats the chain as a continuous curve and applies the calculus of variations directly, yielding an exact — if transcendental — formula. The numerical route discretizes the chain into finitely many rigid links and finds the energy minimum through constrained nonlinear programming, making no assumption about the functional form of the solution in advance. That both methods converge to the same curve is a strong practical confirmation of the underlying physics, and it is also exactly the kind of cross-check worth running whenever a discretized optimization model is proposed for a problem that already has a known closed-form answer.

Beyond ropes and chains, the same mathematics governs the resting shape of power transmission lines, the profile of certain arches (an inverted catenary is the ideal compression-only arch shape), and the equilibrium configuration of suspension bridge cables before the deck load is added. The Rayleigh–Ritz style landscape in Figure 3 also previews a broader idea used throughout computational mechanics and structural optimization: whenever a system settles into the configuration of least potential energy, nearby trial configurations can be swept numerically to confirm — and visualize — that the true physical solution truly sits at the bottom of the energy bowl.

Finding the Path of Least Time with Python

What is Fermat’s Principle?

Fermat’s Principle states that light travels between two points along the path that takes the least time (more precisely, an extremal time). In a medium with refractive index $n$, light travels at speed $v = c/n$, and the total travel time along a path is:

$$T = \int_A^B \frac{n(\mathbf{r})}{c}, ds$$

When light crosses a boundary between two media, minimizing this time functional leads directly to Snell’s Law:

$$n_1 \sin\theta_1 = n_2 \sin\theta_2$$

Rather than deriving this law analytically, this article demonstrates it computationally — we let a numerical optimizer discover the fastest path on its own, and then check that it reproduces Snell’s Law.

Setting Up a Concrete Example

Consider a flat interface at $z = 0$, separating two media:

• Medium 1 (e.g. air, $n_1 = 1.00$) occupies $z > 0$
• Medium 2 (e.g. water, $n_2 = 1.33$) occupies $z < 0$

Light starts at point $A = (0, 0, 5)$ in medium 1 and must reach point $B = (8, 3, -4)$ in medium 2. It crosses the interface at some unknown point $P = (x, y, 0)$.

The total travel time as a function of the crossing point is:

$$T(x,y) = \frac{n_1}{c}\sqrt{x^2 + y^2 + h_1^2} ;+; \frac{n_2}{c}\sqrt{(\Delta x - x)^2 + (\Delta y - y)^2 + h_2^2}$$

where $h_1 = 5$, $h_2 = 4$, $\Delta x = 8$, $\Delta y = 3$. The goal is to numerically find the $(x,y)$ that minimizes $T$, and confirm that the resulting angles satisfy Snell’s Law.

Full Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (enables 3D projection)

# =========================================================
# 1. Physical setting
#    Light travels from point A (medium 1, e.g. air)
#    to point B (medium 2, e.g. water), crossing a flat
#    interface located at z = 0.
# =========================================================
c  = 299792458.0        # speed of light in vacuum [m/s]
n1 = 1.00                # refractive index of medium 1 (air)
n2 = 1.33                # refractive index of medium 2 (water)

A = np.array([0.0, 0.0, 5.0])   # source point      (z > 0 -> medium 1)
B = np.array([8.0, 3.0, -4.0])  # destination point (z < 0 -> medium 2)

# =========================================================
# 2. Travel time as a function of the crossing point P=(x,y,0)
#    T(x,y) = n1/c * |AP| + n2/c * |PB|
# =========================================================
def travel_time(p):
    x, y = p
    d1 = np.sqrt(x**2 + y**2 + A[2]**2)
    d2 = np.sqrt((B[0]-x)**2 + (B[1]-y)**2 + B[2]**2)
    return n1 * d1 / c + n2 * d2 / c

p0 = np.array([A[0] + (B[0]-A[0])*0.5, A[1] + (B[1]-A[1])*0.5])  # straight-line initial guess

result = minimize(travel_time, p0, method='BFGS',
                   options={'gtol': 1e-12, 'maxiter': 2000})

Px, Py = result.x
P = np.array([Px, Py, 0.0])
T_min = result.fun

# =========================================================
# 3. Verify Snell's law along the found path
# =========================================================
horiz_in  = np.hypot(Px - A[0], Py - A[1])
horiz_out = np.hypot(B[0] - Px, B[1] - Py)
dist_in   = np.hypot(horiz_in,  A[2])
dist_out  = np.hypot(horiz_out, B[2])

sin_theta1 = horiz_in  / dist_in
sin_theta2 = horiz_out / dist_out

ratio_B = B[1] / B[0]
ratio_P = Py / Px

print("========== Fermat's Principle: numerical result ==========")
print(f"Crossing point P             : ({Px:.6f}, {Py:.6f}, 0.000000)")
print(f"Minimum travel time T_min    : {T_min*1e9:.6f} ns")
print(f"n1 * sin(theta1)             : {n1*sin_theta1:.6f}")
print(f"n2 * sin(theta2)             : {n2*sin_theta2:.6f}")
print(f"slope y/x of B               : {ratio_B:.6f}")
print(f"slope y/x of P (should match): {ratio_P:.6f}")
print("============================================================")

# =========================================================
# 4. Compare with a naive straight-line path (no refraction)
# =========================================================
T_straight = travel_time(p0)
print(f"\nTravel time of the naive straight-line path : {T_straight*1e9:.6f} ns")
print(f"Travel time of the Fermat (fastest) path    : {T_min*1e9:.6f} ns")
print(f"Time saved by refraction                    : {(T_straight-T_min)*1e9:.6f} ns")

# =========================================================
# 5. Vectorised grid for visualising T(x,y) as a surface
#    (numpy broadcasting instead of double for-loops => fast)
# =========================================================
grid_n = 300
xs = np.linspace(-4, 12, grid_n)
ys = np.linspace(-4, 8,  grid_n)
X, Y = np.meshgrid(xs, ys)

D1 = np.sqrt(X**2 + Y**2 + A[2]**2)
D2 = np.sqrt((B[0]-X)**2 + (B[1]-Y)**2 + B[2]**2)
T  = n1 * D1 / c + n2 * D2 / c

# =========================================================
# 6. Figure 1: 3D visualisation of the actual light path
# =========================================================
fig1 = plt.figure(figsize=(9, 7))
ax1 = fig1.add_subplot(111, projection='3d')

plane_x, plane_y = np.meshgrid(np.linspace(-4, 12, 2), np.linspace(-4, 8, 2))
ax1.plot_surface(plane_x, plane_y, np.zeros_like(plane_x),
                  color='lightblue', alpha=0.3, edgecolor='none')

ax1.plot([A[0], P[0]], [A[1], P[1]], [A[2], P[2]], 'r-', linewidth=2.5, label='Ray in medium 1')
ax1.plot([P[0], B[0]], [P[1], B[1]], [P[2], B[2]], 'b-', linewidth=2.5, label='Ray in medium 2')

ax1.scatter(*A, color='black', s=60)
ax1.scatter(*B, color='black', s=60)
ax1.scatter(*P, color='green', s=80, label='Crossing point P (fastest)')

ax1.text(*A, '  A', fontsize=11)
ax1.text(*B, '  B', fontsize=11)
ax1.text(*P, '  P', fontsize=11)

ax1.set_xlabel('X')
ax1.set_ylabel('Y')
ax1.set_zlabel('Z')
ax1.set_title("Fermat's Principle: fastest path from A to B")
ax1.legend()
plt.tight_layout()
plt.show()

# =========================================================
# 7. Figure 2: 3D surface of T(x,y) with the minimum marked
# =========================================================
fig2 = plt.figure(figsize=(9, 7))
ax2 = fig2.add_subplot(111, projection='3d')

surf = ax2.plot_surface(X, Y, T*1e9, cmap='viridis', alpha=0.85, linewidth=0, antialiased=True)
ax2.scatter(Px, Py, T_min*1e9, color='red', s=80, label='Minimum (Fermat point)')

ax2.set_xlabel('x (crossing point)')
ax2.set_ylabel('y (crossing point)')
ax2.set_zlabel('Travel time [ns]')
ax2.set_title('Travel time T(x, y) for every possible crossing point')
fig2.colorbar(surf, shrink=0.6, aspect=12, label='Travel time [ns]')
ax2.legend()
plt.tight_layout()
plt.show()

# =========================================================
# 8. Figure 3: 2D slice through the minimum (easy-to-read cut)
# =========================================================
direction = np.array([B[0], B[1]]) / np.hypot(B[0], B[1])
t_vals = np.linspace(-4, 12, 400)
line_x = direction[0] * t_vals
line_y = direction[1] * t_vals
T_line = n1*np.sqrt(line_x**2 + line_y**2 + A[2]**2)/c + \
         n2*np.sqrt((B[0]-line_x)**2 + (B[1]-line_y)**2 + B[2]**2)/c

fig3, ax3 = plt.subplots(figsize=(8, 5))
ax3.plot(t_vals, T_line*1e9, color='navy', linewidth=2)
ax3.axvline(np.hypot(Px, Py), color='red', linestyle='--', label='Fermat minimum')
ax3.set_xlabel('Distance along the plane of incidence')
ax3.set_ylabel('Travel time [ns]')
ax3.set_title('Travel time along the plane of incidence (2D cut of the 3D surface)')
ax3.legend()
ax3.grid(alpha=0.3)
plt.tight_layout()
plt.show()

Code Walkthrough

Section 1–2 (Physical setup and time function): We define two points $A$ and $B$ straddling a flat interface at $z=0$, along with the refractive indices of each medium. The function travel_time(p) computes the total time for light to go from $A$ to a candidate crossing point $P=(x,y,0)$, then from $P$ to $B$, using the formula derived above.

Section 2 (Optimization): Instead of a brute-force search, we use scipy.optimize.minimize with the BFGS method — a quasi-Newton algorithm that converges quadratically for smooth, convex problems like this one (the time function is a sum of two convex distance terms, so it has a single global minimum). This converges in a handful of iterations rather than thousands of grid evaluations.

Section 3 (Snell’s Law check): From the optimized crossing point, we compute $\sin\theta_1$ and $\sin\theta_2$ using simple trigonometry (horizontal distance over total distance), then confirm that $n_1\sin\theta_1 = n_2\sin\theta_2$. We also check that $P$ lies exactly on the straight line connecting the $xy$-projections of $A$ and $B$ — this confirms that the optimal ray stays within the plane of incidence, a well-known geometric consequence of Fermat’s Principle.

Section 4 (Comparison): We compare the optimized travel time against a naive straight-line crossing point, showing explicitly that refraction is faster, not just geometrically different.

Section 5 (Vectorized grid — the performance-critical part): To visualize $T(x,y)$ as a full 3D surface, we need to evaluate the time function on a $300 \times 300$ grid (90,000 points). Doing this with nested Python for loops would be extremely slow due to interpreter overhead. Instead, we use NumPy broadcasting with np.meshgrid, computing all 90,000 values in one vectorized array operation — this runs in milliseconds instead of seconds.

Sections 6–8 (Plotting): Three figures are generated: a 3D ray-path diagram, a 3D surface of the time function, and a 2D cross-section for an intuitive read of the minimum.

Results

========== Fermat's Principle: numerical result ==========
Crossing point P             : (5.410675, 2.029003, 0.000000)
Minimum travel time T_min    : 47.062861 ns
n1 * sin(theta1)             : 0.756215
n2 * sin(theta2)             : 0.756341
slope y/x of B               : 0.375000
slope y/x of P (should match): 0.375000
============================================================

Travel time of the naive straight-line path : 47.900133 ns
Travel time of the Fermat (fastest) path    : 47.062861 ns
Time saved by refraction                    : 0.837272 ns

Visualizing the Fastest Path

Figure 1 — The 3D Light Path

This figure shows the actual geometry: point $A$ above the interface, point $B$ below it, and the bent ray path through the crossing point $P$ found by the optimizer. The pale blue plane represents the interface between the two media. Notice the ray bends toward the normal when entering the denser medium (water) — exactly as Snell’s Law predicts.

Figure 2 — The Time Surface (3D)

This is the most instructive plot: it renders $T(x,y)$ as a full 3D bowl-shaped surface over every conceivable crossing point, not just the correct one. The red marker sits exactly at the bottom of the bowl — visually confirming that the point found by scipy.optimize.minimize truly is the global minimum of the travel-time function, which is the entire content of Fermat’s Principle.

Figure 3 — 2D Cross-Section Through the Minimum

Since the 3D bowl can be hard to read precisely, this figure slices the surface along the plane of incidence, producing an ordinary 2D curve. The dashed red line marks the minimum — the same point found numerically in Figure 2, now easy to verify by eye.

Why This Matters

What makes this example powerful is that we never told the program about Snell’s Law. We only told it: “minimize the travel time.” The bent ray, the exact angles, and the well-known refraction formula all emerged automatically from a generic numerical optimizer — a nice demonstration of how a simple variational principle in physics can be rediscovered purely through computation.

Finding the Fastest Path Down

Imagine a ball rolling under gravity from point A to point B. Which path gets it there in the least time? Intuitively, a straight line looks like the shortest route — but shortest in distance isn’t the same as fastest in time. The answer, famously posed by Johann Bernoulli in 1696, is a curve called the cycloid. This is the birth story of the calculus of variations, and today we’ll solve it numerically and visually with Python.

The Problem, Mathematically

We drop a particle from rest at point $A = (0,0)$ and let it slide (frictionlessly, under gravity $g$) along some curve $y(x)$ to point $B = (x_1, y_1)$, where $y$ is measured downward as positive.

By energy conservation, the speed at height $y$ is:

$$
v = \sqrt{2gy}
$$

The time to traverse a small arc length $ds = \sqrt{1 + y’(x)^2}, dx$ is $dt = ds / v$. So the total descent time is the functional:

$$
T[y] = \int_0^{x_1} \frac{\sqrt{1 + y’(x)^2}}{\sqrt{2gy(x)}} , dx
$$

We want to find the function $y(x)$ that minimizes $T[y]$. Applying the Euler–Lagrange equation to this functional leads to the differential equation whose solution is a cycloid — the curve traced by a point on the rim of a rolling circle:

$$
x(\theta) = R(\theta - \sin\theta), \qquad y(\theta) = R(1 - \cos\theta)
$$

where $R$ is the rolling circle’s radius and $\theta$ is the rotation angle. Given the endpoint $B=(x_1,y_1)$, $R$ and the final angle $\theta_1$ are found by solving:

$$
\frac{y_1}{x_1} = \frac{1 - \cos\theta_1}{\theta_1 - \sin\theta_1}
$$

and the minimal descent time has the remarkably clean closed form:

$$
T_{\text{cycloid}} = \theta_1 \sqrt{\frac{R}{g}}
$$

Our Concrete Example

We’ll drop a particle from $A = (0, 0)$ to $B = (3.0,\ 2.0)$ meters (3 m horizontally, 2 m of drop), with $g = 9.81\ \text{m/s}^2$. We’ll compare three paths:

1. A straight line — the “obvious” but wrong answer.
2. A family of quadratic Bézier curves, whose shape we sweep over a grid of control points to search for the fastest one within that family.
3. The true cycloid solution.

This lets us visually confirm that the cycloid beats every other candidate curve, and lets us plot the whole “time landscape” as a 3D surface — a nice, concrete picture of what calculus of variations is actually doing (searching an infinite-dimensional space of curves for a minimum).

Full Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161

# ============================================================
# Brachistochrone Problem — Numerical & Visual Solution
# Run this cell as-is in Google Colaboratory.
# ============================================================

import numpy as np
from scipy.integrate import quad
from scipy.optimize import brentq
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (enables 3D projection)

# ------------------------------------------------------------
# 1. Problem setup
# ------------------------------------------------------------
g = 9.81          # gravitational acceleration [m/s^2]
x1, y1 = 3.0, 2.0 # target point B (y measured DOWNWARD as positive)

# ------------------------------------------------------------
# 2. Generic descent-time integrator (works for ANY parametric curve)
#
#    A curve is given as x(t), y(t) for t in [0, 1], with A at t=0
#    and B at t=1. We substitute t = u^2 to analytically cancel the
#    integrable 1/sqrt(t) singularity that occurs near the start
#    point (where speed -> 0). This keeps scipy.integrate.quad fast,
#    warning-free, and numerically stable for every curve we test.
# ------------------------------------------------------------
def descent_time(xt, yt, dxt, dyt):
    """
    xt, yt   : functions t -> x(t), y(t)
    dxt, dyt : functions t -> dx/dt, dy/dt
    Returns total descent time (seconds).
    """
    def integrand(u):
        t = u ** 2
        y = max(yt(t), 1e-14)               # avoid divide-by-zero at t=0
        speed = np.sqrt(2 * g * y)
        ds_dt = np.sqrt(dxt(t) ** 2 + dyt(t) ** 2)
        return 2 * u * ds_dt / speed        # extra factor from dt = 2u du
    value, _ = quad(integrand, 0, 1, limit=200)
    return value

# ------------------------------------------------------------
# 3. Curve 1: Straight line from A to B
# ------------------------------------------------------------
line_x  = lambda t: x1 * t
line_y  = lambda t: y1 * t
line_dx = lambda t: x1 + 0 * t
line_dy = lambda t: y1 + 0 * t
T_line = descent_time(line_x, line_y, line_dx, line_dy)

# ------------------------------------------------------------
# 4. Curve 2: Quadratic Bezier family, swept over a grid of
#    control points (cx, cy). This is a brute-force search over a
#    2-parameter family of curves — a concrete, visual stand-in for
#    "searching curve-space" as in the calculus of variations.
# ------------------------------------------------------------
def bezier_funcs(cx, cy):
    x_t  = lambda t: 2 * (1 - t) * t * cx + t**2 * x1
    y_t  = lambda t: 2 * (1 - t) * t * cy + t**2 * y1
    dx_t = lambda t: 2 * cx * (1 - 2 * t) + 2 * t * x1
    dy_t = lambda t: 2 * cy * (1 - 2 * t) + 2 * t * y1
    return x_t, y_t, dx_t, dy_t

n_grid = 25
cx_vals = np.linspace(0.05 * x1, 0.95 * x1, n_grid)
cy_vals = np.linspace(0.05 * y1, 0.95 * y1, n_grid)
CX, CY = np.meshgrid(cx_vals, cy_vals, indexing='ij')
T_grid = np.zeros_like(CX)

for i in range(n_grid):
    for j in range(n_grid):
        xt, yt, dxt, dyt = bezier_funcs(CX[i, j], CY[i, j])
        T_grid[i, j] = descent_time(xt, yt, dxt, dyt)

best_idx = np.unravel_index(np.argmin(T_grid), T_grid.shape)
best_cx, best_cy = CX[best_idx], CY[best_idx]
T_bezier_min = T_grid[best_idx]

# ------------------------------------------------------------
# 5. Curve 3: The true Brachistochrone (cycloid) solution
#
#    Solve  y1/x1 = (1 - cos(theta1)) / (theta1 - sin(theta1))  for theta1
# ------------------------------------------------------------
def cycloid_ratio_eq(theta):
    return (1 - np.cos(theta)) / (theta - np.sin(theta)) - y1 / x1

theta1 = brentq(cycloid_ratio_eq, 1e-6, 2 * np.pi - 1e-6)
R = y1 / (1 - np.cos(theta1))

cyc_x  = lambda t: R * (t * theta1 - np.sin(t * theta1))
cyc_y  = lambda t: R * (1 - np.cos(t * theta1))
cyc_dx = lambda t: R * theta1 * (1 - np.cos(t * theta1))
cyc_dy = lambda t: R * theta1 * np.sin(t * theta1)

T_cycloid_numeric = descent_time(cyc_x, cyc_y, cyc_dx, cyc_dy)
T_cycloid_closed  = theta1 * np.sqrt(R / g)   # analytic formula, for cross-check

# ------------------------------------------------------------
# 6. Console summary
# ------------------------------------------------------------
print("=" * 55)
print("BRACHISTOCHRONE RESULTS")
print("=" * 55)
print(f"Target point B                : ({x1}, {y1}) m")
print(f"Cycloid radius R              : {R:.5f} m")
print(f"Cycloid final angle theta1    : {theta1:.5f} rad")
print("-" * 55)
print(f"Straight line time            : {T_line:.5f} s")
print(f"Best Bezier grid time         : {T_bezier_min:.5f} s "
      f"(cx={best_cx:.3f}, cy={best_cy:.3f})")
print(f"Cycloid time (numeric integral): {T_cycloid_numeric:.5f} s")
print(f"Cycloid time (closed form)     : {T_cycloid_closed:.5f} s")
print("=" * 55)

# ------------------------------------------------------------
# 7. Figure 1 — 2D comparison of the three curves
# ------------------------------------------------------------
t_plot = np.linspace(0, 1, 300)

fig1, ax1 = plt.subplots(figsize=(8, 6))
ax1.plot(line_x(t_plot), line_y(t_plot), '--', color='gray',
         label=f'Straight line  (T = {T_line:.3f} s)')

bx, by, _, _ = bezier_funcs(best_cx, best_cy)
ax1.plot(bx(t_plot), by(t_plot), '-.', color='orange',
         label=f'Best Bezier curve  (T = {T_bezier_min:.3f} s)')

ax1.plot(cyc_x(t_plot), cyc_y(t_plot), '-', color='crimson', linewidth=2.5,
         label=f'Cycloid (Brachistochrone)  (T = {T_cycloid_numeric:.3f} s)')

ax1.scatter([0, x1], [0, y1], color='black', zorder=5)
ax1.annotate('A', (0, 0), textcoords="offset points", xytext=(-10, 10))
ax1.annotate('B', (x1, y1), textcoords="offset points", xytext=(10, -15))

ax1.set_xlabel('x [m]')
ax1.set_ylabel('y [m]  (downward = positive)')
ax1.set_title('Brachistochrone Problem: Path Comparison')
ax1.invert_yaxis()   # so "down" visually points down
ax1.legend()
ax1.grid(alpha=0.3)
plt.tight_layout()
plt.show()

# ------------------------------------------------------------
# 8. Figure 2 — 3D surface: descent time over the Bezier
#    control-point parameter space (cx, cy)
# ------------------------------------------------------------
fig2 = plt.figure(figsize=(9, 7))
ax2 = fig2.add_subplot(111, projection='3d')
surf = ax2.plot_surface(CX, CY, T_grid, cmap='viridis', alpha=0.9,
                         edgecolor='none')
ax2.scatter(best_cx, best_cy, T_bezier_min, color='red', s=60,
            label='Best Bezier curve')
ax2.set_xlabel('Control point cx [m]')
ax2.set_ylabel('Control point cy [m]')
ax2.set_zlabel('Descent time T [s]')
ax2.set_title('Descent-Time Landscape over Curve Shape Parameters')
fig2.colorbar(surf, shrink=0.6, aspect=12, label='Time [s]')
ax2.legend()
plt.tight_layout()
plt.show()

Code Walkthrough

Section 2 — the singularity trick. Every candidate curve starts at rest, so speed $v = \sqrt{2gy} \to 0$ as $t \to 0$, and the raw time integrand blows up like $1/\sqrt{t}$ near the start. This singularity is integrable (finite area), but numerically it can make scipy.integrate.quad slow or noisy. The substitution $t = u^2$ (so $dt = 2u,du$) cancels the $1/\sqrt{t}$ term algebraically, turning every curve’s integral into a smooth, fast-converging one. This is why the same descent_time() function works cleanly for the line, the Bézier curves, and the cycloid without any special-casing.

Section 3 — the straight line. This is the naive baseline: constant velocity direction, but it accelerates too slowly at first because it doesn’t drop steeply enough near the start.

Section 4 — the Bézier search. Instead of trying to derive the optimal curve, we brute-force search a 2-parameter family of quadratic Bézier curves (parametrized by a control point $(c_x, c_y)$) over a $25 \times 25$ grid, computing the descent time for each of the 625 candidate curves. This is fast — each quad() call is a few milliseconds — so the whole grid finishes in a couple of seconds in Colab. Conceptually, this loop is a discretized version of “searching the space of all curves for a minimum,” which is exactly what the calculus of variations does analytically via the Euler–Lagrange equation.

Section 5 — the cycloid. brentq solves the transcendental equation for $\theta_1$ (the ratio equation has a guaranteed sign change between $\theta \to 0^+$, where the ratio diverges, and $\theta \to 2\pi^-$, where it approaches 0). Once $\theta_1$ is known, $R$ follows directly, and we compute the descent time two independent ways — via numerical integration and via the closed-form formula $T = \theta_1\sqrt{R/g}$ — as a sanity check that they agree.

Section 7 — the 2D plot. All three curves are drawn on the same axes, with the y-axis inverted so “downward” reads visually as “down.” You should see the cycloid dip more steeply near A than the straight line (trading extra distance for extra early speed), then flatten out toward B.

Section 8 — the 3D plot. This is the most illuminating part. It plots the descent time $T$ as a surface over the 2D space of Bézier control points $(c_x, c_y)$ — literally the “cost landscape” that an optimizer (or evolution, or your own intuition) would need to descend to find the best curve shape. The red marker shows the grid’s discovered minimum. Because the true cycloid isn’t a member of this particular Bézier family, its time (printed in the console) will typically be slightly lower than the Bézier grid minimum — nicely demonstrating that the cycloid is the true global optimum across all curves, not just within one restricted family.

=======================================================
BRACHISTOCHRONE RESULTS
=======================================================
Target point B                : (3.0, 2.0) m
Cycloid radius R              : 1.00133 m
Cycloid final angle theta1    : 3.06878 rad
-------------------------------------------------------
Straight line time            : 1.15116 s
Best Bezier grid time         : 0.98209 s (cx=0.263, cy=1.600)
Cycloid time (numeric integral): 0.98043 s
Cycloid time (closed form)     : 0.98043 s
=======================================================

Reading the Results

Once you’ve pasted in your run, you should observe:

• $T_{\text{cycloid}} < T_{\text{bezier_min}} < T_{\text{line}}$ — the straight line is the slowest, the best-found Bézier curve is faster, and the true cycloid is fastest of all.
• The two cycloid time values (numeric integral vs. closed-form) should match to about 4–5 decimal places, confirming the numerical integration is accurate.
• On the 3D surface, the bowl-shaped landscape has a single, fairly broad minimum region — which is why gradient-based or grid-based search methods converge to it reliably, and also hints at why this class of variational problem tends to have a unique smooth solution rather than many competing local optima.

This little experiment is a nice hands-on echo of the historical event: Bernoulli’s challenge produced the same answer whether you attack it with pure analysis (Euler–Lagrange calculus) or, as we just did, with a grid search and numerical integration on a laptop three centuries later.

Terrain-Based Optimization with Gradient Methods in Python

Optimization problems often come down to two opposite goals: reaching the highest point of a landscape, or avoiding high ground altogether while traveling from A to B. Both problems can be tackled with the same mathematical tool — the gradient — just pointed in different directions.

In this article we build a synthetic terrain (a 2D elevation field made of overlapping mountains), then solve two classic problems on it:

1. Finding the summit of a mountain using gradient ascent (maximization).
2. Finding a low-cost route across the terrain using gradient descent on a path, where the “cost” of a route is the total elevation it climbs — a continuous relaxation of the shortest-path problem.

Both are implemented in NumPy, fully vectorized, and visualized in 3D and as contour maps.

1. Modeling the Terrain

Real elevation data usually comes from a digital elevation model (DEM), but for a clean, reproducible demo we construct our terrain as a sum of Gaussian “bumps,” each acting like a mountain:

$$
Z(x,y) = \sum_{k=1}^{K} A_k \exp!\left(-\left(\frac{(x-x_{0,k})^2}{2\sigma_{x,k}^2} + \frac{(y-y_{0,k})^2}{2\sigma_{y,k}^2}\right)\right)
$$

Here $A_k$ is the height (amplitude) of mountain $k$, $(x_{0,k}, y_{0,k})$ is its peak location, and $\sigma_{x,k}, \sigma_{y,k}$ control how wide it spreads. Because this function is analytic, we can compute its gradient in closed form — no need for slow numerical differentiation:

$$
\frac{\partial Z}{\partial x} = \sum_k A_k \exp(\cdots)\cdot\left(-\frac{x-x_{0,k}}{\sigma_{x,k}^2}\right), \qquad
\frac{\partial Z}{\partial y} = \sum_k A_k \exp(\cdots)\cdot\left(-\frac{y-y_{0,k}}{\sigma_{y,k}^2}\right)
$$

This analytical gradient is the key to making everything below fast: instead of estimating slopes with finite differences (which requires extra function evaluations per step), we get the exact slope in one pass, and the whole thing works equally well on a single point or on an entire array of points at once thanks to NumPy broadcasting.

2. Full Source Code

The script below covers everything: terrain generation, the peak search, the path optimization, and all plots. It’s designed to run top-to-bottom in a single cell without any additional setup.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # enables 3D projection
import matplotlib.cm as cm

np.random.seed(42)

# ---------------------------------------------------------
# 1. Terrain (elevation map) definition
# ---------------------------------------------------------
# Each mountain is a Gaussian bump: (amplitude, center_x, center_y, spread_x, spread_y)
peaks = [
    (8.0, -1.5,  1.5, 1.2, 1.2),
    (6.0,  1.0, -1.0, 1.0, 1.0),
    (5.0,  2.5,  2.0, 0.8, 0.8),
]

def elevation(x, y):
    """Elevation Z(x, y). Works for scalars or NumPy arrays."""
    z = np.zeros_like(x, dtype=float)
    for A, x0, y0, sx, sy in peaks:
        z += A * np.exp(-(((x - x0) ** 2) / (2 * sx ** 2) +
                            ((y - y0) ** 2) / (2 * sy ** 2)))
    return z

def elevation_grad(x, y):
    """Analytical gradient (dZ/dx, dZ/dy). Works for scalars or arrays."""
    dzdx = np.zeros_like(x, dtype=float)
    dzdy = np.zeros_like(y, dtype=float)
    for A, x0, y0, sx, sy in peaks:
        g = A * np.exp(-(((x - x0) ** 2) / (2 * sx ** 2) +
                           ((y - y0) ** 2) / (2 * sy ** 2)))
        dzdx += g * (-(x - x0) / sx ** 2)
        dzdy += g * (-(y - y0) / sy ** 2)
    return dzdx, dzdy

# Grid for plotting
grid_n = 200
x_lin = np.linspace(-5, 5, grid_n)
y_lin = np.linspace(-5, 5, grid_n)
X, Y = np.meshgrid(x_lin, y_lin)
Z = elevation(X, Y)

# ---------------------------------------------------------
# 2. Visualize the raw terrain (3D + contour)
# ---------------------------------------------------------
fig = plt.figure(figsize=(14, 6))

ax1 = fig.add_subplot(1, 2, 1, projection='3d')
surf = ax1.plot_surface(X, Y, Z, cmap=cm.terrain, linewidth=0, antialiased=True, alpha=0.9)
ax1.set_title("Terrain Elevation (3D Surface)")
ax1.set_xlabel("x"); ax1.set_ylabel("y"); ax1.set_zlabel("Elevation Z")
fig.colorbar(surf, ax=ax1, shrink=0.6, label="Elevation")

ax2 = fig.add_subplot(1, 2, 2)
cont = ax2.contourf(X, Y, Z, levels=30, cmap=cm.terrain)
ax2.set_title("Terrain Elevation (Contour Map)")
ax2.set_xlabel("x"); ax2.set_ylabel("y")
fig.colorbar(cont, ax=ax2, label="Elevation")

plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 3. Finding the peak: Gradient Ascent
# ---------------------------------------------------------
def gradient_ascent(start, lr=0.05, n_iter=300):
    path = np.zeros((n_iter + 1, 2))
    path[0] = start
    p = np.array(start, dtype=float)
    for i in range(n_iter):
        gx, gy = elevation_grad(p[0], p[1])
        p = p + lr * np.array([gx, gy])
        path[i + 1] = p
    return path

start_point = np.array([-4.0, -3.5])
ascent_path = gradient_ascent(start_point, lr=0.05, n_iter=300)

print(f"Start point       : {start_point}")
print(f"Reached peak point: {ascent_path[-1]}")
print(f"Elevation at peak : {elevation(ascent_path[-1,0], ascent_path[-1,1]):.4f}")

fig = plt.figure(figsize=(14, 6))

ax1 = fig.add_subplot(1, 2, 1, projection='3d')
ax1.plot_surface(X, Y, Z, cmap=cm.terrain, linewidth=0, alpha=0.6)
ax1.plot(ascent_path[:, 0], ascent_path[:, 1],
          elevation(ascent_path[:, 0], ascent_path[:, 1]),
          color='red', linewidth=2, marker='o', markersize=2, label='Ascent Path')
ax1.scatter(*ascent_path[-1], elevation(ascent_path[-1, 0], ascent_path[-1, 1]),
             color='black', s=60, label='Peak Found')
ax1.set_title("Gradient Ascent to the Peak (3D)")
ax1.set_xlabel("x"); ax1.set_ylabel("y"); ax1.set_zlabel("Elevation Z")
ax1.legend()

ax2 = fig.add_subplot(1, 2, 2)
ax2.contourf(X, Y, Z, levels=30, cmap=cm.terrain)
ax2.plot(ascent_path[:, 0], ascent_path[:, 1], color='red', linewidth=2, label='Ascent Path')
ax2.scatter(*start_point, color='blue', s=60, label='Start')
ax2.scatter(*ascent_path[-1], color='black', s=60, label='Peak Found')
ax2.set_title("Gradient Ascent to the Peak (Contour)")
ax2.set_xlabel("x"); ax2.set_ylabel("y")
ax2.legend()

plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 4. Elevation-cost shortest path: Gradient Descent approximation
# ---------------------------------------------------------
def optimize_path(p_start, p_end, n_points=50,
                    w_elev=1.0, w_smooth=8.0,
                    lr=0.02, n_iter=800):
    p_start = np.array(p_start, dtype=float)
    p_end   = np.array(p_end, dtype=float)

    # Initial path: a straight line between start and goal
    t = np.linspace(0, 1, n_points).reshape(-1, 1)
    path = (1 - t) * p_start + t * p_end
    init_path = path.copy()

    cost_history = []

    for it in range(n_iter):
        # Gradient of the elevation cost at every path point (vectorized)
        gx, gy = elevation_grad(path[:, 0], path[:, 1])
        grad_elev = np.stack([gx, gy], axis=1)

        # Gradient of the smoothness penalty (discrete Laplacian)
        grad_smooth = np.zeros_like(path)
        grad_smooth[1:-1] = 2 * (2 * path[1:-1] - path[:-2] - path[2:])

        grad = w_elev * grad_elev + w_smooth * grad_smooth

        # Start and goal are fixed; only interior points move
        path[1:-1] -= lr * grad[1:-1]

        elev_cost = elevation(path[:, 0], path[:, 1]).sum()
        smooth_cost = np.sum(np.sum((path[1:] - path[:-1]) ** 2, axis=1))
        cost_history.append(w_elev * elev_cost + w_smooth * smooth_cost)

    return path, init_path, np.array(cost_history)

p_start = (-4.0, -4.0)
p_end   = (4.0, 4.0)
opt_path, init_path, cost_history = optimize_path(p_start, p_end)

print(f"Initial total cost: {cost_history[0]:.2f}")
print(f"Final total cost  : {cost_history[-1]:.2f}")

fig = plt.figure(figsize=(14, 6))

ax1 = fig.add_subplot(1, 2, 1, projection='3d')
ax1.plot_surface(X, Y, Z, cmap=cm.terrain, linewidth=0, alpha=0.6)
ax1.plot(init_path[:, 0], init_path[:, 1],
          elevation(init_path[:, 0], init_path[:, 1]),
          color='blue', linestyle='--', linewidth=2, label='Initial (Straight) Path')
ax1.plot(opt_path[:, 0], opt_path[:, 1],
          elevation(opt_path[:, 0], opt_path[:, 1]),
          color='red', linewidth=2, label='Optimized Path')
ax1.set_title("Elevation-Cost Path Optimization (3D)")
ax1.set_xlabel("x"); ax1.set_ylabel("y"); ax1.set_zlabel("Elevation Z")
ax1.legend()

ax2 = fig.add_subplot(1, 2, 2)
ax2.contourf(X, Y, Z, levels=30, cmap=cm.terrain)
ax2.plot(init_path[:, 0], init_path[:, 1], color='blue', linestyle='--', linewidth=2, label='Initial Path')
ax2.plot(opt_path[:, 0], opt_path[:, 1], color='red', linewidth=2, label='Optimized Path')
ax2.scatter(*p_start, color='cyan', s=60, label='Start')
ax2.scatter(*p_end, color='black', s=60, label='Goal')
ax2.set_title("Elevation-Cost Path Optimization (Contour)")
ax2.set_xlabel("x"); ax2.set_ylabel("y")
ax2.legend()

plt.tight_layout()
plt.show()

# ---------------------------------------------------------
# 5. Convergence of the path cost
# ---------------------------------------------------------
plt.figure(figsize=(8, 5))
plt.plot(cost_history, color='darkred')
plt.title("Convergence of Path Cost (Gradient Descent)")
plt.xlabel("Iteration")
plt.ylabel("Total Cost J")
plt.grid(True)
plt.show()

3. Visualizing the Raw Terrain

Before running any optimization, the script first draws the terrain itself — a 3D surface and a matching contour map — so we can see where the three mountains sit and how steep each one is.

4. Part 1 — Climbing to the Summit: Gradient Ascent

The idea

To find a mountain’s peak, we start at some point and repeatedly move in the direction of steepest increase — the gradient. Each step nudges the current position uphill:

$$
\mathbf{p}_{t+1} = \mathbf{p}_t + \eta \nabla Z(\mathbf{p}_t)
$$

where $\eta$ is the learning rate (step size) and $\nabla Z = (\partial Z/\partial x,\ \partial Z/\partial y)$ is the gradient we derived earlier. This is exactly gradient descent, just with the sign flipped — hence “ascent.”

Code walk-through

• gradient_ascent() initializes a position p at start, then loops n_iter times. At each step it evaluates elevation_grad(p[0], p[1]) to get the local slope, and moves p by lr times that slope.
• The entire trajectory is stored in the path array so we can later plot how the search climbed the mountain step by step.
• start_point = [-4.0, -3.5] places the search near the bottom-left corner, far from any peak, so the climb is visually obvious.
• Because the gradient is computed analytically (not numerically), each iteration only costs one pass through the 3 peaks — extremely cheap, so 300 iterations finish instantly.

Note that gradient ascent finds a local maximum — the peak nearest to the starting point in terms of the slope it follows, not necessarily the tallest mountain on the map. This is an important and realistic limitation: try changing start_point and you’ll likely converge to a different summit.

Start point       : [-4.  -3.5]
Reached peak point: [-3.99581568 -3.49174409]
Elevation at peak : 0.0002

5. Part 2 — Terrain-Cost Shortest Path: Gradient Descent Approximation

The idea

Classic shortest-path algorithms like Dijkstra’s algorithm work on discrete graphs. Here, instead, we treat the path as a continuous, deformable curve — a sequence of $N$ points $\mathbf{p}_1, \dots, \mathbf{p}_N$ between a fixed start and goal — and let gradient descent pull that curve toward a low-elevation, low-cost route. This is the same idea behind “elastic band” or “snake” path planning: the path behaves like a stretched band that is simultaneously pulled downhill by the terrain and kept smooth by a tension term.

The objective function to minimize is:

• The first term is the total elevation the path passes through — minimizing it pushes the route away from mountains and into valleys.
• The second term is a smoothness penalty — without it, each point would independently roll straight downhill and the path would tear itself apart instead of staying connected.

We minimize $J$ with standard gradient descent on every interior point simultaneously:

$$
\mathbf{p}_i \leftarrow \mathbf{p}_i - \eta \frac{\partial J}{\partial \mathbf{p}_i}
$$

The smoothness term’s gradient has a clean closed form — the discrete Laplacian of the path:

Code walk-through

• optimize_path() starts with a straight line between p_start and p_end as the initial guess (init_path), interpolated with n_points=50 points.
• At every iteration:
  • elevation_grad(path[:, 0], path[:, 1]) computes the elevation gradient for all 50 points at once — this is the key vectorization trick. There is no inner Python loop over points; NumPy broadcasts the operation across the whole array.
  • grad_smooth is computed using array slicing (path[:-2], path[1:-1], path[2:]) to implement the discrete Laplacian for every interior point in one shot.
  • The two gradients are combined with weights w_elev and w_smooth, and only the interior points (path[1:-1]) are updated — the start and goal stay fixed.
• cost_history records the value of $J$ at every iteration so we can later plot convergence.
• With n_points=50 and n_iter=800, the total workload is on the order of tens of thousands of simple array operations — this finishes in well under a second even on Colab’s default CPU runtime, because the per-point loop that would exist in a naive implementation has been replaced entirely by array-level operations.

Increasing w_smooth makes the path stiffer (straighter, less willing to detour), while increasing w_elev makes it more averse to climbing, bending harder around the mountains. Try adjusting these two weights to see the trade-off directly.

Initial total cost: 103.08
Final total cost  : 59.41

6. Checking Convergence

To confirm the gradient descent actually improved the route (rather than just moving points around), the script plots the total cost $J$ at every iteration.

You should see the cost drop sharply in the first several dozen iterations as the path pulls away from the mountains, then flatten out as it settles into a smooth, low-elevation route between the fixed endpoints.

7. Performance Notes

Both algorithms here are, by design, very fast: the terrain is defined analytically (a sum of 3 Gaussians), so every gradient evaluation is a handful of exp() calls rather than a lookup into a large grid. The main opportunity for slowdown would be looping over each path point individually in Python — this script avoids that entirely by using NumPy array operations (path[:, 0], path[1:-1], slicing-based Laplacians) so that all 50 path points are updated in a single vectorized step per iteration. If you scale this up to, say, a real digital elevation model with thousands of grid cells and a path of hundreds of points, the same vectorization strategy — computing gradients for the whole path array at once instead of point-by-point — is what keeps the optimization fast.

8. Takeaways

• Gradient ascent is a natural way to find a mountain’s peak, but it only guarantees a local maximum — the result depends on where you start.
• Gradient descent on a path turns “shortest path with terrain cost” into a continuous optimization problem: instead of searching a discrete graph, we let a deformable curve relax downhill while a smoothness term keeps it connected. This is an approximation, not an exact shortest path, but it’s fast, differentiable, and easy to extend (e.g., add obstacle-avoidance terms or path-length penalties).
• Both methods share the same core building block: an analytical gradient of the elevation function, evaluated in a fully vectorized way across arrays of points.
• In a real-world setting, elevation() could be replaced by an interpolated function over actual DEM (Digital Elevation Model) data, and the same gradient ascent / descent machinery would apply directly to hiking route planning, drone path planning, or terrain-aware robotics navigation.

Minimizing Total Distance Cost

The Business Problem

Imagine a logistics company that needs to build exactly two distribution warehouses to serve a scattered network of retail stores or customers. Each customer has a different demand volume (some order a lot, some order a little), and the company wants to minimize total transportation cost — defined as the sum of each customer’s shipping distance multiplied by their demand weight, always shipping from whichever of the two warehouses is closer.

This is a classic problem in operations research known as the multi-source Weber problem, and it sits at the intersection of continuous optimization (where exactly should the warehouses sit?) and combinatorial optimization (which customers get served by which warehouse?).

Mathematical Formulation

Given $n$ customer locations $\mathbf{p}_i \in \mathbb{R}^2$ with demand weights $w_i$, and two facility locations $\mathbf{f}_1, \mathbf{f}_2 \in \mathbb{R}^2$, we want to solve:

where $c_i \in {1, 2}$ denotes which facility serves customer $i$. In practice, once the facility locations are fixed, the optimal assignment is trivial — each customer simply goes to the nearer facility:

$$
c_i = \arg\min_{k \in {1,2}} \left| \mathbf{p}_i - \mathbf{f}_k \right|_2
$$

The hard part is finding the facility positions themselves. For a single facility, the optimal point minimizing the weighted sum of Euclidean distances is called the weighted geometric median, and it’s found using Weiszfeld’s algorithm, an iterative fixed-point method:

For two facilities, we combine this with an alternating scheme very similar to k-means clustering:

1. Assign every customer to its nearest facility.
2. Re-optimize each facility’s position using Weiszfeld’s algorithm on its assigned customers.
3. Repeat until the facility positions stop moving.

Because this alternating scheme can get stuck in a local optimum (just like k-means), we run it from many random starting positions and keep the best result found.

Full Python Implementation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

# ==============================================================
# Two-Facility Location Problem (Minimizing Weighted Distance Cost)
# ==============================================================

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
import matplotlib as mpl

# ---------------- Dark theme ----------------
plt.style.use('dark_background')
mpl.rcParams['figure.facecolor'] = '#0d1117'
mpl.rcParams['axes.facecolor'] = '#0d1117'
mpl.rcParams['savefig.facecolor'] = '#0d1117'

np.random.seed(42)

# ---------------- Problem data ----------------
# 12 customer locations (x, y) in km, and demand weights (shipment volume, tons/day)
customers = np.array([
    [ 2.0,  8.5], [ 3.5,  9.8], [ 1.0,  6.0], [ 4.5,  7.0],
    [ 8.0,  8.0], [ 9.5,  9.0], [ 7.0,  6.5], [10.0,  6.0],
    [ 3.0,  2.0], [ 5.0,  1.5], [ 8.5,  2.5], [10.5,  3.5]
])
weights = np.array([4, 6, 3, 5, 7, 5, 4, 6, 8, 3, 5, 4], dtype=float)

# ---------------- Core functions ----------------
def total_cost(facilities, customers, weights):
    """Total weighted distance cost given facility positions, shape (2,2)."""
    d = np.linalg.norm(customers[:, None, :] - facilities[None, :, :], axis=2)  # (n,2)
    nearest = d.min(axis=1)
    return np.sum(weights * nearest)

def assign_customers(facilities, customers):
    """Assign each customer to its nearest facility (0 or 1)."""
    d = np.linalg.norm(customers[:, None, :] - facilities[None, :, :], axis=2)
    return np.argmin(d, axis=1)

def weiszfeld_update(points, weights, f_init, max_iter=200, tol=1e-8, eps=1e-6):
    """Weighted geometric median via Weiszfeld's algorithm."""
    f = f_init.copy()
    for _ in range(max_iter):
        diff = points - f
        dist = np.maximum(np.linalg.norm(diff, axis=1), eps)
        w = weights / dist
        f_new = (w[:, None] * points).sum(axis=0) / w.sum()
        if np.linalg.norm(f_new - f) < tol:
            f = f_new
            break
        f = f_new
    return f

def two_facility_location(customers, weights, f_init, max_outer=100, tol=1e-7):
    """Alternating optimization: assignment step + Weiszfeld update step."""
    facilities = f_init.copy()
    history = [facilities.copy()]
    cost_history = [total_cost(facilities, customers, weights)]
    for _ in range(max_outer):
        assign = assign_customers(facilities, customers)
        new_facilities = facilities.copy()
        for k in range(2):
            mask = assign == k
            if mask.sum() == 0:
                continue
            new_facilities[k] = weiszfeld_update(customers[mask], weights[mask], facilities[k])
        history.append(new_facilities.copy())
        cost_history.append(total_cost(new_facilities, customers, weights))
        if np.linalg.norm(new_facilities - facilities) < tol:
            facilities = new_facilities
            break
        facilities = new_facilities
    return facilities, np.array(history), np.array(cost_history)

# ---------------- Multi-start optimization ----------------
n_starts = 30
best_cost = np.inf
best_result = None
all_costs = []

x_min, x_max = customers[:, 0].min(), customers[:, 0].max()
y_min, y_max = customers[:, 1].min(), customers[:, 1].max()

for _ in range(n_starts):
    f_init = np.column_stack([
        np.random.uniform(x_min, x_max, 2),
        np.random.uniform(y_min, y_max, 2)
    ])
    fac, hist, chist = two_facility_location(customers, weights, f_init)
    all_costs.append(chist)
    if chist[-1] < best_cost:
        best_cost = chist[-1]
        best_result = (fac, hist, chist)

best_facilities, best_history, best_cost_history = best_result
best_assignment = assign_customers(best_facilities, customers)

print(f"Best total cost found: {best_cost:.4f}")
print(f"Facility 1 location: ({best_facilities[0,0]:.3f}, {best_facilities[0,1]:.3f})")
print(f"Facility 2 location: ({best_facilities[1,0]:.3f}, {best_facilities[1,1]:.3f})")

# ---------------- Visualization 1: 2D map with assignment & optimization path ----------------
fig1, ax1 = plt.subplots(figsize=(9, 7))
colors = ['#00e5ff', '#ff6ec7']

for k in range(2):
    mask = best_assignment == k
    ax1.scatter(customers[mask, 0], customers[mask, 1], s=weights[mask] * 40,
                color=colors[k], alpha=0.85, edgecolor='white', linewidth=0.5,
                label=f'Customers -> Facility {k+1}')

for k in range(2):
    path = best_history[:, k, :]
    ax1.plot(path[:, 0], path[:, 1], '--', color=colors[k], alpha=0.5, linewidth=1.2)
    ax1.scatter(*best_facilities[k], marker='*', s=600, color=colors[k],
                edgecolor='white', linewidth=1.2, zorder=5,
                label=f'Facility {k+1} (final)')
    ax1.scatter(*best_history[0, k], marker='x', s=120, color='gray', zorder=4)

ax1.set_title('Two-Facility Location: Customer Assignment & Optimization Path', fontsize=13, color='white')
ax1.set_xlabel('X coordinate (km)')
ax1.set_ylabel('Y coordinate (km)')
ax1.legend(loc='upper left', fontsize=8, framealpha=0.3)
ax1.grid(alpha=0.2)
plt.tight_layout()
plt.show()

# ---------------- Visualization 2: 3D cost landscape (vectorized, no nested loops) ----------------
grid_n = 80
gx = np.linspace(x_min - 1, x_max + 1, grid_n)
gy = np.linspace(y_min - 1, y_max + 1, grid_n)
GX, GY = np.meshgrid(gx, gy)

fixed_facility = best_facilities[1]
grid_points = np.stack([GX.ravel(), GY.ravel()], axis=1)          # (grid_n^2, 2)

d1 = np.linalg.norm(customers[None, :, :] - grid_points[:, None, :], axis=2)  # (grid_n^2, 12)
d2 = np.linalg.norm(customers - fixed_facility, axis=1)                       # (12,)
d_min = np.minimum(d1, d2[None, :])
Z = (d_min * weights[None, :]).sum(axis=1).reshape(grid_n, grid_n)

fig2 = plt.figure(figsize=(10, 8))
ax2 = fig2.add_subplot(111, projection='3d')
ax2.plot_surface(GX, GY, Z, cmap='plasma', alpha=0.9, linewidth=0, antialiased=True)
ax2.scatter(best_facilities[0, 0], best_facilities[0, 1], best_cost,
            color='cyan', s=120, marker='*', label='Optimal Facility 1')
ax2.set_xlabel('Facility 1 X')
ax2.set_ylabel('Facility 1 Y')
ax2.set_zlabel('Total Weighted Cost')
ax2.set_title('Cost Landscape for Facility 1 Position\n(Facility 2 fixed at optimum)', color='white')
ax2.view_init(elev=35, azim=-60)
plt.tight_layout()
plt.show()

# ---------------- Visualization 3: Convergence curves across all random starts ----------------
fig3, ax3 = plt.subplots(figsize=(9, 6))
for chist in all_costs:
    ax3.plot(chist, color='gray', alpha=0.3, linewidth=1)
ax3.plot(best_cost_history, color='#00ff9d', linewidth=2.5, label='Best run')
ax3.set_xlabel('Iteration')
ax3.set_ylabel('Total Weighted Cost')
ax3.set_title('Convergence of Total Cost Across 30 Random Starts', color='white')
ax3.legend()
ax3.grid(alpha=0.2)
plt.tight_layout()
plt.show()

Console Output

Best total cost found: 168.4827
Facility 1 location: (8.978, 6.448)
Facility 2 location: (3.346, 6.470)

Code Walkthrough

Problem data. Twelve customer locations are defined as (x, y) coordinates in kilometers, each paired with a demand weight representing daily shipment volume. In a real deployment these would come from a sales database or GIS system.

total_cost: computes the objective function directly from its mathematical definition. It calculates the distance from every customer to every facility in one vectorized NumPy operation (no Python-level loop), takes the minimum distance per customer (i.e., “ships from whichever warehouse is closer”), and returns the weighted sum.

assign_customers: the combinatorial half of the problem — given fixed facility positions, it returns which facility (0 or 1) is closest to each customer.

weiszfeld_update: the continuous half of the problem. This implements the Weiszfeld fixed-point iteration shown in the formula above. A small eps floor is added to the distance to avoid division by zero if a facility ever lands exactly on a customer’s coordinates.

two_facility_location: the outer loop that alternates between assignment and position updates, very similar in spirit to Lloyd’s algorithm for k-means, except each “centroid” is a weighted geometric median rather than a simple average — appropriate because we’re minimizing straight-line distance, not squared distance.

Multi-start loop: since alternating optimization can converge to different local optima depending on where it starts, the script runs the whole procedure 30 times from random initial facility positions and keeps whichever run achieved the lowest total cost. This is a standard and inexpensive way to guard against poor local optima in this class of problem.

Speed Notes

The 3D cost-landscape plot in principle requires evaluating the objective function at every point on an 80×80 grid — 6,400 evaluations. A naive implementation would use two nested Python for loops, which is slow. The code above avoids this entirely by reshaping the grid into a single array of (x, y) points and computing all customer-to-grid-point distances in one broadcasted NumPy operation, then taking element-wise minimums against the fixed second facility. This turns a 6,400-iteration Python loop into a handful of vectorized array operations, so the surface renders almost instantly even on a standard Colab CPU runtime.

Understanding the Visualizations

Figure 1 — Customer Assignment & Optimization Path. This 2D map shows all twelve customers, sized in proportion to their demand weight, colored according to which of the two final warehouses serves them. The two stars mark the optimal facility positions, the gray X marks show where that particular optimization run started, and the dashed lines trace how each facility “walked” from its random starting point to its final resting place across iterations. Notice how the facilities settle roughly in the geometric weighted center of their respective clusters, rather than at the simple average — larger demand customers pull the facility more strongly toward them.

Figure 2 — 3D Cost Landscape. This surface shows how the total cost changes as Facility 1 is moved anywhere on the map, while Facility 2 stays fixed at its optimal location. The valley (lowest point, marked with a cyan star) corresponds to the optimal position found by the algorithm. This plot is useful for building intuition: the cost surface is not smooth like a simple bowl — it has a somewhat faceted, ridge-like shape because of the “assign to nearest facility” logic switching abruptly between the two warehouses as Facility 1 moves across customer territory.

Figure 3 — Convergence Across Random Starts. Each gray line is one of the 30 independent optimization runs, showing how quickly its total cost drops as the alternating algorithm iterates. The green line highlights the best-performing run — the one whose final result is reported and drawn in Figures 1 and 2. Most runs converge within just a handful of iterations, and the spread between lines illustrates why multiple random starts matter: a few runs stall at a noticeably higher cost, meaning they got trapped in a local optimum that a single-start approach could easily have missed.

Takeaways

The two-facility location problem is a compact but genuinely useful example of combining discrete decisions (which warehouse serves which customer) with continuous optimization (where exactly to place each warehouse). The alternating Weiszfeld approach used here generalizes cleanly to three, four, or more facilities simply by expanding the facility array — making it a practical starting point for real supply-chain network design problems.

Solving Order Quantity and Safety Stock Together

Every warehouse manager eventually runs into the same tension: order in large batches and you pay too much to hold stock; order in tiny batches and you pay too much in ordering fees; keep too little buffer stock and you risk running out during the lead time. This article works through a concrete numerical example that solves the order quantity and safety stock simultaneously, rather than treating them as two separate problems, and visualizes the full cost landscape in 3D.

The Business Scenario

A distribution center sells a mid-volume SKU with the following characteristics:

• Annual demand: 12,000 units/year
• Ordering cost: $45 per purchase order
• Holding cost: $4 per unit per year
• Shortage (backorder) cost: $15 per unit short
• Average daily demand: ~32.9 units, with a daily demand standard deviation of 5 units
• Supplier lead time: 14 days

The question: how many units should we order each time, and how much safety stock should we carry, to minimize total annual cost?

Mathematical Formulation

Classical Economic Order Quantity (EOQ)

Ignoring uncertainty for a moment, the trade-off between ordering cost and holding cost gives the well-known EOQ formula:

$$
Q^{*} = \sqrt{\frac{2DS}{H}}
$$

where $D$ is annual demand, $S$ is the fixed cost per order, and $H$ is the holding cost per unit per year.

Adding Demand Uncertainty: Safety Stock

Because daily demand is random, the demand realized during the lead time $L$ is also random. If lead-time demand has standard deviation $\sigma_L$, we define a safety factor $z$ (a standard normal quantile) so that:

$$
SS = z,\sigma_L, \qquad R = \bar{d}L + z,\sigma_L
$$

where $SS$ is the safety stock, $R$ is the reorder point, and $\bar{d}$ is average daily demand. A larger $z$ means a higher target service level, but also more holding cost.

Total Annual Cost

Combining ordering cost, cycle-stock holding cost, safety-stock holding cost, and the expected cost of running out during the lead time (the Hadley–Whitin formulation), the total annual cost as a function of both decision variables $Q$ and $z$ is:

where $C_s$ is the shortage cost per unit, and $L(z)$ is the standard normal loss function:

$$
L(z) = \phi(z) - z\big(1-\Phi(z)\big)
$$

with $\phi$ the standard normal PDF and $\Phi$ the standard normal CDF. $L(z)$ represents the expected number of units short (in standard-deviation units) per replenishment cycle.

The goal is to jointly minimize $TC(Q,z)$ over both $Q$ and $z$ — this is what makes the problem genuinely two-dimensional rather than two separate one-dimensional problems.

Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

# ============================================================
# Inventory Cost Minimization: Order Quantity + Safety Stock
# ============================================================
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
from scipy.optimize import minimize
from mpl_toolkits.mplot3d import Axes3D

# ---------------- 1. Problem parameters ----------------
D = 12000.0        # annual demand (units/year)
S = 45.0            # ordering cost per order ($)
H = 4.0             # holding cost per unit per year ($)
Cs = 15.0            # shortage cost per unit short ($)
d_daily = D / 365.0       # average daily demand
sigma_daily = 5.0        # std dev of daily demand
L_days = 14.0             # supplier lead time (days)
sigma_L = sigma_daily * np.sqrt(L_days)  # lead-time demand std dev

# ---------------- 2. Cost model ----------------
def loss_function(z):
    """Standard normal unit loss function L(z)."""
    return norm.pdf(z) - z * (1 - norm.cdf(z))

def total_cost(x):
    """Total annual cost as a function of [Q, z]."""
    Q, z = x
    if Q <= 0:
        return 1e12
    ordering_cost = (D / Q) * S
    cycle_holding_cost = (Q / 2) * H
    safety_holding_cost = z * sigma_L * H
    shortage_cost = (D / Q) * Cs * sigma_L * loss_function(z)
    return ordering_cost + cycle_holding_cost + safety_holding_cost + shortage_cost

def total_cost_vec(Q, z):
    """Vectorized version of total_cost, used for fast grid evaluation."""
    ordering_cost = (D / Q) * S
    cycle_holding_cost = (Q / 2) * H
    safety_holding_cost = z * sigma_L * H
    Lz = norm.pdf(z) - z * (1 - norm.cdf(z))
    shortage_cost = (D / Q) * Cs * sigma_L * Lz
    return ordering_cost + cycle_holding_cost + safety_holding_cost + shortage_cost

# ---------------- 3. Joint optimization of Q and z ----------------
Q0 = np.sqrt(2 * D * S / H)           # classical EOQ as initial guess
x0 = [Q0, 1.65]                        # start near a 95% service level

result = minimize(
    total_cost, x0,
    method='L-BFGS-B',
    bounds=[(10, 3000), (0, 4)]
)

Q_opt, z_opt = result.x
SS_opt = z_opt * sigma_L
ROP_opt = d_daily * L_days + SS_opt
TC_opt = result.fun
service_level = norm.cdf(z_opt) * 100

print("=" * 50)
print("INVENTORY OPTIMIZATION RESULTS")
print("=" * 50)
print(f"Classical EOQ (no uncertainty) : {Q0:8.2f} units")
print(f"Optimal order quantity Q*      : {Q_opt:8.2f} units")
print(f"Optimal safety factor z*       : {z_opt:8.4f}")
print(f"Implied service level          : {service_level:8.2f} %")
print(f"Optimal safety stock SS*       : {SS_opt:8.2f} units")
print(f"Optimal reorder point R*       : {ROP_opt:8.2f} units")
print(f"Minimum total annual cost      : ${TC_opt:8.2f}")
print("=" * 50)

# ---------------- 4. Build cost surface (vectorized, fast) ----------------
Q_range = np.linspace(50, 1500, 120)
z_range = np.linspace(0, 3.5, 120)
Q_grid, Z_grid = np.meshgrid(Q_range, z_range)
TC_grid = total_cost_vec(Q_grid, Z_grid)   # no Python loop -> instant on Colab CPU

# ---------------- 5. 3D visualization of the cost surface ----------------
fig = plt.figure(figsize=(12, 9))
fig.patch.set_facecolor('#0d1117')
ax = fig.add_subplot(111, projection='3d')
ax.set_facecolor('#0d1117')

surf = ax.plot_surface(
    Q_grid, Z_grid, TC_grid,
    cmap='plasma', edgecolor='none', alpha=0.92, antialiased=True
)

ax.scatter(
    [Q_opt], [z_opt], [TC_opt],
    color='cyan', s=100, depthshade=False,
    label=f'Optimum: Q*={Q_opt:.0f}, z*={z_opt:.2f}, TC*=${TC_opt:.0f}'
)

ax.set_xlabel('Order Quantity  Q', color='white', labelpad=14)
ax.set_ylabel('Safety Factor  z', color='white', labelpad=14)
ax.set_zlabel('Total Annual Cost ($)', color='white', labelpad=14)
ax.set_title('Total Inventory Cost Surface  TC(Q, z)', color='white', fontsize=15, pad=22)

ax.tick_params(colors='white')
ax.xaxis.pane.set_facecolor((0.05, 0.05, 0.08, 1.0))
ax.yaxis.pane.set_facecolor((0.05, 0.05, 0.08, 1.0))
ax.zaxis.pane.set_facecolor((0.05, 0.05, 0.08, 1.0))
ax.view_init(elev=25, azim=-60)

cbar = fig.colorbar(surf, shrink=0.55, aspect=12, pad=0.1)
cbar.set_label('Total Cost ($)', color='white')
cbar.ax.yaxis.set_tick_params(color='white')
plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='white')

legend = ax.legend(facecolor='#161b22', loc='upper right')
for text in legend.get_texts():
    text.set_color('white')

plt.tight_layout()
plt.savefig('inventory_cost_surface_3d.png', dpi=150, facecolor=fig.get_facecolor())
plt.show()

# ---------------- 6. 2D cost-component breakdown at the optimal z ----------------
Q_line = np.linspace(50, 1500, 300)
order_c = (D / Q_line) * S
hold_c = (Q_line / 2) * H
safety_c = np.full_like(Q_line, z_opt * sigma_L * H)
short_c = (D / Q_line) * Cs * sigma_L * loss_function(z_opt)
total_c = order_c + hold_c + safety_c + short_c

fig2, ax2 = plt.subplots(figsize=(11, 6.5))
fig2.patch.set_facecolor('#0d1117')
ax2.set_facecolor('#0d1117')

ax2.plot(Q_line, order_c, label='Ordering Cost', color='#ff7f0e', lw=2)
ax2.plot(Q_line, hold_c, label='Cycle Holding Cost', color='#1f77b4', lw=2)
ax2.plot(Q_line, safety_c, label='Safety Stock Holding Cost', color='#2ca02c', lw=2, linestyle='--')
ax2.plot(Q_line, short_c, label='Expected Shortage Cost', color='#d62728', lw=2, linestyle=':')
ax2.plot(Q_line, total_c, label='Total Cost', color='white', lw=3)

ax2.axvline(Q_opt, color='cyan', linestyle='--', alpha=0.7)
ax2.scatter([Q_opt], [TC_opt], color='cyan', s=90, zorder=5)

ax2.set_xlabel('Order Quantity  Q', color='white')
ax2.set_ylabel('Annual Cost ($)', color='white')
ax2.set_title(f'Cost Components vs Order Quantity  (z fixed at z* = {z_opt:.2f})',
              color='white', fontsize=13)
ax2.tick_params(colors='white')
ax2.grid(alpha=0.2)
for spine in ax2.spines.values():
    spine.set_color('white')

leg2 = ax2.legend(facecolor='#161b22')
for text in leg2.get_texts():
    text.set_color('white')

plt.tight_layout()
plt.savefig('inventory_cost_components_2d.png', dpi=150, facecolor=fig2.get_facecolor())
plt.show()

Console Output

==================================================
INVENTORY OPTIMIZATION RESULTS
==================================================
Classical EOQ (no uncertainty) :   519.62 units
Optimal order quantity Q*      :   526.11 units
Optimal safety factor z*       :   2.2671
Implied service level          :    98.83 %
Optimal safety stock SS*       :    42.41 units
Optimal reorder point R*       :   502.69 units
Minimum total annual cost      : $ 2274.07
==================================================

Code Walkthrough

Section 1 – Parameters. All business inputs are declared as plain floats at the top: demand, ordering cost, holding cost, shortage cost, and the demand-variability figures (daily demand std dev and lead time). Keeping these separate from the model logic makes the script easy to re-run with a different SKU’s numbers.

Section 2 – Cost model. loss_function(z) implements the standard normal unit loss function $L(z)$, which converts a safety factor into an expected number of units short per cycle. total_cost(x) is the scalar objective function that scipy.optimize.minimize calls; it unpacks x = [Q, z] and sums the four cost terms described in the math section above. A guard clause returns a very large penalty if Q is non-positive, which keeps the optimizer away from degenerate values. total_cost_vec(Q, z) is a NumPy-vectorized twin of the same function — it accepts full arrays for Q and z and returns an array of costs with no Python-level loop, which is what makes the 3D surface below build instantly instead of looping over thousands of grid points one at a time.

Section 3 – Joint optimization. Rather than solving order quantity and safety stock as two independent problems, minimize searches over both Q and z at once, using the classical EOQ value and a 95%-service-level guess (z = 1.65) as the starting point. L-BFGS-B is used because it supports box constraints (bounds), which keeps Q and z in economically sensible ranges during the search. The results are then translated into safety stock and reorder point using the formulas from the math section.

Section 4 – Cost surface. A 120×120 grid of (Q, z) combinations is built with np.meshgrid, and the entire cost surface is evaluated in one vectorized call to total_cost_vec. This avoids a nested double loop (120×120 = 14,400 evaluations) and keeps the whole computation well under a second.

Section 5 – 3D plot. plot_surface draws the cost landscape, and the optimum found in Section 3 is marked as a single cyan point. The surface makes the trade-off visually obvious: moving along the $Q$-axis away from the optimum increases cost because of the ordering/holding trade-off, while moving along the $z$-axis away from the optimum increases cost because of the safety-stock/shortage trade-off. The dark theme (background, panes, tick colors) matches a typical technical-blog dark layout.

Section 6 – 2D breakdown. Fixing $z$ at its optimal value, this chart decomposes total cost into its four components as a function of $Q$ alone — this is the classic “U-shaped EOQ curve” but now shown alongside the safety-stock and shortage-cost lines so the reader can see how much of the total cost each component actually contributes at the optimum.

The 3D surface is bowl-shaped, with a single global minimum in the interior of the plotted region — that’s the optimum reported by the optimizer. Moving toward small $Q$ makes the surface rise steeply because ordering cost explodes as $D/Q$; moving toward small $z$ makes it rise because expected shortage cost grows. The optimum sits at a moderate order quantity with a fairly high safety factor, since in this example the shortage cost is significant relative to holding cost.

In the 2D breakdown, the white “Total Cost” curve is clearly the sum of the other four curves, and its minimum lines up exactly with the cyan marker. Ordering cost falls steadily as $Q$ grows, cycle holding cost rises linearly, and the safety-stock holding cost stays flat (since $z$ is fixed in this chart) while shortage cost falls as larger, more frequent… actually less frequent orders reduce the number of cycles per year in which a shortage can occur.

Interpretation of the Numbers

For this example, the model finds an optimal order quantity of roughly 526 units per order — close to, but not identical to, the classical EOQ of about 520 units, because uncertainty slightly shifts the ideal batch size. The optimal safety factor comes out to about $z \approx 2.27$, which corresponds to an implied service level of roughly 98.8%. That translates into a safety stock of about 42 units and a reorder point of about 503 units, at a minimum total annual cost of roughly $2,274.

The relatively high service level here is a direct consequence of the numbers chosen: a shortage cost of $15 per unit is considerably higher than the $4 holding cost per unit, so the optimizer prefers to carry a bit more safety stock rather than risk running short. Changing that ratio — for example, lowering the shortage cost or raising the holding cost — would pull the optimal $z$ down and reduce the safety stock accordingly. That sensitivity is exactly what the 3D surface makes visible at a glance: the cost landscape’s shape along the $z$-axis directly reflects how expensive shortages are relative to holding inventory.

Optimizing Labor and Capital in Python

Every firm faces the same fundamental question: how much labor and how much capital should it employ to maximize profit? When production follows a Cobb-Douglas technology, this question turns into a clean, well-behaved optimization problem — one that’s perfect for illustrating both the economics and the numerical methods behind it. In this article, we’ll set up a two-variable profit maximization problem, solve it with a closed-form derivation and a numerical solver, and visualize the profit landscape in 3D.

The Economic Setup

Consider a firm that produces output $Q$ using labor $L$ and capital $K$ according to a Cobb-Douglas production function:

$$
Q(L, K) = A , L^{\alpha} K^{\beta}
$$

where $A$ is total factor productivity, and $\alpha, \beta \in (0,1)$ are the output elasticities of labor and capital. When $\alpha + \beta < 1$, the technology exhibits decreasing returns to scale, which guarantees a well-defined, interior profit-maximizing point (rather than a corner solution or an unbounded profit).

The firm sells output at price $p$, pays wage $w$ per unit of labor, and pays rental rate $r$ per unit of capital. Profit is:

$$
\pi(L, K) = p , A , L^{\alpha} K^{\beta} - wL - rK
$$

The firm’s problem is:

$$
\max_{L > 0,, K > 0} ; \pi(L, K)
$$

First-Order Conditions

Taking partial derivatives and setting them to zero gives the classic marginal-revenue-product-equals-input-price conditions:

$$
\frac{\partial \pi}{\partial L} = p A \alpha L^{\alpha - 1} K^{\beta} - w = 0
$$

$$
\frac{\partial \pi}{\partial K} = p A \beta L^{\alpha} K^{\beta - 1} - r = 0
$$

Dividing the first equation by the second eliminates $p$ and $A$, yielding a simple ratio between the optimal capital-labor ratio and the relative input prices:

$$
\frac{K^*}{L^*} = \frac{\beta w}{\alpha r}
$$

Substituting this back into the first FOC and solving for $L^*$ gives a closed-form solution:

$$
L^* = \left[ \frac{w}{pA\alpha \left(\dfrac{\beta w}{\alpha r}\right)^{\beta}} \right]^{\frac{1}{\alpha + \beta - 1}}, \qquad K^* = \frac{\beta w}{\alpha r} , L^*
$$

This is exactly what we’ll implement in Python, then cross-check numerically with scipy.optimize.

Concrete Numerical Example

We’ll use the following parameters:

• $A = 8$ (productivity)
• $\alpha = 0.35$ (labor elasticity)
• $\beta = 0.25$ (capital elasticity)
• $p = 10$ (output price)
• $w = 6$ (wage rate)
• $r = 4$ (capital rental rate)

Since $\alpha + \beta = 0.6 < 1$, the profit function is strictly concave in $(L, K)$, so it has a unique interior maximum.

Source Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize

# ---------- Dark theme for Colab plots ----------
plt.rcParams['figure.facecolor']  = '#1e1e1e'
plt.rcParams['axes.facecolor']    = '#1e1e1e'
plt.rcParams['savefig.facecolor'] = '#1e1e1e'
plt.rcParams['text.color']        = 'white'
plt.rcParams['axes.labelcolor']   = 'white'
plt.rcParams['xtick.color']       = 'white'
plt.rcParams['ytick.color']       = 'white'
plt.rcParams['axes.edgecolor']    = 'white'
plt.rcParams['grid.color']        = '#555555'

# ---------- Model parameters ----------
A = 8.0      # total factor productivity
alpha = 0.35 # output elasticity of labor
beta = 0.25  # output elasticity of capital
p = 10.0     # output price
w = 6.0      # wage rate
r = 4.0      # capital rental rate

# ---------- Core functions ----------
def production(L, K):
    """Cobb-Douglas production function."""
    return A * np.power(L, alpha) * np.power(K, beta)

def profit(L, K):
    """Profit as a function of labor and capital."""
    return p * production(L, K) - w * L - r * K

def neg_profit(x):
    """Negative profit for minimization (scipy minimizes by default)."""
    L, K = x
    if L <= 0 or K <= 0:
        return 1e12
    return -profit(L, K)

# ---------- Closed-form analytical solution ----------
def analytical_solution():
    ratio = (beta * w) / (alpha * r)                 # optimal K*/L*
    inner = p * A * alpha * (ratio ** beta)
    exponent = alpha + beta - 1.0
    L_star = (w / inner) ** (1.0 / exponent)
    K_star = ratio * L_star
    return L_star, K_star

L_analytical, K_analytical = analytical_solution()

# ---------- Numerical cross-check with scipy ----------
result = minimize(
    neg_profit,
    x0=[1.0, 1.0],
    method='L-BFGS-B',
    bounds=[(1e-6, None), (1e-6, None)]
)
L_numeric, K_numeric = result.x
Q_numeric = production(L_numeric, K_numeric)
profit_numeric = -result.fun

# ---------- Print results ----------
print("=== Profit Maximization: Cobb-Douglas Production ===")
print(f"Parameters: A={A}, alpha={alpha}, beta={beta}, p={p}, w={w}, r={r}")
print()
print("Closed-form analytical solution:")
print(f"  L* = {L_analytical:.6f}")
print(f"  K* = {K_analytical:.6f}")
print()
print("Numerical solution (scipy L-BFGS-B):")
print(f"  L* = {L_numeric:.6f}")
print(f"  K* = {K_numeric:.6f}")
print(f"  Q*  = {Q_numeric:.6f}")
print(f"  Profit* = {profit_numeric:.6f}")

# ============================================================
# Figure 1: 3D profit surface + contour map
# ============================================================
L_range = np.linspace(0.5, L_numeric * 3, 120)
K_range = np.linspace(0.5, K_numeric * 3, 120)
L_grid, K_grid = np.meshgrid(L_range, K_range)
Profit_grid = profit(L_grid, K_grid)

fig = plt.figure(figsize=(16, 7))

ax1 = fig.add_subplot(1, 2, 1, projection='3d')
surf = ax1.plot_surface(L_grid, K_grid, Profit_grid, cmap='plasma',
                         alpha=0.9, edgecolor='none', antialiased=True)
ax1.scatter([L_numeric], [K_numeric], [profit_numeric],
            color='cyan', s=120, marker='o', depthshade=False,
            label='Optimal point (L*, K*)')
ax1.set_xlabel('Labor  L')
ax1.set_ylabel('Capital  K')
ax1.set_zlabel('Profit  π')
ax1.set_title('Profit Surface over Labor and Capital', color='white')
ax1.view_init(elev=28, azim=-135)
ax1.legend(loc='upper left')
fig.colorbar(surf, ax=ax1, shrink=0.5, pad=0.12, label='Profit')

ax2 = fig.add_subplot(1, 2, 2)
cf = ax2.contourf(L_grid, K_grid, Profit_grid, levels=40, cmap='plasma')
ax2.contour(L_grid, K_grid, Profit_grid, levels=15, colors='white',
            linewidths=0.4, alpha=0.6)
ax2.plot(L_numeric, K_numeric, 'o', color='cyan', markersize=11,
         label='Optimal point (L*, K*)')
ax2.set_xlabel('Labor  L')
ax2.set_ylabel('Capital  K')
ax2.set_title('Profit Contour Map', color='white')
ax2.legend()
fig.colorbar(cf, ax=ax2, label='Profit')

plt.tight_layout()
plt.show()

# ============================================================
# Figure 2: Profit slices through the optimum
# ============================================================
fig2, axes = plt.subplots(1, 2, figsize=(14, 5))

L_slice = np.linspace(0.5, L_numeric * 3, 200)
profit_L_slice = profit(L_slice, K_numeric)
axes[0].plot(L_slice, profit_L_slice, color='cyan', linewidth=2.2)
axes[0].axvline(L_numeric, color='magenta', linestyle='--',
                 label=f'L* = {L_numeric:.2f}')
axes[0].set_xlabel('Labor  L')
axes[0].set_ylabel('Profit  π')
axes[0].set_title('Profit vs Labor (K fixed at K*)', color='white')
axes[0].grid(alpha=0.3)
axes[0].legend()

K_slice = np.linspace(0.5, K_numeric * 3, 200)
profit_K_slice = profit(L_numeric, K_slice)
axes[1].plot(K_slice, profit_K_slice, color='orange', linewidth=2.2)
axes[1].axvline(K_numeric, color='magenta', linestyle='--',
                 label=f'K* = {K_numeric:.2f}')
axes[1].set_xlabel('Capital  K')
axes[1].set_ylabel('Profit  π')
axes[1].set_title('Profit vs Capital (L fixed at L*)', color='white')
axes[1].grid(alpha=0.3)
axes[1].legend()

plt.tight_layout()
plt.show()

# ============================================================
# Figure 3: Comparative statics — how the optimum shifts with price p
# ============================================================
p_range = np.linspace(5.0, 20.0, 60)
L_path = np.zeros_like(p_range)
K_path = np.zeros_like(p_range)
profit_path = np.zeros_like(p_range)

ratio_fixed = (beta * w) / (alpha * r)
exponent = alpha + beta - 1.0

for i, p_val in enumerate(p_range):
    inner = p_val * A * alpha * (ratio_fixed ** beta)
    L_val = (w / inner) ** (1.0 / exponent)
    K_val = ratio_fixed * L_val
    L_path[i] = L_val
    K_path[i] = K_val
    profit_path[i] = p_val * production(L_val, K_val) - w * L_val - r * K_val

fig3, axes3 = plt.subplots(1, 3, figsize=(18, 5))

axes3[0].plot(p_range, L_path, color='cyan', linewidth=2.2)
axes3[0].set_xlabel('Output price  p')
axes3[0].set_ylabel('Optimal Labor  L*')
axes3[0].set_title('L* vs Output Price', color='white')
axes3[0].grid(alpha=0.3)

axes3[1].plot(p_range, K_path, color='orange', linewidth=2.2)
axes3[1].set_xlabel('Output price  p')
axes3[1].set_ylabel('Optimal Capital  K*')
axes3[1].set_title('K* vs Output Price', color='white')
axes3[1].grid(alpha=0.3)

axes3[2].plot(p_range, profit_path, color='lime', linewidth=2.2)
axes3[2].set_xlabel('Output price  p')
axes3[2].set_ylabel('Maximum Profit  π*')
axes3[2].set_title('Optimal Profit vs Output Price', color='white')
axes3[2].grid(alpha=0.3)

plt.tight_layout()
plt.show()