tensor/internal/core/jacobian_test.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

import (
	"errors"
	"math"
	"strings"
	"testing"
)

// TestJacobianAnalytic pins the Jacobian of f(x, y) = (x·y, x+y),
// whose exact derivative [[y, x], [1, 1]] the central differences
// must reproduce to 1e-8, at points that exercise the step scaling.
func TestJacobianAnalytic(t *testing.T) {
	f := func(x *Array) (*Array, error) {
		xv, yv := x.FloatAt(0), x.FloatAt(1)
		return FromFloats([]float64{xv * yv, xv + yv}, 2)
	}
	cases := []struct {
		name string
		x, y float64
		want []float64
	}{
		{"first quadrant", 2, 3, []float64{3, 2, 1, 1}},
		{"negative coordinate", -1.5, 4, []float64{4, -1.5, 1, 1}},
		{"origin", 0, 0, []float64{0, 0, 1, 1}},
	}
	for _, c := range cases {
		j, err := Jacobian(f, mustFloats(t, []float64{c.x, c.y}), JacobianOptions{})
		if err != nil {
			t.Fatalf("%s: %v", c.name, err)
		}
		if shape := j.Shape(); len(shape) != 2 || shape[0] != 2 || shape[1] != 2 {
			t.Fatalf("%s: shape %v, want (2, 2)", c.name, shape)
		}
		for i, want := range c.want {
			if got := j.FloatAt(i); math.Abs(got-want) > 1e-8 {
				t.Errorf("%s: entry %d = %v, want %v", c.name, i, got, want)
			}
		}
	}
}

// TestJacobianNonlinear pins a nonlinear scalar-input case,
// f(x) = (x², x³) at x = 1.7, whose Jacobian is the column
// (2x, 3x²) = (3.4, 8.67).
func TestJacobianNonlinear(t *testing.T) {
	f := func(x *Array) (*Array, error) {
		xv := x.FloatAt(0)
		return FromFloats([]float64{xv * xv, xv * xv * xv}, 2)
	}
	j, err := Jacobian(f, mustFloats(t, []float64{1.7}), JacobianOptions{})
	if err != nil {
		t.Fatalf("Jacobian: %v", err)
	}
	// The x³ column sits at the default step's rounding floor
	// (~2·10⁻⁸ of cancellation noise), so the pin is 2e-8 there.
	tols := []float64{1e-8, 2e-8}
	for i, want := range []float64{3.4, 8.67} {
		if got := j.FloatAt(i); math.Abs(got-want) > tols[i] {
			t.Errorf("entry %d = %v, want %v", i, got, want)
		}
	}
}

// TestJacobianCustomStep checks that a caller-supplied step is
// honoured: on a linear field any step is exact, so the result pins
// the matrix rather than the step heuristic.
func TestJacobianCustomStep(t *testing.T) {
	f := func(x *Array) (*Array, error) {
		return FromFloats([]float64{2*x.FloatAt(0) - x.FloatAt(1)}, 1)
	}
	j, err := Jacobian(f, mustFloats(t, []float64{5, -7}), JacobianOptions{Step: 1e-6})
	if err != nil {
		t.Fatalf("Jacobian: %v", err)
	}
	if shape := j.Shape(); len(shape) != 2 || shape[0] != 1 || shape[1] != 2 {
		t.Fatalf("shape %v, want (1, 2)", shape)
	}
	for i, want := range []float64{2, -1} {
		// The wide 1e-6 step leaves ~2e-9 of cancellation noise on
		// values of size ~17; the pin still says which matrix it is.
		if got := j.FloatAt(i); math.Abs(got-want) > 1e-8 {
			t.Errorf("entry %d = %v, want %v", i, got, want)
		}
	}
}

// TestJacobianRejects pins the contracts: complex points, empty
// points, non-vector outputs, outputs whose length changes between
// columns, and errors from f all come back as errors.
func TestJacobianRejects(t *testing.T) {
	complexPoint, _ := FromComplexes([]complex128{1, 2}, 2)
	if _, err := Jacobian(vectorIdentity, complexPoint, JacobianOptions{}); err == nil {
		t.Error("expected an error for a complex point")
	} else if !strings.Contains(err.Error(), "complex") {
		t.Errorf("complex point: error %q lacks \"complex\"", err)
	}
	if _, err := Jacobian(vectorIdentity, mustFloats(t, nil), JacobianOptions{}); err == nil {
		t.Error("expected an error for an empty point")
	}

	matrixOut := func(x *Array) (*Array, error) {
		return FromFloats([]float64{1, 0, 0, 1}, 2, 2)
	}
	if _, err := Jacobian(matrixOut, mustFloats(t, []float64{1}), JacobianOptions{}); err == nil {
		t.Error("expected an error for a matrix-shaped output")
	} else if !strings.Contains(err.Error(), "vector") {
		t.Errorf("matrix output: error %q lacks \"vector\"", err)
	}

	emptyOut := func(x *Array) (*Array, error) {
		return FromFloats([]float64{}, 0)
	}
	if _, err := Jacobian(emptyOut, mustFloats(t, []float64{1}), JacobianOptions{}); err == nil {
		t.Error("expected an error for an empty output")
	}

	calls := 0
	changingLength := func(x *Array) (*Array, error) {
		calls++
		if calls == 1 {
			return FromFloats([]float64{1, 2}, 2)
		}
		return FromFloats([]float64{1, 2, 3}, 3)
	}
	if _, err := Jacobian(changingLength, mustFloats(t, []float64{1, 2}), JacobianOptions{}); err == nil {
		t.Error("expected an error for an output length that changes between columns")
	}

	sentinel := errf("f blew up")
	failing := func(x *Array) (*Array, error) {
		if x.FloatAt(0) != 1 {
			return nil, sentinel
		}
		return FromFloats([]float64{1}, 1)
	}
	_, err := Jacobian(failing, mustFloats(t, []float64{1}), JacobianOptions{})
	if err == nil {
		t.Error("expected the probe failure to propagate")
	} else if !strings.Contains(err.Error(), "blew up") {
		t.Errorf("probe failure: error %q lacks the cause", err)
	} else if !errors.Is(err, sentinel) {
		// The wrap keeps the chain open, so the sentinel stays reachable
		// through the entry point's context: a revert to a plain %v
		// formatting fails exactly here.
		t.Errorf("probe failure: error %q does not unwrap to the cause", err)
	}
}

// vectorIdentity is a well-behaved f for the input-side rejections.
func vectorIdentity(x *Array) (*Array, error) {
	return FromFloats([]float64{x.FloatAt(0), x.FloatAt(1)}, 2)
}