vectors.go

// Copyright 2022, NLP Odyssey Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gomaddness

import (
	"math"
	"sort"
)

// Vectors is a collection of Vector objects.
//
// All Vector objects MUST be of the same size.
//
// Given a collection of N vectors, each of size M, this can also be seen as a
// matrix with dimensions N×M (N rows, M columns).
type Vectors[F Float] []Vector[F]

// ColumnWiseVariance returns a new Vector containing the column-wise variance
// computed across all vectors.
func (vs Vectors[F]) ColumnWiseVariance() Vector[F] {
	vecSize := len(vs[0])
	sum := make(Vector[F], vecSize)
	sumOfSquares := make(Vector[F], vecSize)

	for _, v := range vs {
		sum.Add(v)
		sumOfSquares.AddSquares(v)
	}

	n := F(len(vs))
	mean := sum.DivScalar(n)
	squareMean := mean.Square()
	return sumOfSquares.DivScalar(n).Sub(squareMean)
}

// SplitByThreshold splits the collection of vectors by a threshold value,
// compared with the elements at the given index.
//
// Two new collections are returned:
//   - lt: vectors whose value at "index" is lower than "threshold"
//   - gte: vectors whose value at "index" is greater than or equal to "threshold"
func (vs Vectors[F]) SplitByThreshold(index int, threshold F) (lt, gte Vectors[F]) {
	lt = make(Vectors[F], 0)
	gte = make(Vectors[F], 0)
	for _, v := range vs {
		if v[index] < threshold {
			lt = append(lt, v)
			continue
		}
		gte = append(gte, v)
	}
	return
}

// Copy returns a shallow copy of the collection.
func (vs Vectors[F]) Copy() Vectors[F] {
	c := make(Vectors[F], len(vs))
	copy(c, vs)
	return c
}

// SortByColumn sorts vs in place, according to the value at each vector's
// index (column), in ascending order.
// It returns vs.
func (vs Vectors[F]) SortByColumn(index int) Vectors[F] {
	sort.SliceStable(vs, func(i, j int) bool {
		return vs[i][index] < vs[j][index]
	})
	return vs
}

// Reverse reverses vs in place.
// It returns vs.
func (vs Vectors[F]) Reverse() Vectors[F] {
	for l, r := 0, len(vs)-1; l < r; l, r = l+1, r-1 {
		vs[l], vs[r] = vs[r], vs[l]
	}
	return vs
}

// CumulativeSSE returns a new Vector, computing the error sum of squares
// (SSE) over the set of Vectors.
func (vs Vectors[F]) CumulativeSSE() Vector[F] {
	n := len(vs)
	d := len(vs[0])

	out := make(Vector[F], n) // all values are zero by default
	sum := make(Vector[F], d)
	sumOfSquares := make(Vector[F], d)

	_ = sum[len(vs[0])-1]
	for j, x := range vs[0] {
		sum[j] = x
		sumOfSquares[j] = x * x
	}

	for i := 1; i < n; i++ {
		_ = sum[len(vs[i])-1]
		for j, x := range vs[i] {
			sum[j] += x
			sumOfSquares[j] += x * x
			out[i] += sumOfSquares[j] - (sum[j] * sum[j] / F(i+1))
		}
	}

	return out
}

// OptimalSplitThreshold finds the optimal split threshold within the vectors,
// computed over values at the given split-index column.
func (vs Vectors[F]) OptimalSplitThreshold(splitIndex int) (threshold, loss F) {
	sorted := vs.Copy().SortByColumn(splitIndex)
	if len(sorted) == 1 {
		// TODO: check this corner case
		sorted = append(sorted, sorted[0])
	}

	losses := sorted.CumulativeSSE() // head
	tail := sorted.Copy().Reverse().CumulativeSSE().Reverse()

	losses[:len(losses)-1].Add(tail[1:])

	n := losses.ArgMin()
	if n < len(sorted)-1 {
		threshold = (sorted[n][splitIndex] + sorted[n+1][splitIndex]) / 2
	} else {
		// TODO: check this corner case
		threshold = F(math.Nextafter32(float32(sorted[n][splitIndex]), float32(math.Inf(+1))))
	}
	return threshold, losses[n]
}

// Mean calculates and returns the mean vector.
func (vs Vectors[F]) Mean() Vector[F] {
	sum := vs[0].Copy()
	for _, v := range vs[1:] {
		sum.Add(v)
	}
	return sum.DivScalar(F(len(vs)))
}