Skip to content

Commit 7937730

Browse files
authored
Add BFGS and Nelder-Mead multivariate optimization algorithms (#331)
- BFGS: Quasi-Newton method with automatic differentiation - Nelder-Mead: Derivative-free simplex method - Update optimize.livemd with multivariate section - Add efficient_frontier.livemd portfolio optimization example
1 parent 036c654 commit 7937730

6 files changed

Lines changed: 1545 additions & 25 deletions

File tree

lib/scholar/optimize/bfgs.ex

Lines changed: 289 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,289 @@
1+
defmodule Scholar.Optimize.BFGS do
2+
@moduledoc """
3+
BFGS (Broyden-Fletcher-Goldfarb-Shanno) algorithm for multivariate function minimization.
4+
5+
BFGS is a quasi-Newton optimization method that approximates the inverse Hessian
6+
matrix using gradient information. It is well-suited for smooth, differentiable
7+
objective functions and typically converges much faster than derivative-free methods
8+
like Nelder-Mead.
9+
10+
## Algorithm
11+
12+
At each iteration, the algorithm:
13+
1. Computes the gradient using automatic differentiation
14+
2. Determines a search direction from the inverse Hessian approximation
15+
3. Performs a line search to find an acceptable step length
16+
4. Updates the inverse Hessian approximation using the BFGS formula
17+
18+
## Convergence
19+
20+
The algorithm converges when the gradient norm falls below the specified tolerance.
21+
22+
## References
23+
24+
* Nocedal, J. and Wright, S. J. (2006). "Numerical Optimization"
25+
* Fletcher, R. (1987). "Practical Methods of Optimization"
26+
"""
27+
28+
import Nx.Defn
29+
30+
@derive {Nx.Container, containers: [:x, :fun, :converged, :iterations, :fun_evals, :grad_evals]}
31+
defstruct [:x, :fun, :converged, :iterations, :fun_evals, :grad_evals]
32+
33+
@type t :: %__MODULE__{
34+
x: Nx.Tensor.t(),
35+
fun: Nx.Tensor.t(),
36+
converged: Nx.Tensor.t(),
37+
iterations: Nx.Tensor.t(),
38+
fun_evals: Nx.Tensor.t(),
39+
grad_evals: Nx.Tensor.t()
40+
}
41+
42+
# Line search parameter (Armijo condition)
43+
@c1 1.0e-4
44+
45+
opts = [
46+
gtol: [
47+
type: {:custom, Scholar.Options, :positive_number, []},
48+
default: 1.0e-5,
49+
doc: """
50+
Gradient norm tolerance for convergence. The algorithm stops when
51+
the infinity norm of the gradient is below this threshold.
52+
"""
53+
],
54+
maxiter: [
55+
type: :pos_integer,
56+
default: 500,
57+
doc: "Maximum number of iterations."
58+
]
59+
]
60+
61+
@opts_schema NimbleOptions.new!(opts)
62+
63+
@doc """
64+
Minimizes a multivariate function using the BFGS algorithm.
65+
66+
BFGS is a gradient-based method that uses automatic differentiation to compute
67+
gradients and approximates the inverse Hessian matrix for fast convergence.
68+
69+
## Arguments
70+
71+
* `x0` - Initial guess as a 1-D tensor of shape `{n}`.
72+
* `fun` - The objective function to minimize. Must be a defn-compatible function
73+
that takes a 1-D tensor of shape `{n}` and returns a scalar tensor.
74+
* `opts` - Options (see below).
75+
76+
## Options
77+
78+
#{NimbleOptions.docs(@opts_schema)}
79+
80+
## Returns
81+
82+
A `Scholar.Optimize.BFGS` struct with the optimization result:
83+
84+
* `:x` - The optimal point found (shape `{n}`)
85+
* `:fun` - The function value at the optimal point
86+
* `:converged` - Whether the optimization converged (1 if true, 0 if false)
87+
* `:iterations` - Number of iterations performed
88+
* `:fun_evals` - Number of function evaluations
89+
* `:grad_evals` - Number of gradient evaluations
90+
91+
## Examples
92+
93+
iex> # Minimize a simple quadratic: f(x) = x1^2 + x2^2
94+
iex> fun = fn x -> Nx.sum(Nx.pow(x, 2)) end
95+
iex> x0 = Nx.tensor([1.0, 2.0])
96+
iex> result = Scholar.Optimize.BFGS.minimize(x0, fun)
97+
iex> Nx.to_number(result.converged)
98+
1
99+
iex> Nx.all_close(result.x, Nx.tensor([0.0, 0.0]), atol: 1.0e-4) |> Nx.to_number()
100+
1
101+
102+
For higher precision, use f64 tensors:
103+
104+
iex> fun = fn x -> Nx.sum(Nx.pow(x, 2)) end
105+
iex> x0 = Nx.tensor([1.0, 2.0], type: :f64)
106+
iex> result = Scholar.Optimize.BFGS.minimize(x0, fun, gtol: 1.0e-10)
107+
iex> Nx.to_number(result.converged)
108+
1
109+
iex> Nx.all_close(result.x, Nx.tensor([0.0, 0.0]), atol: 1.0e-8) |> Nx.to_number()
110+
1
111+
112+
Minimizing the Rosenbrock function:
113+
114+
iex> # Rosenbrock: f(x,y) = (1-x)^2 + 100*(y-x^2)^2, minimum at (1, 1)
115+
iex> rosenbrock = fn x ->
116+
...> x0 = x[0]
117+
...> x1 = x[1]
118+
...> term1 = Nx.pow(Nx.subtract(1, x0), 2)
119+
...> term2 = Nx.multiply(100, Nx.pow(Nx.subtract(x1, Nx.pow(x0, 2)), 2))
120+
...> Nx.add(term1, term2)
121+
...> end
122+
iex> x0 = Nx.tensor([0.0, 0.0], type: :f64)
123+
iex> result = Scholar.Optimize.BFGS.minimize(x0, rosenbrock, gtol: 1.0e-8, maxiter: 1000)
124+
iex> Nx.to_number(result.converged)
125+
1
126+
iex> Nx.all_close(result.x, Nx.tensor([1.0, 1.0]), atol: 1.0e-4) |> Nx.to_number()
127+
1
128+
"""
129+
deftransform minimize(x0, fun, opts \\ []) do
130+
opts = NimbleOptions.validate!(opts, @opts_schema)
131+
minimize_n(x0, fun, opts[:gtol], opts[:maxiter])
132+
end
133+
134+
defnp minimize_n(x0, fun, gtol, maxiter) do
135+
x0 = Nx.flatten(x0)
136+
{n} = Nx.shape(x0)
137+
138+
# Initial function value and gradient
139+
{f0, g0} = value_and_grad(x0, fun)
140+
141+
# Initialize inverse Hessian as identity matrix
142+
h_inv = Nx.eye(n, type: Nx.type(x0))
143+
144+
# Initial state
145+
initial_state = %{
146+
x: x0,
147+
f: f0,
148+
g: g0,
149+
h_inv: h_inv,
150+
iter: Nx.u32(0),
151+
f_evals: Nx.u32(1),
152+
g_evals: Nx.u32(1)
153+
}
154+
155+
# Main optimization loop
156+
{final_state, _} =
157+
while {state = initial_state, {gtol, maxiter}},
158+
not converged?(state, gtol) and state.iter < maxiter do
159+
new_state = bfgs_step(state, fun)
160+
{new_state, {gtol, maxiter}}
161+
end
162+
163+
# Check convergence
164+
converged = converged?(final_state, gtol)
165+
166+
%__MODULE__{
167+
x: final_state.x,
168+
fun: final_state.f,
169+
converged: converged,
170+
iterations: final_state.iter,
171+
fun_evals: final_state.f_evals,
172+
grad_evals: final_state.g_evals
173+
}
174+
end
175+
176+
# Check convergence: infinity norm of gradient < gtol
177+
defnp converged?(state, gtol) do
178+
grad_norm = Nx.reduce_max(Nx.abs(state.g))
179+
Nx.less(grad_norm, gtol)
180+
end
181+
182+
# Perform one BFGS iteration
183+
defnp bfgs_step(state, fun) do
184+
x = state.x
185+
f = state.f
186+
g = state.g
187+
h_inv = state.h_inv
188+
189+
# Compute search direction: p = -H_inv * g
190+
p = Nx.negate(Nx.dot(h_inv, g))
191+
192+
# Line search to find step length alpha
193+
{alpha, f_new, g_new, ls_f_evals, ls_g_evals} =
194+
line_search(x, f, g, p, fun)
195+
196+
# Compute step and gradient change
197+
s = Nx.multiply(alpha, p)
198+
x_new = Nx.add(x, s)
199+
y = Nx.subtract(g_new, g)
200+
201+
# Update inverse Hessian using BFGS formula
202+
h_inv_new = update_inverse_hessian(h_inv, s, y)
203+
204+
%{
205+
state
206+
| x: x_new,
207+
f: f_new,
208+
g: g_new,
209+
h_inv: h_inv_new,
210+
iter: Nx.add(state.iter, 1),
211+
f_evals: Nx.add(state.f_evals, ls_f_evals),
212+
g_evals: Nx.add(state.g_evals, ls_g_evals)
213+
}
214+
end
215+
216+
# Backtracking line search with Armijo condition
217+
defnp line_search(x, f, g, p, fun) do
218+
# Directional derivative at alpha=0
219+
slope = Nx.dot(g, p)
220+
221+
# Try step sizes: 1, 0.5, 0.25, 0.125, ...
222+
# Unrolled at compile time for fusion.
223+
{alpha, _} =
224+
while {alpha = Nx.tensor(1.0, type: Nx.type(x)), {x, f, p, slope}},
225+
_i <- 0..9//1,
226+
unroll: true do
227+
new_alpha = backtrack_step(x, f, p, slope, fun, alpha)
228+
{new_alpha, {x, f, p, slope}}
229+
end
230+
231+
# Evaluate function and gradient at final alpha
232+
x_final = Nx.add(x, Nx.multiply(alpha, p))
233+
{f_final, g_final} = value_and_grad(x_final, fun)
234+
235+
{alpha, f_final, g_final, Nx.u32(11), Nx.u32(1)}
236+
end
237+
238+
# Single backtracking step
239+
defnp backtrack_step(x, f, p, slope, fun, alpha) do
240+
x_trial = Nx.add(x, Nx.multiply(alpha, p))
241+
f_trial = fun.(x_trial)
242+
243+
# Check Armijo condition
244+
armijo_ok =
245+
Nx.less_equal(f_trial, Nx.add(f, Nx.multiply(@c1, Nx.multiply(alpha, slope))))
246+
247+
# If satisfied, keep alpha; otherwise halve it
248+
Nx.select(armijo_ok, alpha, Nx.multiply(0.5, alpha))
249+
end
250+
251+
# BFGS inverse Hessian update
252+
# H_new = (I - rho*s*y') * H * (I - rho*y*s') + rho*s*s'
253+
defnp update_inverse_hessian(h_inv, s, y) do
254+
{n} = Nx.shape(s)
255+
256+
# rho = 1 / (y' * s)
257+
ys = Nx.dot(y, s)
258+
259+
# Skip update if ys is too small (would cause numerical issues)
260+
skip_update = Nx.less(Nx.abs(ys), 1.0e-10)
261+
262+
rho = Nx.select(skip_update, 0.0, Nx.divide(1.0, ys))
263+
264+
# Compute update terms
265+
# I - rho*s*y' and I - rho*y*s'
266+
eye = Nx.eye(n, type: Nx.type(s))
267+
268+
# s and y as column vectors for outer products
269+
s_col = Nx.reshape(s, {n, 1})
270+
y_col = Nx.reshape(y, {n, 1})
271+
s_row = Nx.reshape(s, {1, n})
272+
y_row = Nx.reshape(y, {1, n})
273+
274+
# (I - rho*s*y')
275+
left = Nx.subtract(eye, Nx.multiply(rho, Nx.dot(s_col, y_row)))
276+
277+
# (I - rho*y*s')
278+
right = Nx.subtract(eye, Nx.multiply(rho, Nx.dot(y_col, s_row)))
279+
280+
# rho*s*s'
281+
ss_term = Nx.multiply(rho, Nx.dot(s_col, s_row))
282+
283+
# H_new = left * H * right + ss_term
284+
h_new = Nx.add(Nx.dot(Nx.dot(left, h_inv), right), ss_term)
285+
286+
# Use old H if update skipped
287+
Nx.select(skip_update, h_inv, h_new)
288+
end
289+
end

0 commit comments

Comments
 (0)