feat: Improve docstrings and variable naming (Pringled#3)

* Renamed alpha to lambda_param * Updated docstrings * Updated typing * Renamed relevances to scores

Author: Pringled

src/pyversity/core.py

@@ -8,29 +8,29 @@
 
 def diversify(
  strategy: Strategy,
- relevances: np.ndarray,
  embeddings: np.ndarray,
+ scores: np.ndarray,
  k: int,
  **kwargs: Any,
 ) -> tuple[np.ndarray, np.ndarray]:
  """
  Diversify a retrieval result using a selected strategy.
 
  :param strategy: The diversification strategy to apply. Supported strategies are: MMR, MSD, COVER, and DPP.
- :param relevances: Array of relevance scores for the items.
  :param embeddings: Array of embeddings for the items.
+ :param scores: Array of relevance scores for the items.
  :param k: The number of items to select in the diversified result.
  :param **kwargs: Additional keyword arguments passed to the specific strategy function.
  :return: A tuple containing an array of indices of the selected items
  and an array of corresponding relevance scores for the selected items.
  :raises ValueError: If the provided strategy is not recognized.
  """
  if strategy == Strategy.MMR:
- return mmr(relevances, embeddings, k, **kwargs)
+ return mmr(scores, embeddings, k, **kwargs)
  if strategy == Strategy.MSD:
- return msd(relevances, embeddings, k, **kwargs)
+ return msd(scores, embeddings, k, **kwargs)
  if strategy == Strategy.COVER:
- return cover(relevances, embeddings, k, **kwargs)
+ return cover(scores, embeddings, k, **kwargs)
  if strategy == Strategy.DPP:
- return dpp(relevances, embeddings, k, **kwargs)
+ return dpp(scores, embeddings, k, **kwargs)
  raise ValueError(f"Unknown strategy: {strategy}")

src/pyversity/datatypes.py

@@ -2,12 +2,16 @@
 
 
 class Strategy(str, Enum):
+ """Supported diversification strategies."""
+
  MMR = "mmr"
  MSD = "msd"
  COVER = "cover"
  DPP = "dpp"
 
 
 class Metric(str, Enum):
+ """Supported similarity metrics."""
+
  COSINE = "cosine"
  DOT = "dot"

src/pyversity/strategies/cover.py

@@ -5,8 +5,8 @@
 
 
 def cover(
- relevances: np.ndarray,
  embeddings: np.ndarray,
+ scores: np.ndarray,
  k: int,
  theta: float = 0.5,
  gamma: float = 0.5,
@@ -19,8 +19,8 @@ def cover(
  This strategy chooses `k` items by combining pure relevance with
  diversity-driven coverage using a concave submodular formulation.
 
- :param relevances: 1D array of relevance scores for each item.
  :param embeddings: 2D array of shape (n_samples, n_features).
+ :param scores: 1D array of relevance scores for each item.
  :param k: Number of items to select.
  :param theta: Trade-off between relevance and coverage in [0, 1].
  1.0 = pure relevance, 0.0 = pure coverage.
@@ -31,47 +31,55 @@ def cover(
  :raises ValueError: If theta is not in [0, 1].
  :raises ValueError: If gamma is not in (0, 1].
  """
+ # Validate parameters
  if not (0.0 <= float(theta) <= 1.0):
  raise ValueError("theta must be in [0, 1]")
  if not (0.0 < float(gamma) <= 1.0):
  raise ValueError("gamma must be in (0, 1]")
 
- relevance_scores, feature_matrix, top_k, early_exit = prepare_inputs(relevances, embeddings, k)
+ # Prepare inputs
+ relevance_scores, feature_matrix, top_k, early_exit = prepare_inputs(scores, embeddings, k)
  if early_exit:
+ # Nothing to select: return empty arrays
  return np.empty(0, np.int32), np.empty(0, np.float32)
 
  if metric == Metric.COSINE and normalize:
+ # Normalize feature vectors to unit length for cosine similarity
  feature_matrix = normalize_rows(feature_matrix)
 
- # Pure relevance: short-circuit
  if float(theta) == 1.0:
+ # Pure relevance: select top-k by relevance scores
  topk = np.argsort(-relevance_scores)[:top_k].astype(np.int32)
  gains = relevance_scores[topk].astype(np.float32, copy=False)
  return topk, gains
 
- # Nonnegative similarities for coverage to avoid concave-power NaNs
+ # Compute non-negative similarities for coverage to avoid concave-power NaNs
  similarity_matrix = pairwise_similarity(feature_matrix, metric)
- transposed_similarity = similarity_matrix.T
+ transposed_similarity_matrix = similarity_matrix.T
 
- n = similarity_matrix.shape[0]
- accumulated_coverage = np.zeros(n, dtype=np.float32)
- selected_mask = np.zeros(n, dtype=bool)
+ # Initialize selection state
+ accumulated_coverage = np.zeros(similarity_matrix.shape[0], dtype=np.float32)
+ selected_mask = np.zeros(similarity_matrix.shape[0], dtype=bool)
  selected_indices = np.empty(top_k, dtype=np.int32)
  marginal_gains = np.empty(top_k, dtype=np.float32)
 
- for t in range(top_k):
+ for step in range(top_k):
+ # Compute coverage gains using concave transformation
  concave_before = np.power(accumulated_coverage, gamma)
- concave_after = np.power(transposed_similarity + accumulated_coverage[None, :], gamma)
+ concave_after = np.power(transposed_similarity_matrix + accumulated_coverage[None, :], gamma)
  coverage_gains = (concave_after - concave_before[None, :]).sum(axis=1)
 
+ # Combine relevance and coverage gains
  candidate_scores = theta * relevance_scores + (1.0 - theta) * coverage_gains
  candidate_scores[selected_mask] = -np.inf
 
- chosen = int(np.argmax(candidate_scores))
- selected_indices[t] = chosen
- marginal_gains[t] = float(candidate_scores[chosen])
- selected_mask[chosen] = True
+ # Select item with highest combined score
+ best_index = int(np.argmax(candidate_scores))
+ selected_indices[step] = best_index
+ marginal_gains[step] = float(candidate_scores[best_index])
+ selected_mask[best_index] = True
 
- accumulated_coverage += similarity_matrix[:, chosen]
+ # Update accumulated coverage
+ accumulated_coverage += similarity_matrix[:, best_index]
 
  return selected_indices, marginal_gains

src/pyversity/strategies/dpp.py

@@ -12,8 +12,8 @@ def _exp_zscore_weights(relevance: np.ndarray, beta: float) -> np.ndarray:
 
 
 def dpp(
- relevances: np.ndarray,
  embeddings: np.ndarray,
+ scores: np.ndarray,
  k: int,
  beta: float = 1.0,
 ) -> tuple[np.ndarray, np.ndarray]:
@@ -24,58 +24,67 @@ def dpp(
  maximizing the determinant of a kernel matrix that balances item relevance
  and pairwise similarity.
 
- :param relevances: 1D array of relevance scores for each item.
  :param embeddings: 2D array of shape (n_samples, n_features).
+ :param scores: 1D array of relevance scores for each item.
  :param k: Number of items to select.
  :param beta: Controls the influence of relevance scores in the DPP kernel.
  Higher values increase the emphasis on relevance.
  :return: Tuple of selected indices and their marginal gains.
  """
- relevance_scores, feature_matrix, top_k, early_exit = prepare_inputs(relevances, embeddings, k)
+ # Prepare inputs
+ relevance_scores, feature_matrix, top_k, early_exit = prepare_inputs(scores, embeddings, k)
  if early_exit:
+ # Nothing to select: return empty arrays
  return np.empty(0, np.int32), np.empty(0, np.float32)
 
+ # Normalize feature vectors to unit length for cosine similarity
  feature_matrix = normalize_rows(feature_matrix)
 
  num_items = feature_matrix.shape[0]
  weights = _exp_zscore_weights(relevance_scores, beta)
 
- # Diagonal of L plus jitter is the initial residual variance.
+ # Initial residual variance is the weighted self-similarity
  residual_variance = (weights * weights + float(EPS32)).astype(np.float32, copy=False)
 
- # Columns will store orthogonalized update components.
+ # Initialize selection state
  component_matrix = np.zeros((num_items, top_k), dtype=np.float32)
-
  selected_indices = np.empty(top_k, dtype=np.int32)
  marginal_gains = np.empty(top_k, dtype=np.float32)
  selected_mask = np.zeros(num_items, dtype=bool)
 
- t = 0
- for t in range(top_k):
+ step = 0
+ for step in range(top_k):
+ # Select item with highest residual variance
  residual_variance[selected_mask] = -np.inf
  best_index = int(np.argmax(residual_variance))
- best_gain = float(residual_variance[best_index])
+ best_score = float(residual_variance[best_index])
 
- selected_indices[t] = best_index
- marginal_gains[t] = best_gain
+ selected_indices[step] = best_index
+ marginal_gains[step] = best_score
  selected_mask[best_index] = True
 
- if t == top_k - 1 or best_gain <= 0.0:
- t += 1
+ if step == top_k - 1 or best_score <= 0.0:
+ # No more items to select or no positive gain
+ step += 1
  break
 
+ # Update residual variance using the new component
  weighted_similarity_to_best = (weights * (feature_matrix @ feature_matrix[best_index])) * weights[best_index]
 
- if t > 0:
- projected_component: np.ndarray = component_matrix[:, :t] @ component_matrix[best_index, :t]
+ if step > 0:
+ # Project out the component in the span of previously selected items
+ projected_component: np.ndarray = component_matrix[:, :step] @ component_matrix[best_index, :step]
  else:
+ # No previous components, so projection is zero
  projected_component = np.zeros(num_items, dtype=np.float32)
 
- sqrt_best_gain = np.float32(np.sqrt(best_gain))
- update_component = (weighted_similarity_to_best - projected_component) / (sqrt_best_gain + EPS32)
+ # Compute update component
+ sqrt_best_score = np.float32(np.sqrt(best_score))
+ update_component = (weighted_similarity_to_best - projected_component) / (sqrt_best_score + EPS32)
 
- component_matrix[:, t] = update_component
+ # Update component matrix and residual variance
+ component_matrix[:, step] = update_component
  residual_variance -= update_component * update_component
  np.maximum(residual_variance, 0.0, out=residual_variance)
 
- return selected_indices[:t], marginal_gains[:t]
+ return selected_indices[:step], marginal_gains[:step]

src/pyversity/strategies/mmr.py

@@ -5,10 +5,10 @@
 
 
 def mmr(
- relevances: np.ndarray,
  embeddings: np.ndarray,
+ scores: np.ndarray,
  k: int,
- alpha: float = 0.5,
+ lambda_param: float = 0.5,
  metric: Metric = Metric.COSINE,
  normalize: bool = True,
 ) -> tuple[np.ndarray, np.ndarray]:
@@ -19,21 +19,21 @@ def mmr(
  iteratively choosing items that maximize a combination of their relevance
  and their dissimilarity to already selected items.
 
- :param relevances: 1D array of relevance scores for each item.
  :param embeddings: 2D array of shape (n_samples, n_features).
+ :param scores: 1D array of relevance scores for each item.
  :param k: Number of items to select.
- :param alpha: Trade-off parameter in [0, 1].
+ :param lambda_param: Trade-off parameter in [0, 1].
  1.0 = pure relevance, 0.0 = pure diversity.
  :param metric: Similarity metric to use. Default is Metric.COSINE.
  :param normalize: Whether to normalize embeddings before computing similarity.
  :return: Tuple of selected indices and their marginal gains.
  """
  return greedy_select(
  "mmr",
- relevances,
+ scores,
  embeddings,
  k,
  metric=metric,
  normalize=normalize,
- alpha=alpha,
+ lambda_param=lambda_param,
  )

src/pyversity/strategies/msd.py

@@ -5,10 +5,10 @@
 
 
 def msd(
- relevances: np.ndarray,
  embeddings: np.ndarray,
+ scores: np.ndarray,
  k: int,
- alpha: float = 0.5,
+ lambda_param: float = 0.5,
  metric: Metric = Metric.COSINE,
  normalize: bool = True,
 ) -> tuple[np.ndarray, np.ndarray]:
@@ -18,10 +18,11 @@ def msd(
  This strategy selects `k` items that balance relevance and diversity by
  iteratively choosing items that maximize a combination of their relevance
  and their total distance to already selected items.
- :param relevances: 1D array of relevance scores for each item.
+
  :param embeddings: 2D array of shape (n_samples, n_features).
+ :param scores: 1D array of relevance scores for each item.
  :param k: Number of items to select.
- :param alpha: Trade-off parameter in [0, 1].
+ :param lambda_param: Trade-off parameter in [0, 1].
  1.0 = pure relevance, 0.0 = pure diversity.
 
  :param metric: Similarity metric to use. Default is Metric.COSINE.
@@ -30,10 +31,10 @@ def msd(
  """
  return greedy_select(
  "msd",
- relevances,
+ scores,
  embeddings,
  k,
  metric=metric,
  normalize=normalize,
- alpha=alpha,
+ lambda_param=lambda_param,
  )