diff --git a/README.md b/README.md index 37a9125..3dabfe2 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,14 @@ We provide examples demonstrating how to use copairs for: - [calculating mAP to assess phenotypic consistency of perturbations](https://github.com/cytomining/copairs/blob/main/docs/examples/phenotypic_consistency.ipynb) - [estimating null size for mAP p-value calculation](https://github.com/cytomining/copairs/blob/main/docs/examples/null_size.ipynb) +### Worker configuration + +Batched similarity, null-distribution, and mAP aggregation calculations use at +most 8 worker threads by default, further limited by the CPUs available to the +process and the number of tasks. Pass `max_workers` to the AP APIs for per-call +control, or set the `COPAIRS_MAX_WORKERS` environment variable to a positive +integer to change the default across these worker pools. + ## Citation If you find this work useful for your research, please cite our [paper](https://doi.org/10.1038/s41467-025-60306-2): diff --git a/src/copairs/compute.py b/src/copairs/compute.py index b7dbdc7..c0199b0 100644 --- a/src/copairs/compute.py +++ b/src/copairs/compute.py @@ -11,11 +11,49 @@ from scipy.spatial.distance import _METRICS_NAMES as SCIPY_METRICS_NAMES from scipy.spatial.distance import cdist +_DEFAULT_MAX_WORKERS = 8 + + +def _available_cpu_count() -> int: + """Return CPUs available to this process, with a portable fallback.""" + try: + return len(os.sched_getaffinity(0)) or 1 + except (AttributeError, NotImplementedError, OSError): + return os.cpu_count() or 1 + + +def _resolve_max_workers(num_items: int, max_workers: Optional[int]) -> int: + """Resolve the worker budget, bounded by the number of tasks.""" + if max_workers is None: + configured = os.environ.get("COPAIRS_MAX_WORKERS") + if configured is not None: + try: + max_workers = int(configured) + except ValueError as error: + raise ValueError( + "COPAIRS_MAX_WORKERS must be a positive integer" + ) from error + else: + # Pair batches allocate gathered feature arrays and run native NumPy + # kernels. A bounded budget limits simultaneous memory pressure and + # avoids scaling Python threads to every CPU on large hosts. + max_workers = min(_DEFAULT_MAX_WORKERS, _available_cpu_count()) + + if isinstance(max_workers, bool) or not isinstance(max_workers, int): + raise TypeError("max_workers must be a positive integer") + if max_workers < 1: + raise ValueError("max_workers must be a positive integer") + if num_items < 1: + return 0 + return min(num_items, max_workers) + def parallel_map( par_func: Callable[[int], None], items: np.ndarray, progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ) -> None: """Execute a function in parallel over a list of items. @@ -30,33 +68,41 @@ def parallel_map( (an item index or value). items : np.ndarray An array or list of items to process. + progress_bar : bool + Whether to display task completion with tqdm. + max_workers : int, optional + Maximum number of worker threads. By default, the worker count is the + smaller of the task count, available CPUs, and 8. The + ``COPAIRS_MAX_WORKERS`` environment variable can override the default. """ - # Total number of items to process num_items = len(items) + pool_size = _resolve_max_workers(num_items, max_workers) + if pool_size == 0: + return - # Determine the number of threads to use, limited by CPU count - pool_size = min(num_items, os.cpu_count()) - - # Calculate chunk size for dividing work among threads - chunksize = num_items // pool_size - - # Use a thread pool to execute the function in parallel - with ThreadPool(pool_size) as pool: - # Map the function to items with unordered execution for better efficiency - tasks = pool.imap_unordered(par_func, items, chunksize=chunksize) - + def consume(tasks) -> None: if progress_bar: - # Display progress using tqdm from tqdm.autonotebook import tqdm - tasks = tqdm(tasks, total=len(items), leave=False) + tasks = tqdm(tasks, total=num_items, leave=False) for _ in tasks: pass + if pool_size == 1: + consume(map(par_func, items)) + return + + chunksize = max(1, num_items // pool_size) + with ThreadPool(pool_size) as pool: + tasks = pool.imap_unordered(par_func, items, chunksize=chunksize) + consume(tasks) + def batch_processing( pairwise_op: Callable[[np.ndarray, np.ndarray], np.ndarray], progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ): """ Add batch processing support to pairwise operations. @@ -72,6 +118,8 @@ def batch_processing( between two arrays of features. progress_bar : bool Whether or not to show tqdm's progress bar. + max_workers : int, optional + Maximum number of worker threads used by the wrapped operation. Returns ------- @@ -79,8 +127,17 @@ def batch_processing( A wrapped function that processes pairwise operations in batches. """ + configured_max_workers = max_workers + + def batched_fn( + feats: np.ndarray, + pair_ix: np.ndarray, + batch_size: int, + *, + max_workers: Optional[int] = None, + ): + worker_budget = configured_max_workers if max_workers is None else max_workers - def batched_fn(feats: np.ndarray, pair_ix: np.ndarray, batch_size: int): # Total number of pairs to process num_pairs = len(pair_ix) @@ -97,7 +154,10 @@ def par_func(i): # Use multithreading to process the batches in parallel parallel_map( - par_func, np.arange(0, num_pairs, batch_size), progress_bar=progress_bar + par_func, + np.arange(0, num_pairs, batch_size), + progress_bar=progress_bar, + max_workers=worker_budget, ) return result @@ -269,7 +329,10 @@ def _cdist_diag_sim( def get_similarity_fn( - distance: Union[str, Callable], progress_bar: bool = True + distance: Union[str, Callable], + progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ) -> Callable: """Retrieve a similarity function based on a distance string identifier or custom callable. @@ -296,6 +359,8 @@ def get_similarity_fn( callable function. progress_bar : bool Whether or not to show tqdm's progress bar. + max_workers : int, optional + Maximum number of worker threads used for batched similarities. Returns ------- @@ -342,7 +407,9 @@ def get_similarity_fn( raise ValueError("Distance must be either a string or a callable object.") # Wrap the distance function for efficient batch processing - return batch_processing(similarity_fn, progress_bar=progress_bar) + return batch_processing( + similarity_fn, progress_bar=progress_bar, max_workers=max_workers + ) def random_binary_matrix(n, m, k, rng): @@ -532,6 +599,8 @@ def get_null_dists( seed: int, cache_dir: Optional[Union[str, Path]] = None, progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ) -> np.ndarray: """Generate null distributions for each configuration of positive and total pairs. @@ -546,6 +615,8 @@ def get_null_dists( Random seed for reproducibility. progress_bar : bool Whether or not to show tqdm's progress bar. + max_workers : int, optional + Maximum number of worker threads used to generate null distributions. Returns ------- @@ -572,7 +643,12 @@ def par_func(i): _cache_write(path, null_dist) null_dists[i] = null_dist - parallel_map(par_func, np.arange(num_confs), progress_bar) + parallel_map( + par_func, + np.arange(num_confs), + progress_bar, + max_workers=max_workers, + ) return null_dists @@ -583,6 +659,8 @@ def p_values( null_size: int, seed: int, progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ): """Calculate p-values for an array of Average Precision (AP) scores using a null distribution. @@ -600,6 +678,8 @@ def p_values( distribution. progress_bar : bool Whether or not to show tqdm's progress bar. + max_workers : int, optional + Maximum number of worker threads used to generate null distributions. Returns ------- @@ -610,7 +690,13 @@ def p_values( confs, rev_ix = np.unique(null_confs, axis=0, return_inverse=True) # Generate null distributions for each unique configuration - null_dists = get_null_dists(confs, null_size, seed, progress_bar=progress_bar) + null_dists = get_null_dists( + confs, + null_size, + seed, + progress_bar=progress_bar, + max_workers=max_workers, + ) # Sort null distributions for efficient p-value computation null_dists.sort(axis=1) diff --git a/src/copairs/map/average_precision.py b/src/copairs/map/average_precision.py index 9a23f60..d2f53f4 100644 --- a/src/copairs/map/average_precision.py +++ b/src/copairs/map/average_precision.py @@ -1,7 +1,7 @@ """Functions to compute average precision.""" import logging -from typing import List +from typing import List, Optional import numpy as np import pandas as pd @@ -92,6 +92,8 @@ def average_precision( batch_size: int = 20000, distance: str = "cosine", progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ) -> pd.DataFrame: """Calculate average precision (AP) scores for pairs of profiles based on their similarity. @@ -142,6 +144,10 @@ def average_precision( distance : str The distance function used for computing similarities. Default is "cosine". + progress_bar : bool + Whether or not to show tqdm's progress bar. + max_workers : int, optional + Maximum number of worker threads used for similarity calculations. Returns ------- @@ -175,7 +181,9 @@ def average_precision( validate_pipeline_input(meta, feats, columns) # Get the distance function for similarity calculations (e.g., cosine) - similarity_fn = compute.get_similarity_fn(distance, progress_bar=progress_bar) + similarity_fn = compute.get_similarity_fn( + distance, progress_bar=progress_bar, max_workers=max_workers + ) # Reset metadata index for consistent indexing meta = meta.reset_index(drop=True).copy() @@ -235,7 +243,12 @@ def average_precision( def p_values( - dframe: pd.DataFrame, null_size: int, seed: int, progress_bar: bool = True + dframe: pd.DataFrame, + null_size: int, + seed: int, + progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ) -> np.ndarray: """Compute p-values for average precision scores based on a null distribution. @@ -257,6 +270,8 @@ def p_values( Random seed for reproducibility of the null distribution. progress_bar : bool Whether or not to show tqdm's progress bar. + max_workers : int, optional + Maximum number of worker threads used to generate null distributions. Returns ------- @@ -275,7 +290,14 @@ def p_values( null_confs = dframe.loc[mask, ["n_pos_pairs", "n_total_pairs"]].values # Compute p-values for profiles with valid configurations using the null distribution - pvals[mask] = compute.p_values(scores, null_confs, null_size, seed, progress_bar) + pvals[mask] = compute.p_values( + scores, + null_confs, + null_size, + seed, + progress_bar, + max_workers=max_workers, + ) # Return the array of p-values, including NaN for invalid profiles return pvals diff --git a/src/copairs/map/map.py b/src/copairs/map/map.py index a5b9cac..a5ea0a0 100644 --- a/src/copairs/map/map.py +++ b/src/copairs/map/map.py @@ -1,7 +1,6 @@ """Functions to compute mean average precision.""" import logging -from os import cpu_count from typing import List, Union, Optional from pathlib import Path from concurrent.futures import ThreadPoolExecutor @@ -46,8 +45,9 @@ def get_map_pvalue( Random seed for reproducibility. progress_bar : bool Whether or not to show tqdm's progress bar. - max_workers : int - Number of workers used. Default defined by tqdm's `thread_map`. + max_workers : int, optional + Maximum workers used for null generation and mAP aggregation. The + default uses ``COPAIRS_MAX_WORKERS`` or at most 8 available CPUs. cache_dir : str or Path Location to save the cache. @@ -73,7 +73,12 @@ def get_map_pvalue( # Generate null distributions for each unique configuration null_dists = compute.get_null_dists( - null_confs, null_size, seed=seed, cache_dir=cache_dir, progress_bar=progress_bar + null_confs, + null_size, + seed=seed, + cache_dir=cache_dir, + progress_bar=progress_bar, + max_workers=max_workers, ) ap_scores["null_ix"] = rev_ix @@ -100,18 +105,28 @@ def get_p_value(params): "mean_normalized_average_precision", ] - # Compute p-values for each group using the null distributions + # Compute p-values for each group using the null distributions. Resolve this + # pool separately so the same explicit/environment budget constrains both + # native null generation and mAP aggregation. params = map_scores[["mean_average_precision", "indices"]] + aggregation_workers = compute._resolve_max_workers(len(params), max_workers) - if progress_bar: + if aggregation_workers == 0: + p_values = [] + elif progress_bar: from tqdm.contrib.concurrent import thread_map p_values = thread_map( - get_p_value, params.values, leave=False, max_workers=max_workers + get_p_value, + params.values, + leave=False, + max_workers=aggregation_workers, ) else: p_values = silent_thread_map( - get_p_value, params.values, max_workers=max_workers + get_p_value, + params.values, + max_workers=aggregation_workers, ) map_scores["p_value"] = p_values @@ -150,8 +165,9 @@ def mean_average_precision( Random seed for reproducibility. progress_bar : bool Whether or not to show tqdm's progress bar. - max_workers : int - Number of workers used. Default defined by tqdm's `thread_map`. + max_workers : int, optional + Maximum workers used for null generation and mAP aggregation. The + default uses ``COPAIRS_MAX_WORKERS`` or at most 8 available CPUs. cache_dir : str or Path Location to save the cache. @@ -238,8 +254,9 @@ def mean_average_precision_hierarchical( grouped structure. progress_bar : bool Whether or not to show tqdm's progress bar. - max_workers : int - Number of workers used. Default defined by tqdm's `thread_map`. + max_workers : int, optional + Maximum workers used for null generation and mAP aggregation. The + default uses ``COPAIRS_MAX_WORKERS`` or at most 8 available CPUs. cache_dir : str or Path Location to save the cache. @@ -297,7 +314,7 @@ def silent_thread_map(fn, *iterables, **kwargs): **kwargs : dict Additional keyword arguments. Accepts: - max_workers : int, optional - Maximum number of workers [default: min(32, cpu_count() + 4)]. + Maximum number of workers [default: at most 8 available CPUs]. - chunksize : int, optional Size of chunks for each worker [default: 1]. """ @@ -305,7 +322,7 @@ def silent_thread_map(fn, *iterables, **kwargs): # (github.com/tqdm/tqdm/blob/0ed5d7f18fa3153834cbac0aa57e8092b217cc16/tqdm/contrib/concurrent.py#L29). kwargs = kwargs.copy() - max_workers = kwargs.pop("max_workers", min(32, cpu_count() + 4)) + max_workers = kwargs.pop("max_workers", min(8, compute._available_cpu_count())) chunksize = kwargs.pop("chunksize", 1) with ThreadPoolExecutor(max_workers=max_workers) as ex: return list(ex.map(fn, *iterables, chunksize=chunksize, **kwargs)) diff --git a/src/copairs/map/multilabel.py b/src/copairs/map/multilabel.py index d321dac..7785383 100644 --- a/src/copairs/map/multilabel.py +++ b/src/copairs/map/multilabel.py @@ -1,7 +1,7 @@ """Functions to compute mAP with multilabel support.""" import logging -from typing import List +from typing import List, Optional import numpy as np import pandas as pd @@ -79,10 +79,13 @@ def average_precision( batch_size=20000, distance="cosine", progress_bar: bool = True, + *, + max_workers: Optional[int] = None, ) -> pd.DataFrame: """ Compute average precision with multilabel support. + ``max_workers`` limits worker threads used for similarity calculations. Returns normalized_average_precision in addition to average_precision. See Also @@ -92,7 +95,9 @@ def average_precision( columns = flatten_str_list(pos_sameby, pos_diffby, neg_sameby, neg_diffby) meta, columns = evaluate_and_filter(meta, columns) validate_pipeline_input(meta, feats, columns) - distance_fn = compute.get_similarity_fn(distance, progress_bar=progress_bar) + distance_fn = compute.get_similarity_fn( + distance, progress_bar=progress_bar, max_workers=max_workers + ) # Critical!, otherwise the indexing wont work meta = meta.reset_index(drop=True).copy() diff --git a/tests/test_compute.py b/tests/test_compute.py index 6624ab8..e1e6f9b 100644 --- a/tests/test_compute.py +++ b/tests/test_compute.py @@ -12,6 +12,107 @@ rng = np.random.default_rng(SEED) +def test_resolve_max_workers(monkeypatch): + """Worker resolution is bounded, configurable, and task-aware.""" + monkeypatch.delenv("COPAIRS_MAX_WORKERS", raising=False) + monkeypatch.setattr(compute.os, "cpu_count", lambda: 384) + monkeypatch.setattr( + compute.os, "sched_getaffinity", lambda _: set(range(64)), raising=False + ) + assert compute._resolve_max_workers(1000, None) == 8 + assert compute._resolve_max_workers(3, None) == 3 + assert compute._resolve_max_workers(1000, 7) == 7 + + monkeypatch.setenv("COPAIRS_MAX_WORKERS", "5") + assert compute._resolve_max_workers(1000, None) == 5 + assert compute._resolve_max_workers(2, None) == 2 + assert compute._resolve_max_workers(0, None) == 0 + assert compute._resolve_max_workers(1000, 4) == 4 + + +def test_resolve_max_workers_uses_affinity_with_fallback(monkeypatch): + """Defaults respect process affinity and portably fall back to CPU count.""" + monkeypatch.delenv("COPAIRS_MAX_WORKERS", raising=False) + monkeypatch.setattr( + compute.os, "sched_getaffinity", lambda _: set(range(4)), raising=False + ) + monkeypatch.setattr(compute.os, "cpu_count", lambda: 384) + assert compute._resolve_max_workers(100, None) == 4 + + def unavailable(_): + raise OSError("affinity unavailable") + + monkeypatch.setattr(compute.os, "sched_getaffinity", unavailable) + monkeypatch.setattr(compute.os, "cpu_count", lambda: 6) + assert compute._resolve_max_workers(100, None) == 6 + + +@pytest.mark.parametrize("num_items", [0, 2]) +@pytest.mark.parametrize("value", ["0", "-1", "invalid"]) +def test_resolve_max_workers_rejects_invalid_environment(monkeypatch, num_items, value): + """The process-wide worker override is validated even without tasks.""" + monkeypatch.setenv("COPAIRS_MAX_WORKERS", value) + with pytest.raises(ValueError, match="positive integer"): + compute._resolve_max_workers(num_items, None) + + +@pytest.mark.parametrize("num_items", [0, 2]) +@pytest.mark.parametrize("value", [0, -1, 1.5, True]) +def test_resolve_max_workers_rejects_invalid_argument(num_items, value): + """Explicit worker budgets are validated even without tasks.""" + error = TypeError if isinstance(value, (float, bool)) else ValueError + with pytest.raises(error, match="positive integer"): + compute._resolve_max_workers(num_items, value) + + +def test_parallel_map_empty_validates_worker_budget(monkeypatch): + """Empty task collections do not bypass worker-budget validation.""" + with pytest.raises(ValueError, match="positive integer"): + compute.parallel_map(lambda _: None, [], progress_bar=False, max_workers=0) + + monkeypatch.setenv("COPAIRS_MAX_WORKERS", "invalid") + with pytest.raises(ValueError, match="positive integer"): + compute.parallel_map(lambda _: None, [], progress_bar=False) + + +def test_parallel_map_serial_honors_progress_bar(monkeypatch): + """Serial execution still reports task progress.""" + progress_calls = [] + + def fake_tqdm(tasks, **kwargs): + progress_calls.append(kwargs) + return tasks + + monkeypatch.setattr("tqdm.autonotebook.tqdm", fake_tqdm) + monkeypatch.setattr( + compute, + "ThreadPool", + lambda *_: pytest.fail("serial execution created a thread pool"), + ) + visited = [] + compute.parallel_map( + visited.append, + np.arange(3), + progress_bar=True, + max_workers=1, + ) + + assert visited == [0, 1, 2] + assert progress_calls == [{"total": 3, "leave": False}] + + +def test_batched_similarity_output_parity_across_worker_budgets(): + """Worker count does not affect batched similarity output.""" + feats = rng.normal(size=(12, 7)) + pairs = rng.integers(0, len(feats), size=(31, 2)) + similarity_fn = compute.get_similarity_fn("cosine", progress_bar=False) + + serial = similarity_fn(feats, pairs, 4, max_workers=1) + parallel = similarity_fn(feats, pairs, 4, max_workers=4) + + np.testing.assert_array_equal(serial, parallel) + + def corrcoef_naive(feats, pairs): """Compute correlation coefficient between pairs of features.""" corr = np.empty((len(pairs),)) @@ -202,6 +303,26 @@ def test_hamming(): assert np.allclose(hamming_gt, hamming) +def test_null_distribution_output_parity_across_worker_budgets(tmp_path): + """Native null-distribution output is unchanged by the worker budget.""" + confs = np.asarray([[1, 5], [2, 7], [3, 9]]) + kwargs = { + "confs": confs, + "null_size": 100, + "seed": 42, + "progress_bar": False, + } + + serial = compute.get_null_dists( + **kwargs, cache_dir=tmp_path / "serial", max_workers=1 + ) + parallel = compute.get_null_dists( + **kwargs, cache_dir=tmp_path / "parallel", max_workers=3 + ) + + np.testing.assert_array_equal(serial, parallel) + + def test_null_dist_cached(): """Test that null_dist_cached creates and uses cache.""" with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/test_map.py b/tests/test_map.py index 364395e..72abd40 100644 --- a/tests/test_map.py +++ b/tests/test_map.py @@ -6,7 +6,8 @@ from sklearn.metrics import average_precision_score from copairs import compute -from copairs.map import average_precision +from copairs.map import map as map_module +from copairs.map import average_precision, mean_average_precision from tests.helpers import simulate_random_dframe from copairs.matching import UnpairedException from copairs.map.multilabel import average_precision as multilabel_average_precision @@ -246,6 +247,137 @@ def test_progress_bar_consistency(): assert with_pb.equals(no_pb), "The progress_bar argument changed results" +def test_average_precision_output_parity_across_worker_budgets(): + """AP output is unchanged by the similarity worker budget.""" + length = 20 + vocab_size = {"p": 5, "w": 3, "l": 4} + rng = np.random.default_rng(SEED) + meta = simulate_random_dframe(length, vocab_size, ["l"], ["p"], rng) + feats = rng.uniform(size=(len(meta), 8)) + kwargs = { + "meta": meta, + "feats": feats, + "pos_sameby": ["l"], + "pos_diffby": ["p"], + "neg_sameby": [], + "neg_diffby": ["l"], + "batch_size": 3, + "progress_bar": False, + } + + serial = average_precision(**kwargs, max_workers=1) + parallel = average_precision(**kwargs, max_workers=4) + + pd.testing.assert_frame_equal(serial, parallel, check_exact=True) + + +def test_multilabel_worker_budget_output_parity(): + """Multilabel AP output is unchanged by the similarity worker budget.""" + rng = np.random.default_rng(SEED) + meta = simulate_random_dframe(20, {"p": 3, "w": 5, "l": 4}, ["l"], [], rng) + meta = meta.groupby(["p", "w"])["l"].unique().reset_index() + feats = rng.uniform(size=(len(meta), 8)) + kwargs = { + "meta": meta, + "feats": feats, + "pos_sameby": ["l"], + "pos_diffby": [], + "neg_sameby": [], + "neg_diffby": ["l"], + "multilabel_col": "l", + "batch_size": 3, + "progress_bar": False, + } + + serial = multilabel_average_precision(**kwargs, max_workers=1) + parallel = multilabel_average_precision(**kwargs, max_workers=4) + + pd.testing.assert_frame_equal(serial, parallel, check_exact=True) + + +def test_mean_average_precision_propagates_worker_budget(monkeypatch, tmp_path): + """The public mAP pipeline applies its budget to native null generation.""" + ap_scores = pd.DataFrame( + { + "treatment": ["a", "a", "b", "b"], + "average_precision": [0.8, 0.7, 0.6, 0.5], + "normalized_average_precision": [0.7, 0.6, 0.5, 0.4], + "n_pos_pairs": [1, 1, 2, 2], + "n_total_pairs": [5, 5, 7, 7], + } + ) + observed_budgets = [] + aggregation_budgets = [] + get_null_dists = compute.get_null_dists + silent_thread_map = map_module.silent_thread_map + + def recording_get_null_dists(*args, **kwargs): + observed_budgets.append(kwargs.get("max_workers")) + return get_null_dists(*args, **kwargs) + + def recording_silent_thread_map(*args, **kwargs): + aggregation_budgets.append(kwargs.get("max_workers")) + return silent_thread_map(*args, **kwargs) + + monkeypatch.setattr(compute, "get_null_dists", recording_get_null_dists) + monkeypatch.setattr(map_module, "silent_thread_map", recording_silent_thread_map) + mean_average_precision( + ap_scores, + sameby=["treatment"], + null_size=20, + threshold=0.05, + seed=42, + progress_bar=False, + max_workers=1, + cache_dir=tmp_path, + ) + + assert observed_budgets == [1] + assert aggregation_budgets == [1] + + +def test_map_worker_environment_constrains_both_pools(monkeypatch, tmp_path): + """The environment budget constrains null generation and mAP aggregation.""" + ap_scores = pd.DataFrame( + { + "treatment": ["a", "a", "b", "b"], + "average_precision": [0.8, 0.7, 0.6, 0.5], + "normalized_average_precision": [0.7, 0.6, 0.5, 0.4], + "n_pos_pairs": [1, 1, 2, 2], + "n_total_pairs": [5, 5, 7, 7], + } + ) + resolved_budgets = [] + aggregation_budgets = [] + resolve_max_workers = compute._resolve_max_workers + silent_thread_map = map_module.silent_thread_map + + def recording_resolve_max_workers(num_items, max_workers): + resolved = resolve_max_workers(num_items, max_workers) + resolved_budgets.append((num_items, max_workers, resolved)) + return resolved + + def recording_silent_thread_map(*args, **kwargs): + aggregation_budgets.append(kwargs.get("max_workers")) + return silent_thread_map(*args, **kwargs) + + monkeypatch.setenv("COPAIRS_MAX_WORKERS", "1") + monkeypatch.setattr(compute, "_resolve_max_workers", recording_resolve_max_workers) + monkeypatch.setattr(map_module, "silent_thread_map", recording_silent_thread_map) + mean_average_precision( + ap_scores, + sameby=["treatment"], + null_size=20, + threshold=0.05, + seed=42, + progress_bar=False, + cache_dir=tmp_path, + ) + + assert resolved_budgets == [(2, None, 1), (2, None, 1)] + assert aggregation_budgets == [1] + + def test_multilabel_has_normalized_ap(): """Test that multilabel AP includes normalized_average_precision column.""" length = 10