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)
|
|||
|
|
}
|
|||
|
|
}
|