Hoist weighted-distance weight vector out of the MMR loop

Building the per-dimension weight vector inside compute_weighted_distance
rebuilt it on every call in the O(n²) MMR loop. Split out build_dimension_weights
+ compute_weighted_distance_vec so apply_mmr builds the vector once and reuses it,
roughly halving per-call cost. compute_weighted_distance keeps its signature as a
thin wrapper for the remaining callers.

Author: marcelveldt

music_assistant/providers/sonic_similarity/similarity.py

@@ -11,7 +11,7 @@
 
 import numpy as np
 
-from .vectors import compute_weighted_distance
+from .vectors import build_dimension_weights, compute_weighted_distance_vec
 
 
 class Candidate(NamedTuple):
@@ -111,9 +111,11 @@ def apply_mmr(
     seed_arr = np.array(seed_vec, dtype=np.float64)
 
     if weights is not None:
+        # Build the per-dimension weight vector once; _similarity runs O(n²) times.
+        dim_weights = build_dimension_weights(weights)
 
         def _similarity(a: np.ndarray, b: np.ndarray) -> float:
-            d = compute_weighted_distance(a, b, weights)
+            d = compute_weighted_distance_vec(a, b, dim_weights)
             return 1.0 / (1.0 + d)
 
         relevance: dict[str, float] = {

music_assistant/providers/sonic_similarity/vectors.py

@@ -174,26 +174,60 @@ def compute_group_distances(
     return result
 
 
+def build_dimension_weights(weights: dict[str, float]) -> np.ndarray:
+    """
+    Expand per-group weights into a per-dimension weight vector.
+
+    Callers in a hot loop (e.g. MMR re-ranking) can build this once and reuse it
+    across many compute_weighted_distance_vec calls instead of rebuilding it per call.
+
+    :param weights: Per-group weight overrides keyed by FEATURE_GROUPS name. Groups
+        absent from the dict default to weight 1.0.
+    :returns: Float64 weight array of length VECTOR_DIMENSIONS.
+    """
+    dim_weights = np.ones(VECTOR_DIMENSIONS, dtype=np.float64)
+    for group, (start, end) in FEATURE_GROUPS.items():
+        if group in weights:
+            dim_weights[start:end] = weights[group]
+    return dim_weights
+
+
 def compute_weighted_distance(
     sig_a: list[float] | np.ndarray,
     sig_b: list[float] | np.ndarray,
     weights: dict[str, float],
 ) -> float:
-    """Compute per-group weighted Euclidean distance between two feature vectors.
+    """
+    Compute per-group weighted Euclidean distance between two feature vectors.
 
     :param sig_a: First feature vector (list or numpy array).
     :param sig_b: Second feature vector (list or numpy array).
     :param weights: Per-group weight overrides keyed by FEATURE_GROUPS name.
     :returns: Weighted normalized distance as a float.
     """
+    return compute_weighted_distance_vec(sig_a, sig_b, build_dimension_weights(weights))
+
+
+def compute_weighted_distance_vec(
+    sig_a: list[float] | np.ndarray,
+    sig_b: list[float] | np.ndarray,
+    dim_weights: np.ndarray,
+) -> float:
+    """
+    Compute weighted Euclidean distance from a precomputed per-dimension weight vector.
+
+    :param sig_a: First feature vector (list or numpy array).
+    :param sig_b: Second feature vector (list or numpy array).
+    :param dim_weights: Per-dimension weights as built by build_dimension_weights.
+    :returns: Weighted normalized distance as a float.
+    """
+    total_weighted_dims = float(dim_weights.sum())
+    if total_weighted_dims == 0.0:
+        return 0.0
     # np.asarray is a no-op when the caller already holds a float64 array (the
     # MMR hot path), avoiding the list round-trip the previous version forced.
     a = np.asarray(sig_a, dtype=np.float64)
     b = np.asarray(sig_b, dtype=np.float64)
-    dim_weights = _expand_group_weights(weights)
-    total_weighted_dims = float(dim_weights.sum())
-    if total_weighted_dims == 0.0:
-        return 0.0
     diff = a - b
     weighted_sq_sum = float(np.dot(dim_weights, diff * diff))
     return math.sqrt(weighted_sq_sum / total_weighted_dims)
@@ -222,18 +256,3 @@ def build_debug_breakdown(
             for k, v in compute_group_distances(seed_normalized, cand_normalized).items()
         },
     }
-
-
-def _expand_group_weights(weights: dict[str, float]) -> np.ndarray:
-    """
-    Expand per-group weights into a per-dimension weight vector.
-
-    :param weights: Per-group weight overrides keyed by FEATURE_GROUPS name. Groups
-        absent from the dict default to weight 1.0.
-    :returns: Float64 weight array of length VECTOR_DIMENSIONS.
-    """
-    dim_weights = np.ones(VECTOR_DIMENSIONS, dtype=np.float64)
-    for group, (start, end) in FEATURE_GROUPS.items():
-        if group in weights:
-            dim_weights[start:end] = weights[group]
-    return dim_weights

tests/providers/sonic_similarity/test_vector_assembly.py

@@ -12,8 +12,10 @@
     FEATURE_GROUPS,
     VECTOR_DIMENSIONS,
     assemble_vector,
+    build_dimension_weights,
     compute_corpus_stats,
     compute_weighted_distance,
+    compute_weighted_distance_vec,
     encode_key_mode,
     normalize_features,
 )
@@ -454,3 +456,24 @@ def test_all_zero_weights_returns_zero(self) -> None:
         b = self._make_vector(1.0)
         all_zero = dict.fromkeys(FEATURE_GROUPS, 0.0)
         assert compute_weighted_distance(a, b, all_zero) == 0.0
+
+    def test_precomputed_weights_match_dict_path(self) -> None:
+        """compute_weighted_distance_vec with a prebuilt vector equals the dict-based path."""
+        rng = np.random.default_rng(7)
+        a = rng.standard_normal(VECTOR_DIMENSIONS)
+        b = rng.standard_normal(VECTOR_DIMENSIONS)
+        weights = {"rhythm": 2.0, "dynamics": 0.0}
+        dim_weights = build_dimension_weights(weights)
+        assert compute_weighted_distance_vec(a, b, dim_weights) == pytest.approx(
+            compute_weighted_distance(a, b, weights)
+        )
+
+    def test_build_dimension_weights_expands_groups(self) -> None:
+        """Each group's weight lands on its dimensions; absent groups default to 1.0."""
+        dim_weights = build_dimension_weights({"timbre": 3.0})
+        assert len(dim_weights) == VECTOR_DIMENSIONS
+        start, end = FEATURE_GROUPS["timbre"]
+        assert list(dim_weights[start:end]) == [3.0] * (end - start)
+        # rhythm was not overridden, so it keeps the 1.0 default
+        r_start, r_end = FEATURE_GROUPS["rhythm"]
+        assert list(dim_weights[r_start:r_end]) == [1.0] * (r_end - r_start)