cores/embedding/vectors.py

"""Vector operations — normalize, cosine, euclidean, batch."""

from __future__ import annotations

import numpy as np


def normalize(v: np.ndarray) -> np.ndarray:
    """L2-normalize a vector."""
    n = np.linalg.norm(v)
    return v / n if n > 0 else v


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity between two vectors."""
    na = np.linalg.norm(a)
    nb = np.linalg.norm(b)
    if na == 0 or nb == 0:
        return 0.0
    return float(np.dot(a, b) / (na * nb))


def euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    """Euclidean distance between two vectors."""
    return float(np.linalg.norm(a - b))


def batch_cosine_similarity(query: np.ndarray, matrix: np.ndarray) -> np.ndarray:
    """Cosine similarity between a query vector and a matrix of vectors.

    Args:
        query: 1-D array of shape (d,)
        matrix: 2-D array of shape (n, d)

    Returns:
        1-D array of shape (n,) with similarity scores.
    """
    query_norm = np.linalg.norm(query)
    if query_norm == 0:
        return np.zeros(matrix.shape[0])
    matrix_norms = np.linalg.norm(matrix, axis=1)
    # Avoid division by zero
    safe_norms = np.where(matrix_norms == 0, 1.0, matrix_norms)
    return (matrix @ query) / (safe_norms * query_norm)