diff --git a/cyteonto/cyteonto.py b/cyteonto/cyteonto.py index 5f96e94..3358f3b 100644 --- a/cyteonto/cyteonto.py +++ b/cyteonto/cyteonto.py @@ -6,7 +6,7 @@ import uuid from datetime import datetime, timezone from pathlib import Path -from typing import Any, Sequence +from typing import Any, Literal, Sequence import numpy as np import pandas as pd # type: ignore @@ -152,7 +152,7 @@ def __init__( self.max_description_concurrency = max_description_concurrency self.use_pubmed_tool = use_pubmed_tool self.reasoning = reasoning - self.mapping = OntologyMapping(self.paths.ontology_csv) + self.mapping = OntologyMapping(self.paths.ontology_mapping_csv) self.usage = AgentUsage(agentName="CyteOnto", modelName=llm.model) self.model_pair_usage = ModelPairUsage(llm=llm, embedding=self.embedding) @@ -487,19 +487,28 @@ async def _embed_user_labels( # matching + scoring def _match( - self, query_embeddings: np.ndarray, min_similarity: float = 0.1 + self, + query_embeddings: np.ndarray, + min_similarity: float = 0.1, + # labels: list[str] | None = None ) -> list[tuple[str | None, float]]: onto_emb, onto_ids = self._load_ontology_embeddings() if query_embeddings.ndim == 1: query_embeddings = query_embeddings.reshape(1, -1) sims = cosine_similarity(query_embeddings, onto_emb) out: list[tuple[str | None, float]] = [] - for row in sims: + for i, row in enumerate(sims): idx = int(np.argmax(row)) + # nearest_5 = [ + # (str(onto_ids[int(j)]), float(row[int(j)])) for j in np.argsort(row)[::-1][:5] + # ] score = float(row[idx]) if score < min_similarity: out.append((None, score)) else: + # if labels is not None: + # logger.info(f"ontology match for {labels[i]!r}: {onto_ids[idx]!r} {score:.4f}") + # logger.info(f"nearest 5: {nearest_5}") out.append((str(onto_ids[idx]), score)) return out @@ -610,6 +619,7 @@ async def compare( metric_params: dict[str, Any] | None = None, min_match_similarity: float = 0.1, use_cache: bool = True, + compound_scoring: Literal["max", "hungarian_mean"] = "max", ) -> pd.DataFrame: """Compare author labels with one or more algorithm label sets. @@ -629,12 +639,21 @@ async def compare( unmatched against the ontology. use_cache: When ``True``, reuse on-disk embeddings and descriptions if the labels match. + compound_scoring: How to reduce the m×n part score matrix for compound + pairs. ``"max"`` takes the maximum entry; ``"hungarian_mean"`` uses + Hungarian assignment mean with a coverage penalty when m != n. Returns: DataFrame with one row per (algorithm, pair_index) combination. The ``run_id`` column is populated with the id used for this call (the caller-provided value or the auto-generated UUID). """ + if compound_scoring not in ("max", "hungarian_mean"): + raise ValueError( + f"compound_scoring must be 'max' or 'hungarian_mean', " + f"got {compound_scoring!r}" + ) + if run_id is None: run_id = f"run-{uuid.uuid4()}" logger.info( @@ -657,6 +676,17 @@ async def compare( raise ValueError("Algorithm name 'author' is reserved") seen.add(name) + author_labels = [ + lbl.lower() if not _is_empty(lbl) else lbl for lbl in author_labels + ] + algo_items = [ + ( + name, + [lbl.lower() if not _is_empty(lbl) else lbl for lbl in labels], + ) + for name, labels in algo_items + ] + logger.info( f"compare(run_id='{run_id}', author={len(author_labels)}, " f"algorithms={len(algo_items)}, metric='{metric}')" @@ -676,7 +706,12 @@ async def compare( identifier="author", use_cache=use_cache, ) - author_matches = self._match(author_emb, min_similarity=min_match_similarity) + + author_matches = self._match( + author_emb, + min_similarity=min_match_similarity, + # labels=author_parts_list + ) author_lookup = dict(zip(author_parts_list, author_matches)) author_empty = [_is_empty(lbl) for lbl in author_labels] similarity = self._ensure_similarity() @@ -700,7 +735,11 @@ async def compare( identifier=algo_name, use_cache=use_cache, ) - algo_matches = self._match(algo_emb, min_similarity=min_match_similarity) + algo_matches = self._match( + algo_emb, + min_similarity=min_match_similarity, + # labels=algo_parts_list + ) algo_lookup = dict(zip(algo_parts_list, algo_matches)) algo_empty = [_is_empty(lbl) for lbl in algo_labels] @@ -729,10 +768,13 @@ async def compare( if author_empty[i]: g_parts = algo_parts_map[g_lbl] g_part_sims = [algo_lookup[gp][1] for gp in g_parts] - g_part_ids = [ - algo_lookup[gp][0] for gp in g_parts if algo_lookup[gp][0] - ] - g_id_out = ";".join(dict.fromkeys(g_part_ids)) + g_ids = [algo_lookup[gp][0] or "" for gp in g_parts] + g_id_out = ";".join(g_ids) + g_emb_out: str | float = ( + ";".join(str(round(s, 4)) for s in g_part_sims) + if len(g_parts) > 1 + else round(g_part_sims[0], 4) + ) rows.append( { "run_id": run_id, @@ -744,12 +786,11 @@ async def compare( "author_ontology_name": "", "author_embedding_similarity": 0.0, "algorithm_ontology_id": g_id_out, - "algorithm_ontology_name": self._ontology_names_for_ids( - g_id_out - ), - "algorithm_embedding_similarity": round( - sum(g_part_sims) / len(g_part_sims), 4 + "algorithm_ontology_name": ";".join( + self._ontology_name_for_id(oid) if oid else "" + for oid in g_ids ), + "algorithm_embedding_similarity": g_emb_out, "cytescore_similarity": 0.0, "similarity_method": "empty", } @@ -757,12 +798,13 @@ async def compare( else: a_parts = author_parts_map[a_lbl] a_part_sims = [author_lookup[ap][1] for ap in a_parts] - a_part_ids = [ - author_lookup[ap][0] - for ap in a_parts - if author_lookup[ap][0] - ] - a_id_out = ";".join(dict.fromkeys(a_part_ids)) + a_ids = [author_lookup[ap][0] or "" for ap in a_parts] + a_id_out = ";".join(a_ids) + a_emb_out: str | float = ( + ";".join(str(round(s, 4)) for s in a_part_sims) + if len(a_parts) > 1 + else round(a_part_sims[0], 4) + ) rows.append( { "run_id": run_id, @@ -771,12 +813,11 @@ async def compare( "author_label": a_lbl, "algorithm_label": g_lbl, "author_ontology_id": a_id_out, - "author_ontology_name": self._ontology_names_for_ids( - a_id_out - ), - "author_embedding_similarity": round( - sum(a_part_sims) / len(a_part_sims), 4 + "author_ontology_name": ";".join( + self._ontology_name_for_id(oid) if oid else "" + for oid in a_ids ), + "author_embedding_similarity": a_emb_out, "algorithm_ontology_id": "", "algorithm_ontology_name": "", "algorithm_embedding_similarity": 0.0, @@ -792,10 +833,31 @@ async def compare( a_part_sims = [author_lookup[ap][1] for ap in a_parts] g_part_sims = [algo_lookup[gp][1] for gp in g_parts] - a_sim_mean = sum(a_part_sims) / len(a_part_sims) - g_sim_mean = sum(g_part_sims) / len(g_part_sims) is_compound = m > 1 or n > 1 + a_ids = [author_lookup[ap][0] or "" for ap in a_parts] + g_ids = [algo_lookup[gp][0] or "" for gp in g_parts] + a_id_out = ";".join(a_ids) + g_id_out = ";".join(g_ids) + a_name_out = ";".join( + self._ontology_name_for_id(oid) if oid else "" for oid in a_ids + ) + g_name_out = ";".join( + self._ontology_name_for_id(oid) if oid else "" for oid in g_ids + ) + if len(a_parts) > 1: + a_emb_out: str | float = ";".join( + str(round(s, 4)) for s in a_part_sims + ) + else: + a_emb_out = round(a_part_sims[0], 4) + if len(g_parts) > 1: + g_emb_out: str | float = ";".join( + str(round(s, 4)) for s in g_part_sims + ) + else: + g_emb_out = round(g_part_sims[0], 4) + if not is_compound: a_id, _ = author_lookup[a_parts[0]] g_id, _ = algo_lookup[g_parts[0]] @@ -806,19 +868,12 @@ async def compare( if a_id and g_id else 0.0 ) - a_id_out = a_id or "" - g_id_out = g_id or "" method = self._method_for(a_id, g_id, hier) else: - # Hungarian match mean on an m x n score matrix S. Pick one-to-one - # assignments that maximize total score, then average those k pairs. - # When m != n, multiply by coverage = min(m,n) / max(m,n). - # - # 2x2 same compound (diagonal ~1, off-diagonal ~0): - # cytescore_similarity ~ 1.0 - # - # 3x2 partial overlap (author AT2+plasma+NK, algo AT2+T cell): - # match_mean ~ 0.53, coverage 2/3, cytescore_similarity ~ 0.35 + # Build an m x n part score matrix S, then reduce by compound_scoring: + # - max: highest entry in S + # - hungarian_mean: one-to-one assignment mean; when m != n, + # multiply by coverage = min(m,n) / max(m,n) score_matrix = np.zeros((m, n), dtype=np.float64) for ai, ap in enumerate(a_parts): a_id, _ = author_lookup[ap] @@ -835,19 +890,11 @@ async def compare( else 0.0 ) - hier, assignments = _hungarian_match_mean(score_matrix) + if compound_scoring == "max": + hier = float(score_matrix.max()) if score_matrix.size else 0.0 + else: + hier, _ = _hungarian_match_mean(score_matrix) method = "cytescore_compound" - a_matched_ids: list[str] = [] - g_matched_ids: list[str] = [] - for ai, aj in assignments: - a_id = author_lookup[a_parts[ai]][0] - g_id = algo_lookup[g_parts[aj]][0] - if a_id: - a_matched_ids.append(a_id) - if g_id: - g_matched_ids.append(g_id) - a_id_out = ";".join(dict.fromkeys(a_matched_ids)) - g_id_out = ";".join(dict.fromkeys(g_matched_ids)) rows.append( { @@ -857,13 +904,11 @@ async def compare( "author_label": a_lbl, "algorithm_label": g_lbl, "author_ontology_id": a_id_out, - "author_ontology_name": self._ontology_names_for_ids(a_id_out), - "author_embedding_similarity": round(a_sim_mean, 4), + "author_ontology_name": a_name_out, + "author_embedding_similarity": a_emb_out, "algorithm_ontology_id": g_id_out, - "algorithm_ontology_name": self._ontology_names_for_ids( - g_id_out - ), - "algorithm_embedding_similarity": round(g_sim_mean, 4), + "algorithm_ontology_name": g_name_out, + "algorithm_embedding_similarity": g_emb_out, "cytescore_similarity": round(hier, 4), "similarity_method": method, } @@ -887,6 +932,7 @@ async def compare_anndata( metric: str = "cosine_kernel", metric_params: dict[str, Any] | None = None, use_cache: bool = True, + compound_scoring: Literal["max", "hungarian_mean"] = "max", ) -> pd.DataFrame: """Pull algorithm labels out of ``adata.obs`` and delegate to ``compare``. @@ -897,6 +943,10 @@ async def compare_anndata( if len(names) != len(target_columns): raise ValueError("algorithm_names length must match target_columns length") + author_labels = [ + lbl.lower() if not _is_empty(lbl) else lbl for lbl in author_labels + ] + collected: list[tuple[str, list[str]]] = [] for adata in anndata_objects: if author_column not in adata.obs: @@ -908,7 +958,16 @@ async def compare_anndata( if col not in adata.obs: logger.warning(f"Column '{col}' missing from AnnData; skipping") continue - collected.append((name, adata.obs[col].astype(str).tolist())) + algo_labels = adata.obs[col].astype(str).tolist() + collected.append( + ( + name, + [ + lbl.lower() if not _is_empty(lbl) else lbl + for lbl in algo_labels + ], + ) + ) return await self.compare( author_labels=author_labels, @@ -917,6 +976,7 @@ async def compare_anndata( metric=metric, metric_params=metric_params, use_cache=use_cache, + compound_scoring=compound_scoring, ) # cache utilities diff --git a/cyteonto/describe.py b/cyteonto/describe.py index 380c103..c6320bb 100644 --- a/cyteonto/describe.py +++ b/cyteonto/describe.py @@ -131,59 +131,66 @@ def _build_prompt(label: str, use_pubmed: bool) -> str: if use_pubmed else "Do not call any tools." ) - return dedent( - f""" - Task: produce a structured description of one cell type. - - Input label: - {label} - - The input label may be a single name, an abbreviation, or multiple - synonyms separated by semicolons. Treat them as referring to the same - cell type. If the label is ambiguous, choose the most likely - interpretation in a standard mammalian reference context and commit - to it throughout the response. - - {tool_block} - - Fill every field below. Write each field in third person, factual, - present tense. Do not add text outside the structured object. Do not - repeat the label text verbatim inside the other fields unless it is - the only correct wording. - - initialLabel - Copy the input label verbatim, including punctuation and - semicolons. Do not normalise case, expand abbreviations, or - strip synonyms. - - descriptiveName - One noun phrase, at most 12 words. A precise name that identifies - the cell type without relying on the initial label. Prefer lineage - and location when available (for example "tissue-resident alveolar - macrophage"). - - function - One to three sentences, at most 60 words total. State the primary - biological role of the cell. Mention key effector mechanisms only - when they are defining for this cell type. - - diseaseRelevance - One to three sentences, at most 60 words total. State conditions - in which this cell type is pathologically involved, depleted, - expanded, or is a therapeutic target. If the cell type has no - well-established disease relevance, write "Not established". - - developmentalStage - One sentence, at most 30 words. State the lineage, progenitor, or - developmental window the cell belongs to. Use "Terminally - differentiated" when applicable. Use "Not established" when the - stage is unknown. - - Do not include marker genes, tissue lists, species ranges, or - citations inside any field. Do not hedge with phrases such as - "it is believed that" or "some studies suggest". State facts directly. - """ - ).strip() + return ( + dedent( + """ + Task: produce a structured description of one cell type. + + Input label: + {label} + + The input label may be a single name, an abbreviation, or multiple + synonyms separated by semicolons. Treat them as referring to the same + cell type. If the label is ambiguous, choose the most likely + interpretation in a standard mammalian reference context and commit + to it throughout the response. + + {tool_block} + + Fill every field below. Write each field in third person, factual, + present tense. Do not add text outside the structured object. Do not + repeat the label text verbatim inside the other fields unless it is + the only correct wording. + + initialLabel + Copy the input label verbatim, including punctuation and + semicolons. Do not normalise case, expand abbreviations, or + strip synonyms. + + descriptiveName + One noun phrase, at most 12 words. A precise name that identifies + the cell type without relying on the initial label. Prefer lineage + and location when available (for example "tissue-resident alveolar + macrophage"). + + function + One to three sentences, at most 60 words total. State the primary + biological role of the cell. Mention key effector mechanisms only + when they are defining for this cell type. + + diseaseRelevance + One to three sentences, at most 60 words total. State conditions + in which this cell type is pathologically involved, depleted, + expanded, or is a therapeutic target. If the cell type has no + well-established disease relevance, write "Not established". + + developmentalStage + One sentence, at most 30 words. State the lineage, progenitor, or + developmental window the cell belongs to. Use "Terminally + differentiated" when applicable. Use "Not established" when the + stage is unknown. + + Do not include marker genes, tissue lists, species ranges, or + citations inside any field. Do not hedge with phrases such as + "it is believed that" or "some studies suggest". State facts directly. + """ + ) + .format( + label=label, + tool_block=tool_block, + ) + .strip() + ) def _build_descriptor_agent( @@ -381,51 +388,56 @@ def _normalize_decomposition( return LabelDecomposition.single(label) if not output.isCompound: return LabelDecomposition.single(label) - cleaned = [p.strip() for p in output.parts if p and p.strip()] + cleaned = [p.strip().lower() for p in output.parts if p and p.strip()] if len(cleaned) < 2: return LabelDecomposition.single(label) return LabelDecomposition(initialLabel=label, isCompound=True, parts=cleaned) def _build_decompose_prompt(label: str) -> str: - return dedent( - f""" - Task: decide whether a cell-type annotation label refers to one cell - type or multiple distinct cell types. - - Input label: - {label} - - Context: labels come from single-cell RNA-seq cluster annotations. - Some labels name a single cell type; others name a mixture such as a - doublet or mixed population. - - Rules: - - Semicolon-separated text lists synonyms for one cell type. Do not - split those. - - Commas and hyphens inside a single cell-type name do not indicate - multiple types. - - When the label names multiple distinct cell types, set isCompound to - true and list each type in parts as a short sanitized name suitable - for ontology lookup. Remove mixture qualifiers such as "doublet", - "mixed population", or "contamination" from part names. - - When the label names one cell type, set isCompound to false and set - parts to a list containing only the original label verbatim. - - initialLabel - Copy the input label exactly as given. - - isCompound - true when the label names multiple distinct cell types; false - otherwise. - - parts - When isCompound is false, must contain exactly one element equal - to initialLabel. - When isCompound is true, must contain two or more sanitized - cell-type names. + return ( + dedent( + """ + Task: decide whether a cell-type annotation label refers to one cell + type or multiple distinct cell types. + + Input label: + {label} + + Context: labels come from single-cell RNA-seq cluster annotations. + Some labels name a single cell type; others name a mixture such as a + doublet or mixed population. + + Rules: + - Semicolon-separated text lists synonyms for one cell type. Do not + split those. + - Commas and hyphens inside a single cell-type name do not indicate + multiple types. + - When the label names multiple distinct cell types, set isCompound to + true and list each type in parts as a short sanitized name suitable + for ontology lookup. Remove mixture qualifiers such as "doublet", + "mixed population", or "contamination" from part names. + - When isCompound is true, return every part name in lowercase. + - When the label names one cell type, set isCompound to false and set + parts to a list containing only the original label verbatim. + + initialLabel + Copy the input label exactly as given. + + isCompound + true when the label names multiple distinct cell types; false + otherwise. + + parts + When isCompound is false, must contain exactly one element equal + to initialLabel. + When isCompound is true, must contain two or more sanitized + cell-type names. """ - ).strip() + ) + .format(label=label) + .strip() + ) def _build_decompose_agent(base_agent: Agent, reasoning: bool = False) -> Agent: diff --git a/cyteonto/ontology.py b/cyteonto/ontology.py index 95bde3e..ae988ff 100644 --- a/cyteonto/ontology.py +++ b/cyteonto/ontology.py @@ -26,12 +26,39 @@ def load(self) -> bool: return True try: df = pd.read_csv(self.csv_path) + if "label_normalized" not in df.columns: + df["label_normalized"] = df["label"].astype(str).str.lower() + else: + df["label_normalized"] = df["label_normalized"].astype(str).str.lower() + dup_mask = df.duplicated( + subset=["ontology_id", "label_normalized"], keep=False + ) + if dup_mask.any(): + for (oid, norm), grp in df[dup_mask].groupby( + ["ontology_id", "label_normalized"] + ): + originals = grp["label"].astype(str).tolist() + logger.warning( + f"Normalized label collision for {oid} {norm!r}: " + f"original labels {originals}; keeping {originals[0]!r}" + ) + before = len(df) + df = df.drop_duplicates( + subset=["ontology_id", "label_normalized"], keep="first" + ) + dropped = before - len(df) + if dropped: + logger.info( + f"Dropped {dropped} duplicate ontology rows " + "(same ontology_id and label_normalized)" + ) self._df = df for _, row in df.iterrows(): - label = str(row["label"]) + norm = str(row["label_normalized"]) + original = str(row["label"]) oid = str(row["ontology_id"]) - self._id_to_labels.setdefault(oid, []).append(label) - self._label_to_id.setdefault(label, oid) + self._id_to_labels.setdefault(oid, []).append(original) + self._label_to_id.setdefault(norm, oid) self._loaded = True logger.info(f"Loaded {len(df)} ontology mappings from {self.csv_path}") return True @@ -49,7 +76,7 @@ def df(self) -> pd.DataFrame: def label_to_id(self, label: str) -> str | None: if not self._loaded: self.load() - return self._label_to_id.get(label) + return self._label_to_id.get(label.lower()) def labels_for_id(self, ontology_id: str) -> list[str]: if not self._loaded: @@ -61,8 +88,12 @@ def ids_and_joined_labels(self) -> tuple[list[str], list[str]]: if not self._loaded: self.load() assert self._df is not None - grouped = self._df.groupby("ontology_id")["label"].apply(";".join).reset_index() - return grouped["ontology_id"].tolist(), grouped["label"].tolist() + grouped = ( + self._df.groupby("ontology_id")["label_normalized"] + .apply(lambda s: ";".join(dict.fromkeys(s.astype(str)))) + .reset_index() + ) + return grouped["ontology_id"].tolist(), grouped["label_normalized"].tolist() class OntologySimilarity: diff --git a/cyteonto/paths.py b/cyteonto/paths.py index f403d6d..363ee97 100644 --- a/cyteonto/paths.py +++ b/cyteonto/paths.py @@ -19,6 +19,7 @@ class PathConfig: Layout produced under ``data_dir``:: cell_ontology/cell_to_cell_ontology.csv + cell_ontology/cell_to_cell_ontology_enriched.csv cell_ontology/cl.owl embedding/cell_ontology/embeddings__.npz embedding/descriptions/descriptions_.json @@ -43,6 +44,18 @@ def __init__( def ontology_csv(self) -> Path: return self.data_dir / "cell_ontology" / "cell_to_cell_ontology.csv" + @property + def ontology_enriched_csv(self) -> Path: + return self.data_dir / "cell_ontology" / "cell_to_cell_ontology_enriched.csv" + + @property + def ontology_mapping_csv(self) -> Path: + """CSV used for lookups: enriched file when present, else shipped original.""" + enriched = self.ontology_enriched_csv + if enriched.exists(): + return enriched + return self.ontology_csv + @property def ontology_owl(self) -> Path: return self.data_dir / "cell_ontology" / "cl.owl" @@ -113,5 +126,6 @@ def run_kind_dirs(self, run_id: str, kind: str) -> list[Path]: def core_files_present(self) -> dict[str, bool]: return { "ontology_csv": self.ontology_csv.exists(), + "ontology_enriched_csv": self.ontology_enriched_csv.exists(), "ontology_owl": self.ontology_owl.exists(), } diff --git a/cyteonto/setup.py b/cyteonto/setup.py index 9e8acab..32ef05e 100644 --- a/cyteonto/setup.py +++ b/cyteonto/setup.py @@ -9,6 +9,7 @@ import sys from pathlib import Path +import pandas as pd # type: ignore import requests from tqdm.auto import tqdm # type: ignore @@ -45,6 +46,9 @@ BACKUP_EMBD_KEY = BACKUP_EMBEDDING.to_artifact_key() ONTOLOGY_CSV_URL: str = f"{BASE_URL}/cell_ontology/cell_to_cell_ontology.csv" +ONTOLOGY_ENRICHED_CSV_URL: str = ( + f"{BASE_URL}/cell_ontology/cell_to_cell_ontology_enriched.csv" +) ONTOLOGY_OWL_URL: str = f"{BASE_URL}/cell_ontology/cl.owl" @@ -59,24 +63,32 @@ def _embeddings_url(llm_key: ModelArtifactKey, embd_key: ModelArtifactKey) -> st ) -def _ontology_artifact_targets(paths: PathConfig) -> list[tuple[str, Path]]: - """Primary and backup description + embedding files on the CDN.""" - pairs: list[tuple[ModelArtifactKey, ModelArtifactKey]] = [ - (PRIMARY_LLM_KEY, PRIMARY_EMBD_KEY), - (BACKUP_LLM_KEY, BACKUP_EMBD_KEY), +def _primary_ontology_artifact_targets( + paths: PathConfig, +) -> list[tuple[str, Path]]: + return [ + ( + _descriptions_url(PRIMARY_LLM_KEY), + paths.ontology_descriptions(PRIMARY_LLM_KEY), + ), + ( + _embeddings_url(PRIMARY_LLM_KEY, PRIMARY_EMBD_KEY), + paths.ontology_embeddings(PRIMARY_LLM_KEY, PRIMARY_EMBD_KEY), + ), + ] + + +def _backup_ontology_artifact_targets(paths: PathConfig) -> list[tuple[str, Path]]: + return [ + ( + _descriptions_url(BACKUP_LLM_KEY), + paths.ontology_descriptions(BACKUP_LLM_KEY), + ), + ( + _embeddings_url(BACKUP_LLM_KEY, BACKUP_EMBD_KEY), + paths.ontology_embeddings(BACKUP_LLM_KEY, BACKUP_EMBD_KEY), + ), ] - targets: list[tuple[str, Path]] = [] - for llm_key, embd_key in pairs: - targets.append( - (_descriptions_url(llm_key), paths.ontology_descriptions(llm_key)) - ) - targets.append( - ( - _embeddings_url(llm_key, embd_key), - paths.ontology_embeddings(llm_key, embd_key), - ) - ) - return targets def _download(url: str, destination: Path, *, force: bool) -> bool: @@ -130,24 +142,79 @@ def main() -> int: paths = PathConfig(data_dir=args.data_dir) - targets = [ + required_targets = [ (ONTOLOGY_CSV_URL, paths.ontology_csv), (ONTOLOGY_OWL_URL, paths.ontology_owl), - *_ontology_artifact_targets(paths), + *_primary_ontology_artifact_targets(paths), ] + optional_targets = _backup_ontology_artifact_targets(paths) - failures: list[str] = [] - for url, dest in targets: + required_failures: list[str] = [] + for url, dest in required_targets: try: _download(url, dest, force=args.force) except Exception as exc: logger.error(f"Failed to download {dest.name}: {exc}") - failures.append(dest.name) + required_failures.append(dest.name) - if failures: - logger.error(f"{len(failures)} download(s) failed: {failures}") + if required_failures: + logger.error( + f"{len(required_failures)} required download(s) failed: {required_failures}" + ) return 1 + optional_failures: list[str] = [] + for url, dest in optional_targets: + try: + _download(url, dest, force=args.force) + except Exception as exc: + logger.warning(f"Optional backup download failed for {dest.name}: {exc}") + optional_failures.append(dest.name) + + if optional_failures: + logger.warning( + f"{len(optional_failures)} optional backup download(s) failed: " + f"{optional_failures}; continuing with primary model pair only" + ) + + enriched_path = paths.ontology_enriched_csv + try: + _download(ONTOLOGY_ENRICHED_CSV_URL, enriched_path, force=args.force) + except Exception as exc: + logger.warning( + f"Enriched ontology CSV not available from CDN ({exc}); " + "will build locally from the shipped original if needed" + ) + + csv_path = paths.ontology_csv + if not enriched_path.exists() and csv_path.exists(): + logger.info("Building enriched ontology CSV locally from shipped original") + df = pd.read_csv(csv_path) + df["label_normalized"] = df["label"].astype(str).str.lower() + dup_mask = df.duplicated(subset=["ontology_id", "label_normalized"], keep=False) + if dup_mask.any(): + for (oid, norm), grp in df[dup_mask].groupby( + ["ontology_id", "label_normalized"] + ): + originals = grp["label"].astype(str).tolist() + logger.warning( + f"Normalized label collision for {oid} {norm!r}: " + f"original labels {originals}; keeping {originals[0]!r}" + ) + before = len(df) + df = df.drop_duplicates( + subset=["ontology_id", "label_normalized"], keep="first" + ) + dropped = before - len(df) + if dropped: + logger.info( + f"Dropped {dropped} duplicate ontology rows " + "(same ontology_id and label_normalized)" + ) + enriched_path.parent.mkdir(parents=True, exist_ok=True) + df.to_csv(enriched_path, index=False) + logger.info(f"Wrote enriched ontology CSV: {enriched_path}") + logger.info(f"Setup complete. Data tree is ready under {paths.data_dir}") return 0 diff --git a/modal_app/models.py b/modal_app/models.py index 651b4c9..92a29bf 100644 --- a/modal_app/models.py +++ b/modal_app/models.py @@ -48,6 +48,7 @@ class CompareRequest(BaseModel): metric: str = "cosine_kernel" metricParams: dict[str, Any] | None = None minMatchSimilarity: float = Field(default=0.1, ge=0.0, le=1.0) + compoundScoring: Literal["max", "hungarian_mean"] = "max" useCache: bool = True diff --git a/modal_app/worker.py b/modal_app/worker.py index bd1cc81..0a352c7 100644 --- a/modal_app/worker.py +++ b/modal_app/worker.py @@ -217,6 +217,7 @@ async def run_compare_job(run_id: str, payload: dict[str, Any], volume) -> None: metric_params=payload.get("metricParams"), min_match_similarity=payload["minMatchSimilarity"], use_cache=payload["useCache"], + compound_scoring=payload.get("compoundScoring", "max"), ) csv_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/notebooks/quick_tutorial.ipynb b/notebooks/quick_tutorial.ipynb index 8dbb9d2..92e85c3 100644 --- a/notebooks/quick_tutorial.ipynb +++ b/notebooks/quick_tutorial.ipynb @@ -383,14 +383,269 @@ }, { "cell_type": "markdown", - "id": "d2a9d5a1", + "id": "84651030", "metadata": {}, "source": [ - "## Compound label scoring\n", + "# Compound Label Scoring\n", + "\n", + "## Default Compound Comparison is Max-Score\n", "\n", "Mixture labels (doublets, mixed populations) are decomposed by the LLM into cell-type parts before scoring. Each input pair still produces **one row**.\n", "\n", - "For compound pairs (`similarity_method == \"cytescore_compound\"`):\n", + "For compound pairs (`similarity_method == \"cytescore_compound\"`):\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "719ef730", + "metadata": {}, + "outputs": [], + "source": [ + "compound_scenarios = [\n", + " (\"2x1_coverage\", \"AT2 cell-plasma cell doublet\", \"Plasma cell\"),\n", + " (\"2x2_two_overlap\", \"AT2 cell-plasma cell doublet\", \"AT2 cell / Plasma cell\"),\n", + " (\"2x2_one_overlap\", \"AT2 cell-plasma cell doublet\", \"AT2 cell / T cell\"),\n", + " (\"2x2_zero_overlap\", \"AT2 cell-plasma cell doublet\", \"T cell / B cell\"),\n", + " (\"3x2_one_overlap\", \"AT2 cell / plasma cell / NK cell\", \"AT2 cell / T cell\"),\n", + " (\"simple_control\", \"T cell\", \"T cell\"),\n", + "]\n", + "\n", + "compound_author_labels = [author for _, author, _ in compound_scenarios]\n", + "compound_algorithm_labels = [algo for _, _, algo in compound_scenarios]\n", + "\n", + "compound_df_max = await cyto.compare(\n", + " author_labels=compound_author_labels,\n", + " algorithms={\"scenarios\": compound_algorithm_labels},\n", + " run_id=\"quick_tutorial_compound\",\n", + " metric=\"cosine_kernel\",\n", + ")\n", + "compound_df_max.insert(0, \"scenario\", [name for name, _, _ in compound_scenarios])" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "95fbf3c0", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
scenariorun_idalgorithmpair_indexauthor_labelalgorithm_labelauthor_ontology_idauthor_ontology_nameauthor_embedding_similarityalgorithm_ontology_idalgorithm_ontology_namealgorithm_embedding_similaritycytescore_similaritysimilarity_method
02x1_coveragequick_tutorial_compoundscenarios0at2 cell-plasma cell doubletplasma cellCL:0002063;CL:0000946pulmonary alveolar type 2 cell;antibody secret...0.9066;0.9281CL:0000786plasma cell0.91810.8975cytescore_compound
12x2_two_overlapquick_tutorial_compoundscenarios1at2 cell-plasma cell doubletat2 cell / plasma cellCL:0002063;CL:0000946pulmonary alveolar type 2 cell;antibody secret...0.9066;0.9281CL:0002063;CL:0000786pulmonary alveolar type 2 cell;plasma cell0.8747;0.91811.0000cytescore_compound
22x2_one_overlapquick_tutorial_compoundscenarios2at2 cell-plasma cell doubletat2 cell / t cellCL:0002063;CL:0000946pulmonary alveolar type 2 cell;antibody secret...0.9066;0.9281CL:0002063;CL:0000084pulmonary alveolar type 2 cell;T cell0.8747;0.88931.0000cytescore_compound
32x2_zero_overlapquick_tutorial_compoundscenarios3at2 cell-plasma cell doublett cell / b cellCL:0002063;CL:0000946pulmonary alveolar type 2 cell;antibody secret...0.9066;0.9281CL:0000084;CL:0000945T cell;lymphocyte of B lineage0.8893;0.89530.4046cytescore_compound
43x2_one_overlapquick_tutorial_compoundscenarios4at2 cell / plasma cell / nk cellat2 cell / t cellCL:0002063;CL:0000946;CL:0000623pulmonary alveolar type 2 cell;antibody secret...0.9024;0.9281;0.8382CL:0002063;CL:0000084pulmonary alveolar type 2 cell;T cell0.8747;0.88931.0000cytescore_compound
5simple_controlquick_tutorial_compoundscenarios5t cellt cellCL:0000084T cell0.8688CL:0000084T cell0.88931.0000cytescore
\n", + "
" + ], + "text/plain": [ + " scenario run_id algorithm pair_index \\\n", + "0 2x1_coverage quick_tutorial_compound scenarios 0 \n", + "1 2x2_two_overlap quick_tutorial_compound scenarios 1 \n", + "2 2x2_one_overlap quick_tutorial_compound scenarios 2 \n", + "3 2x2_zero_overlap quick_tutorial_compound scenarios 3 \n", + "4 3x2_one_overlap quick_tutorial_compound scenarios 4 \n", + "5 simple_control quick_tutorial_compound scenarios 5 \n", + "\n", + " author_label algorithm_label \\\n", + "0 at2 cell-plasma cell doublet plasma cell \n", + "1 at2 cell-plasma cell doublet at2 cell / plasma cell \n", + "2 at2 cell-plasma cell doublet at2 cell / t cell \n", + "3 at2 cell-plasma cell doublet t cell / b cell \n", + "4 at2 cell / plasma cell / nk cell at2 cell / t cell \n", + "5 t cell t cell \n", + "\n", + " author_ontology_id \\\n", + "0 CL:0002063;CL:0000946 \n", + "1 CL:0002063;CL:0000946 \n", + "2 CL:0002063;CL:0000946 \n", + "3 CL:0002063;CL:0000946 \n", + "4 CL:0002063;CL:0000946;CL:0000623 \n", + "5 CL:0000084 \n", + "\n", + " author_ontology_name \\\n", + "0 pulmonary alveolar type 2 cell;antibody secret... \n", + "1 pulmonary alveolar type 2 cell;antibody secret... \n", + "2 pulmonary alveolar type 2 cell;antibody secret... \n", + "3 pulmonary alveolar type 2 cell;antibody secret... \n", + "4 pulmonary alveolar type 2 cell;antibody secret... \n", + "5 T cell \n", + "\n", + " author_embedding_similarity algorithm_ontology_id \\\n", + "0 0.9066;0.9281 CL:0000786 \n", + "1 0.9066;0.9281 CL:0002063;CL:0000786 \n", + "2 0.9066;0.9281 CL:0002063;CL:0000084 \n", + "3 0.9066;0.9281 CL:0000084;CL:0000945 \n", + "4 0.9024;0.9281;0.8382 CL:0002063;CL:0000084 \n", + "5 0.8688 CL:0000084 \n", + "\n", + " algorithm_ontology_name algorithm_embedding_similarity \\\n", + "0 plasma cell 0.9181 \n", + "1 pulmonary alveolar type 2 cell;plasma cell 0.8747;0.9181 \n", + "2 pulmonary alveolar type 2 cell;T cell 0.8747;0.8893 \n", + "3 T cell;lymphocyte of B lineage 0.8893;0.8953 \n", + "4 pulmonary alveolar type 2 cell;T cell 0.8747;0.8893 \n", + "5 T cell 0.8893 \n", + "\n", + " cytescore_similarity similarity_method \n", + "0 0.8975 cytescore_compound \n", + "1 1.0000 cytescore_compound \n", + "2 1.0000 cytescore_compound \n", + "3 0.4046 cytescore_compound \n", + "4 1.0000 cytescore_compound \n", + "5 1.0000 cytescore " + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "compound_df_max" + ] + }, + { + "cell_type": "markdown", + "id": "d2a9d5a1", + "metadata": {}, + "source": [ + "## Compound label scoring - Hungarian Matrix Scoring\n", "\n", "1. Build an **m×n** cytescore matrix `S` between author parts (rows) and algorithm parts (columns).\n", "2. Run **max-weight bipartite matching** (Hungarian algorithm): pick `k = min(m, n)` one-to-one pairs that maximize total score.\n", @@ -578,18 +833,6 @@ "metadata": {}, "outputs": [], "source": [ - "compound_scenarios = [\n", - " (\"2x1_coverage\", \"AT2 cell-plasma cell doublet\", \"Plasma cell\"),\n", - " (\"2x2_two_overlap\", \"AT2 cell-plasma cell doublet\", \"AT2 cell / Plasma cell\"),\n", - " (\"2x2_one_overlap\", \"AT2 cell-plasma cell doublet\", \"AT2 cell / T cell\"),\n", - " (\"2x2_zero_overlap\", \"AT2 cell-plasma cell doublet\", \"T cell / B cell\"),\n", - " (\"3x2_one_overlap\", \"AT2 cell / plasma cell / NK cell\", \"AT2 cell / T cell\"),\n", - " (\"simple_control\", \"T cell\", \"T cell\"),\n", - "]\n", - "\n", - "compound_author_labels = [author for _, author, _ in compound_scenarios]\n", - "compound_algorithm_labels = [algo for _, _, algo in compound_scenarios]\n", - "\n", "compound_df = await cyto.compare(\n", " author_labels=compound_author_labels,\n", " algorithms={\"scenarios\": compound_algorithm_labels},\n", diff --git a/tests/test_cyteonto.py b/tests/test_cyteonto.py index b80a6b6..a039294 100644 --- a/tests/test_cyteonto.py +++ b/tests/test_cyteonto.py @@ -271,7 +271,18 @@ def test_compound_keeps_cleaned_parts(self): LabelDecomposition(initialLabel="A/B", isCompound=True, parts=["A", " B "]), ) assert out.isCompound is True - assert out.parts == ["A", "B"] + assert out.parts == ["a", "b"] + + def test_compound_lowercases_parts(self): + out = _normalize_decomposition( + "a/b", + LabelDecomposition( + initialLabel="a/b", + isCompound=True, + parts=["AT2 cell", "Plasma cell"], + ), + ) + assert out.parts == ["at2 cell", "plasma cell"] class TestDecomposeLabels: @@ -301,7 +312,7 @@ async def test_compound_input(self, mock_base_agent): mock_base_agent, "AT2 cell–plasma cell doublet" ) assert dec.isCompound is True - assert dec.parts == ["AT2 cell", "plasma cell"] + assert dec.parts == ["at2 cell", "plasma cell"] assert usage.requests == 1 @pytest.mark.asyncio @@ -430,8 +441,8 @@ async def test_hungarian_2x1_with_coverage(self): _mock_mapping_for_compare(inst) inst._resolve_label_parts = AsyncMock( side_effect=[ - {"A1/A2": ["A1", "A2"]}, - {"G1": ["G1"]}, + {"a1/a2": ["a1", "a2"]}, + {"g1": ["g1"]}, ] ) inst._embed_user_labels = AsyncMock( @@ -452,26 +463,76 @@ async def test_hungarian_2x1_with_coverage(self): author_labels=["A1/A2"], algorithms={"algo0": ["G1"]}, run_id="run-test", + compound_scoring="hungarian_mean", ) row = df.iloc[0] - assert row["author_label"] == "A1/A2" - assert row["algorithm_label"] == "G1" + assert row["author_label"] == "a1/a2" + assert row["algorithm_label"] == "g1" assert row["cytescore_similarity"] == 0.3 assert row["similarity_method"] == "cytescore_compound" - assert row["author_ontology_id"] == "CL:0000001" + assert row["author_ontology_id"] == "CL:0000001;CL:0000002" assert row["algorithm_ontology_id"] == "CL:0000001" - assert row["author_ontology_name"] == "T cell" + assert row["author_ontology_name"] == "T cell;NK cell" + assert row["author_embedding_similarity"] == "0.9;0.8" + assert row["algorithm_embedding_similarity"] == 0.9 assert sim.similarity.call_count == 2 + @pytest.mark.asyncio + async def test_max_2x1_default(self): + inst = object.__new__(CyteOnto) + _mock_mapping_for_compare(inst) + inst._resolve_label_parts = AsyncMock( + side_effect=[ + {"a1/a2": ["a1", "a2"]}, + {"g1": ["g1"]}, + ] + ) + inst._embed_user_labels = AsyncMock( + return_value=np.zeros((3, 2), dtype=np.float32) + ) + inst._match = Mock( + return_value=[ + ("CL:0000001", 0.9), + ("CL:0000002", 0.8), + ("CL:0000003", 0.7), + ] + ) + sim = Mock() + sim.similarity.side_effect = [0.6, 0.2] + inst._ensure_similarity = Mock(return_value=sim) + + df = await inst.compare( + author_labels=["A1/A2"], + algorithms={"algo0": ["G1"]}, + run_id="run-test", + ) + + row = df.iloc[0] + assert row["cytescore_similarity"] == 0.6 + assert row["similarity_method"] == "cytescore_compound" + + @pytest.mark.asyncio + async def test_invalid_compound_scoring_raises(self): + inst = object.__new__(CyteOnto) + _mock_mapping_for_compare(inst) + + with pytest.raises(ValueError, match="compound_scoring"): + await inst.compare( + author_labels=["T cell"], + algorithms={"algo0": ["B cell"]}, + run_id="run-test", + compound_scoring="mean", # type: ignore[arg-type] + ) + @pytest.mark.asyncio async def test_hungarian_2x2_same_compound(self): inst = object.__new__(CyteOnto) _mock_mapping_for_compare(inst) inst._resolve_label_parts = AsyncMock( side_effect=[ - {"A/B": ["A", "B"]}, - {"A/B": ["A", "B"]}, + {"a/b": ["a", "b"]}, + {"a/b": ["a", "b"]}, ] ) inst._embed_user_labels = AsyncMock( @@ -491,6 +552,7 @@ async def test_hungarian_2x2_same_compound(self): author_labels=["A/B"], algorithms={"algo0": ["A/B"]}, run_id="run-test", + compound_scoring="hungarian_mean", ) row = df.iloc[0] @@ -498,6 +560,8 @@ async def test_hungarian_2x2_same_compound(self): assert row["similarity_method"] == "cytescore_compound" assert row["author_ontology_id"] == "CL:0000001;CL:0000002" assert row["algorithm_ontology_id"] == "CL:0000001;CL:0000002" + assert row["author_embedding_similarity"] == "0.9;0.8" + assert row["algorithm_embedding_similarity"] == "0.85;0.75" assert row["pair_index"] == 0 @pytest.mark.asyncio @@ -505,17 +569,17 @@ async def test_compound_pair_index_not_shadowed(self): """Repeated author strings must keep distinct pair_index values.""" inst = object.__new__(CyteOnto) _mock_mapping_for_compare(inst) - doublet = "AT2 cell-plasma cell doublet" + doublet = "at2 cell-plasma cell doublet" inst._resolve_label_parts = AsyncMock( side_effect=[ { - doublet: ["AT2 cell", "plasma cell"], - "T cell": ["T cell"], + doublet: ["at2 cell", "plasma cell"], + "t cell": ["t cell"], }, { - "Plasma cell": ["Plasma cell"], - "AT2 cell / Plasma cell": ["AT2 cell", "Plasma cell"], - "T cell": ["T cell"], + "plasma cell": ["plasma cell"], + "at2 cell / plasma cell": ["at2 cell", "plasma cell"], + "t cell": ["t cell"], }, ] ) @@ -543,9 +607,9 @@ async def test_compound_pair_index_not_shadowed(self): assert df["pair_index"].tolist() == [0, 1, 2] assert df["algorithm_label"].tolist() == [ - "Plasma cell", - "AT2 cell / Plasma cell", - "T cell", + "plasma cell", + "at2 cell / plasma cell", + "t cell", ] @pytest.mark.asyncio @@ -554,8 +618,8 @@ async def test_non_compound_pair_unchanged(self): _mock_mapping_for_compare(inst) inst._resolve_label_parts = AsyncMock( side_effect=[ - {"T cell": ["T cell"]}, - {"B cell": ["B cell"]}, + {"t cell": ["t cell"]}, + {"b cell": ["b cell"]}, ] ) inst._embed_user_labels = AsyncMock( diff --git a/tests/test_ontology.py b/tests/test_ontology.py index 5f8b9f3..c4fb380 100644 --- a/tests/test_ontology.py +++ b/tests/test_ontology.py @@ -20,7 +20,7 @@ def test_load_missing_file(self, temp_dir): def test_label_to_id_and_back(self, sample_ontology_csv_file): mapping = OntologyMapping(sample_ontology_csv_file) mapping.load() - assert mapping.label_to_id("T cell") == "CL:0000001" + assert mapping.label_to_id("t cell") == "CL:0000001" assert mapping.label_to_id("not a cell") is None assert mapping.labels_for_id("CL:0000002") == ["B cell"] assert mapping.labels_for_id("CL:9999999") == [] @@ -39,8 +39,22 @@ def test_ids_and_joined_labels_groups_duplicates(self, temp_dir): joined_by_id = dict(zip(ids, joined)) assert set(ids) == {"CL:0000001", "CL:0000002"} - assert joined_by_id["CL:0000001"] == "T cell;T lymphocyte" - assert joined_by_id["CL:0000002"] == "B cell" + assert joined_by_id["CL:0000001"] == "t cell;t lymphocyte" + assert joined_by_id["CL:0000002"] == "b cell" + + def test_load_lowercases_mixed_case_labels(self, temp_dir): + data = { + "ontology_id": ["CL:0000786", "CL:0000623"], + "label": ["Plasma cell", "NK cell"], + } + csv_path = temp_dir / "mixed.csv" + pd.DataFrame(data).to_csv(csv_path, index=False) + + mapping = OntologyMapping(csv_path) + mapping.load() + + assert mapping.label_to_id("plasma cell") == "CL:0000786" + assert mapping.labels_for_id("CL:0000623") == ["NK cell"] class TestOntologySimilarityHelpers: diff --git a/tests/test_setup.py b/tests/test_setup.py new file mode 100644 index 0000000..78a2b8b --- /dev/null +++ b/tests/test_setup.py @@ -0,0 +1,82 @@ +"""Tests for cyteonto.setup.""" + +import sys + +import pandas as pd # type: ignore + +from cyteonto.paths import PathConfig +from cyteonto.setup import ( + BACKUP_EMBD_KEY, + BACKUP_LLM_KEY, + ONTOLOGY_ENRICHED_CSV_URL, + PRIMARY_LLM_KEY, + main, +) + + +def test_setup_builds_deduped_enriched_csv_when_not_on_cdn(temp_dir, monkeypatch): + paths = PathConfig(str(temp_dir)) + csv_path = paths.ontology_csv + enriched_path = paths.ontology_enriched_csv + csv_path.parent.mkdir(parents=True, exist_ok=True) + paths.ontology_owl.parent.mkdir(parents=True, exist_ok=True) + pd.DataFrame( + { + "ontology_id": ["CL:1", "CL:1", "CL:2"], + "label": ["Plasma cell", "PLASMA CELL", "NK cell"], + } + ).to_csv(csv_path, index=False) + + def fake_download(url: str, destination, *, force: bool) -> bool: + destination.parent.mkdir(parents=True, exist_ok=True) + if destination == csv_path: + return False + if url == ONTOLOGY_ENRICHED_CSV_URL: + raise OSError("not on CDN yet") + destination.touch() + return True + + monkeypatch.setattr("cyteonto.setup._download", fake_download) + monkeypatch.setattr(sys, "argv", ["setup.py", "--data-dir", str(temp_dir)]) + + assert main() == 0 + + original = pd.read_csv(csv_path) + assert original["label"].tolist() == ["Plasma cell", "PLASMA CELL", "NK cell"] + + enriched = pd.read_csv(enriched_path) + assert len(enriched) == 2 + assert enriched["label"].tolist() == ["Plasma cell", "NK cell"] + assert enriched["label_normalized"].tolist() == ["plasma cell", "nk cell"] + + +def test_setup_continues_when_backup_downloads_fail(temp_dir, monkeypatch): + paths = PathConfig(str(temp_dir)) + csv_path = paths.ontology_csv + csv_path.parent.mkdir(parents=True, exist_ok=True) + paths.ontology_owl.parent.mkdir(parents=True, exist_ok=True) + pd.DataFrame({"ontology_id": ["CL:1"], "label": ["T cell"]}).to_csv( + csv_path, index=False + ) + backup_desc = paths.ontology_descriptions(BACKUP_LLM_KEY) + backup_emb = paths.ontology_embeddings(BACKUP_LLM_KEY, BACKUP_EMBD_KEY) + primary_desc = paths.ontology_descriptions(PRIMARY_LLM_KEY) + + def fake_download(url: str, destination, *, force: bool) -> bool: + destination.parent.mkdir(parents=True, exist_ok=True) + if destination in (backup_desc, backup_emb): + raise OSError("backup unavailable") + if destination == csv_path: + return False + if url == ONTOLOGY_ENRICHED_CSV_URL: + raise OSError("not on CDN yet") + destination.touch() + return True + + monkeypatch.setattr("cyteonto.setup._download", fake_download) + monkeypatch.setattr(sys, "argv", ["setup.py", "--data-dir", str(temp_dir)]) + + assert main() == 0 + assert primary_desc.exists() + assert not backup_desc.exists() + assert not backup_emb.exists()