diff --git a/.gitignore b/.gitignore index 7d598aa..a6cf19c 100644 --- a/.gitignore +++ b/.gitignore @@ -14,9 +14,14 @@ models/ *.py[cod] .venv/ +.venv310/ /data/ /dumps/ /pgdata/ *.dump *.sql -downloads/ \ No newline at end of file +downloads/ + +tmp/ +out/ +.snakemake/ \ No newline at end of file diff --git a/database/Snakemake b/database/Snakemake new file mode 100644 index 0000000..f22bb63 --- /dev/null +++ b/database/Snakemake @@ -0,0 +1,240 @@ +"""RetroMol database-construction pipeline. + +Run from the repo root with: + + snakemake -s database/Snakemake --configfile database/config.yaml --cores 4 + +Fill in database/config.yaml's `sources` URLs before running. The pipeline: + + 1. create_db - empty DuckDB database + 2. download_* - fetch NPAtlas SDF + MIBiG JSON/GBK archives + 3. parse_npatlas } run RetroMol on NPAtlas compounds + 4. load_npatlas_compounds} turn results into "compound" db entries + 5. extract_mibig_compounds + parse_mibig_compounds } run RetroMol on MIBiG's compounds + 6. load_mibig_compounds } turn results into "compound" db entries (linked to MIBiG's URL) + 7. parse_mibig_gbks - antiSMASH GBKs -> linear module readouts (PARAS-annotated) + 8. load_mibig_bgcs - turn readouts into "bgc" db entries + +Steps 4, 6, and 8 all mutate the same DuckDB file, so they're chained through marker +files (rather than each declaring the database itself as `output`) to force +Snakemake to serialize them -- DuckDB doesn't support concurrent writers. +""" + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(workflow.basedir) / "scripts")) + +WORKDIR = Path(config["paths"]["workdir"]) +DB_PATH = Path(config["paths"]["database"]) +MARKERS = WORKDIR / "markers" + +RXN_RULES = config["paths"].get("reaction_rules") +MXN_RULES = config["paths"].get("matching_rules") +PARAS_MODEL_PATH = config["paths"].get("paras_model") + +PARAS_THRESHOLD = config["paras"]["threshold"] +PARAS_KEEP_TOP = config["paras"]["keep_top"] + +PARSE_COMPOUNDS_WORKERS = config["compute"]["parse_compounds_workers"] +PARSE_GBKS_WORKERS = config["compute"]["parse_gbks_workers"] + + +rule all: + input: + MARKERS / "bgcs_loaded.done" + + +# --------------------------------------------------------------------------- +# Step 1: empty database +# --------------------------------------------------------------------------- + +rule create_db: + output: + marker=touch(MARKERS / "db_created.done") + run: + import create_db + create_db.run(db_path=DB_PATH, overwrite=True) + + +# --------------------------------------------------------------------------- +# Step 2: downloads +# --------------------------------------------------------------------------- + +rule download_npatlas: + output: + raw=WORKDIR / "npatlas" / "download.raw", + extract_dir=directory(WORKDIR / "npatlas" / "extracted") + params: + url=config["sources"]["npatlas_sdf_url"] + run: + import download_sources + download_sources.run(url=params.url, download_path=output.raw, extract_dir=output.extract_dir) + + +rule resolve_npatlas_sdf: + input: + extract_dir=WORKDIR / "npatlas" / "extracted" + output: + sdf=WORKDIR / "npatlas" / "npatlas.sdf" + run: + import shutil + candidates = sorted(Path(input.extract_dir).rglob("*.sdf")) + if not candidates: + raise FileNotFoundError(f"no .sdf file found under {input.extract_dir}") + shutil.copy2(candidates[0], output.sdf) + + +rule download_mibig_json: + output: + raw=WORKDIR / "mibig_json" / "download.raw", + extract_dir=directory(WORKDIR / "mibig_json" / "extracted") + params: + url=config["sources"]["mibig_json_url"] + run: + import download_sources + download_sources.run(url=params.url, download_path=output.raw, extract_dir=output.extract_dir) + + +rule download_mibig_gbk: + output: + raw=WORKDIR / "mibig_gbk" / "download.raw", + extract_dir=directory(WORKDIR / "mibig_gbk" / "extracted") + params: + url=config["sources"]["mibig_gbk_url"] + run: + import download_sources + download_sources.run(url=params.url, download_path=output.raw, extract_dir=output.extract_dir) + + +# --------------------------------------------------------------------------- +# Steps 3-4: NPAtlas compounds +# --------------------------------------------------------------------------- + +rule parse_npatlas: + input: + sdf=WORKDIR / "npatlas" / "npatlas.sdf" + output: + results=WORKDIR / "npatlas" / "results.jsonl" + threads: PARSE_COMPOUNDS_WORKERS + run: + import parse_compounds + parse_compounds.run( + input_path=input.sdf, + input_format="sdf", + output_path=output.results, + reaction_rules_path=RXN_RULES, + matching_rules_path=MXN_RULES, + workers=threads, + ) + + +rule load_npatlas_compounds: + input: + results=WORKDIR / "npatlas" / "results.jsonl", + db_created=MARKERS / "db_created.done" + output: + marker=touch(MARKERS / "npatlas_loaded.done") + run: + import load_compounds + load_compounds.run( + results_path=input.results, + db_path=DB_PATH, + source="npatlas", + reaction_rules_path=RXN_RULES, + matching_rules_path=MXN_RULES, + ) + + +# --------------------------------------------------------------------------- +# Steps 5-6: MIBiG compounds +# --------------------------------------------------------------------------- + +rule extract_mibig_compounds: + input: + extract_dir=WORKDIR / "mibig_json" / "extracted" + output: + compounds=WORKDIR / "mibig_json" / "compounds.jsonl" + run: + import extract_mibig_compounds + extract_mibig_compounds.run(mibig_json_dir=input.extract_dir, output_path=output.compounds) + + +rule parse_mibig_compounds: + input: + compounds=WORKDIR / "mibig_json" / "compounds.jsonl" + output: + results=WORKDIR / "mibig_json" / "results.jsonl" + threads: PARSE_COMPOUNDS_WORKERS + run: + import parse_compounds + parse_compounds.run( + input_path=input.compounds, + input_format="jsonl", + output_path=output.results, + reaction_rules_path=RXN_RULES, + matching_rules_path=MXN_RULES, + workers=threads, + ) + + +rule load_mibig_compounds: + input: + results=WORKDIR / "mibig_json" / "results.jsonl", + # Depends on the GBK-derived version map, not just the npatlas-load marker, + # since MIBiG URLs need an accession's version (only present in the GBKs' + # ACCESSION/VERSION line, see common.split_accession_version). + versions=WORKDIR / "mibig_gbk" / "versions.json", + prev=MARKERS / "npatlas_loaded.done" + output: + marker=touch(MARKERS / "mibig_compounds_loaded.done") + run: + import load_compounds + load_compounds.run( + results_path=input.results, + db_path=DB_PATH, + source="mibig", + reaction_rules_path=RXN_RULES, + matching_rules_path=MXN_RULES, + mibig_versions_path=input.versions, + ) + + +# --------------------------------------------------------------------------- +# Steps 7-8: MIBiG BGCs +# --------------------------------------------------------------------------- + +rule parse_mibig_gbks: + input: + gbk_dir=WORKDIR / "mibig_gbk" / "extracted" + output: + readouts=WORKDIR / "mibig_gbk" / "readouts.jsonl", + versions=WORKDIR / "mibig_gbk" / "versions.json" + threads: PARSE_GBKS_WORKERS + run: + import parse_gbks + parse_gbks.run( + gbk_dir=input.gbk_dir, + readouts_output_path=output.readouts, + versions_output_path=output.versions, + paras_threshold=PARAS_THRESHOLD, + paras_keep_top=PARAS_KEEP_TOP, + paras_model_path=PARAS_MODEL_PATH, + workers=threads, + ) + + +rule load_mibig_bgcs: + input: + readouts=WORKDIR / "mibig_gbk" / "readouts.jsonl", + prev=MARKERS / "mibig_compounds_loaded.done" + output: + marker=touch(MARKERS / "bgcs_loaded.done") + run: + import load_bgcs + load_bgcs.run( + readouts_path=input.readouts, + db_path=DB_PATH, + reaction_rules_path=RXN_RULES, + matching_rules_path=MXN_RULES, + ) diff --git a/database/config.yaml b/database/config.yaml new file mode 100644 index 0000000..ead4c85 --- /dev/null +++ b/database/config.yaml @@ -0,0 +1,30 @@ +sources: + # Direct download links. MIBiG ships as tar.gz archives (one file per BGC inside); + # NPAtlas ships as a single SDF (optionally gzipped) -- download_sources.py handles + # extraction based on the URL's extension either way. + npatlas_sdf_url: "https://www.npatlas.org/static/downloads/NPAtlas_download.sdf" + mibig_json_url: "https://dl.secondarymetabolites.org/mibig/mibig_json_4.0.tar.gz" + mibig_gbk_url: "https://dl.secondarymetabolites.org/mibig/mibig_gbk_4.0.tar.gz" + +paths: + # Final DuckDB database produced by the pipeline. + database: "database/output/retromol.duckdb" + + # Scratch space for downloads and intermediate per-step results. + workdir: "database/work" + + # null -> ParasModel downloads/caches its own default model. + paras_model: null + + # null -> RuleSet.load_default()'s bundled reaction/matching rules. + reaction_rules: null + matching_rules: null + +paras: + threshold: 0.1 + keep_top: 3 + +compute: + # Both are embarrassingly parallel over independent compounds/files. + parse_compounds_workers: 4 + parse_gbks_workers: 2 diff --git a/database/scripts/common.py b/database/scripts/common.py new file mode 100644 index 0000000..14f3d4b --- /dev/null +++ b/database/scripts/common.py @@ -0,0 +1,177 @@ +"""Shared helpers for the database-construction pipeline. + +The fingerprint recipe here (vocabulary = every matching rule's name + pseudonyms, +Fingerprinter with n_bits=FINGERPRINT_SIZE, n_hashes=2) must match +gui/src/server/routes/discovery.py's `_build_context` exactly -- that's what the +webapp uses to encode a query at search time. Deviating here would silently make +every fingerprint stored by this pipeline incomparable to a live query. +""" + +from multiprocessing import Pool +from pathlib import Path +from typing import Any, Iterable, Iterator + +from rdkit import RDLogger + +from retromol.io.streaming import ResultEvent, _init_worker, _process_compound, _task_buffered_iterator +from retromol.model.result import Result +from retromol.model.rules import MatchingRule, RuleSet +from retromol_database.duckdb import FINGERPRINT_SIZE +from retromol_fingerprint.fingerprint import Fingerprinter, Vocabulary + +# Silences RDKit's kekulization/valence/etc. warnings in *this* (single) process -- +# every pipeline script imports common, so this alone covers create_db.py, +# load_compounds.py, load_bgcs.py, and parse_gbks.py's main process. It does NOT +# reliably reach parse_compounds.py's multiprocessing workers: those are spawned +# fresh by retromol.io.streaming.run_retromol_stream's own Pool, with its own +# initializer, and whether a fresh worker re-runs this module-level call at all +# depends on whether the *original* entry point was this script directly (works) +# or something else re-importing it, like Snakemake's generated run: script +# (doesn't) -- see run_retromol_stream_quiet below for the reliable fix. +RDLogger.DisableLog("rdApp.*") + +# Group-level PKS pseudo-tokens a BGC's PKS module resolves to (see +# retromol_antismash.modules.PKSExtenderUnit / module_primary_sequence_tokens). +# Not matching-rule names themselves, but already part of the fingerprint +# vocabulary as pseudonyms of every rule at that reduction level. +PK_GROUP_TOKENS = ("PK_A", "PK_B", "PK_C", "PK_D") + +MIBIG_URL_TEMPLATE = "https://mibig.secondarymetabolites.org/repository/{accession}.{version}/index.html#r1c1" +NPATLAS_URL_TEMPLATE = "https://www.npatlas.org/explore/compounds/{npaid}" + + +def load_ruleset( + reaction_rules_path: str | Path | None, + matching_rules_path: str | Path | None, + match_stereochemistry: bool = False, +) -> RuleSet: + """Load a RuleSet, falling back to RetroMol's bundled default rules when a path is None/empty.""" + return RuleSet.load_from_files( + reaction_rules_path=reaction_rules_path or None, + matching_rules_path=matching_rules_path or None, + match_stereochemistry=match_stereochemistry, + ) + + +def build_fingerprint_context(ruleset: RuleSet) -> tuple[dict[str, MatchingRule], Fingerprinter]: + """Build the (name_to_rule, fingerprinter) pair used to fingerprint primary sequences.""" + vocab_tokens: set[str] = set() + name_to_rule: dict[str, MatchingRule] = {} + + for rule in ruleset.matching_rules: + vocab_tokens.add(rule.name) + vocab_tokens.update(rule.pseudonyms) + name_to_rule.setdefault(rule.name, rule) + + vocab = Vocabulary(vocab_tokens) + fingerprinter = Fingerprinter(vocab, n_bits=FINGERPRINT_SIZE, n_hashes=2) + return name_to_rule, fingerprinter + + +def per_monomer_tokens(name: str, name_to_rule: dict[str, MatchingRule]) -> list[str]: + """Fingerprinting token list for one compound primary-sequence block (mirrors discovery.py's `_per_monomer_tokens`).""" + if name in PK_GROUP_TOKENS: + return [name, "PK"] + + rule = name_to_rule.get(name) + if rule is None: + return [] + + tokens = {rule.name} + tokens.update(rule.pseudonyms) + return list(tokens) + + +def find_key_ci(record: dict[str, Any], substrings: list[str]) -> str | None: + """Find the first key in `record` whose lowercased form contains any of `substrings`.""" + for key in record: + lowered = str(key).lower() + if any(sub in lowered for sub in substrings): + return key + return None + + +def mibig_url(accession: str | None, version: str | None) -> str | None: + if not accession or not version: + return None + return MIBIG_URL_TEMPLATE.format(accession=accession, version=version) + + +def npatlas_url(npaid: str | None) -> str | None: + if not npaid: + return None + return NPATLAS_URL_TEMPLATE.format(npaid=npaid) + + +def primary_sequences_from_result(result: Result, min_length: int = 2) -> list[list[str]]: + """ + Every candidate primary sequence for a parsed compound, read directly off + `result.linear_readout.paths` -- no backbone reconstruction involved, that's a + display-only concern this pipeline has no use for. `result.linear_readout` is + already computed by retromol.pipelines.parsing.run_retromol (it's just a field + on Result), each path is one candidate ordering of monomers through the + molecule, and each becomes its own db entry (see load_compounds.py). An + unidentified node is named "X", the same convention used everywhere else in + RetroMol. + + `result.linear_readout` includes single-node paths for tailoring events that + don't connect to any chain (e.g. a lone "glycosylation" or "methylation") -- + real for the molecule, but not a "sequence" in any useful sense, so those are + dropped by the `min_length` floor. + + :param result: a parsed RetroMol Result + :param min_length: drop paths shorter than this (default 2) + :return: one name list per path meeting `min_length` + """ + return [ + [node.identity.matched_rule.name if node.is_identified else "X" for node in path] + for path in result.linear_readout.paths + if len(path) >= min_length + ] + + +def _init_worker_quiet(ruleset: RuleSet) -> None: + """Worker-process initializer: set up the ruleset global exactly like retromol.io.streaming's own + _init_worker does, then also disable RDKit logging -- unlike a module-level RDLogger call, this is + guaranteed to run once per worker process no matter how the pool was launched.""" + _init_worker(ruleset) + RDLogger.DisableLog("rdApp.*") + + +def run_retromol_stream_quiet( + ruleset: RuleSet, + row_iter: Iterable[dict[str, Any]], + smiles_col: str = "smiles", + workers: int = 1, + batch_size: int = 2000, + pool_chunksize: int = 50, + maxtasksperchild: int = 2000, +) -> Iterator[ResultEvent]: + """ + Drop-in replacement for retromol.io.streaming.run_retromol_stream that also + disables RDKit's C-level logging inside every worker process. Reuses that + module's own batching/worker-task functions -- only the Pool's initializer + differs (see _init_worker_quiet). + """ + with Pool( + processes=workers, + initializer=_init_worker_quiet, + initargs=(ruleset,), + maxtasksperchild=maxtasksperchild, + ) as pool: + for task_batch in _task_buffered_iterator(row_iter, smiles_col=smiles_col, batch_size=batch_size): + for serialized, err in pool.imap_unordered(_process_compound, task_batch, chunksize=pool_chunksize): + yield ResultEvent(serialized, err) + + +def split_accession_version(record_id: str) -> tuple[str, str | None]: + """ + Split a GenBank-style "ACCESSION.VERSION" id (e.g. "BGC0000001.5") in two. + + :param record_id: the record id, as set on retromol_antismash Region/LinearReadout.id + :return: (accession, version) -- version is None if record_id has no "." suffix + """ + if "." in record_id: + accession, version = record_id.rsplit(".", 1) + return accession, version + return record_id, None diff --git a/database/scripts/create_db.py b/database/scripts/create_db.py new file mode 100644 index 0000000..1f96d06 --- /dev/null +++ b/database/scripts/create_db.py @@ -0,0 +1,26 @@ +"""Step 1: create an empty RetroMol database.""" + +import argparse +from pathlib import Path + +from retromol_database.duckdb import RetroMolDuckDB + + +def run(db_path: str | Path, overwrite: bool = True) -> None: + db_path = Path(db_path) + db_path.parent.mkdir(parents=True, exist_ok=True) + db = RetroMolDuckDB.create(db_path, overwrite=overwrite) + db.close() + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--db-path", required=True) + ap.add_argument("--overwrite", action="store_true") + args = ap.parse_args() + + run(db_path=args.db_path, overwrite=args.overwrite) + + +if __name__ == "__main__": + main() diff --git a/database/scripts/download_sources.py b/database/scripts/download_sources.py new file mode 100644 index 0000000..1237223 --- /dev/null +++ b/database/scripts/download_sources.py @@ -0,0 +1,88 @@ +"""Step 2: download a source file/archive and optionally extract it. + +Handles the two shapes the pipeline needs: a single (possibly gzipped) file, like +the NPAtlas SDF, and a tar.gz/zip archive containing many files, like MIBiG's JSON +and GenBank bundles. +""" + +import argparse +import gzip +import shutil +import tarfile +import zipfile +from pathlib import Path +from urllib.parse import urlsplit + +import requests + + +def download(url: str, dest: str | Path) -> None: + dest = Path(dest) + dest.parent.mkdir(parents=True, exist_ok=True) + with requests.get(url, stream=True, timeout=300) as resp: + resp.raise_for_status() + with open(dest, "wb") as fh: + for chunk in resp.iter_content(chunk_size=1024 * 1024): + if chunk: + fh.write(chunk) + + +def _url_filename(url: str) -> str: + """Best-effort original filename from a URL's path (ignores query strings).""" + name = Path(urlsplit(url).path).name + return name or "download" + + +def extract(path: str | Path, out_dir: str | Path, source_name: str | None = None) -> None: + """ + Extract/decompress a downloaded file, or copy it through unchanged. + + :param path: the downloaded file on disk -- often a fixed, extension-less name + (e.g. "download.raw") the caller chose for the raw download, so it can't be + used to decide *how* to extract or to name a plain copy/decompressed output. + :param out_dir: directory to extract/copy into. + :param source_name: the original filename (e.g. from the source URL), used both + to pick the extraction strategy and, for the plain-copy/gzip cases, to name + the resulting file so its extension survives (e.g. "*.sdf" stays findable). + Falls back to `path`'s own name if not given. + """ + path = Path(path) + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + display_name = source_name or path.name + name = display_name.lower() + + if name.endswith((".tar.gz", ".tgz", ".tar")): + mode = "r:gz" if name.endswith((".tar.gz", ".tgz")) else "r" + with tarfile.open(path, mode) as tf: + tf.extractall(out_dir) + elif name.endswith(".zip"): + with zipfile.ZipFile(path) as zf: + zf.extractall(out_dir) + elif name.endswith(".gz"): + out_path = out_dir / Path(display_name).with_suffix("").name + with gzip.open(path, "rb") as fin, open(out_path, "wb") as fout: + shutil.copyfileobj(fin, fout) + else: + shutil.copy2(path, out_dir / display_name) + + +def run(url: str, download_path: str | Path, extract_dir: str | Path | None = None) -> None: + download(url, download_path) + if extract_dir is not None: + extract(download_path, extract_dir, source_name=_url_filename(url)) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--url", required=True) + ap.add_argument("--download-path", required=True, help="where the raw downloaded file is stored") + ap.add_argument("--extract-dir", default=None, help="if set, extract/decompress the download into this directory") + args = ap.parse_args() + + run(url=args.url, download_path=args.download_path, extract_dir=args.extract_dir) + + +if __name__ == "__main__": + main() diff --git a/database/scripts/extract_mibig_compounds.py b/database/scripts/extract_mibig_compounds.py new file mode 100644 index 0000000..e35fbc6 --- /dev/null +++ b/database/scripts/extract_mibig_compounds.py @@ -0,0 +1,98 @@ +"""Step 5 (prep): flatten MIBiG's per-BGC JSON files into one compounds JSONL. + +MIBiG's JSON schema has shifted over releases (older releases nest everything +under a "cluster" key; newer ones are flatter), so field lookup here is +deliberately tolerant -- entries that don't have a resolvable accession/SMILES +are skipped (and counted) rather than raising. +""" + +import argparse +import json +import logging +from pathlib import Path +from typing import Any, Iterator + +log = logging.getLogger(__name__) + + +def _entry_root(data: dict[str, Any]) -> dict[str, Any]: + """Older MIBiG JSON nests everything under "cluster"; newer releases are flat.""" + cluster = data.get("cluster") + return cluster if isinstance(cluster, dict) else data + + +def _accession(data: dict[str, Any], root: dict[str, Any]) -> str | None: + return root.get("mibig_accession") or root.get("accession") or data.get("accession") + + +def _iter_compound_records(path: Path) -> Iterator[dict[str, Any]]: + with open(path) as fh: + data = json.load(fh) + + if not isinstance(data, dict): + return + + root = _entry_root(data) + accession = _accession(data, root) + compounds = root.get("compounds") + + if not accession or not isinstance(compounds, list): + return + + for idx, compound in enumerate(compounds): + if not isinstance(compound, dict): + continue + + smiles = compound.get("chem_struct") or compound.get("smiles") or compound.get("structure") + if not smiles: + continue + + name = compound.get("compound") or compound.get("name") or f"{accession} compound {idx + 1}" + + yield { + "id": f"{accession}:{idx}", + "smiles": smiles, + "name": name, + "mibig_accession": accession, + } + + +def run(mibig_json_dir: str | Path, output_path: str | Path) -> None: + mibig_json_dir = Path(mibig_json_dir) + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + json_files = sorted(mibig_json_dir.rglob("*.json")) + written = 0 + skipped = 0 + + with open(output_path, "w") as out: + for path in json_files: + try: + had_any = False + for record in _iter_compound_records(path): + out.write(json.dumps(record) + "\n") + written += 1 + had_any = True + if not had_any: + skipped += 1 + except Exception: + log.exception("failed to parse MIBiG JSON file: %s", path) + skipped += 1 + + log.info("extract_mibig_compounds: wrote %d compound records, skipped %d files", written, skipped) + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--mibig-json-dir", required=True) + ap.add_argument("--output", required=True) + args = ap.parse_args() + + run(mibig_json_dir=args.mibig_json_dir, output_path=args.output) + + +if __name__ == "__main__": + main() diff --git a/database/scripts/load_bgcs.py b/database/scripts/load_bgcs.py new file mode 100644 index 0000000..33ffb12 --- /dev/null +++ b/database/scripts/load_bgcs.py @@ -0,0 +1,95 @@ +"""Step 8: turn parsed BGC readouts into database entries. + +Each antiSMASH region becomes one "bgc" entry. bgc_primary_sequence maps every +module's predicted substrate onto the same matching-rule vocabulary compound +primary sequences use (PKS modules resolve to a reduction-level pseudonym like +"PK_A", NRPS modules resolve by matching PARAS' predicted substrate structurally +against the ruleset) -- see retromol_antismash.modules for why the two module +types resolve differently. That shared vocabulary is what makes a BGC entry's +fingerprint and primary sequence comparable to a compound's. +""" + +import argparse +import json +import logging +from pathlib import Path + +from common import build_fingerprint_context, load_ruleset +from retromol_antismash.modules import LinearReadout, bgc_primary_sequence +from retromol_database.duckdb import RetroMolDuckDB + +log = logging.getLogger(__name__) + + +def run( + readouts_path: str | Path, + db_path: str | Path, + reaction_rules_path: str | Path | None, + matching_rules_path: str | Path | None, + match_stereochemistry: bool = False, + include_raw_gbk: bool = True, +) -> None: + ruleset = load_ruleset(reaction_rules_path, matching_rules_path, match_stereochemistry) + _, fingerprinter = build_fingerprint_context(ruleset) + + added = 0 + skipped = 0 + + db = RetroMolDuckDB.open(db_path) + try: + with open(readouts_path) as fh: + for line in fh: + line = line.strip() + if not line: + continue + + entry = json.loads(line) + readout = LinearReadout.from_dict(entry["readout"]) + names, tokens = bgc_primary_sequence(readout, ruleset) + + if not names: + skipped += 1 + continue + + fp = fingerprinter.encode(tokens) + name = f"{entry['accession']} ({readout.id})" if entry.get("accession") else readout.id + + db.add_entry( + name=name, + url=entry.get("url"), + raw=entry.get("raw_gbk") if include_raw_gbk else None, + entry_type="bgc", + primary_sequence=names, + fingerprint=fp, + ) + added += 1 + finally: + db.close() + + log.info("load_bgcs: added=%d skipped=%d", added, skipped) + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--readouts", required=True) + ap.add_argument("--db-path", required=True) + ap.add_argument("--rxn-rules", default=None) + ap.add_argument("--mxn-rules", default=None) + ap.add_argument("--match-stereochemistry", action="store_true") + ap.add_argument("--no-raw-gbk", action="store_true") + args = ap.parse_args() + + run( + readouts_path=args.readouts, + db_path=args.db_path, + reaction_rules_path=args.rxn_rules, + matching_rules_path=args.mxn_rules, + match_stereochemistry=args.match_stereochemistry, + include_raw_gbk=not args.no_raw_gbk, + ) + + +if __name__ == "__main__": + main() diff --git a/database/scripts/load_compounds.py b/database/scripts/load_compounds.py new file mode 100644 index 0000000..4337c49 --- /dev/null +++ b/database/scripts/load_compounds.py @@ -0,0 +1,153 @@ +"""Steps 4 & 6: turn parsed compound results into database entries. + +For each RetroMol Result, every candidate primary sequence -- one per path in +result.linear_readout.paths, read directly off the Result (see +common.primary_sequences_from_result) -- becomes its own "compound" entry. +Ambiguous parses intentionally produce multiple queryable entries, the same +convention the webapp uses for an uploaded compound with more than one candidate +reading. `raw` is always the original input SMILES, the same for every entry a +given compound produces. +""" + +import argparse +import json +import logging +from pathlib import Path +from typing import Literal + +from tqdm import tqdm + +from common import ( + build_fingerprint_context, + find_key_ci, + load_ruleset, + mibig_url, + npatlas_url, + per_monomer_tokens, + primary_sequences_from_result, +) +from retromol.model.result import Result +from retromol_database.duckdb import RetroMolDuckDB + +log = logging.getLogger(__name__) + + +def _npatlas_name_and_url(props: dict) -> tuple[str | None, str | None]: + npaid_key = find_key_ci(props, ["npaid"]) + npaid = props.get(npaid_key) if npaid_key else None + + name_key = find_key_ci(props, ["original_name", "compound_name", "name"]) + name = props.get(name_key) if name_key else None + + return name, npatlas_url(npaid) + + +def _mibig_name_and_url(props: dict, versions: dict[str, str]) -> tuple[str | None, str | None]: + accession = props.get("mibig_accession") + version = versions.get(accession) if accession else None + name = props.get("name") + return name, mibig_url(accession, version) + + +def run( + results_path: str | Path, + db_path: str | Path, + source: Literal["npatlas", "mibig"], + reaction_rules_path: str | Path | None, + matching_rules_path: str | Path | None, + match_stereochemistry: bool = False, + mibig_versions_path: str | Path | None = None, + log_every: int = 1000, +) -> None: + ruleset = load_ruleset(reaction_rules_path, matching_rules_path, match_stereochemistry) + name_to_rule, fingerprinter = build_fingerprint_context(ruleset) + + versions: dict[str, str] = {} + if source == "mibig" and mibig_versions_path is not None: + with open(mibig_versions_path) as fh: + versions = json.load(fh) + + compounds = 0 + added = 0 + skipped = 0 + + db = RetroMolDuckDB.open(db_path) + try: + with open(results_path) as fh: + with tqdm(desc=f"load_compounds[{source}]", unit="cmpd") as pbar: + for line in fh: + line = line.strip() + if not line: + continue + + result = Result.from_dict(json.loads(line)) + props = result.submission.props or {} + + if source == "npatlas": + name, url = _npatlas_name_and_url(props) + else: + name, url = _mibig_name_and_url(props, versions) + + name = name or result.submission.name or result.submission.inchikey + + for names in primary_sequences_from_result(result): + if not names: + skipped += 1 + continue + + tokens = [per_monomer_tokens(n, name_to_rule) for n in names] + fp = fingerprinter.encode(tokens) + + db.add_entry( + name=name, + url=url, + raw=result.submission.smiles, + entry_type="compound", + primary_sequence=names, + fingerprint=fp, + ) + added += 1 + + compounds += 1 + pbar.update(1) + pbar.set_postfix(added=added, skipped=skipped) + + if log_every > 0 and compounds % log_every == 0: + log.info( + "load_compounds[%s]: processed %d compounds (added=%d skipped=%d)", + source, compounds, added, skipped, + ) + finally: + db.close() + + log.info("load_compounds[%s]: compounds=%d added=%d skipped=%d", source, compounds, added, skipped) + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--results", required=True) + ap.add_argument("--db-path", required=True) + ap.add_argument("--source", choices=["npatlas", "mibig"], required=True) + ap.add_argument("--rxn-rules", default=None) + ap.add_argument("--mxn-rules", default=None) + ap.add_argument("--match-stereochemistry", action="store_true") + ap.add_argument("--mibig-versions", default=None, help="required when --source=mibig") + ap.add_argument("--log-every", type=int, default=1000, help="log a progress line every N compounds (0 to disable)") + args = ap.parse_args() + + run( + results_path=args.results, + db_path=args.db_path, + source=args.source, + reaction_rules_path=args.rxn_rules, + matching_rules_path=args.mxn_rules, + match_stereochemistry=args.match_stereochemistry, + mibig_versions_path=args.mibig_versions, + log_every=args.log_every, + ) + + +if __name__ == "__main__": + main() diff --git a/database/scripts/parse_compounds.py b/database/scripts/parse_compounds.py new file mode 100644 index 0000000..6d4b33a --- /dev/null +++ b/database/scripts/parse_compounds.py @@ -0,0 +1,106 @@ +"""Steps 3 & 5: run RetroMol over a batch of compounds (NPAtlas SDF or MIBiG compounds JSONL). + +Shared between both compound sources -- only the input format differs. This is the +slow step (one retrosynthetic analysis per compound), so it's parallelized with +`workers` worker processes via common.run_retromol_stream_quiet -- the same +worker/batching machinery the `retromol` CLI's batch mode uses +(retromol.io.streaming.run_retromol_stream), but with RDKit's C-level logging also +disabled inside every worker (see common.py for why that needs its own initializer). +""" + +import argparse +import json +import logging +from pathlib import Path +from typing import Literal + +from tqdm import tqdm + +from common import load_ruleset, run_retromol_stream_quiet +from retromol.io.streaming import stream_json_records, stream_sdf_records + +log = logging.getLogger(__name__) + + +def run( + input_path: str | Path, + input_format: Literal["sdf", "jsonl"], + output_path: str | Path, + reaction_rules_path: str | Path | None, + matching_rules_path: str | Path | None, + match_stereochemistry: bool = False, + smiles_col: str = "smiles", + workers: int = 1, + batch_size: int = 2000, + log_every: int = 1000, +) -> None: + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + ruleset = load_ruleset(reaction_rules_path, matching_rules_path, match_stereochemistry) + + if input_format == "sdf": + row_iter = stream_sdf_records(str(input_path)) + else: + row_iter = stream_json_records(str(input_path), jsonl=True) + + successes = 0 + errors = 0 + + with open(output_path, "w", buffering=1) as out: + with tqdm(desc="parse_compounds", unit="cmpd") as pbar: + for evt in run_retromol_stream_quiet( + ruleset=ruleset, + row_iter=row_iter, + smiles_col=smiles_col, + workers=workers, + batch_size=batch_size, + ): + if evt.error is not None: + errors += 1 + elif evt.result is not None: + out.write(json.dumps(evt.result) + "\n") + successes += 1 + + pbar.update(1) + pbar.set_postfix(ok=successes, err=errors) + + total = successes + errors + if log_every > 0 and total % log_every == 0: + log.info("parse_compounds: parsed %d (successes=%d errors=%d)", total, successes, errors) + + log.info("parse_compounds: successes=%d errors=%d -> %s", successes, errors, output_path) + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--input", required=True) + ap.add_argument("--input-format", choices=["sdf", "jsonl"], required=True) + ap.add_argument("--output", required=True) + ap.add_argument("--rxn-rules", default=None) + ap.add_argument("--mxn-rules", default=None) + ap.add_argument("--match-stereochemistry", action="store_true") + ap.add_argument("--smiles-col", default="smiles") + ap.add_argument("--workers", type=int, default=1) + ap.add_argument("--batch-size", type=int, default=2000) + ap.add_argument("--log-every", type=int, default=1000, help="log a progress line every N compounds (0 to disable)") + args = ap.parse_args() + + run( + input_path=args.input, + input_format=args.input_format, + output_path=args.output, + reaction_rules_path=args.rxn_rules, + matching_rules_path=args.mxn_rules, + match_stereochemistry=args.match_stereochemistry, + smiles_col=args.smiles_col, + workers=args.workers, + batch_size=args.batch_size, + log_every=args.log_every, + ) + + +if __name__ == "__main__": + main() diff --git a/database/scripts/parse_gbks.py b/database/scripts/parse_gbks.py new file mode 100644 index 0000000..0ee1e07 --- /dev/null +++ b/database/scripts/parse_gbks.py @@ -0,0 +1,162 @@ +"""Step 7: parse antiSMASH GenBank files into linear module readouts. + +For each region: PARAS predicts NRPS A-domain substrate specificities +(retromol_antismash.inference.registry.annotate_region), then +retromol_antismash.modules.linear_readout collects PKS/NRPS modules in +biosynthetic order. One file per worker process (PARAS model loading is the +expensive per-process setup cost, so it's paid once per worker, not once per file). + +Emits two outputs: +- readouts JSONL: one line per antiSMASH region, with its LinearReadout, the raw + GenBank text of its source file, and the MIBiG accession/version/URL parsed out + of the region id (e.g. "BGC0000001.5"). +- a small accession -> version JSON map, reused by load_compounds.py to link MIBiG + compound entries to the same BGC page. +""" + +import argparse +import json +import logging +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path + +from rdkit import RDLogger + +from common import mibig_url, split_accession_version +from retromol_antismash.inference.model_paras import ParasModel +from retromol_antismash.inference.registry import annotate_region +from retromol_antismash.io import AntiSmashOptions, parse_antismash_gbk +from retromol_antismash.modules import linear_readout + +log = logging.getLogger(__name__) + +GBK_GLOBS = ("*.gbk", "*.gb", "*.gbff") + +_G_PARAS_MODEL: ParasModel | None = None + + +def _init_worker(paras_threshold: float, paras_keep_top: int, paras_model_path: str | None, paras_cache_dir: str) -> None: + global _G_PARAS_MODEL + # Belt-and-braces: importing this module (to resolve _init_worker as this pool's + # initializer) already re-runs common.py's own RDLogger.DisableLog at the top of + # this file's import chain, but this initializer is what the pool guarantees runs + # once per worker no matter what -- see common.run_retromol_stream_quiet's + # docstring for why that guarantee matters more than it might seem. + RDLogger.DisableLog("rdApp.*") + _G_PARAS_MODEL = ParasModel( + threshold=paras_threshold, + keep_top=paras_keep_top, + model_path=paras_model_path, + cache_dir=paras_cache_dir, + ) + + +def _process_file(path_str: str) -> tuple[list[dict], dict[str, str], str | None]: + """Parse one GenBank file. Returns (entries, accession->version, error message).""" + path = Path(path_str) + try: + raw_gbk = path.read_text() + regions = parse_antismash_gbk(path, AntiSmashOptions()) + + entries: list[dict] = [] + versions: dict[str, str] = {} + + for region in regions: + annotate_region(region, domain_models=[_G_PARAS_MODEL]) + readout = linear_readout(region) + + accession, version = split_accession_version(region.id) + if version is not None: + versions[accession] = version + + entries.append({ + "accession": accession, + "version": version, + "url": mibig_url(accession, version), + "file_name": region.file_name, + "raw_gbk": raw_gbk, + "readout": readout.to_dict(), + }) + + return entries, versions, None + except Exception as e: + return [], {}, f"{path}: {e}" + + +def run( + gbk_dir: str | Path, + readouts_output_path: str | Path, + versions_output_path: str | Path, + paras_threshold: float = 0.1, + paras_keep_top: int = 3, + paras_model_path: str | Path | None = None, + paras_cache_dir: str | Path = "paras_cache", + workers: int = 1, +) -> None: + gbk_dir = Path(gbk_dir) + readouts_output_path = Path(readouts_output_path) + versions_output_path = Path(versions_output_path) + readouts_output_path.parent.mkdir(parents=True, exist_ok=True) + versions_output_path.parent.mkdir(parents=True, exist_ok=True) + + paths = sorted({p for pattern in GBK_GLOBS for p in gbk_dir.rglob(pattern)}) + + all_versions: dict[str, str] = {} + n_entries = 0 + n_errors = 0 + + init_args = ( + paras_threshold, + paras_keep_top, + str(paras_model_path) if paras_model_path else None, + str(paras_cache_dir), + ) + + with open(readouts_output_path, "w", buffering=1) as out: + with ProcessPoolExecutor(max_workers=workers, initializer=_init_worker, initargs=init_args) as pool: + futures = [pool.submit(_process_file, str(p)) for p in paths] + for fut in as_completed(futures): + entries, versions, error = fut.result() + if error is not None: + log.error("parse_gbks: %s", error) + n_errors += 1 + continue + for entry in entries: + out.write(json.dumps(entry) + "\n") + n_entries += 1 + all_versions.update(versions) + + with open(versions_output_path, "w") as fh: + json.dump(all_versions, fh, indent=2, sort_keys=True) + + log.info("parse_gbks: files=%d regions=%d errors=%d", len(paths), n_entries, n_errors) + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--gbk-dir", required=True) + ap.add_argument("--readouts-output", required=True) + ap.add_argument("--versions-output", required=True) + ap.add_argument("--paras-threshold", type=float, default=0.1) + ap.add_argument("--paras-keep-top", type=int, default=3) + ap.add_argument("--paras-model-path", default=None) + ap.add_argument("--paras-cache-dir", default="paras_cache") + ap.add_argument("--workers", type=int, default=1) + args = ap.parse_args() + + run( + gbk_dir=args.gbk_dir, + readouts_output_path=args.readouts_output, + versions_output_path=args.versions_output, + paras_threshold=args.paras_threshold, + paras_keep_top=args.paras_keep_top, + paras_model_path=args.paras_model_path, + paras_cache_dir=args.paras_cache_dir, + workers=args.workers, + ) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 66206b5..113d56a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,8 +15,8 @@ dependencies = [ "biopython", "pyhmmer", "tqdm", - "numpy<2", # scikit-learn 1.2.0 requires numpy<2 - "scikit-learn==1.2.0", + "numpy", + "scikit-learn", "scipy", "pandas", "pyyaml", diff --git a/src/retromol/data/mxn.yml b/src/retromol/data/mxn.yml index 58abc23..3f67042 100644 --- a/src/retromol/data/mxn.yml +++ b/src/retromol/data/mxn.yml @@ -437,12 +437,12 @@ smiles: 'C/C(C(O)=O)=C\SO' display_smiles: 'C/C(C(S[2*])=O)=C\[1*]' stereochemistry: true - pseudonyms: ["PK_C", "PK"] + pseudonyms: ["C2", "PK_C", "PK"] - name: C^Z2 smiles: 'C/C(C(O)=O)=C/SO' display_smiles: 'C/C(C(S[2*])=O)=C/[1*]' stereochemistry: true - pseudonyms: ["PK_C", "PK"] + pseudonyms: ["C2", "PK_C", "PK"] - name: C2 smiles: "CC(=CSO)C(=O)O" display_smiles: "CC(=C[1*])C(=O)S[2*]" @@ -452,12 +452,12 @@ smiles: 'CC/C(C(O)=O)=C\SO' display_smiles: 'CC/C(C(S[2*])=O)=C\[1*]' stereochemistry: true - pseudonyms: ["PK_C", "PK"] + pseudonyms: ["C4", "PK_C", "PK"] - name: C^Z4 smiles: 'CC/C(C(O)=O)=C/SO' display_smiles: 'CC/C(C(S[2*])=O)=C/[1*]' stereochemistry: true - pseudonyms: ["PK_C", "PK"] + pseudonyms: ["C4", "PK_C", "PK"] - name: C4 smiles: "CCC(=CSO)C(=O)O" display_smiles: "CCC(=C[1*])C(=O)S[2*]" @@ -467,12 +467,12 @@ smiles: 'O=C(/C(CO)=C/SO)O' display_smiles: 'O=C(/C(CO)=C/[1*])S[2*]' stereochemistry: true - pseudonyms: ["PK_C", "PK"] + pseudonyms: ["C7", "PK_C", "PK"] - name: C^Z7 smiles: 'O=C(/C(CO)=C\SO)O' display_smiles: 'O=C(/C(CO)=C\[1*])S[2*]' stereochemistry: true - pseudonyms: ["PK_C", "PK"] + pseudonyms: ["C7", "PK_C", "PK"] - name: C7 smiles: "O=C(O)C(=CSO)CO" display_smiles: "O=C(S[2*])C(=C[1*])CO"