diff --git a/src/copairs/compute.py b/src/copairs/compute.py index b7dbdc7..617f9f1 100644 --- a/src/copairs/compute.py +++ b/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. diff --git a/src/copairs/map/average_precision.py b/src/copairs/map/average_precision.py index 9a23f60..672f9e0 100644 --- a/src/copairs/map/average_precision.py +++ b/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...") diff --git a/src/copairs/map/multilabel.py b/src/copairs/map/multilabel.py index d321dac..0ba0c3c 100644 --- a/src/copairs/map/multilabel.py +++ b/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) diff --git a/tests/test_compute.py b/tests/test_compute.py index 6624ab8..162f13f 100644 --- a/tests/test_compute.py +++ b/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 diff --git a/tests/test_map.py b/tests/test_map.py index 364395e..e377077 100644 --- a/tests/test_map.py +++ b/tests/test_map.py @@ -14,6 +14,17 @@ SEED = 0 +def _with_layout(feats: np.ndarray, layout: str) -> np.ndarray: + """Return feature values with the requested memory layout.""" + if layout == "fortran": + return np.asfortranarray(feats) + if layout == "strided": + backing = np.empty((len(feats), feats.shape[1] * 2), dtype=feats.dtype) + backing[:, ::2] = feats + return backing[:, ::2] + return feats + + def binary2indices(arr: np.ndarray) -> np.ndarray: """Convert a binary matrix to a list of indices.""" return np.nonzero(arr)[1].reshape(arr.shape[0], -1) @@ -218,6 +229,208 @@ def test_raise_nan_error(): ) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("layout", ["fortran", "strided"]) +def test_cosine_fast_path_exact_regular_parity_with_ties(dtype, layout): + """The cosine fast path preserves positive-first ranking for exact ties.""" + meta = pd.DataFrame( + { + "compound": ["a", "a", "b", "b", "c", "c"], + "plate": ["p1", "p2", "p1", "p2", "p1", "p2"], + } + ) + # Every pair has exactly the same cosine similarity, so the expected AP of + # one verifies that positive pairs still precede tied negative pairs. + feats = np.tile( + np.asarray([0.125, -1.5, 2.25, 0.75, -0.0625], dtype=dtype), + (len(meta), 1), + ) + feats = _with_layout(feats, layout) + kwargs = { + "meta": meta, + "feats": feats, + "pos_sameby": ["compound"], + "pos_diffby": [], + "neg_sameby": [], + "neg_diffby": ["compound"], + "progress_bar": False, + "batch_size": 2, + } + + generic = average_precision(distance=compute.pairwise_cosine, **kwargs) + optimized = average_precision(distance="cosine", **kwargs) + + pd.testing.assert_frame_equal(optimized, generic, check_exact=True) + np.testing.assert_array_equal(optimized["average_precision"], 1.0) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("layout", ["fortran", "strided"]) +def test_cosine_fast_path_exact_multilabel_parity_with_ties(dtype, layout): + """Multilabel cosine retrieval preserves exact grouped tie semantics.""" + meta = pd.DataFrame( + { + "compound": ["a", "b", "c", "d"], + "target": [["x"], ["x", "y"], ["y"], ["z"]], + } + ) + feats = np.tile( + np.asarray([0.125, -1.5, 2.25, 0.75, -0.0625], dtype=dtype), + (len(meta), 1), + ) + feats = _with_layout(feats, layout) + kwargs = { + "meta": meta, + "feats": feats, + "pos_sameby": ["target"], + "pos_diffby": [], + "neg_sameby": [], + "neg_diffby": ["target"], + "multilabel_col": "target", + "progress_bar": False, + "batch_size": 2, + } + + generic = multilabel_average_precision(distance=compute.pairwise_cosine, **kwargs) + optimized = multilabel_average_precision(distance="cosine", **kwargs) + + pd.testing.assert_frame_equal(optimized, generic, check_exact=True) + np.testing.assert_array_equal(optimized["average_precision"], 1.0) + + +def test_cosine_fast_path_regular_skips_unpaired_nonfinite_rows(): + """Regular cosine preparation only normalizes profiles found in pairs.""" + meta = pd.DataFrame( + { + "compound": ["a", "a", "b", "b", "zero", "inf"], + "cohort": ["main", "main", "main", "main", "zero", "inf"], + } + ) + feats = np.asarray( + [ + [1.0, 0.5], + [0.75, 1.0], + [-0.5, 1.0], + [1.0, -0.25], + [0.0, 0.0], + [np.inf, 1.0], + ] + ) + kwargs = { + "meta": meta, + "feats": feats, + "pos_sameby": ["compound"], + "pos_diffby": [], + "neg_sameby": ["cohort"], + "neg_diffby": ["compound"], + "progress_bar": False, + "batch_size": 2, + } + + with np.errstate(all="raise"): + generic = average_precision(distance=compute.pairwise_cosine, **kwargs) + optimized = average_precision(distance="cosine", **kwargs) + + pd.testing.assert_frame_equal(optimized, generic, check_exact=True) + + +def test_cosine_fast_path_multilabel_skips_unpaired_nonfinite_rows(): + """Multilabel preparation only normalizes profiles found in pairs.""" + meta = pd.DataFrame( + { + "target": [["x"], ["x"], ["y"], ["y"], ["zero"], ["inf"]], + "cohort": ["main", "main", "main", "main", "zero", "inf"], + } + ) + feats = np.asarray( + [ + [1.0, 0.5], + [0.75, 1.0], + [-0.5, 1.0], + [1.0, -0.25], + [0.0, 0.0], + [np.inf, 1.0], + ] + ) + kwargs = { + "meta": meta, + "feats": feats, + "pos_sameby": ["target"], + "pos_diffby": [], + "neg_sameby": ["cohort"], + "neg_diffby": ["target"], + "multilabel_col": "target", + "progress_bar": False, + "batch_size": 2, + } + + with np.errstate(all="raise"): + generic = multilabel_average_precision( + distance=compute.pairwise_cosine, **kwargs + ) + optimized = multilabel_average_precision(distance="cosine", **kwargs) + + pd.testing.assert_frame_equal(optimized, generic, check_exact=True) + + +def test_generic_string_and_callable_similarity_fallback(monkeypatch): + """Non-cosine strings and custom callables retain generic batch processing.""" + meta = pd.DataFrame( + {"compound": ["a", "a", "b", "b"], "plate": ["p1", "p2", "p1", "p2"]} + ) + feats = np.eye(len(meta)) + kwargs = { + "meta": meta, + "feats": feats, + "pos_sameby": ["compound"], + "pos_diffby": [], + "neg_sameby": [], + "neg_diffby": ["compound"], + "progress_bar": False, + } + + def fail_fast_path(*args, **kwargs): + raise AssertionError("cosine fast path must not handle fallback metrics") + + monkeypatch.setattr(compute, "cosine_pairs", fail_fast_path) + string_result = average_precision(distance="euclidean", **kwargs) + + calls = [] + + def custom_distance(x_sample, y_sample): + calls.append(len(x_sample)) + return compute.pairwise_euclidean(x_sample, y_sample) + + callable_result = average_precision(distance=custom_distance, **kwargs) + + pd.testing.assert_frame_equal(string_result, callable_result, check_exact=True) + assert calls + + +def test_cosine_fast_path_forwards_progress_setting(monkeypatch): + """Each cosine pair set retains the requested progress-bar setting.""" + progress_settings = [] + + def sequential_map(par_func, items, progress_bar=True): + progress_settings.append(progress_bar) + for item in items: + par_func(item) + + monkeypatch.setattr(compute, "parallel_map", sequential_map) + meta = pd.DataFrame({"compound": ["a", "a", "b", "b"]}) + average_precision( + meta, + np.eye(len(meta)), + pos_sameby=["compound"], + pos_diffby=[], + neg_sameby=[], + neg_diffby=["compound"], + progress_bar=False, + ) + + assert progress_settings == [False, False] + + def test_progress_bar_consistency(): """Test that the progress_bar argument does not change results.""" length = 10