vibesonic/internal/similarity/index.go

// Package similarity builds a small in-process nearest-neighbor index
// over numeric audio-feature vectors and uses it for two things:
// ranking tracks by actual sonic closeness, and walking a musically
// reasonable path between two tracks through feature space. There is
// no external vector database — the whole index is just normalized
// vectors held in memory, rebuilt from internal/database's
// audio_features table whenever a caller needs it.
package similarity

import (
	"math"
	"sort"
)

// Item is one track's raw (un-normalized) feature vector, keyed by
// song ID.
type Item struct {
	ID     string
	Vector []float64
}

// Index holds z-score-normalized vectors for a set of tracks. Each
// dimension is normalized independently using the mean/stddev of that
// dimension across exactly the items the index was built from, so
// dimensions with a wide natural range (like tempo, in BPM) don't
// dominate Euclidean distance over dimensions with a narrow one (like a
// normalized chroma bin).
type Index struct {
	ids  []string
	norm [][]float64
}

// Build normalizes and indexes items. An item whose Vector length
// disagrees with the first item's is dropped rather than causing a
// panic or silently misaligning dimensions — a defensive check, not
// something expected to trigger with correctly-produced feature rows.
func Build(items []Item) *Index {
	idx := &Index{}
	if len(items) == 0 {
		return idx
	}
	dims := len(items[0].Vector)

	var kept []Item
	for _, it := range items {
		if len(it.Vector) == dims {
			kept = append(kept, it)
		}
	}
	if len(kept) == 0 {
		return idx
	}

	means := make([]float64, dims)
	for _, it := range kept {
		for d, v := range it.Vector {
			means[d] += v
		}
	}
	for d := range means {
		means[d] /= float64(len(kept))
	}

	stds := make([]float64, dims)
	for _, it := range kept {
		for d, v := range it.Vector {
			diff := v - means[d]
			stds[d] += diff * diff
		}
	}
	for d := range stds {
		stds[d] = math.Sqrt(stds[d] / float64(len(kept)))
	}

	idx.ids = make([]string, len(kept))
	idx.norm = make([][]float64, len(kept))
	for i, it := range kept {
		nv := make([]float64, dims)
		for d, v := range it.Vector {
			if stds[d] > 1e-9 {
				nv[d] = (v - means[d]) / stds[d]
			}
		}
		idx.ids[i] = it.ID
		idx.norm[i] = nv
	}
	return idx
}

// Match is one similarity result: another track and its distance from
// the query track in normalized feature space (0 = identical, larger =
// less similar).
type Match struct {
	ID       string
	Distance float64
}

// Nearest returns up to k tracks closest to id, nearest first,
// excluding id itself. Returns nil if id isn't in the index.
func (idx *Index) Nearest(id string, k int) []Match {
	pos := idx.indexOf(id)
	if pos < 0 {
		return nil
	}
	matches := make([]Match, 0, len(idx.ids)-1)
	for i, other := range idx.ids {
		if i == pos {
			continue
		}
		matches = append(matches, Match{ID: other, Distance: euclidean(idx.norm[pos], idx.norm[i])})
	}
	sort.Slice(matches, func(a, b int) bool { return matches[a].Distance < matches[b].Distance })
	if k > 0 && len(matches) > k {
		matches = matches[:k]
	}
	return matches
}

// Path walks from fromID to toID in `steps` hops (including both
// endpoints), picking at each intermediate step the not-yet-used track
// closest to the corresponding point on the straight line between
// fromID's and toID's feature vectors. That keeps the sequence
// perceptually gradual — each hop is a small, real step in sonic
// character rather than a random jump — without ever inventing a track
// that isn't actually in the library. Returns nil if either endpoint
// isn't indexed.
func (idx *Index) Path(fromID, toID string, steps int) []string {
	fromPos := idx.indexOf(fromID)
	toPos := idx.indexOf(toID)
	if fromPos < 0 || toPos < 0 {
		return nil
	}
	if steps < 2 {
		steps = 2
	}

	path := []string{fromID}
	used := map[string]bool{fromID: true, toID: true}
	fromVec, toVec := idx.norm[fromPos], idx.norm[toPos]

	for step := 1; step < steps-1; step++ {
		t := float64(step) / float64(steps-1)
		target := lerp(fromVec, toVec, t)

		best := -1
		bestDist := math.Inf(1)
		for i, id := range idx.ids {
			if used[id] {
				continue
			}
			d := euclidean(idx.norm[i], target)
			if d < bestDist {
				bestDist = d
				best = i
			}
		}
		if best < 0 {
			break // library exhausted; fall through with what we have
		}
		path = append(path, idx.ids[best])
		used[idx.ids[best]] = true
	}

	path = append(path, toID)
	return path
}

// DistanceBetween returns the normalized-feature-space Euclidean
// distance between two indexed tracks, or +Inf if either isn't
// indexed.
func (idx *Index) DistanceBetween(idA, idB string) float64 {
	a, b := idx.indexOf(idA), idx.indexOf(idB)
	if a < 0 || b < 0 {
		return math.Inf(1)
	}
	return euclidean(idx.norm[a], idx.norm[b])
}

func (idx *Index) indexOf(id string) int {
	for i, v := range idx.ids {
		if v == id {
			return i
		}
	}
	return -1
}

func euclidean(a, b []float64) float64 {
	var sum float64
	for i := range a {
		d := a[i] - b[i]
		sum += d * d
	}
	return math.Sqrt(sum)
}

func lerp(a, b []float64, t float64) []float64 {
	out := make([]float64, len(a))
	for i := range a {
		out[i] = a[i] + (b[i]-a[i])*t
	}
	return out
}