296 lines
11 KiB
Go
296 lines
11 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package optim
|
||
|
||
import (
|
||
"math"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
"testing"
|
||
)
|
||
|
||
// TestLevenbergMarquardt fits y = a·e^{−bx} + c to exact data and
|
||
// checks that the fitted parameters match the generating values
|
||
// within the convergence tolerance of the LM fitter.
|
||
func TestLevenbergMarquardt(t *testing.T) {
|
||
xData := []float64{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
|
||
yData := make([]float64, len(xData))
|
||
for i, x := range xData {
|
||
yData[i] = 3*math.Exp(-0.5*float64(x)) + 0.5
|
||
}
|
||
yArr := mustFloats(t, yData, len(yData))
|
||
residual := func(p *core.Array) (*core.Array, error) {
|
||
a, b, c := p.FloatAt(0), p.FloatAt(1), p.FloatAt(2)
|
||
out := core.New(core.Float, len(xData))
|
||
for i := range len(xData) {
|
||
out.RawFloats()[i] = yArr.FloatAt(i) - (a*math.Exp(-b*float64(xData[i])) + c)
|
||
}
|
||
return out, nil
|
||
}
|
||
p0 := mustFloats(t, []float64{2, 0.3, 0.1}, 3)
|
||
pOpt, chi2, err := LevenbergMarquardt(residual, p0, LMOptions{})
|
||
if err != nil {
|
||
t.Fatalf("LevenbergMarquardt: %v", err)
|
||
}
|
||
if chi2 > 1e-4 {
|
||
t.Fatalf("χ² = %v, want < 1e-4 for exact data", chi2)
|
||
}
|
||
a := pOpt.FloatAt(0)
|
||
b := pOpt.FloatAt(1)
|
||
if math.Abs(a-3) > 0.01 || math.Abs(b-0.5) > 0.01 {
|
||
t.Fatalf("a = %.8f, b = %.8f, want ≈ 3, 0.5", a, b)
|
||
}
|
||
}
|
||
|
||
// TestLevenbergMarquardtErrors pins the error contract.
|
||
func TestLevenbergMarquardtErrors(t *testing.T) {
|
||
f := func(p *core.Array) (*core.Array, error) { return nil, nil }
|
||
cx, _ := core.FromComplexes([]complex128{1}, 1)
|
||
if _, _, err := LevenbergMarquardt(f, cx, LMOptions{}); err == nil {
|
||
t.Fatal("expected an error for complex parameters")
|
||
}
|
||
empty, _ := core.FromFloats(nil, 0)
|
||
if _, _, err := LevenbergMarquardt(f, empty, LMOptions{}); err == nil {
|
||
t.Fatal("expected an error for an empty parameter vector")
|
||
}
|
||
}
|
||
|
||
// TestLevenbergMarquardtAnalyticJacobian fits the same exponential
|
||
// model twice, once with the analytic Jacobian and once with the
|
||
// central-difference fallback: both must land on the same optimum,
|
||
// the analytic one reaching it without the extra evaluations.
|
||
func TestLevenbergMarquardtAnalyticJacobian(t *testing.T) {
|
||
xData := []float64{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
|
||
yData := make([]float64, len(xData))
|
||
for i, x := range xData {
|
||
yData[i] = 3*math.Exp(-0.5*float64(x)) + 0.5
|
||
}
|
||
residual := func(p *core.Array) (*core.Array, error) {
|
||
a, b, c := p.FloatAt(0), p.FloatAt(1), p.FloatAt(2)
|
||
out := core.New(core.Float, len(xData))
|
||
for i := range len(xData) {
|
||
out.RawFloats()[i] = yData[i] - (a*math.Exp(-b*float64(xData[i])) + c)
|
||
}
|
||
return out, nil
|
||
}
|
||
jacobian := func(p *core.Array) (*core.Array, error) {
|
||
a, b := p.FloatAt(0), p.FloatAt(1)
|
||
out := core.New(core.Float, len(xData), 3)
|
||
for i := range len(xData) {
|
||
e := math.Exp(-b * float64(xData[i]))
|
||
out.RawFloats()[i*3+0] = -e
|
||
out.RawFloats()[i*3+1] = a * float64(xData[i]) * e
|
||
out.RawFloats()[i*3+2] = -1
|
||
}
|
||
return out, nil
|
||
}
|
||
p0 := mustFloats(t, []float64{2, 0.3, 0.1}, 3)
|
||
pA, chi2A, err := LevenbergMarquardt(residual, p0, LMOptions{Jacobian: jacobian})
|
||
if err != nil {
|
||
t.Fatalf("LevenbergMarquardt analytic: %v", err)
|
||
}
|
||
pD, chi2D, err := LevenbergMarquardt(residual, p0, LMOptions{})
|
||
if err != nil {
|
||
t.Fatalf("LevenbergMarquardt differences: %v", err)
|
||
}
|
||
for j := range 3 {
|
||
if math.Abs(pA.FloatAt(j)-pD.FloatAt(j)) > 1e-6 {
|
||
t.Fatalf("parameter %d: analytic %.12g, differences %.12g", j, pA.FloatAt(j), pD.FloatAt(j))
|
||
}
|
||
}
|
||
if chi2A > 1e-4 {
|
||
t.Fatalf("analytic χ² = %v, want < 1e-4", chi2A)
|
||
}
|
||
if math.Abs(chi2A-chi2D) > 1e-8 {
|
||
t.Fatalf("χ² disagree: analytic %v, differences %v", chi2A, chi2D)
|
||
}
|
||
}
|
||
|
||
// TestLevenbergMarquardtAnalyticExact solves a determined linear
|
||
// system, where the Gauss-Newton step is exact and the analytic
|
||
// Jacobian reaches the solution the equations dictate.
|
||
func TestLevenbergMarquardtAnalyticExact(t *testing.T) {
|
||
residual := func(p *core.Array) (*core.Array, error) {
|
||
return core.FromFloats([]float64{
|
||
p.FloatAt(0) + 2*p.FloatAt(1) - 3,
|
||
2*p.FloatAt(0) + p.FloatAt(1) - 4,
|
||
}, 2)
|
||
}
|
||
jacobian := func(p *core.Array) (*core.Array, error) {
|
||
return core.FromFloats([]float64{1, 2, 2, 1}, 2, 2)
|
||
}
|
||
p0 := mustFloats(t, []float64{0, 0}, 2)
|
||
pOpt, chi2, err := LevenbergMarquardt(residual, p0, LMOptions{Jacobian: jacobian, Tolerance: 1e-14})
|
||
if err != nil {
|
||
t.Fatalf("LevenbergMarquardt: %v", err)
|
||
}
|
||
if math.Abs(pOpt.FloatAt(0)-5.0/3) > 1e-8 || math.Abs(pOpt.FloatAt(1)-2.0/3) > 1e-8 {
|
||
t.Fatalf("p = (%.12g, %.12g), want (5/3, 2/3)", pOpt.FloatAt(0), pOpt.FloatAt(1))
|
||
}
|
||
if chi2 > 1e-16 {
|
||
t.Fatalf("χ² = %v, want 0", chi2)
|
||
}
|
||
}
|
||
|
||
// TestLevenbergMarquardtJacobianErrors pins the analytic Jacobian's
|
||
// own contract: shape mismatches and callback errors surface.
|
||
func TestLevenbergMarquardtJacobianErrors(t *testing.T) {
|
||
residual := func(p *core.Array) (*core.Array, error) {
|
||
return core.FromFloats([]float64{p.FloatAt(0) - 1, p.FloatAt(1) - 2}, 2)
|
||
}
|
||
rank1 := func(p *core.Array) (*core.Array, error) {
|
||
return core.FromFloats([]float64{1, 1}, 2)
|
||
}
|
||
if _, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2),
|
||
LMOptions{Jacobian: rank1}); err == nil {
|
||
t.Fatal("expected an error for a rank-1 Jacobian")
|
||
}
|
||
wrongDims := func(p *core.Array) (*core.Array, error) {
|
||
return core.FromFloats([]float64{1, 1, 1}, 3, 1)
|
||
}
|
||
if _, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2),
|
||
LMOptions{Jacobian: wrongDims}); err == nil {
|
||
t.Fatal("expected an error for a Jacobian of the wrong dimensions")
|
||
}
|
||
boom := func(p *core.Array) (*core.Array, error) {
|
||
return nil, base.Errf("jacobian failed")
|
||
}
|
||
if _, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2),
|
||
LMOptions{Jacobian: boom}); err == nil {
|
||
t.Fatal("expected the Jacobian error to propagate")
|
||
}
|
||
}
|
||
|
||
// TestLevenbergMarquardtDtypes pins the parameter-vector promotion:
|
||
// Int and Float32 starts must behave exactly like Float64 ones. The
|
||
// old RawFloats() copy silently started those fits from zeros.
|
||
func TestLevenbergMarquardtDtypes(t *testing.T) {
|
||
residual := func(p *core.Array) (*core.Array, error) {
|
||
return core.FromFloats([]float64{
|
||
p.FloatAt(0) + 2*p.FloatAt(1) - 3,
|
||
2*p.FloatAt(0) + p.FloatAt(1) - 4,
|
||
}, 2)
|
||
}
|
||
jacobian := func(*core.Array) (*core.Array, error) {
|
||
return core.FromFloats([]float64{1, 2, 2, 1}, 2, 2)
|
||
}
|
||
ints, err := core.FromInts([]int64{0, 0}, 2)
|
||
if err != nil {
|
||
t.Fatalf("FromInts: %v", err)
|
||
}
|
||
thirtyTwo, err := core.FromFloat32s([]float32{0, 0}, 2)
|
||
if err != nil {
|
||
t.Fatalf("FromFloat32s: %v", err)
|
||
}
|
||
for name, p0 := range map[string]*core.Array{"int": ints, "float32": thirtyTwo} {
|
||
pOpt, chi2, err := LevenbergMarquardt(residual, p0, LMOptions{Jacobian: jacobian, Tolerance: 1e-14})
|
||
if err != nil {
|
||
t.Fatalf("LevenbergMarquardt(%s start): %v", name, err)
|
||
}
|
||
if math.Abs(pOpt.FloatAt(0)-5.0/3) > 1e-8 || math.Abs(pOpt.FloatAt(1)-2.0/3) > 1e-8 {
|
||
t.Fatalf("LevenbergMarquardt(%s start): p = (%.12g, %.12g), want (5/3, 2/3)",
|
||
name, pOpt.FloatAt(0), pOpt.FloatAt(1))
|
||
}
|
||
if chi2 > 1e-16 {
|
||
t.Fatalf("LevenbergMarquardt(%s start): χ² = %v, want 0", name, chi2)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestLevenbergMarquardtResidualLengthChange pins the residual
|
||
// contract: a callback whose output length changes mid-fit is an
|
||
// error naming the mismatch, not a shape panic.
|
||
func TestLevenbergMarquardtResidualLengthChange(t *testing.T) {
|
||
calls := 0
|
||
residual := func(*core.Array) (*core.Array, error) {
|
||
calls++
|
||
if calls == 1 {
|
||
return core.FromFloats([]float64{1, 2, 3}, 3)
|
||
}
|
||
return core.FromFloats([]float64{1, 2}, 2)
|
||
}
|
||
_, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2), LMOptions{})
|
||
if err == nil {
|
||
t.Fatal("residual length change mid-fit: want an error")
|
||
}
|
||
}
|
||
|
||
// TestLevenbergMarquardtStencilMatchesAnalyticJacobian pins the reused
|
||
// difference stencil against the analytic route. The stencil carries
|
||
// the offset on one parameter at a time and puts it back as soon as the
|
||
// column is differenced, so each column's Jacobian is the central
|
||
// difference of that parameter alone and the two routes land on the
|
||
// same point. A stencil that leaves its offset behind differences every
|
||
// later column at a point that is also displaced earlier: a different
|
||
// Jacobian, and a fit that parts company with the analytic route.
|
||
func TestLevenbergMarquardtStencilMatchesAnalyticJacobian(t *testing.T) {
|
||
// A decay with an offset and a smooth nuisance term, so every
|
||
// parameter carries curvature and the columns are coupled.
|
||
const n = 40
|
||
xs := make([]float64, n)
|
||
ys := make([]float64, n)
|
||
for i := range n {
|
||
xs[i] = float64(i) * 0.25
|
||
ys[i] = 1.5*math.Exp(-0.8*xs[i]) + 0.7 + 0.05*math.Sin(2*xs[i])
|
||
}
|
||
residual := func(p *core.Array) (*core.Array, error) {
|
||
out := core.New(core.Float, n)
|
||
v := out.RawFloats()
|
||
for i := range n {
|
||
v[i] = ys[i] - (p.FloatAt(0)*math.Exp(-p.FloatAt(1)*xs[i]) +
|
||
p.FloatAt(2) + p.FloatAt(3)*math.Sin(2*xs[i]))
|
||
}
|
||
return out, nil
|
||
}
|
||
jacobian := func(p *core.Array) (*core.Array, error) {
|
||
v := make([]float64, n*4)
|
||
for i := range n {
|
||
e := math.Exp(-p.FloatAt(1) * xs[i])
|
||
v[i*4+0] = -e
|
||
v[i*4+1] = p.FloatAt(0) * xs[i] * e
|
||
v[i*4+2] = -1
|
||
v[i*4+3] = -math.Sin(2 * xs[i])
|
||
}
|
||
return core.FromFloats(v, n, 4)
|
||
}
|
||
start, err := core.FromFloats([]float64{1.2, 0.9, 0.4, 0.04}, 4)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
fd, fdChi, err := LevenbergMarquardt(residual, start, LMOptions{})
|
||
if err != nil {
|
||
t.Fatalf("LevenbergMarquardt with a difference stencil: %v", err)
|
||
}
|
||
analytic, analyticChi, err := LevenbergMarquardt(residual, start, LMOptions{Jacobian: jacobian})
|
||
if err != nil {
|
||
t.Fatalf("LevenbergMarquardt with an analytic Jacobian: %v", err)
|
||
}
|
||
if fd.Len() != analytic.Len() {
|
||
t.Fatalf("the two routes fit %d and %d parameters", fd.Len(), analytic.Len())
|
||
}
|
||
for i := range fd.Len() {
|
||
if d := math.Abs(fd.FloatAt(i) - analytic.FloatAt(i)); d > 1e-12 {
|
||
t.Fatalf("parameter %d: the stencil route gives %.17g, the analytic one %.17g (differ by %g)",
|
||
i, fd.FloatAt(i), analytic.FloatAt(i), d)
|
||
}
|
||
}
|
||
if math.Abs(fdChi-analyticChi) > 1e-12 {
|
||
t.Fatalf("the two routes report chi2 %.17g and %.17g", fdChi, analyticChi)
|
||
}
|
||
// The stencil is carried across the fit's iterations, so the same
|
||
// start must reproduce the same bits.
|
||
again, againChi, err := LevenbergMarquardt(residual, start, LMOptions{})
|
||
if err != nil {
|
||
t.Fatalf("the repeated stencil fit: %v", err)
|
||
}
|
||
for i := range again.Len() {
|
||
if math.Float64bits(again.FloatAt(i)) != math.Float64bits(fd.FloatAt(i)) {
|
||
t.Fatalf("parameter %d moved between stencil fits: %.17g against %.17g", i, again.FloatAt(i), fd.FloatAt(i))
|
||
}
|
||
}
|
||
if math.Float64bits(againChi) != math.Float64bits(fdChi) {
|
||
t.Fatalf("chi2 moved between stencil fits: %.17g against %.17g", againChi, fdChi)
|
||
}
|
||
}
|