Precompute cosine normalization for indexed pairs

src/copairs/compute.py

@@ -165,6 +165,79 @@ def pairwise_cosine(x_sample: np.ndarray, y_sample: np.ndarray) -> np.ndarray:
     return c_sim
 
 
+def prepare_cosine(feats: np.ndarray, profile_ix: np.ndarray) -> np.ndarray:
+    """Normalize referenced feature rows for repeated cosine pair evaluation.
+
+    Parameters
+    ----------
+    feats : np.ndarray
+        A 2D feature matrix with profiles in rows.
+    profile_ix : np.ndarray
+        Row indices that will be referenced by subsequent pair evaluation.
+
+    Returns
+    -------
+    np.ndarray
+        A feature-sized matrix whose referenced rows are row-normalized.
+    """
+    feats = np.asarray(feats)
+    profile_ix = np.asarray(profile_ix)
+    # Advanced indexing gathers a C-order matrix, matching the generic batched
+    # cosine path regardless of the source matrix's memory layout.
+    referenced_feats = feats[profile_ix]
+    normalized_rows = (
+        referenced_feats / np.linalg.norm(referenced_feats, axis=1)[:, np.newaxis]
+    )
+    normalized_feats = np.empty(feats.shape, dtype=normalized_rows.dtype)
+    normalized_feats[profile_ix] = normalized_rows
+    return normalized_feats
+
+
+def cosine_pairs(
+    normalized_feats: np.ndarray,
+    pair_ix: np.ndarray,
+    batch_size: int,
+    progress_bar: bool = True,
+) -> np.ndarray:
+    """Compute indexed dot products from row-normalized features in batches.
+
+    Parameters
+    ----------
+    normalized_feats : np.ndarray
+        A feature matrix normalized with :func:`prepare_cosine`.
+    pair_ix : np.ndarray
+        A two-column array containing row-index pairs.
+    batch_size : int
+        Number of pairs to process per batch.
+    progress_bar : bool
+        Whether or not to show tqdm's progress bar.
+
+    Returns
+    -------
+    np.ndarray
+        A float32 array of cosine similarities for the indexed pairs.
+    """
+    num_pairs = len(pair_ix)
+    result = np.empty(num_pairs, dtype=np.float32)
+    if num_pairs == 0:
+        return result
+
+    def par_func(i: int) -> None:
+        pairs = pair_ix[i : i + batch_size]
+        x_sample = normalized_feats[pairs[:, 0]]
+        y_sample = normalized_feats[pairs[:, 1]]
+        # Match pairwise_cosine's multiply-then-sum arithmetic so exact ties and
+        # floating-point ranking behavior do not change.
+        result[i : i + len(pairs)] = np.sum(x_sample * y_sample, axis=1)
+
+    parallel_map(
+        par_func,
+        np.arange(0, num_pairs, batch_size),
+        progress_bar=progress_bar,
+    )
+    return result
+
+
 def pairwise_abs_cosine(x_sample: np.ndarray, y_sample: np.ndarray) -> np.ndarray:
     """Compute the absolute cosine similarity for paired rows of two matrices.
 

src/copairs/map/average_precision.py

@@ -194,13 +194,35 @@ def average_precision(
     if len(neg_pairs) == 0:
         raise UnpairedException("Unable to find negative pairs.")
 
-    # Compute similarities for positive pairs
-    logger.info("Computing positive similarities...")
-    pos_sims = similarity_fn(feats, pos_pairs, batch_size)
-
-    # Compute similarities for negative pairs
-    logger.info("Computing negative similarities...")
-    neg_sims = similarity_fn(feats, neg_pairs, batch_size)
+    # Cosine retrieval reuses profile rows across many pairs. For the exact
+    # built-in cosine string, normalize each referenced row once and evaluate
+    # indexed pair dot products. Other strings and callables retain the generic
+    # batched similarity path.
+    if isinstance(distance, str) and distance == "cosine":
+        profile_ix = np.unique(np.concatenate((pos_pairs, neg_pairs)))
+        normalized_feats = compute.prepare_cosine(feats, profile_ix)
+
+        logger.info("Computing positive similarities...")
+        pos_sims = compute.cosine_pairs(
+            normalized_feats,
+            pos_pairs,
+            batch_size,
+            progress_bar=progress_bar,
+        )
+
+        logger.info("Computing negative similarities...")
+        neg_sims = compute.cosine_pairs(
+            normalized_feats,
+            neg_pairs,
+            batch_size,
+            progress_bar=progress_bar,
+        )
+    else:
+        logger.info("Computing positive similarities...")
+        pos_sims = similarity_fn(feats, pos_pairs, batch_size)
+
+        logger.info("Computing negative similarities...")
+        neg_sims = similarity_fn(feats, neg_pairs, batch_size)
 
     # Build rank lists for calculating average precision
     logger.info("Building rank lists...")

src/copairs/map/multilabel.py

@@ -115,11 +115,31 @@ def average_precision(
     logger.info("Dropping dups in negative pairs...")
     neg_pairs = np.unique(neg_pairs, axis=0)
 
-    logger.info("Computing positive similarities...")
-    pos_sims = distance_fn(feats, pos_pairs, batch_size)
+    if isinstance(distance, str) and distance == "cosine":
+        profile_ix = np.unique(np.concatenate((pos_pairs, neg_pairs)))
+        normalized_feats = compute.prepare_cosine(feats, profile_ix)
+
+        logger.info("Computing positive similarities...")
+        pos_sims = compute.cosine_pairs(
+            normalized_feats,
+            pos_pairs,
+            batch_size,
+            progress_bar=progress_bar,
+        )
+
+        logger.info("Computing negative similarities...")
+        neg_sims = compute.cosine_pairs(
+            normalized_feats,
+            neg_pairs,
+            batch_size,
+            progress_bar=progress_bar,
+        )
+    else:
+        logger.info("Computing positive similarities...")
+        pos_sims = distance_fn(feats, pos_pairs, batch_size)
 
-    logger.info("Computing negative similarities...")
-    neg_sims = distance_fn(feats, neg_pairs, batch_size)
+        logger.info("Computing negative similarities...")
+        neg_sims = distance_fn(feats, neg_pairs, batch_size)
 
     logger.info("Computing AP per label...")
     negs_for = _create_neg_query_solver(neg_pairs, neg_sims)

tests/test_compute.py

@@ -112,6 +112,96 @@ def test_cosine():
     assert np.allclose(cosine_gt, cosine)
 
 
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("layout", ["c", "fortran", "strided"])
+def test_cosine_pairs_exactly_match_generic_batches(dtype, layout):
+    """Prepared pair dots match gathered cosine across dtypes and layouts."""
+    local_rng = np.random.default_rng(SEED)
+    feats = local_rng.normal(size=(17, 11)).astype(dtype)
+    if layout == "fortran":
+        feats = np.asfortranarray(feats)
+    elif layout == "strided":
+        backing = np.empty((len(feats), feats.shape[1] * 2), dtype=dtype)
+        backing[:, ::2] = feats
+        feats = backing[:, ::2]
+    pairs = np.asarray(
+        [(i, (i * 7 + 3) % len(feats)) for i in range(len(feats))],
+        dtype=np.uint32,
+    )
+    generic = compute.get_similarity_fn("cosine", progress_bar=False)(
+        feats, pairs, batch_size=3
+    )
+    normalized = compute.prepare_cosine(feats, np.unique(pairs))
+    actual = compute.cosine_pairs(normalized, pairs, batch_size=3, progress_bar=False)
+
+    assert actual.dtype == np.float32
+    np.testing.assert_array_equal(actual, generic)
+
+
+def test_prepare_cosine_skips_unreferenced_nonfinite_rows_under_strict_errstate():
+    """Only pair-referenced profiles participate in normalization."""
+    feats = np.asarray([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.0, 0.0], [np.inf, 1.0]])
+    pairs = np.asarray([[0, 1], [1, 2], [2, 0]], dtype=np.uint32)
+
+    with np.errstate(all="raise"):
+        normalized = compute.prepare_cosine(feats, np.unique(pairs))
+        actual = compute.cosine_pairs(
+            normalized, pairs, batch_size=2, progress_bar=False
+        )
+        expected = compute.pairwise_cosine(
+            feats[pairs[:, 0]], feats[pairs[:, 1]]
+        ).astype(np.float32)
+
+    np.testing.assert_array_equal(actual, expected)
+
+
+@pytest.mark.parametrize("bad_row", [[0.0, 0.0], [np.inf, 1.0]])
+def test_prepare_cosine_preserves_referenced_nonfinite_errstate(bad_row):
+    """Referenced zero and infinite rows retain strict NumPy error behavior."""
+    feats = np.asarray([[1.0, 0.0], bad_row])
+
+    with np.errstate(all="raise"), pytest.raises(FloatingPointError):
+        compute.prepare_cosine(feats, np.asarray([1]))
+
+
+def test_cosine_pairs_preserve_zero_norm_and_nonfinite_results():
+    """Pre-normalization retains generic cosine behavior for zero-norm rows."""
+    feats = np.asarray([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float64)
+    pairs = np.asarray([[0, 1], [1, 2], [1, 1]], dtype=np.uint32)
+
+    with np.errstate(invalid="ignore"):
+        generic = compute.get_similarity_fn("cosine", progress_bar=False)(
+            feats, pairs, batch_size=1
+        )
+        actual = compute.cosine_pairs(
+            compute.prepare_cosine(feats, np.unique(pairs)),
+            pairs,
+            batch_size=2,
+            progress_bar=False,
+        )
+
+    np.testing.assert_array_equal(actual, generic)
+    assert np.isnan(actual[0])
+
+
+def test_cosine_pairs_returns_typed_empty_result(monkeypatch):
+    """An empty pair array returns without starting parallel workers."""
+
+    def fail_parallel_map(*args, **kwargs):
+        raise AssertionError("parallel_map must not run for empty pairs")
+
+    monkeypatch.setattr(compute, "parallel_map", fail_parallel_map)
+    actual = compute.cosine_pairs(
+        np.empty((3, 2)),
+        np.empty((0, 2), dtype=np.uint32),
+        batch_size=2,
+        progress_bar=False,
+    )
+
+    assert actual.shape == (0,)
+    assert actual.dtype == np.float32
+
+
 def test_euclidean():
     """Test euclidean similarity computation."""
     n_samples = 10