Snakemake pipeline database construction

.gitignore

@@ -14,9 +14,14 @@ models/
 
 *.py[cod]
 .venv/
+.venv310/
 /data/
 /dumps/
 /pgdata/
 *.dump
 *.sql
-downloads/
+downloads/
+
+tmp/
+out/
+.snakemake/

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,
+        )

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