feat: initial release

Assisted-by: GLM 5.3 Flash
2026-09-03 10:00:00 +02:00
commit af4ee19703
617 changed files with 191195 additions and 0 deletions
@@ -0,0 +1,341 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package optim
+
+import (
+	"fmt"
+	"math"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+	"sourcedock.dev/petrbalvin/tensor/internal/engine"
+)
+
+// Tests and benchmarks for the opt-in parallel sweep of the
+// central-difference Jacobian, LMOptions.ParallelJacobian and
+// RootSystemOptions.ParallelJacobian. The contract under test: the
+// default false never calls the callback from more than one goroutine
+// and answers bit for bit what it answered before; the true value may
+// overlap the callback's evaluations across columns and still answers
+// bit for bit the same numbers, because every column is differenced by
+// the same stencil against the same point.
+
+// concurrencyTracker records the largest number of callbacks it has
+// seen inside at once. The counter and the peak are atomic: the tracker
+// itself must not be the thing that serialises the calls it measures.
+type concurrencyTracker struct {
+	cur  atomic.Int64
+	peak atomic.Int64
+}
+
+func (c *concurrencyTracker) enter() {
+	n := c.cur.Add(1)
+	for {
+		p := c.peak.Load()
+		if n <= p || c.peak.CompareAndSwap(p, n) {
+			return
+		}
+	}
+}
+
+func (c *concurrencyTracker) leave() { c.cur.Add(-1) }
+
+// lmParallelProblem builds an eight-parameter two-exponential plus
+// sinusoid fit with a known generating model, hard enough that the fit
+// walks a number of damped steps before it converges.
+func lmParallelProblem() (residual func(*core.Array) (*core.Array, error), start []float64) {
+	const nObs = 48
+	t := make([]float64, nObs)
+	y := make([]float64, nObs)
+	truth := []float64{2.0, 0.7, 1.5, 1.9, 0.8, 3.1, 0.4, 0.05}
+	for i := range nObs {
+		t[i] = float64(i) / 6
+		y[i] = truth[0]*math.Exp(-truth[1]*t[i]) + truth[2]*math.Exp(-truth[3]*t[i]) +
+			truth[4]*math.Sin(truth[5]*t[i]+truth[6]) + truth[7]*t[i]
+	}
+	residual = func(p *core.Array) (*core.Array, error) {
+		out := core.New(core.Float, nObs)
+		vals := out.RawFloats()
+		for i := range nObs {
+			vals[i] = p.FloatAt(0)*math.Exp(-p.FloatAt(1)*t[i]) + p.FloatAt(2)*math.Exp(-p.FloatAt(3)*t[i]) +
+				p.FloatAt(4)*math.Sin(p.FloatAt(5)*t[i]+p.FloatAt(6)) + p.FloatAt(7)*t[i] - y[i]
+		}
+		return out, nil
+	}
+	return residual, []float64{1.6, 0.5, 1.2, 1.6, 0.6, 2.7, 0.2, 0.02}
+}
+
+// rootParallelProblem builds eight coupled quadratic equations whose
+// root sits near xᵢ = i+1, hard enough that Newton walks several rounds.
+func rootParallelProblem() (f func(*core.Array) (*core.Array, error), start []float64) {
+	const n = 8
+	f = func(x *core.Array) (*core.Array, error) {
+		out := core.New(core.Float, n)
+		vals := out.RawFloats()
+		for i := range n {
+			vals[i] = x.FloatAt(i)*x.FloatAt(i) - float64(i+1)*float64(i+1) + 0.1*x.FloatAt((i+1)%n)
+		}
+		return out, nil
+	}
+	start = make([]float64, n)
+	for i := range start {
+		start[i] = float64(i+1) + 0.5
+	}
+	return f, start
+}
+
+func TestParallelJacobianBitIdenticalFit(t *testing.T) {
+	residual, start := lmParallelProblem()
+	run := func(parallel bool) (*core.Array, float64) {
+		t.Helper()
+		p, chi2, err := LevenbergMarquardt(residual, mustFloats(t, start), LMOptions{
+			MaxIterations:    100,
+			ParallelJacobian: parallel,
+		})
+		if err != nil {
+			t.Fatalf("LevenbergMarquardt(ParallelJacobian=%v): %v", parallel, err)
+		}
+		return p, chi2
+	}
+	pSerial, chi2Serial := run(false)
+	pParallel, chi2Parallel := run(true)
+	if math.Float64bits(chi2Serial) != math.Float64bits(chi2Parallel) {
+		t.Fatalf("chi2 differs: serial %v, parallel %v", chi2Serial, chi2Parallel)
+	}
+	serial, parallel := pSerial.RawFloats(), pParallel.RawFloats()
+	for i := range serial {
+		if math.Float64bits(serial[i]) != math.Float64bits(parallel[i]) {
+			t.Fatalf("parameter %d differs: serial %.17g, parallel %.17g", i, serial[i], parallel[i])
+		}
+	}
+}
+
+func TestParallelJacobianBitIdenticalRoot(t *testing.T) {
+	f, start := rootParallelProblem()
+	run := func(parallel, broyden bool) (*core.Array, float64) {
+		t.Helper()
+		x, res, err := FindRootSystem(f, mustFloats(t, start), RootSystemOptions{
+			MaxIterations:    100,
+			UseBroyden:       broyden,
+			ParallelJacobian: parallel,
+		})
+		if err != nil {
+			t.Fatalf("FindRootSystem(UseBroyden=%v, ParallelJacobian=%v): %v", broyden, parallel, err)
+		}
+		return x, res
+	}
+	for _, broyden := range []bool{false, true} {
+		xSerial, resSerial := run(false, broyden)
+		xParallel, resParallel := run(true, broyden)
+		if math.Float64bits(resSerial) != math.Float64bits(resParallel) {
+			t.Fatalf("UseBroyden=%v: residual norm differs: serial %v, parallel %v", broyden, resSerial, resParallel)
+		}
+		sSerial, sParallel := xSerial.RawFloats(), xParallel.RawFloats()
+		for i := range sSerial {
+			if math.Float64bits(sSerial[i]) != math.Float64bits(sParallel[i]) {
+				t.Fatalf("UseBroyden=%v: root coordinate %d differs: serial %.17g, parallel %.17g",
+					broyden, i, sSerial[i], sParallel[i])
+			}
+		}
+	}
+}
+
+// TestParallelJacobianBitIdenticalBroydenRebuilds drives the rebuild
+// path: the duplicated equation x² = 1 has a rank-one Jacobian, so a
+// Broyden run from the hopeless start rebuilds the numerical Jacobian
+// every round until the budget refuses it. Serial and parallel must
+// refuse with the same words, which pins every rebuilt Jacobian's bits.
+func TestParallelJacobianBitIdenticalBroydenRebuilds(t *testing.T) {
+	rankOne := func(x *core.Array) (*core.Array, error) {
+		v := x.FloatAt(0)*x.FloatAt(0) - 1
+		out := core.New(core.Float, 2)
+		vals := out.RawFloats()
+		vals[0], vals[1] = v, v
+		return out, nil
+	}
+	run := func(parallel bool) string {
+		t.Helper()
+		_, _, err := FindRootSystem(rankOne, mustFloats(t, []float64{1.5, 1.5}), RootSystemOptions{
+			MaxIterations:    12,
+			UseBroyden:       true,
+			ParallelJacobian: parallel,
+		})
+		if err == nil {
+			t.Fatalf("FindRootSystem(ParallelJacobian=%v): want the budget refusal", parallel)
+		}
+		return err.Error()
+	}
+	if serial, parallel := run(false), run(true); serial != parallel {
+		t.Fatalf("the refusal differs: serial %q, parallel %q", serial, parallel)
+	}
+}
+
+// TestParallelJacobianCallbackConcurrency pins the consent boundary:
+// with the option off the residual callback never runs inside more
+// than one goroutine at once, and with it on the sweeps do overlap.
+func TestParallelJacobianCallbackConcurrency(t *testing.T) {
+	prev := engine.SetNumWorkers(4)
+	defer engine.SetNumWorkers(prev)
+
+	const nP = 8
+	const nObs = 8
+	peak := func(t *testing.T, parallel bool) int64 {
+		t.Helper()
+		var tr concurrencyTracker
+		residual := func(p *core.Array) (*core.Array, error) {
+			tr.enter()
+			defer tr.leave()
+			if parallel {
+				// The caller holds its ground until a second one is
+				// inside, so the overlap the option promises is
+				// demonstrated by construction instead of left to the
+				// scheduler to interleave two calls this short. A
+				// sweep that collapsed to one goroutine would wait
+				// out the grace alone and fail the peak below.
+				deadline := time.Now().Add(250 * time.Millisecond)
+				for tr.peak.Load() < 2 && time.Now().Before(deadline) {
+					runtime.Gosched()
+				}
+			}
+			s := 0.0
+			for k := range 800 {
+				s += math.Sin(float64(k)*7e-4+p.FloatAt(k&(nP-1))) * math.Cos(float64(k)*3e-4)
+			}
+			out := core.New(core.Float, nObs)
+			vals := out.RawFloats()
+			for i := range nObs {
+				vals[i] = s + p.FloatAt(i&(nP-1))
+			}
+			return out, nil
+		}
+		start := make([]float64, nP)
+		if _, _, err := LevenbergMarquardt(residual, mustFloats(t, start), LMOptions{
+			MaxIterations:    2,
+			AllowBudgetExit:  true,
+			ParallelJacobian: parallel,
+		}); err != nil {
+			t.Fatalf("LevenbergMarquardt(ParallelJacobian=%v): %v", parallel, err)
+		}
+		return tr.peak.Load()
+	}
+	t.Run("serial", func(t *testing.T) {
+		if got := peak(t, false); got != 1 {
+			t.Fatalf("the serial run reached %d concurrent callback calls, want exactly 1", got)
+		}
+	})
+	t.Run("parallel", func(t *testing.T) {
+		if got := peak(t, true); got < 2 {
+			t.Fatalf("the parallel run reached %d concurrent callback calls, want more than 1", got)
+		}
+	})
+}
+
+// TestParallelJacobianErrorMatchesSerial fails one column's stencil
+// and requires the parallel sweep to report the same error the serial
+// walk reports: the lowest failing column, not whichever worker got
+// there first.
+func TestParallelJacobianErrorMatchesSerial(t *testing.T) {
+	const nP = 6
+	const nObs = 10
+	// Coordinate 3 is nonzero exactly while column 3's stencils run,
+	// so only that column's residual evaluations fail.
+	residual := func(p *core.Array) (*core.Array, error) {
+		if p.FloatAt(3) != 0 {
+			return nil, fmt.Errorf("the stencil touched column 3")
+		}
+		out := core.New(core.Float, nObs)
+		for i := range nObs {
+			out.RawFloats()[i] = p.FloatAt(0) + float64(i)
+		}
+		return out, nil
+	}
+	run := func(parallel bool) string {
+		t.Helper()
+		start := make([]float64, nP)
+		_, _, err := LevenbergMarquardt(residual, mustFloats(t, start), LMOptions{
+			MaxIterations:    10,
+			ParallelJacobian: parallel,
+		})
+		if err == nil {
+			t.Fatalf("LevenbergMarquardt(ParallelJacobian=%v): want the stencil error", parallel)
+		}
+		return err.Error()
+	}
+	serial := run(false)
+	parallel := run(true)
+	if serial != parallel {
+		t.Fatalf("the error differs: serial %q, parallel %q", serial, parallel)
+	}
+}
+
+// benchJacResidual builds a sixteen-parameter least squares problem
+// whose residual costs roughly twenty microseconds of real arithmetic
+// per call, and whose data sit far from every model the fit can reach,
+// so a run spends a fixed number of rounds, almost all of it inside
+// the Jacobian sweep, and never converges.
+func benchJacResidual(b *testing.B) (residual func(*core.Array) (*core.Array, error), start *core.Array) {
+	b.Helper()
+	const nP = 16
+	const nObs = 64
+	const spin = 512
+	t := make([]float64, nObs)
+	y := make([]float64, nObs)
+	for i := range nObs {
+		t[i] = float64(i) / 8
+		y[i] = 1e3 + float64(i)
+	}
+	residual = func(p *core.Array) (*core.Array, error) {
+		s := 0.0
+		for k := range spin {
+			s += math.Sin(float64(k)*7e-4+p.FloatAt(k&(nP-1))) * math.Cos(float64(k)*3e-4)
+		}
+		out := core.New(core.Float, nObs)
+		vals := out.RawFloats()
+		for i := range nObs {
+			v := s
+			for k := range nP / 2 {
+				v += p.FloatAt(2*k) * math.Exp(-p.FloatAt(2*k+1)*t[i])
+			}
+			vals[i] = v - y[i]
+		}
+		return out, nil
+	}
+	p0 := make([]float64, nP)
+	for i := range p0 {
+		p0[i] = 0.5
+	}
+	arr, err := core.FromFloats(p0, nP)
+	if err != nil {
+		b.Fatal(err)
+	}
+	return residual, arr
+}
+
+// BenchmarkLevenbergMarquardtJacobianSweep times the central-difference
+// Jacobian sweep serial against parallel in one binary, the same fit
+// walked both ways. The worker count mirrors the GOMAXPROCS=8 the run
+// is read under: eight workers, one per thread.
+func BenchmarkLevenbergMarquardtJacobianSweep(b *testing.B) {
+	restore := engine.SetNumWorkers(8)
+	defer engine.SetNumWorkers(restore)
+	residual, start := benchJacResidual(b)
+	for _, parallel := range []bool{false, true} {
+		name := "serial"
+		if parallel {
+			name = "parallel"
+		}
+		b.Run(name, func(b *testing.B) {
+			opts := LMOptions{MaxIterations: 12, AllowBudgetExit: true, ParallelJacobian: parallel}
+			b.ReportAllocs()
+			for b.Loop() {
+				if _, _, err := LevenbergMarquardt(residual, start, opts); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}