values_test.go

package ndarray

import (
	_ "embed"
	"encoding/json"
	"fmt"
	"math"
	"strconv"
	"strings"
	"testing"
)

// values.json holds NumPy's results for every operation below on every dtype,
// generated by testdata/values.py. Numbers are strings so that NaN, the
// infinities and 64-bit integers survive JSON; complex numbers are [re, im].
//
//go:embed testdata/values.json
var valuesJSON []byte

type npResult struct {
	DType string            `json:"dtype"`
	Shape []int             `json:"shape"`
	Data  []json.RawMessage `json:"data"`
	Error string            `json:"error"`
}

type npValues struct {
	NumPy, Machine string
	Cases          map[string]map[string]json.RawMessage
}

// parseStore builds a storage slice of dtype dt from NumPy's encoded values.
func parseStore(t *testing.T, dt DType, raw []json.RawMessage) any {
	t.Helper()
	cs := make([]complex128, len(raw))
	for i, r := range raw {
		var pair []string
		if json.Unmarshal(r, &pair) == nil {
			cs[i] = complex(parseNum(t, pair[0]), parseNum(t, pair[1]))
			continue
		}
		var s string
		if err := json.Unmarshal(r, &s); err != nil {
			t.Fatal(err)
		}
		switch {
		case s == "true" || s == "false":
			cs[i] = complex(map[bool]float64{true: 1}[s == "true"], 0)
		case dt.IsInteger():
			// 64-bit integers do not fit a float64: store them directly.
			st := makeStore(dt, len(raw))
			for j, r := range raw {
				var v string
				_ = json.Unmarshal(r, &v)
				if dt.IsSigned() {
					n, _ := strconv.ParseInt(v, 10, 64)
					convertStore(sliceAt(st, j), []int64{n})
				} else {
					n, _ := strconv.ParseUint(v, 10, 64)
					convertStore(sliceAt(st, j), []uint64{n})
				}
			}
			return st
		default:
			cs[i] = complex(parseNum(t, s), 0)
		}
	}
	st := makeStore(dt, len(raw))
	convertStore(st, cs)
	return st
}

// sliceAt returns element j of a storage slice as a one-element subslice.
func sliceAt(s any, j int) any {
	switch s := s.(type) {
	case []int8:
		return s[j : j+1]
	case []int16:
		return s[j : j+1]
	case []int32:
		return s[j : j+1]
	case []int64:
		return s[j : j+1]
	case []uint8:
		return s[j : j+1]
	case []uint16:
		return s[j : j+1]
	case []uint32:
		return s[j : j+1]
	case []uint64:
		return s[j : j+1]
	}
	panic("sliceAt: not an integer slice")
}

func parseNum(t *testing.T, s string) float64 {
	t.Helper()
	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		t.Fatalf("parse %q: %v", s, err)
	}
	return f
}

// exactOps are correctly rounded in both libraries, so they must agree to the
// bit; the others may differ by a few units in the last place.
var exactOps = map[string]bool{
	"add": true, "sub": true, "maximum": true, "minimum": true, "eq": true, "ne": true,
	"lt": true, "le": true, "gt": true, "ge": true, "abs": true, "neg": true, "floor": true,
	"max0": true, "min1": true, "argmax1": true, "argmin0": true, "nonzero": true, "where": true,
}

// operations maps each fixture key to the operation under test.
func operations(x, y, m *Array) map[string]func() (*Array, error) {
	ok := func(f func() *Array) func() (*Array, error) {
		return func() (*Array, error) { return f(), nil }
	}
	ops := map[string]func() (*Array, error){
		"add": func() (*Array, error) { return x.Add(y) }, "sub": func() (*Array, error) { return x.Sub(y) },
		"mul": func() (*Array, error) { return x.Mul(y) }, "div": func() (*Array, error) { return x.Div(y) },
		"maximum": func() (*Array, error) { return x.Maximum(y) }, "minimum": func() (*Array, error) { return x.Minimum(y) },
		"eq": func() (*Array, error) { return x.Equal(y) }, "ne": func() (*Array, error) { return x.NotEqual(y) },
		"lt": func() (*Array, error) { return x.Less(y) }, "le": func() (*Array, error) { return x.LessEqual(y) },
		"gt": func() (*Array, error) { return x.Greater(y) }, "ge": func() (*Array, error) { return x.GreaterEqual(y) },
		"sqrt": ok(x.Sqrt), "exp": ok(x.Exp), "log": ok(x.Log), "sin": ok(x.Sin), "abs": ok(x.Abs),
		"neg": ok(x.Neg), "floor": ok(x.Floor), "square": ok(x.Square),
		"power":    ok(func() *Array { return x.Power(2.5) }),
		"sum0":     func() (*Array, error) { return m.SumAxis(0, false) },
		"sum1":     func() (*Array, error) { return m.SumAxis(1, false) },
		"prod1":    func() (*Array, error) { return m.ProdAxis(1, false) },
		"max0":     func() (*Array, error) { return m.MaxAxis(0, false) },
		"min1":     func() (*Array, error) { return m.MinAxis(1, false) },
		"argmax1":  func() (*Array, error) { return m.ArgMaxAxis(1, false) },
		"argmin0":  func() (*Array, error) { return m.ArgMinAxis(0, false) },
		"cumsum1":  func() (*Array, error) { return m.CumSum(1) },
		"cumprod0": func() (*Array, error) { return m.CumProd(0) },
		"mean1":    func() (*Array, error) { return m.MeanAxis(1, false) },
		"matmul":   func() (*Array, error) { return m.MatMul(m.Transpose()) },
		"nonzero":  ok(x.Nonzero),
		"where": func() (*Array, error) {
			if x.dtype == Bool {
				return Where(x, x, y)
			}
			c, _ := x.Greater(y)
			return Where(c, x, y)
		},
	}
	for _, to := range allDTypes() {
		ops["astype_"+to.String()] = ok(func() *Array { return x.AsType(to) })
	}
	return ops
}

// run calls f, turning a panic into an error.
func run(f func() (*Array, error)) (r *Array, err error) {
	defer func() {
		if p := recover(); p != nil {
			err = fmt.Errorf("panic: %v", p)
		}
	}()
	return f()
}

// TestValuesMatchNumPy runs every operation on every dtype and compares the
// dtype, shape and values with NumPy's.
func TestValuesMatchNumPy(t *testing.T) {
	var v npValues
	if err := json.Unmarshal(valuesJSON, &v); err != nil {
		t.Fatal(err)
	}
	checked := 0
	for _, dt := range allDTypes() {
		c := v.Cases[dt.String()]
		var xr, yr []json.RawMessage
		_ = json.Unmarshal(c["x"], &xr)
		_ = json.Unmarshal(c["y"], &yr)
		x := fromStore(parseStore(t, dt, xr), []int{len(xr)})
		y := fromStore(parseStore(t, dt, yr), []int{len(yr)})
		m, _ := x.Reshape(2, 3)
		for name, op := range operations(x, y, m) {
			var want npResult
			if err := json.Unmarshal(c[name], &want); err != nil {
				t.Fatalf("%v %s: %v", dt, name, err)
			}
			got, err := run(op)
			if want.Error != "" {
				if err == nil {
					t.Errorf("%v %s: got %v, NumPy raises %s", dt, name, got, want.Error)
				}
				continue
			}
			if err != nil {
				t.Errorf("%v %s: %v", dt, name, err)
				continue
			}
			wdt, _ := ParseDType(noFloat16(want.DType))
			if got.DType() != wdt || !sameShape(got.Shape(), want.Shape) {
				t.Errorf("%v %s: %v %v, NumPy %s %v", dt, name, got.DType(), got.Shape(), want.DType, want.Shape)
				continue
			}
			ws := parseStore(t, wdt, want.Data)
			gs := got.contiguousStore()
			for i := range storeLen(ws) {
				if strings.HasPrefix(name, "astype_") && outOfRange(x, i, wdt) {
					continue // NumPy leaves it undefined; see AsType
				}
				g, w := elemComplex(gs, i), elemComplex(ws, i)
				if !closeTo(g, w, exactOps[name] || wdt.IsInteger() || wdt == Bool, want.DType) {
					t.Errorf("%v %s[%d] = %v, NumPy %v", dt, name, i, elemString(gs, i), elemString(ws, i))
				}
				checked++
			}
		}
	}
	t.Logf("%d values checked against NumPy %s (%s)", checked, v.NumPy, v.Machine)
}

// outOfRange reports whether element i of x, converted to the integer dtype
// to, is a NaN, an infinity or beyond to's range: a cast NumPy does not
// define.
func outOfRange(x *Array, i int, to DType) bool {
	if !to.IsInteger() || !(x.dtype.IsFloat() || x.dtype.IsComplex()) {
		return false
	}
	f := real(elemComplex(x.store(), i))
	if math.IsNaN(f) || math.IsInf(f, 0) {
		return true
	}
	lo, hi := 0.0, math.Ldexp(1, 8*to.Size())
	if to.IsSigned() {
		lo, hi = -hi/2, hi/2
	}
	return f < lo || f >= hi
}

// closeTo compares two values: to the bit when exact is set (NaN equal to
// NaN, and the sign of a zero counted); otherwise within 8 machine epsilons of
// the result's magnitude, eps being that of NumPy's result dtype (float16 where
// NumPy computes in half precision and this package in float32, which is the
// documented difference). The magnitude is the complex modulus, so a component
// near zero is judged against the size of the number, not its own.
func closeTo(g, w complex128, exact bool, npDType string) bool {
	if cmplxIsNaN(g) || cmplxIsNaN(w) {
		return cmplxIsNaN(g) == cmplxIsNaN(w) &&
			math.IsNaN(real(g)) == math.IsNaN(real(w)) && math.IsNaN(imag(g)) == math.IsNaN(imag(w))
	}
	if exact {
		return g == w && math.Signbit(real(g)) == math.Signbit(real(w)) &&
			math.Signbit(imag(g)) == math.Signbit(imag(w))
	}
	if g == w {
		return true
	}
	eps := 0x1p-52
	switch npDType {
	case "float16":
		eps = 0x1p-10
	case "float32", "complex64":
		eps = 0x1p-23
	}
	mag := math.Max(cabs(g), cabs(w))
	return cabs(g-w) <= 8*eps*mag
}

func cmplxIsNaN(z complex128) bool { return math.IsNaN(real(z)) || math.IsNaN(imag(z)) }

func cabs(z complex128) float64 { return math.Hypot(real(z), imag(z)) }