diff --git a/plans/2026-06-12-tunable-search-multi-hit.md b/plans/2026-06-12-tunable-search-multi-hit.md new file mode 100644 index 0000000..e5a1d17 --- /dev/null +++ b/plans/2026-06-12-tunable-search-multi-hit.md @@ -0,0 +1,326 @@ +# Plan: Tunable, multi-hit `Search` with PYOPSIN structure anchoring + +**Date:** 2026-06-12 +**Author:** Ali A Eftekhari + Claude Code +**Target module:** `src/provesid/search.py` (and helpers in `src/provesid/tools.py`) +**Status:** Proposed + +--- + +## 1. Problem statement + +`Search.search()` returns exactly **one** merged result per query, and that result is +not always correct — most often for **name** queries. Two structural causes: + +1. **Each source is collapsed to a single candidate before consensus runs.** + In `_candidates_from_name` (and the other resolvers) we take `rows[0]` (CompTox/ + PubChem/ChEMBL) or the first synonym hit (ChEBI), so the consensus engine + (`_compute_consensus`) only ever sees **one candidate per source**. If the correct + compound is hit #2 in a source while a wrong compound is hit #1, the correct + compound is invisible to scoring. A single wrong #1 hit can dominate the whole + answer. + +2. **The scoring/fuzzy machinery is hard-coded and inflexible.** rapidfuzz cut-off, + scorer choice, the `0.7` ZeroPM magic constant, the consensus-compatibility + threshold (`0.35`), and `_PRIORITY_TOLERANCE` (`0.05`) are all baked in. There is + no way for the caller to tune precision/recall, and no way to use a structure + oracle (PYOPSIN) to anchor an ambiguous name to a real molecule. + +### Design decisions (confirmed with user, 2026-06-12) + +- **Multi-hit output = distinct compounds, ranked.** When more than one hit is + requested, each output row is one *distinct candidate compound* (deduped across + sources by InChIKey skeleton / canonical structure), ranked by confidence, with a + new `hit_rank` column. Default remains **one row per query**. +- **PYOPSIN is opt-in, default off** (`use_opsin=False`) because it needs a Java + runtime. When enabled, IUPAC names are converted to SMILES offline and used as a + high-confidence structure anchor. +- **Full root-cause refactor**: collect **top-K candidates per source**, pool them, + cluster by structure, and let scoring choose among *all* candidates rather than one + pre-chosen winner per source. + +--- + +## 2. Current architecture (as-is) + +``` +search(queries) + └─ _resolve_single(q) + └─ _resolve_name(q) # (and _resolve_cas/_smiles/...) + ├─ _candidates_from_name(q) # {source: ONE candidate or None} + │ chebi -> rows[0] + │ comptox-> get_by_name or search_by_name[0] + │ pubchem-> search_by_name[0] + │ zeropm -> id_table -> one candidate + │ chembl -> search_by_name[0] + ├─ (fuzzy fallback) _fuzzy_name_candidates(q) + └─ _finalise_result(result, candidates, method, details) + ├─ _compute_consensus(candidates) # picks 1 anchor source + ├─ apply compatible candidates -> ONE merged result dict + ├─ normalize_structure / salts / confidence + └─ returns ONE row +``` + +Key data structure today: `candidates: Dict[str, Optional[dict]]` — **one candidate +per source key**. This is the bottleneck to break. + +Relevant source-method capabilities already available (currently under-used): +- `chebi.search_by_name(name, exact)`, `search_by_synonym(name, exact)` → `List[dict]` +- `comptox.search_by_name(name, exact, limit)` → `List[dict]`; `get_by_name` → one +- `pubchem.search_by_name(name, exact, limit)` → `List[dict]` +- `zeropm.query_similar_name(name, number_of_results=5, score_cutoff=80)` (fuzzy!) +- `chembl.search_by_name(name, limit)` → `List[dict]` + +--- + +## 3. Target architecture (to-be) + +Introduce a **candidate pool**: a flat list of candidate records (not one-per-source), +each tagged with its originating source and a per-source match score. Consensus and +ranking then operate over the pool and over **structure clusters**. + +``` +search(queries, n_hits=1, ...) + └─ _resolve_single(q) -> List[dict] # now returns a LIST of ranked hits + └─ _resolve_name(q) + ├─ pool = _candidate_pool_from_name(q) # list[candidate], top-K/source + │ + (opt) PYOPSIN anchor candidate (structure) + ├─ clusters = _cluster_candidates(pool) # group by structure identity + ├─ ranked = _rank_clusters(clusters, q) # confidence per cluster + └─ _finalise_hits(ranked, n_hits) # 1..N merged rows + hit_rank + └─ flatten list-of-lists -> DataFrame (n_hits column added) +``` + +### 3.1 Candidate record (extended) + +Reuse the existing `_make_candidate(...)` shape, adding two transient fields used only +during ranking (dropped from final output unless surfaced): + +- `query_match_score: float` — how well this candidate matches the *query itself* + (e.g. name similarity for name queries, 1.0 for exact CAS/InChIKey, Tanimoto for + structure fallback). This is the missing signal today. +- `_origin_rank: int` — the rank position the candidate had within its source's result + list (0-based). Used as a tie-breaker and for diagnostics. + +### 3.2 Structure clustering / dedup + +`_cluster_candidates(pool)` groups candidates that refer to the **same compound**: + +1. Primary key: full InChIKey when present. +2. Fallback key: 14-char InChIKey **skeleton** (connectivity) — merges + stereo/charge/isotope variants when `cluster_by_skeleton=True` (default True; + matches the spirit of `inchikey_skeleton`). +3. Fallback key: canonical SMILES. +4. Candidates with no structure at all form singleton clusters keyed by normalized + name (so name-only sources still contribute, but score lower). + +Each cluster collects: member candidates, the set of contributing sources, and the +best `query_match_score` among members. + +### 3.3 Cluster ranking → confidence + +Per cluster, compute a confidence in `[0, 1]` combining: + +- **base** from `_BASE_CONFIDENCE[match_method]` (existing table), +- **query agreement** = `query_match_score` (name/Tanimoto/exact signal), +- **cross-source consensus** = number/agreement of distinct sources backing the + cluster (re-using `_candidate_similarity` pairwise within and across clusters), and +- **PYOPSIN bonus**: when an OPSIN-derived structure anchor exists and a cluster + matches it (by InChIKey or skeleton), that cluster gets a strong confidence boost — + this is what fixes wrong name hits, because OPSIN gives a near-ground-truth structure + for parseable IUPAC names. + +Formula (generalizes the current `_compute_confidence`): + +``` +confidence = clamp( + base + * (w_q * query_match_score + (1 - w_q)) # query agreement term + * (0.5 + 0.5 * consensus_score) # existing consensus modulation + + opsin_match_bonus # additive, capped +) +``` + +Weights (`w_q`, `opsin_match_bonus`, etc.) become tunable attributes (see §4). + +The **best cluster** (rank 0) is the new "single best match" — equivalent to today's +output when `n_hits=1`, but chosen from the *full* pool rather than from pre-collapsed +per-source winners. + +--- + +## 4. New `Search` attributes (constructor) + +All keyword-only, all with defaults preserving current behavior. + +| Attribute | Type | Default | Purpose | +|---|---|---|---| +| `n_hits` | `int \| "all"` | `1` | Default number of ranked hits to return per query. Can be overridden per-call in `search()`. | +| `min_confidence` | `float` | `0.0` | Drop hits below this confidence before truncating to `n_hits`. | +| `use_opsin` | `bool` | `False` | Enable PYOPSIN IUPAC→structure anchoring for name queries (needs Java). | +| `opsin_jar_fpath` | `str` | `"default"` | Passed to `PYOPSIN`. | +| `top_k_per_source` | `int` | `5` | How many candidates to pull from each source before pooling. | +| `cluster_by_skeleton` | `bool` | `True` | Merge stereo/charge variants when clustering. | +| `fuzzy_score_cutoff` | `float` | `80.0` | rapidfuzz / ZeroPM score cut-off (0–100). Replaces hard-coded `80`/`0.7`. | +| `fuzzy_scorer` | `str` | `"WRatio"` | rapidfuzz scorer name (`WRatio`,`token_sort_ratio`,`partial_ratio`,…). | +| `consensus_compat_threshold` | `float` | `0.35` | Was hard-coded in `_candidate_compatible_with_consensus`. | +| `query_weight` (`w_q`) | `float` | `0.5` | Weight of query-agreement vs base in confidence. | +| `return_alternatives` | `bool` | `False` | When `n_hits=1`, still attach runner-up summaries in a new `alternatives` column (for inspection without exploding rows). | + +`search()` gains per-call overrides: + +```python +def search(self, queries, *, column=None, n_hits=None, min_confidence=None): ... +``` +(falling back to the instance attributes when `None`). + +Validation: `n_hits` must be `>= 1` or the literal `"all"`; raise `ValueError` +otherwise. `fuzzy_scorer` validated against a whitelist mapped to rapidfuzz functions. + +--- + +## 5. Output schema changes + +Add three columns to `OUTPUT_COLUMNS`: + +- `hit_rank` (int, 0 = best) — position of this row among the hits for its query. +- `n_source_support` (int) — how many distinct databases back this hit's cluster. +- `opsin_smiles` (str, nullable) — SMILES OPSIN produced for the query name (only + populated when `use_opsin=True` and parse succeeded), for traceability. + +Optional column (only when `return_alternatives=True` and `n_hits==1`): +- `alternatives` (list[dict]) — compact `{name, InChIKey, confidence, source}` for the + next few runner-up clusters. + +**Backward compatibility:** with defaults (`n_hits=1`), output is one row per query as +today, plus the three new always-present columns. Existing column order preserved; +new columns appended. Downstream code selecting by name keeps working. + +--- + +## 6. PYOPSIN integration + +- Add lazy `self._opsin` (a `PYOPSIN` instance) created on first name query when + `use_opsin=True`. Wrap construction/calls in try/except; on any failure (no Java, + py2opsin import error) log a warning **once** and disable OPSIN for the session + (`self._opsin_available = False`) so we don't spam. +- In `_candidate_pool_from_name`: + 1. `smiles = self._opsin.get_smiles(name)` (offline, fast). + 2. If non-empty, run it through `normalize_structure` to get canonical SMILES + + InChIKey, build an **anchor candidate** (`source="OPSIN"`, `query_match_score=1.0`, + `match_method="opsin"`), and add it to the pool. + 3. Additionally use the OPSIN InChIKey to do a direct **structure lookup** across + sources (`_candidate_pool_from_inchikey`-style), pulling in the *correct* compound + even when its name spelling differs — this is the main accuracy win. +- Add `"opsin": 0.97` to `_BASE_CONFIDENCE` (just below exact InChIKey). +- OPSIN also usable for `smiles`/`inchi` cross-checks later, but scope here = names. + +--- + +## 7. Implementation steps + +Ordered, each independently testable. + +1. **Introduce candidate-pool helpers in `tools.py`** (new, additive — don't break + existing `_compute_consensus` callers in `tools.py`'s `ids_from_*`): + - `_cluster_candidates(pool, by_skeleton) -> List[Cluster]` + - `_cluster_confidence(cluster, *, base, weights, opsin_anchor) -> float` + - `_rank_clusters(clusters, query, ...) -> List[(confidence, cluster)]` + - Keep these pure/stateless so they're unit-testable without DB clients. + +2. **Refactor name resolution to pool-based** in `search.py`: + - New `_candidate_pool_from_name(name) -> List[candidate]` that pulls `top_k_per_source` + from each source (using existing `limit`/`exact` params), tags each with + `query_match_score = _text_similarity(name, cand.name/synonyms)` and `_origin_rank`. + - Fold the existing fuzzy logic in: when exact pool is weak, widen with + `exact=False` + `fuzzy_score_cutoff`/`fuzzy_scorer` and ZeroPM + `query_similar_name(number_of_results=top_k_per_source, score_cutoff=fuzzy_score_cutoff)`. + - Add OPSIN anchor (§6). + +3. **Generalize `_finalise_result` → `_finalise_hits`**: + - Build merged result dict per cluster (reuse `_apply_candidate_to_result` over the + cluster's members, anchored on the cluster's best-structure member). + - Run `normalize_structure`, salt stripping, confidence per cluster. + - Set `hit_rank`, `n_source_support`, `opsin_smiles`. + - Return `List[dict]` (ranked, filtered by `min_confidence`, truncated to `n_hits`). + - Keep a thin `_finalise_result` wrapper returning `[0]` for resolvers not yet + migrated, so we can land name first. + +4. **Update `_resolve_single` + `search()`** to handle list-returning resolvers: + - `_resolve_single` returns `List[dict]`; `search()` flattens with `tqdm` over + queries and concatenates. Each query's hits stay contiguous and ordered. + - Apply per-call `n_hits`/`min_confidence` overrides; validate. + - DataFrame assembly: extra-column merge (from DataFrame/file input) must now + **broadcast** original row data across the multiple hit-rows of that query — join + on a per-query index rather than positional `reset_index`. (Important edge case: + today's positional concat assumes 1 row per query; multi-hit breaks that.) + +5. **Migrate remaining resolvers** to the pool model for consistency (CAS, SMILES, + InChI, InChIKey, DTXSID, formula). For exact-identifier types the pool is usually + a single cluster, so behavior is unchanged at `n_hits=1`; but `formula` (inherently + many hits) and `smiles` Tanimoto fallback benefit immediately from ranked multi-hit. + +6. **Constructor + validation** for all new attributes (§4); update class docstring and + the module-level usage examples. + +7. **Confidence table + weights**: add `opsin` base; expose weights; keep + `_compute_confidence` working for the single-cluster path. + +--- + +## 8. Testing plan + +Add to the existing test suite (mirror current test layout; check `tests/`). + +- **Unit (no DB, no Java):** + - `_cluster_candidates`: same-InChIKey merge; skeleton merge on/off; SMILES-only and + name-only fallbacks; empty pool. + - `_cluster_confidence` / ranking: ordering monotonic in query score and source + support; OPSIN anchor pushes the matching cluster to rank 0. + - `n_hits` validation; `"all"`; `min_confidence` filtering. +- **Integration (DB clients, marked like existing offline tests):** + - A known ambiguous name where today's code picks the wrong compound (capture a + real example from the user — see §10) → assert correct compound is rank 0 *or* + appears within top-`n_hits`. + - `n_hits="all"` returns >1 row for an ambiguous name; `hit_rank` contiguous 0..N. + - DataFrame input: extra columns correctly broadcast across multi-hit rows. +- **OPSIN (marked, skip if Java/py2opsin missing — reuse the existing `opsin` pytest + marker):** IUPAC name → correct InChIKey anchor; graceful no-op when OPSIN + unavailable (no exception, falls back to name matching). +- **Regression:** default `Search("name").search([...])` still yields exactly one row + per query; existing columns unchanged; `Search("cas")` outputs identical to before. + +--- + +## 9. Risks & mitigations + +| Risk | Mitigation | +|---|---| +| Multi-hit rows break downstream code expecting 1 row/query | Default `n_hits=1` unchanged; new columns appended, not inserted. Document the broadcast behavior for DataFrame inputs. | +| Pulling `top_k_per_source` from 5 sources slows batch runs | `top_k_per_source` default modest (5); clustering is O(pool²) but pool is tiny per query. Profile on a batch; keep `show_progress`. | +| OPSIN/Java absent in CI or user env | Opt-in (`use_opsin=False`); one-time warning + session disable; tests skip via marker. | +| Skeleton clustering over-merges distinct stereoisomers | `cluster_by_skeleton=True` default but configurable; `n_source_support` + full InChIKey preserved in output for the user to disambiguate. | +| Tunable weights make confidence non-comparable across configs | Document that `confidence` is only comparable within a fixed configuration; keep default weights == current behavior at `n_hits=1` as far as possible. | + +--- + +## 10. Open questions for the user (non-blocking) + +1. **A concrete failing example.** Please provide 1–3 name queries where the current + `search` picks the *wrong* compound (and what the right answer is). These become + regression tests and let us calibrate default weights against reality. +2. **Default `top_k_per_source`** — is 5 reasonable, or do you want it higher for + recall (at some speed cost)? +3. **`alternatives` column** — useful, or noise? (Off by default either way.) +4. Should the pool refactor also expose a **raw, unranked** debug accessor (e.g. + `Search.debug_pool(query)` returning every candidate + scores) to help you tune + weights interactively? Recommended; cheap to add. + +--- + +## 11. Out of scope (future work) + +- Vectorized Parquet + fingerprint Tanimoto search (already noted as future in + `_tanimoto_candidates`). +- Learning/auto-tuning the confidence weights from a labeled set. +- Applying OPSIN to non-name identifier types as a cross-validation oracle. diff --git a/plans/README.md b/plans/README.md new file mode 100644 index 0000000..bddb5e4 --- /dev/null +++ b/plans/README.md @@ -0,0 +1,2 @@ +# Plans +The folder contains the plan I have made in collaboration with Claude Code. These are temporarily stored here and will be removed or archived later. \ No newline at end of file diff --git a/scripts/validate_name_search.py b/scripts/validate_name_search.py new file mode 100644 index 0000000..52d791d --- /dev/null +++ b/scripts/validate_name_search.py @@ -0,0 +1,188 @@ +"""Validate name-based resolution against CASRN ground truth (CompTox). + +Strategy (suggested by the maintainer): + + CompTox carries both CASRN and PREFERRED_NAME for every chemical. A CASRN + lookup is (almost) always correct, so it serves as ground truth. We: + + 1. Sample N CompTox rows that have CASRN, PREFERRED_NAME and INCHIKEY. + 2. Treat each row's (DTXSID, INCHIKEY) as the truth for that CASRN. + 3. Resolve the *same* chemical by its PREFERRED_NAME through a CompTox-only + ``Search`` and check whether we recover the same compound. + 4. Report accuracy and dump the mismatches — those are the real + "wrong name hit" cases, ready to become regression tests. + +The Search is restricted to the CompTox source so the experiment isolates the +name-matching / ranking logic against the very table we are trying to +reproduce ("can you reproduce the CASRN table from the names?"). + +Usage:: + + uv run python scripts/validate_name_search.py --n 300 + uv run python scripts/validate_name_search.py --n 300 --tuned + uv run python scripts/validate_name_search.py --n 300 --out mismatches.csv +""" + +from __future__ import annotations + +import argparse +import sys + +import pandas as pd + +from provesid.comptox import CompToxID +from provesid.search import Search + + +def _skeleton(ik): + return ik[:14] if isinstance(ik, str) and len(ik) >= 14 else ik + + +def sample_rows(comptox: CompToxID, n: int) -> pd.DataFrame: + """Deterministically sample n rows spread across the CompTox table. + + Uses a ``rowid``-stride so the sample is reproducible without RNG and is + spread across the whole table rather than clustered alphabetically. + """ + total = comptox.conn.execute("SELECT COUNT(*) FROM chemicals").fetchone()[0] + step = max(1, total // (n * 4)) # over-sample; we filter nulls/dupes below + cur = comptox.conn.execute( + f""" + SELECT DTXSID, CASRN, PREFERRED_NAME, IUPAC_NAME, INCHIKEY, IDENTIFIER + FROM chemicals + WHERE CASRN IS NOT NULL AND CASRN != '' + AND PREFERRED_NAME IS NOT NULL AND PREFERRED_NAME != '' + AND INCHIKEY IS NOT NULL AND INCHIKEY != '' + AND (rowid % {step}) = 1 + LIMIT {n * 4} + """ + ) + rows = [dict(r) for r in cur.fetchall()] + return pd.DataFrame(rows) + + +def pick_query_name(row, name_source: str): + """Choose the query name for a row given the requested name source. + + Returns ``None`` when no suitable name exists (row is then skipped). + """ + if name_source == "preferred": + return row["PREFERRED_NAME"] + if name_source == "iupac": + return row.get("IUPAC_NAME") or None + if name_source == "synonym": + ident = row.get("IDENTIFIER") or "" + parts = [p.strip() for p in ident.split("|") if p.strip()] + pref = (row["PREFERRED_NAME"] or "").strip().lower() + cas = (row["CASRN"] or "").strip() + # Skip the CAS token and the preferred name; take a different synonym. + for p in parts: + if p == cas or p.lower() == pref: + continue + if any(ch.isalpha() for ch in p): # looks like a name, not a code + return p + return None + raise ValueError(name_source) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--n", type=int, default=300, help="number of chemicals to test") + ap.add_argument("--tuned", action="store_true", + help="enable fuzzy + OPSIN + larger top_k tuning") + ap.add_argument("--all-sources", action="store_true", + help="resolve through all 5 sources (real scenario where " + "cross-source disagreement produces wrong hits)") + ap.add_argument("--name-source", choices=["preferred", "iupac", "synonym"], + default="preferred", + help="which name to query by (default: preferred)") + ap.add_argument("--out", default=None, help="write mismatches to this CSV") + args = ap.parse_args(argv) + + comptox = CompToxID() + df = sample_rows(comptox, args.n) + + # Keep only CASRNs that map to exactly one chemical (the documented caveat). + cas_counts = df.groupby("CASRN")["DTXSID"].nunique() + unique_cas = set(cas_counts[cas_counts == 1].index) + df = df[df["CASRN"].isin(unique_cas)].drop_duplicates("CASRN") + + # Resolve the query name per the chosen source; drop rows without one. + df["_query"] = df.apply(lambda r: pick_query_name(r, args.name_source), axis=1) + df = df[df["_query"].notna()].drop_duplicates("_query").head(args.n) + df = df.reset_index(drop=True) + print(f"Testing {len(df)} chemicals | name-source={args.name_source} | " + f"sources={'all' if args.all_sources else 'comptox-only'} | " + f"{'tuned' if args.tuned else 'default'} config\n") + + kwargs = dict(identifier_type="name", show_progress=True) + if args.all_sources: + # Pass no clients so Search lazily builds all 5 from local data. + pass + else: + # CompTox-only: pass comptox, force the rest off. + kwargs.update(comptox=comptox, chebi=None, pubchem=None, + zeropm=None, chembl=None) + if args.tuned: + kwargs.update(fuzzy=True, use_opsin=True, top_k_per_source=10, + fuzzy_score_cutoff=70.0) + s = Search(**kwargs) + + res = s.search(df["_query"].tolist()) + + truth_dtxsid = df["DTXSID"].tolist() + truth_ik = df["INCHIKEY"].tolist() + + # Correctness is structure identity (skeleton), which is source-independent. + # Three outcomes per query: + # correct – returned a structure whose skeleton matches the truth + # wrong_hit – returned a *different* structure (the precision failure) + # not_found – returned no structure at all (a recall gap, not a wrong hit) + correct = wrong = not_found = 0 + wrong_hits = [] + for i, row in res.iterrows(): + got_ik = row.get("InChIKey") + truth_skel = _skeleton(truth_ik[i]) + if got_ik is None or (isinstance(got_ik, float)): + not_found += 1 + continue + if _skeleton(got_ik) == truth_skel: + correct += 1 + else: + wrong += 1 + wrong_hits.append({ + "query_name": df["_query"][i], + "preferred_name": df["PREFERRED_NAME"][i], + "truth_CASRN": df["CASRN"][i], + "truth_DTXSID": truth_dtxsid[i], + "truth_InChIKey": truth_ik[i], + "got_name": row.get("name"), + "got_InChIKey": got_ik, + "got_DTXSID": row.get("DTXSID"), + "confidence": row.get("confidence"), + "match_method": row.get("match_method"), + }) + + n = len(df) + print(f"\n correct (structure reproduced): {correct}/{n} ({correct/n:.1%})") + print(f" WRONG HIT (different structure): {wrong}/{n} ({wrong/n:.1%})") + print(f" not found (no structure): {not_found}/{n} ({not_found/n:.1%})") + found = correct + wrong + if found: + print(f" precision when a hit is returned: {correct}/{found} ({correct/found:.1%})") + + if wrong_hits: + wdf = pd.DataFrame(wrong_hits) + print(f"\nWRONG-HIT cases (returned the wrong compound) — regression candidates:\n") + with pd.option_context("display.max_colwidth", 34, "display.width", 220): + print(wdf[["query_name", "preferred_name", "got_name", + "match_method", "confidence"]].head(25).to_string(index=False)) + if args.out: + wdf.to_csv(args.out, index=False) + print(f"\nWrote {len(wdf)} wrong-hit cases to {args.out}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/provesid/search.py b/src/provesid/search.py index 52193d5..25a73ba 100644 --- a/src/provesid/search.py +++ b/src/provesid/search.py @@ -49,6 +49,7 @@ from .chebi import ChebiSDF from .chembl import CheMBL from .comptox import CompToxID +from .opsin import PYOPSIN from .pubchem import PubChemID from .zeropm import ZeroPM from .tools import ( @@ -111,6 +112,7 @@ "inchi": 0.95, "exact_cas": 0.90, "dtxsid": 0.90, + "opsin": 0.97, "exact_name": 0.80, "inchikey_skeleton": 0.75, "tanimoto": 0.0, # filled dynamically from Tanimoto score @@ -144,6 +146,9 @@ "match_score", "consensus_source", "source_match_scores", + "hit_rank", + "n_source_support", + "opsin_smiles", ] # Name-normalization: prefixes to strip before fuzzy matching. @@ -333,6 +338,18 @@ class Search: - **Salt/solvent stripping**: RDKit SaltRemover + largest-fragment picker (enabled with ``strip_salts=True``); ``parent_smiles`` and ``parent_inchikey`` populated in results. + - **Candidate pooling + structure clustering**: the top-``k`` candidates from + every source are pooled and clustered into distinct compounds by InChIKey + (skeleton-aware), so a wrong top hit from one source no longer dominates. + - **Query-aware ranking**: clusters are ranked by a confidence that combines + the method base, how well each candidate matches the *query itself* + (name similarity / Tanimoto), and cross-source support. + - **Multi-hit output**: ``n_hits`` returns the best ``N`` (or ``"all"``) + distinct compounds per query, ranked with a ``hit_rank`` column. Default + is one row per query. + - **PYOPSIN structure anchoring** (opt-in, ``use_opsin=True``): IUPAC names + are converted to SMILES offline and used as a high-confidence anchor; + requires a Java runtime. - **Traceability**: ``source_details`` field records which sources were queried, whether they matched, and which output fields they contributed. @@ -346,6 +363,16 @@ class Search: show_progress (bool): Display tqdm progress bar during batch queries. salt_smarts (list[str]): Additional SMARTS patterns to remove during salt stripping. + n_hits (int | str): Default hits to return per query (int or ``"all"``). + min_confidence (float): Confidence floor applied before truncation. + use_opsin (bool): Enable PYOPSIN IUPAC→structure anchoring (needs Java). + top_k_per_source (int): Candidates pulled per source before pooling. + cluster_by_skeleton (bool): Merge stereo/charge variants when clustering. + fuzzy_score_cutoff (float): Fuzzy score cut-off in [0, 100]. + fuzzy_scorer (str): rapidfuzz scorer name. + consensus_compat_threshold (float): Min similarity to merge with anchor. + query_weight (float): Weight of query agreement in the confidence score. + return_alternatives (bool): Attach runner-up summaries when ``n_hits=1``. Example:: @@ -357,6 +384,13 @@ class Search: s_fuzzy = Search("name", fuzzy=True) df = s_fuzzy.search(["asprin", "paracetamol"]) + + # Inspect every plausible interpretation of an ambiguous name + df = Search("name").search("xylene", n_hits="all") + print(df[["hit_rank", "name", "InChIKey", "confidence"]]) + + # Anchor IUPAC names to a real structure via OPSIN (needs Java) + df = Search("name", use_opsin=True).search("2-(acetyloxy)benzoic acid") """ SUPPORTED_TYPES: frozenset = frozenset( @@ -372,6 +406,12 @@ class Search: "chembl": "ChEMBL", } + # rapidfuzz scorer whitelist (name -> scorer callable resolved lazily). + _FUZZY_SCORERS: frozenset = frozenset( + ["WRatio", "ratio", "partial_ratio", "token_sort_ratio", + "token_set_ratio", "QRatio"] + ) + def __init__( self, identifier_type: str = "cas", @@ -382,6 +422,17 @@ def __init__( inchikey_skeleton: bool = False, show_progress: bool = True, salt_smarts: Optional[List[str]] = None, + n_hits: Union[int, str] = 1, + min_confidence: float = 0.0, + use_opsin: bool = False, + opsin_jar_fpath: str = "default", + top_k_per_source: int = 5, + cluster_by_skeleton: bool = True, + fuzzy_score_cutoff: float = 80.0, + fuzzy_scorer: str = "WRatio", + consensus_compat_threshold: float = 0.35, + query_weight: float = 0.5, + return_alternatives: bool = False, data_dir: Optional[Union[str, Path]] = None, redownload: bool = False, chebi: Optional[ChebiSDF] = None, @@ -409,6 +460,34 @@ def __init__( show_progress: Display a tqdm progress bar during batch queries. salt_smarts: Additional SMARTS patterns passed to :func:`strip_salts` when ``strip_salts=True``. + n_hits: Default number of ranked hits to return per query. Either a + positive integer or the literal ``"all"``. Defaults to ``1`` + (one row per query). Can be overridden per-call in + :meth:`search`. + min_confidence: Drop hits whose confidence is below this value + before truncating to ``n_hits``. Defaults to ``0.0``. + use_opsin: Enable PYOPSIN IUPAC-name → structure anchoring for name + queries. Requires a Java runtime; falls back to plain name + matching (with a one-time warning) when unavailable. Defaults + to ``False``. + opsin_jar_fpath: ``jar_fpath`` passed to :class:`~provesid.PYOPSIN`. + top_k_per_source: Number of candidate rows pulled from each source + before pooling / clustering. Defaults to ``5``. + cluster_by_skeleton: Merge stereo/charge/isotope variants when + clustering candidates by structure (14-char InChIKey skeleton). + Defaults to ``True``. + fuzzy_score_cutoff: rapidfuzz / ZeroPM fuzzy score cut-off in + [0, 100]. Defaults to ``80.0``. + fuzzy_scorer: rapidfuzz scorer name; one of ``WRatio``, ``ratio``, + ``partial_ratio``, ``token_sort_ratio``, ``token_set_ratio``, + ``QRatio``. Defaults to ``"WRatio"``. + consensus_compat_threshold: Minimum candidate similarity for a + candidate to be merged with the consensus anchor. Defaults to + ``0.35``. + query_weight: Weight (in [0, 1]) of the query-agreement term versus + the method base in the confidence formula. Defaults to ``0.5``. + return_alternatives: When ``n_hits == 1``, attach compact runner-up + summaries in an ``alternatives`` column. Defaults to ``False``. data_dir: Optional shared data root used when lazily initialising source clients. redownload: If True, lazily initialised source clients force a @@ -437,9 +516,32 @@ def __init__( self.inchikey_skeleton = inchikey_skeleton self.show_progress = show_progress self.salt_smarts: List[str] = list(salt_smarts or []) + + # Multi-hit / tuning attributes + self.n_hits = self._validate_n_hits(n_hits) + self.min_confidence = float(min_confidence) + self.use_opsin = bool(use_opsin) + self.opsin_jar_fpath = opsin_jar_fpath + self.top_k_per_source = max(1, int(top_k_per_source)) + self.cluster_by_skeleton = bool(cluster_by_skeleton) + self.fuzzy_score_cutoff = float(fuzzy_score_cutoff) + if fuzzy_scorer not in self._FUZZY_SCORERS: + raise ValueError( + f"fuzzy_scorer must be one of {sorted(self._FUZZY_SCORERS)}, " + f"got {fuzzy_scorer!r}" + ) + self.fuzzy_scorer = fuzzy_scorer + self.consensus_compat_threshold = float(consensus_compat_threshold) + self.query_weight = float(query_weight) + self.return_alternatives = bool(return_alternatives) + self.data_dir = str(data_dir) if data_dir is not None else None self.redownload = redownload + # OPSIN client — created lazily; disabled for the session on failure. + self._opsin: Optional[PYOPSIN] = None + self._opsin_available: bool = use_opsin + # Client references — may be None until _ensure_clients() is called. self._chebi = chebi self._comptox = comptox @@ -483,6 +585,79 @@ def _ensure_clients(self) -> None: self._clients_initialized = True + @staticmethod + def _validate_n_hits(n_hits: Union[int, str]) -> Union[int, str]: + """Validate and normalise the ``n_hits`` argument. + + Args: + n_hits: Either a positive integer or the literal ``"all"``. + + Returns: + ``"all"`` or a positive ``int``. + + Raises: + ValueError: If ``n_hits`` is neither ``"all"`` nor a positive int. + """ + if isinstance(n_hits, str): + if n_hits.lower() == "all": + return "all" + raise ValueError(f"n_hits string must be 'all', got {n_hits!r}") + if isinstance(n_hits, bool) or not isinstance(n_hits, int) or n_hits < 1: + raise ValueError(f"n_hits must be a positive int or 'all', got {n_hits!r}") + return n_hits + + def _get_opsin(self) -> Optional[PYOPSIN]: + """Lazily create the PYOPSIN client; disable for the session on failure. + + Returns: + A :class:`~provesid.PYOPSIN` instance, or ``None`` when OPSIN is + disabled or unavailable (e.g. no Java runtime). + """ + if not self._opsin_available: + return None + if self._opsin is None: + try: + self._opsin = PYOPSIN(jar_fpath=self.opsin_jar_fpath) + except Exception as exc: # pragma: no cover - environment dependent + log.warning( + "PYOPSIN unavailable (%s); disabling OPSIN anchoring for this " + "session.", exc, + ) + self._opsin_available = False + return None + return self._opsin + + def _opsin_anchor(self, name: str) -> Optional[Dict[str, Any]]: + """Convert an IUPAC name to a normalised structure anchor via PYOPSIN. + + Args: + name: Chemical (IUPAC) name. + + Returns: + Dict with keys ``smiles``, ``canonical_smiles``, ``inchikey`` when + OPSIN parsed the name, else ``None``. + """ + opsin = self._get_opsin() + if opsin is None: + return None + try: + smiles = opsin.get_smiles(name) + except Exception as exc: # pragma: no cover - environment dependent + log.warning( + "PYOPSIN parse failed for %r (%s); disabling OPSIN for session.", + name, exc, + ) + self._opsin_available = False + return None + if _is_missing(smiles) or not str(smiles).strip(): + return None + norm = normalize_structure(str(smiles)) + return { + "smiles": str(smiles), + "canonical_smiles": norm["canonical_smiles"] or str(smiles), + "inchikey": norm["inchikey"], + } + # ── Public entry point ──────────────────────────────────────────────────── def search( @@ -490,6 +665,8 @@ def search( queries: Union[str, List[str], pd.DataFrame, Path], *, column: Optional[str] = None, + n_hits: Optional[Union[int, str]] = None, + min_confidence: Optional[float] = None, ) -> pd.DataFrame: """Resolve one or more chemical identifiers and return a DataFrame. @@ -500,21 +677,27 @@ def search( - A list of strings — one row per query. - A :class:`pandas.DataFrame` — the column given by ``column`` is used as the query list. All other columns are preserved - in the output (left-joined on the original index). + in the output (broadcast across the hit rows of each query). - A file path (:class:`pathlib.Path` or string ending in ``.csv`` / ``.parquet``) — read into a DataFrame first; ``column`` must be provided. column: Column name to read from a DataFrame or file input. Required when ``queries`` is a DataFrame or file path. + n_hits: Per-call override of the instance ``n_hits`` (positive int + or ``"all"``). When ``None`` the instance default is used. + min_confidence: Per-call override of the instance + ``min_confidence``. When ``None`` the instance default is used. Returns: - DataFrame with columns defined in :data:`OUTPUT_COLUMNS`. Rows - correspond to input queries in the same order. + DataFrame with columns defined in :data:`OUTPUT_COLUMNS`. When + ``n_hits == 1`` (the default) there is one row per query; otherwise + up to ``n_hits`` ranked rows per query, ordered by descending + confidence with a ``hit_rank`` column (0 = best). Raises: ValueError: If a DataFrame/file input is given but ``column`` is - not specified. + not specified, or if ``n_hits`` is invalid. FileNotFoundError: If the given file path does not exist. Example:: @@ -522,9 +705,20 @@ def search( s = Search("cas") df = s.search(["50-00-0", "64-17-5"]) df = s.search(Path("compounds.csv"), column="CAS") + + # Return every plausible interpretation of an ambiguous name + s_name = Search("name") + df = s_name.search("xylene", n_hits="all") """ self._ensure_clients() + effective_n_hits = ( + self.n_hits if n_hits is None else self._validate_n_hits(n_hits) + ) + effective_min_conf = ( + self.min_confidence if min_confidence is None else float(min_confidence) + ) + query_list, extra_df = self._coerce_queries(queries, column) iterator = ( @@ -533,21 +727,33 @@ def search( else query_list ) - rows = [self._resolve_single(q) for q in iterator] + # Each query yields a list of ranked hit dicts. Track the source query + # index so DataFrame/file extra columns can be broadcast across hits. + rows: List[Dict[str, Any]] = [] + origin_index: List[int] = [] + for q_idx, q in enumerate(iterator): + hits = self._resolve_single(q, effective_n_hits, effective_min_conf) + for hit in hits: + rows.append(hit) + origin_index.append(q_idx) result_df = pd.DataFrame(rows) # Ensure all output columns are present (fill missing with None) for col in OUTPUT_COLUMNS: if col not in result_df.columns: result_df[col] = None - result_df = result_df[OUTPUT_COLUMNS] + ordered = list(OUTPUT_COLUMNS) + if self.return_alternatives and "alternatives" in result_df.columns: + ordered = ordered + ["alternatives"] + result_df = result_df[ordered] - # Merge extra columns from the original DataFrame if provided - if extra_df is not None: + # Broadcast extra columns from the original DataFrame across hit rows. + if extra_df is not None and origin_index: extra_cols = [c for c in extra_df.columns if c not in result_df.columns] if extra_cols: + broadcast = extra_df[extra_cols].iloc[origin_index].reset_index(drop=True) result_df = pd.concat( - [result_df, extra_df[extra_cols].reset_index(drop=True)], + [result_df.reset_index(drop=True), broadcast], axis=1, ) @@ -605,14 +811,26 @@ def _coerce_queries( # ── Single-query dispatcher ─────────────────────────────────────────────── - def _resolve_single(self, query: str) -> Dict[str, Any]: - """Dispatch one query to the appropriate resolver and return a result dict. + def _resolve_single( + self, + query: str, + n_hits: Union[int, str], + min_confidence: float, + ) -> List[Dict[str, Any]]: + """Dispatch one query to the appropriate resolver and return ranked hits. + + Each resolver returns ``(base_template, pool, opsin_anchor)``; this + method clusters the pool, ranks the clusters, and truncates to + ``n_hits``. Args: query: A single identifier string. + n_hits: Number of ranked hits to return (positive int or ``"all"``). + min_confidence: Drop hits below this confidence before truncation. Returns: - Fully populated result dict matching :data:`OUTPUT_COLUMNS`. + List of result dicts matching :data:`OUTPUT_COLUMNS` (length 1 when + ``n_hits == 1``). """ dispatch = { "cas": self._resolve_cas, @@ -623,7 +841,8 @@ def _resolve_single(self, query: str) -> Dict[str, Any]: "dtxsid": self._resolve_dtxsid, "formula": self._resolve_formula, } - return dispatch[self.identifier_type](query) + base_template, pool, opsin_anchor = dispatch[self.identifier_type](query) + return self._finalise_hits(base_template, pool, n_hits, min_confidence, opsin_anchor) # ── Empty result template ───────────────────────────────────────────────── @@ -661,11 +880,180 @@ def _empty_result(self, query: str, foundby: str) -> Dict[str, Any]: "match_score": 0.0, "consensus_source": None, "source_match_scores": {}, + "hit_rank": 0, + "n_source_support": 0, + "opsin_smiles": None, } + # ── Pool construction helpers ───────────────────────────────────────────── + + @staticmethod + def _tag_candidate( + cand: Dict[str, Any], + source_key: str, + origin_rank: int, + match_method: str, + query_match_score: float, + ) -> Dict[str, Any]: + """Annotate a candidate record with pool/ranking metadata (in place). + + Args: + cand: Candidate record from a ``_candidate_from_*`` helper. + source_key: Originating source key (e.g. ``"chebi"``, ``"opsin"``). + origin_rank: Rank position within the source's result list (0-based). + match_method: How the candidate was found (key into + :data:`_BASE_CONFIDENCE`). + query_match_score: How well the candidate matches the query in + [0, 1]. + + Returns: + The same candidate dict, mutated with the transient ``_`` keys. + """ + cand["_source_key"] = source_key + cand["_origin_rank"] = int(origin_rank) + cand["_match_method"] = match_method + cand["query_match_score"] = float(query_match_score) + return cand + + def _pool_from_candidates_dict( + self, + candidates: Dict[str, Optional[Dict[str, Any]]], + match_method: str, + *, + default_score: float = 1.0, + ) -> List[Dict[str, Any]]: + """Convert a per-source ``{key: candidate}`` dict into a tagged pool. + + Args: + candidates: Mapping of source key → candidate (or ``None``). + match_method: Match method to tag each candidate with. + default_score: ``query_match_score`` assigned to every candidate + (1.0 for exact-identifier matches; a similarity for fuzzy/ + Tanimoto matches). + + Returns: + List of tagged candidate records. + """ + pool: List[Dict[str, Any]] = [] + for key in self._SOURCE_KEYS: + cand = candidates.get(key) + if cand is None: + continue + self._tag_candidate(cand, key, 0, match_method, default_score) + pool.append(cand) + return pool + + def _name_score(self, query: str, cand: Dict[str, Any]) -> float: + """Best similarity between the query name and a candidate's names. + + Compares the query against the candidate ``name``, ``IUPAC_name`` and + each individual synonym using the configured fuzzy scorer (rapidfuzz) + when available, falling back to :func:`_text_similarity`. + + Args: + query: Query name. + cand: Candidate record. + + Returns: + Best similarity in [0, 1]. + """ + names: List[str] = [] + for field in ("name", "IUPAC_name"): + val = cand.get(field) + if not _is_missing(val): + names.append(str(val)) + syn = cand.get("Synonyms") + if not _is_missing(syn): + names.extend(s.strip() for s in str(syn).split(";") if s.strip()) + if not names: + return 0.0 + + if RAPIDFUZZ_AVAILABLE and _fuzz is not None: + scorer = getattr(_fuzz, self.fuzzy_scorer, _fuzz.WRatio) + try: + return max(scorer(query, n) for n in names) / 100.0 + except Exception: + pass + return max(_text_similarity(query, n) for n in names) + + @staticmethod + def _completeness_score(cand: Dict[str, Any]) -> float: + """Fraction of key structural/identifier fields populated in [0, 1]. + + Used as the query-agreement signal for formula matches (which have no + name to compare against). + + Args: + cand: Candidate record. + + Returns: + Completeness fraction in [0, 1]. + """ + fields = ("SMILES", "InChIKey", "InChI", "molecular_mass", "name", "DTXSID") + present = sum(1 for f in fields if not _is_missing(cand.get(f))) + present += 1 if (cand.get("CAS_candidates") or []) else 0 + return present / (len(fields) + 1) + + def _inchikey_pool( + self, inchikey: str, match_method: str, query_match_score: float + ) -> List[Dict[str, Any]]: + """Query every source by InChIKey and return a tagged candidate pool. + + Used by OPSIN anchoring to pull the structurally-correct compound from + each source regardless of name spelling. + + Args: + inchikey: Full InChIKey to look up. + match_method: Match method to tag candidates with. + query_match_score: Query-agreement score for the candidates. + + Returns: + List of tagged candidate records (one per source that matched). + """ + candidates: Dict[str, Optional[Dict[str, Any]]] = {k: None for k in self._SOURCE_KEYS} + + if self._chebi is not None: + try: + row = self._chebi.search_by_inchikey(inchikey) + if row: + candidates["chebi"] = _candidate_from_chebi_row(row) + except Exception as exc: + log.warning("ChEBI OPSIN-InChIKey lookup failed: %s", exc) + if self._comptox is not None: + try: + row = self._comptox.get_by_inchikey(inchikey) + if row: + candidates["comptox"] = _candidate_from_comptox_row(row) + except Exception as exc: + log.warning("CompTox OPSIN-InChIKey lookup failed: %s", exc) + if self._pubchem is not None: + try: + row = self._pubchem.get_by_inchikey(inchikey) + if row: + candidates["pubchem"] = _candidate_from_pubchem_row(row) + except Exception as exc: + log.warning("PubChemID OPSIN-InChIKey lookup failed: %s", exc) + if self._zeropm is not None: + try: + table = self._zeropm.get_id_table_from_inchikey(inchikey) + candidates["zeropm"] = _candidate_from_zeropm_name_table(inchikey, table) + except Exception as exc: + log.warning("ZeroPM OPSIN-InChIKey lookup failed: %s", exc) + if self._chembl is not None: + try: + row = self._chembl.search_by_inchikey(inchikey) + if row: + candidates["chembl"] = _candidate_from_chembl_row(row, self._chembl) + except Exception as exc: + log.warning("ChEMBL OPSIN-InChIKey lookup failed: %s", exc) + + return self._pool_from_candidates_dict( + candidates, match_method, default_score=query_match_score + ) + # ── CAS resolver ───────────────────────────────────────────────────────── - def _resolve_cas(self, cas: str) -> Dict[str, Any]: + def _resolve_cas(self, cas: str) -> Tuple[Dict[str, Any], List[Dict[str, Any]], Optional[Dict[str, Any]]]: """Resolve a CAS Registry Number into a unified identifier record. Queries ChEBI → CompTox → PubChemID → ZeroPM with waterfall priority, @@ -675,9 +1063,10 @@ def _resolve_cas(self, cas: str) -> Dict[str, Any]: cas: CAS Registry Number string. Returns: - Fully populated result dict. + Tuple of (base result template, candidate pool, ``None``). """ result = self._empty_result(cas, "CASRN") + result["match_method"] = "exact_cas" result["CASRN"] = cas candidates: Dict[str, Optional[Dict[str, Any]]] = {k: None for k in self._SOURCE_KEYS} @@ -727,42 +1116,188 @@ def _resolve_cas(self, cas: str) -> Dict[str, Any]: except Exception as exc: log.warning("ChEMBL CAS enrichment failed for %r: %s", cas, exc) - source_details = self._build_source_details(candidates) - return self._finalise_result(result, candidates, "exact_cas", source_details) + pool = self._pool_from_candidates_dict(candidates, "exact_cas") + return result, pool, None # ── Name resolver ───────────────────────────────────────────────────────── - def _resolve_name(self, name: str) -> Dict[str, Any]: - """Resolve a chemical name into a unified identifier record. + def _resolve_name( + self, name: str + ) -> Tuple[Dict[str, Any], List[Dict[str, Any]], Optional[Dict[str, Any]]]: + """Resolve a chemical name into a candidate pool for clustering. - First attempts exact-name matching across all sources, then falls back - to fuzzy matching when ``self.fuzzy`` is True and no exact match is - found. + Pools the top-``k`` candidates from every source (exact name/synonym + matches first; fuzzy-widened when ``self.fuzzy`` is set and exact + matches are weak), then adds a PYOPSIN structure anchor and an + InChIKey-driven structure lookup when ``self.use_opsin`` is enabled. Args: name: Chemical name string (common or IUPAC). Returns: - Fully populated result dict. + Tuple of (base result template, candidate pool, OPSIN anchor dict). """ result = self._empty_result(name, "name") - candidates = self._candidates_from_name(name, exact=True) + result["match_method"] = "exact_name" + pool = self._candidate_pool_from_name(name) + + # OPSIN structure anchoring (opt-in; needs Java). + opsin_anchor: Optional[Dict[str, Any]] = None + if self.use_opsin: + opsin_anchor = self._opsin_anchor(name) + if opsin_anchor is not None: + anchor_cand = _make_candidate( + "OPSIN", + name=name, + smiles=opsin_anchor.get("smiles"), + inchikey=opsin_anchor.get("inchikey"), + ) + self._tag_candidate(anchor_cand, "opsin", 0, "opsin", 1.0) + pool.append(anchor_cand) + # Pull the *correct* compound from each source by the OPSIN + # InChIKey, even when the name spelling differs. + if not _is_missing(opsin_anchor.get("inchikey")): + pool.extend( + self._inchikey_pool(str(opsin_anchor["inchikey"]), "opsin", 1.0) + ) - fuzzy_score: Optional[float] = None - match_method = "exact_name" + return result, pool, opsin_anchor - # Fuzzy fallback - if not _any_candidate(candidates) and self.fuzzy: - candidates, fuzzy_score = self._fuzzy_name_candidates(name) - match_method = "fuzzy_name" + def _candidate_pool_from_name(self, name: str) -> List[Dict[str, Any]]: + """Build a flat, tagged candidate pool from a name query. - source_details = self._build_source_details(candidates) - result = self._finalise_result(result, candidates, match_method, source_details) - if fuzzy_score is not None: - result["confidence"] = self._compute_confidence( - "fuzzy_name", result["match_score"], fuzzy_score=fuzzy_score - ) - return result + Pulls up to ``self.top_k_per_source`` candidates from each source. + When ``self.fuzzy`` is enabled and the exact pass yields no strong + match, the search is widened with non-exact matching and ZeroPM's + fuzzy ``query_similar_name``. + + Args: + name: Chemical name to search. + + Returns: + List of candidate records tagged with ``_source_key``, + ``_origin_rank``, ``_match_method`` and ``query_match_score``. + """ + pool: List[Dict[str, Any]] = [] + k = self.top_k_per_source + + def add(cand, source_key, rank, method): + if cand is None: + return + self._tag_candidate(cand, source_key, rank, method, self._name_score(name, cand)) + pool.append(cand) + + # ── Exact pass ────────────────────────────────────────────────────── + if self._chebi is not None: + try: + rows = self._chebi.search_by_name(name, exact=True) or [] + if not rows: + rows = self._chebi.search_by_synonym(name, exact=True) or [] + for rank, row in enumerate(rows[:k]): + add(_candidate_from_chebi_row(row), "chebi", rank, "exact_name") + except Exception as exc: + log.warning("ChEBI name lookup failed for %r: %s", name, exc) + + if self._comptox is not None: + try: + rows = self._comptox.search_by_name(name, exact=True, limit=k) or [] + if not rows: + row = self._comptox.get_by_name(name) + rows = [row] if row else [] + for rank, row in enumerate(rows[:k]): + add(_candidate_from_comptox_row(row), "comptox", rank, "exact_name") + except Exception as exc: + log.warning("CompTox name lookup failed for %r: %s", name, exc) + + if self._pubchem is not None: + try: + rows = self._pubchem.search_by_name(name, exact=True, limit=k) or [] + for rank, row in enumerate(rows[:k]): + add(_candidate_from_pubchem_row(row), "pubchem", rank, "exact_name") + except Exception as exc: + log.warning("PubChemID name lookup failed for %r: %s", name, exc) + + if self._zeropm is not None: + try: + table = self._zeropm.get_id_table_from_name(name) + add(_candidate_from_zeropm_name_table(name, table), "zeropm", 0, "exact_name") + except Exception as exc: + log.warning("ZeroPM name lookup failed for %r: %s", name, exc) + + if self._chembl is not None: + try: + rows = self._chembl.search_by_name(name, limit=k) or [] + for rank, row in enumerate(rows[:k]): + add(_candidate_from_chembl_row(row, self._chembl), "chembl", rank, "exact_name") + except Exception as exc: + log.warning("ChEMBL name lookup failed for %r: %s", name, exc) + + # ── Fuzzy widening ────────────────────────────────────────────────── + cutoff = self.fuzzy_score_cutoff / 100.0 + strong = any(c["query_match_score"] >= cutoff for c in pool) + if self.fuzzy and not strong: + norm_name = self._normalize_name(name) + + def add_fuzzy(cand, source_key, rank): + if cand is None: + return + score = self._name_score(name, cand) + if score < cutoff: + return + self._tag_candidate(cand, source_key, rank, "fuzzy_name", score) + pool.append(cand) + + if self._chebi is not None: + try: + rows = self._chebi.search_by_name(norm_name, exact=False) or [] + if not rows: + rows = self._chebi.search_by_synonym(norm_name, exact=False) or [] + for rank, row in enumerate(rows[:k]): + add_fuzzy(_candidate_from_chebi_row(row), "chebi", rank) + except Exception as exc: + log.warning("ChEBI fuzzy name lookup failed for %r: %s", name, exc) + + if self._comptox is not None: + try: + rows = self._comptox.search_by_name(norm_name, exact=False, limit=k) or [] + for rank, row in enumerate(rows[:k]): + add_fuzzy(_candidate_from_comptox_row(row), "comptox", rank) + except Exception as exc: + log.warning("CompTox fuzzy name lookup failed for %r: %s", name, exc) + + if self._pubchem is not None: + try: + rows = self._pubchem.search_by_name(norm_name, exact=False, limit=k) or [] + for rank, row in enumerate(rows[:k]): + add_fuzzy(_candidate_from_pubchem_row(row), "pubchem", rank) + except Exception as exc: + log.warning("PubChemID fuzzy name lookup failed for %r: %s", name, exc) + + if self._zeropm is not None: + try: + results = self._zeropm.query_similar_name( + norm_name, + number_of_results=k, + score_cutoff=self.fuzzy_score_cutoff, + ) + if isinstance(results, pd.DataFrame) and not results.empty: + add_fuzzy( + _candidate_from_zeropm_name_table(name, results), "zeropm", 0 + ) + except TypeError: + # Stub / older signature without keyword args. + try: + results = self._zeropm.query_similar_name(norm_name) + if isinstance(results, pd.DataFrame) and not results.empty: + add_fuzzy( + _candidate_from_zeropm_name_table(name, results), "zeropm", 0 + ) + except Exception as exc: + log.warning("ZeroPM fuzzy name lookup failed for %r: %s", name, exc) + except Exception as exc: + log.warning("ZeroPM fuzzy name lookup failed for %r: %s", name, exc) + + return pool def _candidates_from_name( self, name: str, exact: bool = True @@ -826,7 +1361,7 @@ def _candidates_from_name( # ── SMILES resolver ─────────────────────────────────────────────────────── - def _resolve_smiles(self, smiles: str) -> Dict[str, Any]: + def _resolve_smiles(self, smiles: str) -> Tuple[Dict[str, Any], List[Dict[str, Any]], Optional[Dict[str, Any]]]: """Resolve a SMILES string into a unified identifier record. Canonicalises the input, derives an InChIKey, and queries sources by @@ -838,9 +1373,10 @@ def _resolve_smiles(self, smiles: str) -> Dict[str, Any]: smiles: SMILES string. Returns: - Fully populated result dict. + Tuple of (base result template, candidate pool, ``None``). """ result = self._empty_result(smiles, "SMILES") + result["match_method"] = "exact_smiles" result["SMILES"] = smiles norm = normalize_structure(smiles) @@ -901,21 +1437,18 @@ def _resolve_smiles(self, smiles: str) -> Dict[str, Any]: if not _any_candidate(candidates) and self.similarity_threshold > 0: sim_candidates, tanimoto_score = self._tanimoto_candidates(smiles) if _any_candidate(sim_candidates): - candidates = sim_candidates - match_method = "tanimoto" - source_details = self._build_source_details(candidates) - result = self._finalise_result(result, candidates, match_method, source_details) - result["confidence"] = self._compute_confidence( - "tanimoto", result["match_score"], tanimoto=tanimoto_score + score = tanimoto_score if tanimoto_score is not None else 0.0 + pool = self._pool_from_candidates_dict( + sim_candidates, "tanimoto", default_score=score ) - return result + return result, pool, None - source_details = self._build_source_details(candidates) - return self._finalise_result(result, candidates, match_method, source_details) + pool = self._pool_from_candidates_dict(candidates, match_method) + return result, pool, None # ── InChI resolver ──────────────────────────────────────────────────────── - def _resolve_inchi(self, inchi: str) -> Dict[str, Any]: + def _resolve_inchi(self, inchi: str) -> Tuple[Dict[str, Any], List[Dict[str, Any]], Optional[Dict[str, Any]]]: """Resolve an InChI string into a unified identifier record. Converts the InChI to InChIKey via RDKit and delegates to @@ -926,9 +1459,10 @@ def _resolve_inchi(self, inchi: str) -> Dict[str, Any]: inchi: InChI string (must start with ``"InChI="``). Returns: - Fully populated result dict. + Tuple of (base result template, candidate pool, ``None``). """ result = self._empty_result(inchi, "InChI") + result["match_method"] = "inchi" result["InChI"] = inchi # Derive InChIKey and SMILES via RDKit @@ -995,12 +1529,12 @@ def _resolve_inchi(self, inchi: str) -> Dict[str, Any]: except Exception as exc: log.warning("ChEMBL InChI lookup failed: %s", exc) - source_details = self._build_source_details(candidates) - return self._finalise_result(result, candidates, "inchi", source_details) + pool = self._pool_from_candidates_dict(candidates, "inchi") + return result, pool, None # ── InChIKey resolver ───────────────────────────────────────────────────── - def _resolve_inchikey(self, inchikey: str) -> Dict[str, Any]: + def _resolve_inchikey(self, inchikey: str) -> Tuple[Dict[str, Any], List[Dict[str, Any]], Optional[Dict[str, Any]]]: """Resolve an InChIKey into a unified identifier record. Queries all offline sources by InChIKey. Falls back to 14-character @@ -1012,9 +1546,10 @@ def _resolve_inchikey(self, inchikey: str) -> Dict[str, Any]: (``XXXXXXXXXXXXXX-XXXXXXXXXX-X``). Returns: - Fully populated result dict. + Tuple of (base result template, candidate pool, ``None``). """ result = self._empty_result(inchikey, "InChIKey") + result["match_method"] = "exact_inchikey" result["InChIKey"] = inchikey candidates: Dict[str, Optional[Dict[str, Any]]] = {k: None for k in self._SOURCE_KEYS} @@ -1071,12 +1606,12 @@ def _resolve_inchikey(self, inchikey: str) -> Dict[str, Any]: candidates = skel_candidates match_method = "inchikey_skeleton" - source_details = self._build_source_details(candidates) - return self._finalise_result(result, candidates, match_method, source_details) + pool = self._pool_from_candidates_dict(candidates, match_method) + return result, pool, None # ── DTXSID resolver ─────────────────────────────────────────────────────── - def _resolve_dtxsid(self, dtxsid: str) -> Dict[str, Any]: + def _resolve_dtxsid(self, dtxsid: str) -> Tuple[Dict[str, Any], List[Dict[str, Any]], Optional[Dict[str, Any]]]: """Resolve a CompTox DTXSID into a unified identifier record. Queries CompTox as the primary source, then cross-references other @@ -1086,9 +1621,10 @@ def _resolve_dtxsid(self, dtxsid: str) -> Dict[str, Any]: dtxsid: CompTox DTXSID string (e.g., ``"DTXSID7020182"``). Returns: - Fully populated result dict. + Tuple of (base result template, candidate pool, ``None``). """ result = self._empty_result(dtxsid, "DTXSID") + result["match_method"] = "dtxsid" result["DTXSID"] = dtxsid candidates: Dict[str, Optional[Dict[str, Any]]] = {k: None for k in self._SOURCE_KEYS} @@ -1142,57 +1678,67 @@ def _resolve_dtxsid(self, dtxsid: str) -> Dict[str, Any]: except Exception as exc: log.warning("ChEMBL DTXSID cross-ref failed: %s", exc) - source_details = self._build_source_details(candidates) - return self._finalise_result(result, candidates, "dtxsid", source_details) + pool = self._pool_from_candidates_dict(candidates, "dtxsid") + return result, pool, None # ── Formula resolver ────────────────────────────────────────────────────── - def _resolve_formula(self, formula: str) -> Dict[str, Any]: - """Resolve a molecular formula into a unified identifier record. + def _resolve_formula( + self, formula: str + ) -> Tuple[Dict[str, Any], List[Dict[str, Any]], Optional[Dict[str, Any]]]: + """Resolve a molecular formula into a candidate pool. - Returns the most complete match across CompTox, PubChemID, and ChEBI. - Formulas are not unique identifiers, so confidence is capped at 0.30. + Formulas are not unique identifiers, so all sources may return many + rows. The top-``k`` rows per source are pooled and clustered; distinct + compounds are returned ranked by completeness/consensus. Confidence is + capped low (base 0.30) because formula matches are ambiguous. Args: formula: Molecular formula string (e.g., ``"C9H8O4"``). Returns: - Fully populated result dict for the best (most complete) match. + Tuple of (base result template, candidate pool, ``None``). """ result = self._empty_result(formula, "formula") + result["match_method"] = "formula" result["molecular_formula"] = formula - candidates: Dict[str, Optional[Dict[str, Any]]] = {k: None for k in self._SOURCE_KEYS} + pool: List[Dict[str, Any]] = [] + k = self.top_k_per_source + + def add(cand, source_key, rank): + if cand is None: + return + # Completeness drives the query_match_score for formula matches. + score = self._completeness_score(cand) + self._tag_candidate(cand, source_key, rank, "formula", score) + pool.append(cand) if self._chebi is not None: try: - rows = self._chebi.search_by_formula(formula) - if rows: - best = _most_complete_row(rows) - candidates["chebi"] = _candidate_from_chebi_row(best) + rows = self._chebi.search_by_formula(formula) or [] + for rank, row in enumerate(_rank_rows_by_completeness(rows)[:k]): + add(_candidate_from_chebi_row(row), "chebi", rank) except Exception as exc: log.warning("ChEBI formula lookup failed for %r: %s", formula, exc) if self._comptox is not None: try: - rows = self._comptox.search_by_formula(formula) - if rows: - best = _most_complete_row(rows) - candidates["comptox"] = _candidate_from_comptox_row(best) + rows = self._comptox.search_by_formula(formula) or [] + for rank, row in enumerate(_rank_rows_by_completeness(rows)[:k]): + add(_candidate_from_comptox_row(row), "comptox", rank) except Exception as exc: log.warning("CompTox formula lookup failed for %r: %s", formula, exc) if self._pubchem is not None: try: - rows = self._pubchem.search_by_formula(formula) - if rows: - best = _most_complete_row(rows) - candidates["pubchem"] = _candidate_from_pubchem_row(best) + rows = self._pubchem.search_by_formula(formula) or [] + for rank, row in enumerate(_rank_rows_by_completeness(rows)[:k]): + add(_candidate_from_pubchem_row(row), "pubchem", rank) except Exception as exc: log.warning("PubChemID formula lookup failed for %r: %s", formula, exc) - source_details = self._build_source_details(candidates) - return self._finalise_result(result, candidates, "formula", source_details) + return result, pool, None # ── Fuzzy name search ───────────────────────────────────────────────────── @@ -1492,17 +2038,25 @@ def _compute_confidence( *, fuzzy_score: Optional[float] = None, tanimoto: Optional[float] = None, + query_score: Optional[float] = None, ) -> float: """Compute the final confidence score for a result. The base confidence depends on the match method. For fuzzy and - Tanimoto methods, the raw score is used as the base. The base is then - modulated by the cross-source consensus score so that multi-source - agreement boosts confidence. + Tanimoto methods, the raw similarity is used as the base. The base is + modulated by a query-agreement term (weighted by ``self.query_weight``) + and by the cross-source consensus score, so that both a strong query + match and multi-source agreement boost confidence. Formula:: - final = base × (0.5 + 0.5 × consensus_score) + final = base + × (w_q × query_score + (1 − w_q)) + × (0.5 + 0.5 × consensus_score) + + For exact-identifier methods ``query_score`` is 1.0, which collapses the + middle term to 1.0 and reproduces the original + ``base × (0.5 + 0.5 × consensus_score)`` behaviour. Args: match_method: One of the keys in :data:`_BASE_CONFIDENCE`. @@ -1511,6 +2065,9 @@ def _compute_confidence( ``match_method == "fuzzy_name"``. tanimoto: Tanimoto similarity in [0, 1]; used when ``match_method == "tanimoto"``. + query_score: Query-agreement signal in [0, 1] for name/formula + methods. Ignored (treated as 1.0) for fuzzy/Tanimoto where the + similarity already lives in the base. Returns: Confidence value in [0, 1]. @@ -1519,102 +2076,224 @@ def _compute_confidence( return 0.0 base = _BASE_CONFIDENCE.get(match_method, 0.5) + q = 1.0 if query_score is None else max(0.0, min(1.0, query_score)) if match_method == "fuzzy_name": base = fuzzy_score if fuzzy_score is not None else 0.5 + q = 1.0 # similarity already captured in base elif match_method == "tanimoto": base = (tanimoto * 0.85) if tanimoto is not None else 0.5 + q = 1.0 - modulated = base * (0.5 + 0.5 * max(0.0, min(1.0, consensus_score))) + w_q = max(0.0, min(1.0, self.query_weight)) + query_term = w_q * q + (1.0 - w_q) + modulated = base * query_term * (0.5 + 0.5 * max(0.0, min(1.0, consensus_score))) return round(min(1.0, max(0.0, modulated)), 4) # ── Result finalisation ─────────────────────────────────────────────────── - def _finalise_result( + def _finalise_hits( self, - result: Dict[str, Any], - candidates: Dict[str, Optional[Dict[str, Any]]], - match_method: str, - source_details: Dict[str, Any], - ) -> Dict[str, Any]: - """Apply candidates to result, run consensus, and fill derived fields. + base_template: Dict[str, Any], + pool: List[Dict[str, Any]], + n_hits: Union[int, str], + min_confidence: float, + opsin_anchor: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: + """Cluster a candidate pool, rank the clusters, and return ranked hits. - This method: + Args: + base_template: Empty result template (carries query/foundby and any + pre-populated query fields). + pool: Flat list of tagged candidate records. + n_hits: Number of hits to return (positive int or ``"all"``). + min_confidence: Drop hits below this confidence before truncation. + opsin_anchor: Optional OPSIN structure anchor for this query. + + Returns: + List of fully-populated result dicts ordered by descending + confidence with ``hit_rank`` set. Always contains at least one row + (an empty/no-match row when nothing was found). + """ + opsin_smiles = opsin_anchor.get("smiles") if opsin_anchor else None + + if not pool: + # No source matched — still finalise structure/salt fields from any + # pre-populated query fields (e.g. a SMILES/InChI query) via an empty + # cluster, preserving the resolver's default match_method. + empty = self._build_result_for_cluster(base_template, {"members": []}, opsin_smiles) + empty["hit_rank"] = 0 + return [empty] + + clusters = _cluster_candidates(pool, by_skeleton=self.cluster_by_skeleton) + hits = [ + self._build_result_for_cluster(base_template, cluster, opsin_smiles) + for cluster in clusters + ] + + # Rank: OPSIN match first, then confidence, support, query agreement, + # and (lower) origin rank as a final tie-break. + hits.sort( + key=lambda h: ( + 1 if h["_opsin_match"] else 0, + h["confidence"], + h["n_source_support"], + h["_cluster_query_score"], + -h["_min_origin_rank"], + ), + reverse=True, + ) + + filtered = [h for h in hits if h["confidence"] >= min_confidence] + if not filtered: + # Everything was below the floor — represent the query with a single + # no-match row so it is not silently dropped. + empty = self._build_result_for_cluster(base_template, {"members": []}, opsin_smiles) + empty["hit_rank"] = 0 + return [empty] + + if n_hits != "all": + filtered = filtered[: int(n_hits)] + + alternatives = None + if self.return_alternatives and n_hits == 1 and len(hits) > 1: + alternatives = [ + { + "name": h.get("name"), + "InChIKey": h.get("InChIKey"), + "confidence": h.get("confidence"), + "source": h.get("source"), + } + for h in hits[1:6] + ] + + for rank, hit in enumerate(filtered): + hit["hit_rank"] = rank + if alternatives is not None and rank == 0: + hit["alternatives"] = alternatives + + return filtered + + def _build_result_for_cluster( + self, + base_template: Dict[str, Any], + cluster: Dict[str, Any], + opsin_smiles: Optional[str], + ) -> Dict[str, Any]: + """Build one fully-populated result dict from a single structure cluster. - 1. Runs :func:`~provesid.tools._compute_consensus` to pick the anchor. - 2. Applies each compatible candidate to ``result`` via - :func:`~provesid.tools._apply_candidate_to_result`. - 3. Runs :func:`normalize_structure` on the final SMILES to populate - ``canonical_smiles``, ``kekulized_smiles``, ``InChIKey``, etc. - 4. Validates the source-provided InChIKey against the RDKit-derived one. - 5. Runs salt stripping if ``self.strip_salts`` is True. - 6. Computes and stores the ``confidence`` score. + All members of a cluster denote the same compound; this picks the best + member per source, runs the existing consensus/merge machinery over + them, normalises the structure, and computes confidence. Args: - result: Partially populated result dict. - candidates: Per-source candidate records. - match_method: How the match was obtained. - source_details: Pre-built source traceability dict. + base_template: Empty result template to populate. + cluster: A cluster dict with a ``members`` list of tagged candidates. + opsin_smiles: OPSIN SMILES for the query (for the ``opsin_smiles`` + column), if any. Returns: - Fully populated result dict. + A populated result dict, plus transient ``_``-prefixed ranking keys + (dropped before output). """ - consensus_source, source_match_scores, match_score = _compute_consensus(candidates) - consensus_candidate = candidates.get(consensus_source) if consensus_source else None + result = dict(base_template) + members: List[Dict[str, Any]] = cluster["members"] + + opsin_match = any(m.get("_source_key") == "opsin" for m in members) + + # Best member per data source (lowest origin rank, then best query score). + per_source: Dict[str, Dict[str, Any]] = {} + for m in members: + key = m.get("_source_key") + if key in (None, "opsin"): + continue + cur = per_source.get(key) + rank_tuple = (m.get("_origin_rank", 0), -m.get("query_match_score", 0.0)) + if cur is None or rank_tuple < ( + cur.get("_origin_rank", 0), + -cur.get("query_match_score", 0.0), + ): + per_source[key] = m + + # Cluster match method = the strongest method among members; fall back + # to the resolver's default (carried on the template) for empty clusters. + cluster_method = max( + (m.get("_match_method", "unknown") for m in members), + key=lambda mm: _BASE_CONFIDENCE.get(mm, 0.5), + default=base_template.get("match_method", "unknown"), + ) + if opsin_match: + cluster_method = "opsin" + + consensus_source, source_match_scores, match_score = _compute_consensus(per_source) + consensus_candidate = per_source.get(consensus_source) if consensus_source else None for source_key in ["chebi", "comptox", "pubchem", "zeropm"]: - candidate = candidates.get(source_key) - if _candidate_compatible_with_consensus(candidate, consensus_candidate): + candidate = per_source.get(source_key) + if _candidate_compatible_with_consensus( + candidate, consensus_candidate, self.consensus_compat_threshold + ): _apply_candidate_to_result(result, candidate) - # ChEMBL enrichment - chembl_cand = candidates.get("chembl") - if _candidate_compatible_with_consensus(chembl_cand, consensus_candidate): + chembl_cand = per_source.get("chembl") + if _candidate_compatible_with_consensus( + chembl_cand, consensus_candidate, self.consensus_compat_threshold + ): _apply_candidate_to_result(result, chembl_cand) + # OPSIN supplies a structure even when no source row carried one. + if opsin_match and _is_missing(result.get("SMILES")) and not _is_missing(opsin_smiles): + result["SMILES"] = opsin_smiles + # Structure normalisation norm = normalize_structure(result.get("SMILES")) result["canonical_smiles"] = norm["canonical_smiles"] result["kekulized_smiles"] = norm["kekulized_smiles"] result["molecular_mass"] = _pick_first(result.get("molecular_mass"), norm["mol_weight"]) - # InChI / InChIKey — prefer RDKit-derived values; warn on mismatch rdkit_inchi = norm["inchi"] rdkit_ik = norm["inchikey"] - if not _is_missing(result.get("InChIKey")) and not _is_missing(rdkit_ik): if result["InChIKey"] != rdkit_ik: log.debug( "InChIKey mismatch for query %r: source=%r rdkit=%r", - result["query"], - result["InChIKey"], - rdkit_ik, + result["query"], result["InChIKey"], rdkit_ik, ) result["InChI"] = _pick_first(result.get("InChI"), rdkit_inchi) result["InChIKey"] = _pick_first(result.get("InChIKey"), rdkit_ik) - # Align name / IUPAC_name result["name"] = _pick_first(result.get("name"), result.get("IUPAC_name")) result["IUPAC_name"] = _pick_first(result.get("IUPAC_name"), result.get("name")) - - # Source result["source"] = _pick_first( result.get("source"), consensus_candidate.get("source") if consensus_candidate else None, ) - # Consensus & score metadata result["consensus_source"] = ( consensus_candidate.get("source") if consensus_candidate else None ) result["source_match_scores"] = { - (candidates[src].get("source") if candidates.get(src) else src): round(score, 4) + (per_source[src].get("source") if per_source.get(src) else src): round(score, 4) for src, score in source_match_scores.items() } result["match_score"] = round(match_score, 4) - result["source_details"] = source_details - result["match_method"] = match_method - result["confidence"] = self._compute_confidence(match_method, match_score) + result["source_details"] = self._build_source_details(per_source) + result["match_method"] = cluster_method + + cluster_query_score = max( + (m.get("query_match_score", 0.0) for m in members), default=0.0 + ) + fuzzy_score = cluster_query_score if cluster_method == "fuzzy_name" else None + tanimoto = cluster_query_score if cluster_method == "tanimoto" else None + result["confidence"] = self._compute_confidence( + cluster_method, + match_score, + fuzzy_score=fuzzy_score, + tanimoto=tanimoto, + query_score=cluster_query_score, + ) + result["n_source_support"] = len(per_source) + result["opsin_smiles"] = opsin_smiles # Salt stripping if self.strip_salts and not _is_missing(result.get("SMILES")): @@ -1625,6 +2304,12 @@ def _finalise_result( parent_norm = normalize_structure(parent) result["parent_inchikey"] = parent_norm["inchikey"] + # Transient ranking metadata (dropped before DataFrame assembly). + result["_opsin_match"] = opsin_match + result["_cluster_query_score"] = cluster_query_score + result["_min_origin_rank"] = min( + (m.get("_origin_rank", 0) for m in members), default=0 + ) return result @@ -1632,6 +2317,112 @@ def _finalise_result( # Module-level private helpers # ───────────────────────────────────────────────────────────────────────────── +def _rank_rows_by_completeness(rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Sort source rows by number of non-null fields (most complete first). + + Args: + rows: List of raw source record dicts. + + Returns: + New list ordered by descending completeness; stable for ties. + """ + if not rows: + return [] + return sorted( + rows, + key=lambda r: sum(1 for v in r.values() if not _is_missing(v)), + reverse=True, + ) + + +def _candidate_cluster_keys(cand: Dict[str, Any], by_skeleton: bool) -> set: + """Return the set of structure-identity keys a candidate belongs to. + + Two candidates are merged into the same cluster when their key sets + intersect. Priority: full InChIKey (plus 14-char skeleton when + ``by_skeleton``) → canonical SMILES → normalised name. + + Args: + cand: A tagged candidate record. + by_skeleton: Whether to also emit the 14-char InChIKey skeleton key + (merging stereo/charge/isotope variants). + + Returns: + A set of hashable key tuples (empty when the candidate has no usable + identity, signalling a unique singleton cluster). + """ + keys: set = set() + ik = cand.get("InChIKey") + if not _is_missing(ik): + ik_s = str(ik) + keys.add(("ik", ik_s)) + if by_skeleton and len(ik_s) >= 14: + keys.add(("skel", ik_s[:14])) + return keys + + canon = cand.get("canonical_smiles") or cand.get("SMILES") + if not _is_missing(canon): + keys.add(("smi", str(canon))) + return keys + + name = cand.get("name") or cand.get("IUPAC_name") + if not _is_missing(name): + keys.add(("name", str(name).strip().lower())) + return keys + + +def _cluster_candidates( + pool: List[Dict[str, Any]], by_skeleton: bool = True +) -> List[Dict[str, Any]]: + """Group a candidate pool into structure-identity clusters. + + Uses a union-find over candidate identity keys so that, e.g., a candidate + carrying a full InChIKey and one carrying only the matching skeleton end up + in the same cluster. Candidates with no usable identity become singletons. + + Args: + pool: Flat list of tagged candidate records. + by_skeleton: Merge stereo/charge/isotope variants via skeleton keys. + + Returns: + List of cluster dicts, each ``{"members": [candidate, ...]}``. + """ + parent: Dict[Any, Any] = {} + + def find(x: Any) -> Any: + root = x + while parent[root] != root: + root = parent[root] + while parent[x] != root: + parent[x], x = root, parent[x] + return root + + def union(a: Any, b: Any) -> None: + parent.setdefault(a, a) + parent.setdefault(b, b) + ra, rb = find(a), find(b) + if ra != rb: + parent[rb] = ra + + cand_anchor: List[Any] = [] + for idx, cand in enumerate(pool): + keys = _candidate_cluster_keys(cand, by_skeleton) + if not keys: + keys = {("uniq", idx)} # unique singleton + key_list = list(keys) + for k in key_list: + parent.setdefault(k, k) + for k in key_list[1:]: + union(key_list[0], k) + cand_anchor.append(key_list[0]) + + groups: Dict[Any, List[Dict[str, Any]]] = {} + for cand, anchor in zip(pool, cand_anchor): + groups.setdefault(find(anchor), []).append(cand) + + return [{"members": members} for members in groups.values()] + + def _any_candidate(candidates: Dict[str, Optional[Dict[str, Any]]]) -> bool: """Return True if at least one candidate is non-None. diff --git a/tests/test_search_multihit.py b/tests/test_search_multihit.py new file mode 100644 index 0000000..9a71f0a --- /dev/null +++ b/tests/test_search_multihit.py @@ -0,0 +1,230 @@ +"""Tests for the multi-hit / tunable Search features. + +Covers the candidate-pool refactor, structure clustering, ranked multi-hit +output, query-aware ranking (the wrong-name-hit fix), tunable attributes, the +DataFrame broadcast for multi-hit rows, and OPSIN anchoring (skipped when no +Java/py2opsin runtime is available). + +Run with:: + + uv run pytest tests/test_search_multihit.py -v +""" + +from __future__ import annotations + +import pandas as pd +import pytest + +from provesid.search import ( + OUTPUT_COLUMNS, + Search, + _candidate_cluster_keys, + _cluster_candidates, +) +from provesid.tools import _make_candidate + + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers / stubs +# ───────────────────────────────────────────────────────────────────────────── + +def _tag(cand, source_key, rank=0, method="exact_name", score=1.0): + cand["_source_key"] = source_key + cand["_origin_rank"] = rank + cand["_match_method"] = method + cand["query_match_score"] = score + return cand + + +class _PubChemNameStub: + """Returns a fixed list of rows for any name search; no InChIKey hits.""" + + def __init__(self, rows): + self._rows = rows + + def search_by_name(self, name, exact=False, limit=10): + return self._rows + + def get_by_inchikey(self, ik): + return None + + +_ETHANOL = { + "cmpdname": "Wrong Hit", + "iupacname": "ethanol", + "mf": "C2H6O", + "smiles": "CCO", + "inchi": None, + "inchikey": "LFQSCWFLJHTTHZ-UHFFFAOYSA-N", + "mw": 46.07, + "synonyms": ["ethanol"], + "cas_numbers": ["64-17-5"], +} + +_XYLENE = { + "cmpdname": "xylene", + "iupacname": "1,2-dimethylbenzene", + "mf": "C8H10", + "smiles": "Cc1ccccc1C", + "inchi": None, + "inchikey": "CTQNGGLPUBDAKN-UHFFFAOYSA-N", + "mw": 106.16, + "synonyms": ["o-xylene", "xylene"], + "cas_numbers": ["95-47-6"], +} + + +def _name_search(rows, **kwargs): + return Search( + "name", + show_progress=False, + pubchem=_PubChemNameStub(rows), + chebi=None, + comptox=None, + zeropm=None, + chembl=None, + **kwargs, + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# Clustering +# ───────────────────────────────────────────────────────────────────────────── + +class TestClustering: + def test_same_inchikey_merges(self): + a = _tag(_make_candidate("ChEBI", name="aspirin", inchikey="BSYNRYMUTXBXSQ-UHFFFAOYSA-N"), "chebi") + b = _tag(_make_candidate("CompTox", name="Aspirin", inchikey="BSYNRYMUTXBXSQ-UHFFFAOYSA-N"), "comptox") + c = _tag(_make_candidate("PubChemID", name="caffeine", inchikey="RYYVLZVUVIJVGH-UHFFFAOYSA-N"), "pubchem") + clusters = _cluster_candidates([a, b, c], by_skeleton=True) + assert len(clusters) == 2 + assert sorted(len(cl["members"]) for cl in clusters) == [1, 2] + + def test_skeleton_merge_toggle(self): + d = _tag(_make_candidate("ChEBI", name="x", inchikey="ABCDEFGHIJKLMN-AAAAAAAAAA-N"), "chebi") + e = _tag(_make_candidate("CompTox", name="y", inchikey="ABCDEFGHIJKLMN-BBBBBBBBBB-N"), "comptox") + assert len(_cluster_candidates([d, e], by_skeleton=True)) == 1 + assert len(_cluster_candidates([d, e], by_skeleton=False)) == 2 + + def test_name_only_fallback_singletons(self): + a = _tag(_make_candidate("ChEBI", name="foo"), "chebi") + b = _tag(_make_candidate("CompTox", name="bar"), "comptox") + assert len(_cluster_candidates([a, b], by_skeleton=True)) == 2 + + def test_cluster_keys_priority(self): + ik = _make_candidate("X", inchikey="BSYNRYMUTXBXSQ-UHFFFAOYSA-N", smiles="CCO") + keys = _candidate_cluster_keys(ik, by_skeleton=True) + # InChIKey present -> ik + skel keys, no smiles key + assert ("ik", "BSYNRYMUTXBXSQ-UHFFFAOYSA-N") in keys + assert ("skel", "BSYNRYMUTXBXSQ") in keys + assert not any(k[0] == "smi" for k in keys) + + +# ───────────────────────────────────────────────────────────────────────────── +# n_hits validation +# ───────────────────────────────────────────────────────────────────────────── + +class TestNHitsValidation: + @pytest.mark.parametrize("value", [1, 3, "all", "ALL"]) + def test_valid(self, value): + Search("name", n_hits=value, show_progress=False) + + @pytest.mark.parametrize("value", [0, -1, "two", 1.5, True]) + def test_invalid(self, value): + with pytest.raises(ValueError): + Search("name", n_hits=value, show_progress=False) + + def test_invalid_fuzzy_scorer(self): + with pytest.raises(ValueError): + Search("name", fuzzy_scorer="does_not_exist", show_progress=False) + + +# ───────────────────────────────────────────────────────────────────────────── +# Multi-hit output + ranking +# ───────────────────────────────────────────────────────────────────────────── + +class TestMultiHit: + def test_default_single_row(self): + s = _name_search([_ETHANOL, _XYLENE]) + df = s.search("xylene") + assert len(df) == 1 + assert df.iloc[0]["hit_rank"] == 0 + + def test_query_aware_ranking_fixes_wrong_hit(self): + # PubChem returns ethanol as hit #0, xylene as hit #1. The query + # "xylene" must still resolve to xylene at rank 0. + s = _name_search([_ETHANOL, _XYLENE]) + df = s.search("xylene") + assert df.iloc[0]["name"] == "xylene" + + def test_n_hits_all_returns_distinct_compounds(self): + s = _name_search([_ETHANOL, _XYLENE]) + df = s.search("xylene", n_hits="all") + assert len(df) == 2 + assert list(df["hit_rank"]) == [0, 1] + # confidence is non-increasing with rank + assert df.iloc[0]["confidence"] >= df.iloc[1]["confidence"] + + def test_n_hits_int_truncates(self): + s = _name_search([_ETHANOL, _XYLENE]) + assert len(s.search("xylene", n_hits=1)) == 1 + assert len(s.search("xylene", n_hits=5)) == 2 # only 2 distinct compounds + + def test_min_confidence_filters(self): + s = _name_search([_ETHANOL, _XYLENE]) + # high floor keeps only the strong (exact-name) hit + df = s.search("xylene", n_hits="all", min_confidence=0.6) + assert len(df) == 1 + assert df.iloc[0]["name"] == "xylene" + + def test_new_columns_present(self): + s = _name_search([_XYLENE]) + df = s.search("xylene") + for col in ("hit_rank", "n_source_support", "opsin_smiles"): + assert col in df.columns + assert col in OUTPUT_COLUMNS + + def test_dataframe_broadcast(self): + s = _name_search([_ETHANOL, _XYLENE]) + inp = pd.DataFrame({"q": ["xylene"], "batch": ["b1"]}) + out = s.search(inp, column="q", n_hits="all") + assert len(out) == 2 + assert set(out["batch"]) == {"b1"} + + def test_return_alternatives(self): + s = _name_search([_ETHANOL, _XYLENE], return_alternatives=True) + df = s.search("xylene") # n_hits defaults to 1 + assert len(df) == 1 + alts = df.iloc[0]["alternatives"] + assert isinstance(alts, list) and len(alts) >= 1 + assert alts[0]["name"] == "Wrong Hit" + + def test_no_match_still_one_row(self): + s = _name_search([]) + df = s.search("nonexistent-compound") + assert len(df) == 1 + assert df.iloc[0]["confidence"] == 0.0 + assert df.iloc[0]["match_method"] == "exact_name" + + +# ───────────────────────────────────────────────────────────────────────────── +# OPSIN anchoring (requires Java + py2opsin) +# ───────────────────────────────────────────────────────────────────────────── + +@pytest.mark.opsin +class TestOpsinAnchor: + def test_opsin_disabled_by_default(self): + s = _name_search([_XYLENE]) + assert s.use_opsin is False + # OPSIN column present but empty when disabled + df = s.search("xylene") + assert df.iloc[0]["opsin_smiles"] is None + + def test_opsin_anchor_when_available(self): + s = _name_search([_XYLENE], use_opsin=True) + anchor = s._opsin_anchor("1,2-dimethylbenzene") + if anchor is None: + pytest.skip("PYOPSIN/Java unavailable in this environment") + assert anchor["smiles"] + df = s.search("1,2-dimethylbenzene") + assert df.iloc[0]["opsin_smiles"] is not None diff --git a/tests/test_search_precision_regression.py b/tests/test_search_precision_regression.py new file mode 100644 index 0000000..5c3b7a2 --- /dev/null +++ b/tests/test_search_precision_regression.py @@ -0,0 +1,175 @@ +"""Precision regression test for name-based resolution. + +Uses CompTox as ground truth: every chemical carries a CASRN and several +names/synonyms. We resolve a sample of chemicals *by a synonym* through a +CompTox-only :class:`~provesid.Search` (precision-first default config) and +assert that the resolver never returns a **structurally different** compound +(a "wrong hit"). + +Correctness is measured by InChIKey *skeleton* (the connectivity layer), which +is source-independent and tolerant of stereo/charge/tautomer naming variants. +Three outcomes are possible per query: + +- **correct** — returned a structure whose skeleton matches the truth, +- **wrong_hit** — returned a *different* structure (the precision failure we + guard against), or +- **not_found** — returned no structure (a recall gap, explicitly allowed). + +The guard asserts ``wrong_hit == 0``: the resolver may fail to find a compound, +but it must not confidently return the wrong one. This locks in the behaviour +established by the candidate-pool / clustering / query-aware-ranking rework. + +Marked ``integration`` + ``slow``; skipped automatically when the offline +CompTox database is unavailable (e.g. a fresh checkout with no downloaded data). + +Run with:: + + uv run pytest tests/test_search_precision_regression.py -v -m integration +""" + +from __future__ import annotations + +import pytest + +from provesid.search import Search + +# Sample size — kept modest so the test runs in a couple of seconds with the +# CompTox-only resolver while still exercising a representative spread. +SAMPLE_SIZE = 80 + + +def _skeleton(inchikey): + """Return the 14-char InChIKey connectivity skeleton, or the input as-is.""" + if isinstance(inchikey, str) and len(inchikey) >= 14: + return inchikey[:14] + return inchikey + + +def _pick_synonym(row): + """Choose a query synonym for a CompTox row that is not its CAS or name. + + Args: + row: A ``sqlite3.Row`` (mapping) with ``CASRN``, ``PREFERRED_NAME`` and + the pipe-delimited ``IDENTIFIER`` field. + + Returns: + A synonym string distinct from the CASRN and preferred name, or + ``None`` when no suitable name-like token exists. + """ + ident = row["IDENTIFIER"] or "" + pref = (row["PREFERRED_NAME"] or "").strip().lower() + cas = (row["CASRN"] or "").strip() + for part in (p.strip() for p in ident.split("|")): + if not part or part == cas or part.lower() == pref: + continue + if any(ch.isalpha() for ch in part): # name-like, not a bare code + return part + return None + + +@pytest.fixture(scope="module") +def comptox(): + """Yield an offline CompToxID client, skipping if its database is absent.""" + from provesid.comptox import CompToxID + + try: + client = CompToxID() + # Touch the connection so a missing/empty DB skips rather than errors. + client.conn.execute("SELECT 1 FROM chemicals LIMIT 1").fetchone() + except Exception as exc: # pragma: no cover - environment dependent + pytest.skip(f"CompTox offline database unavailable: {exc}") + return client + + +@pytest.fixture(scope="module") +def synonym_sample(comptox): + """Build a deterministic sample of (query synonym, truth skeleton) pairs. + + Samples rows spread across the CompTox table via a ``rowid`` stride (no RNG, + so the sample is reproducible), keeps only CASRNs that map to a single + structure, and selects a synonym distinct from the preferred name. + + Returns: + List of ``(synonym, truth_inchikey)`` tuples of length up to + :data:`SAMPLE_SIZE`. + """ + total = comptox.conn.execute("SELECT COUNT(*) FROM chemicals").fetchone()[0] + step = max(1, total // (SAMPLE_SIZE * 6)) + cur = comptox.conn.execute( + f""" + SELECT CASRN, PREFERRED_NAME, INCHIKEY, IDENTIFIER + FROM chemicals + WHERE CASRN IS NOT NULL AND CASRN != '' + AND PREFERRED_NAME IS NOT NULL AND PREFERRED_NAME != '' + AND INCHIKEY IS NOT NULL AND INCHIKEY != '' + AND IDENTIFIER IS NOT NULL AND IDENTIFIER != '' + AND (rowid % {step}) = 1 + LIMIT {SAMPLE_SIZE * 6} + """ + ) + rows = [dict(r) for r in cur.fetchall()] + + # Drop CASRNs that map to more than one structure (the documented caveat). + cas_structs = {} + for r in rows: + cas_structs.setdefault(r["CASRN"], set()).add(r["INCHIKEY"]) + ambiguous = {c for c, ss in cas_structs.items() if len(ss) > 1} + + sample = [] + seen = set() + for r in rows: + if r["CASRN"] in ambiguous: + continue + syn = _pick_synonym(r) + if syn is None or syn.lower() in seen: + continue + seen.add(syn.lower()) + sample.append((syn, r["INCHIKEY"])) + if len(sample) >= SAMPLE_SIZE: + break + + if len(sample) < 10: # pragma: no cover - environment dependent + pytest.skip("Not enough CompTox synonym samples to run the regression.") + return sample + + +@pytest.mark.integration +@pytest.mark.slow +def test_name_resolution_has_zero_wrong_hits(comptox, synonym_sample): + """The resolver must never return a structurally different compound. + + Resolves each sampled synonym through a CompTox-only, precision-first + ``Search`` and asserts that no query yields a structure whose InChIKey + skeleton differs from the CASRN-derived ground truth. Not-found results + are permitted (recall is not asserted here). + """ + s = Search( + "name", + show_progress=False, + comptox=comptox, + chebi=None, + pubchem=None, + zeropm=None, + chembl=None, + ) + + queries = [syn for syn, _ in synonym_sample] + res = s.search(queries) + + wrong_hits = [] + correct = not_found = 0 + for (syn, truth_ik), (_, row) in zip(synonym_sample, res.iterrows()): + got_ik = row.get("InChIKey") + if not isinstance(got_ik, str): # None / NaN + not_found += 1 + continue + if _skeleton(got_ik) == _skeleton(truth_ik): + correct += 1 + else: + wrong_hits.append((syn, row.get("name"), truth_ik, got_ik)) + + # Recall (correct vs not_found) is informational; precision is the contract. + assert not wrong_hits, ( + f"{len(wrong_hits)} wrong hit(s) out of {len(synonym_sample)} " + f"(correct={correct}, not_found={not_found}). Examples: {wrong_hits[:5]}" + )