Allow for competing assembly graphs

README.md

@@ -51,11 +51,10 @@ The RetroMol CLI has two modes, single and batch:
 * `single`: process a single compound at a time.
 * `batch`: process multiple compounds in a single command.
 
-In either case a the output folder will contain a log file together with the results in either JSON or JSONL format. JSONL is standard output mode for batch mode to allow for easy parsing of large result sets. Batch mode also supports parallel processing.
-
-Any column, field, or property in the input file, either CSV, TSV, SDF, or JSON, is preserved as props in the output JSON or JSONL.
-
-Stereochemistry parsing is supported by supplying the `-c` flag.
+In either case the output folder contains a log file and results. Batch mode
+defaults to gzip-compressed JSONL (`results.jsonl.gz`), supports parallel
+processing, and writes results as they arrive. Single mode writes compact JSON
+(`result.json`); add `--compression gzip` to write `result.json.gz`.
 
 Result JSONs or lines from a JSONL file can be loaded into Python using RetroMol's `Result` class for further downstream analyses.
 

database/Snakemake

@@ -58,7 +58,7 @@ CHEBI_DIR = WORKDIR / "chebi"
 # Toggle whole branches of the pipeline off (see config.yaml's `enabled` comment) --
 # disabling npatlas/mibig skips their download+parse rules entirely, not just db
 # loading, since the "real" inputs below are only declared when enabled: nothing else
-# in the DAG needs npatlas/results.jsonl or mibig_gbk/readouts.jsonl, so Snakemake
+# in the DAG needs npatlas/results.jsonl.gz or mibig_gbk/readouts.jsonl, so Snakemake
 # never schedules the rules that would produce them. Same idea for chebi: disabling it
 # skips its bulk download entirely.
 ENABLED = config.get("enabled", {})
@@ -158,7 +158,7 @@ rule parse_npatlas:
     input:
         sdf=WORKDIR / "npatlas" / "npatlas.sdf"
     output:
-        results=WORKDIR / "npatlas" / "results.jsonl"
+        results=WORKDIR / "npatlas" / "results.jsonl.gz"
     threads: PARSE_COMPOUNDS_WORKERS
     run:
         import parse_compounds
@@ -175,9 +175,9 @@ rule parse_npatlas:
 rule load_npatlas_compounds:
     input:
         db_created=MARKERS / "db_created.done",
-        # Only declared when enabled -- nothing else in the DAG needs npatlas/results.jsonl,
+        # Only declared when enabled -- nothing else in the DAG needs npatlas/results.jsonl.gz,
         # so download_npatlas/resolve_npatlas_sdf/parse_npatlas never run when disabled.
-        **({"results": WORKDIR / "npatlas" / "results.jsonl"} if NPATLAS_ENABLED else {}),
+        **({"results": WORKDIR / "npatlas" / "results.jsonl.gz"} if NPATLAS_ENABLED else {}),
         **({"taxdump_names": TAXDUMP_DIR / "names.dmp", "taxdump_nodes": TAXDUMP_DIR / "nodes.dmp"} if TAXONOMY_ENABLED and NPATLAS_ENABLED else {})
     output:
         marker=touch(MARKERS / "npatlas_loaded.done")
@@ -227,7 +227,7 @@ rule parse_mibig_compounds:
     input:
         compounds=WORKDIR / "mibig_json" / "compounds.jsonl"
     output:
-        results=WORKDIR / "mibig_json" / "results.jsonl"
+        results=WORKDIR / "mibig_json" / "results.jsonl.gz"
     threads: PARSE_COMPOUNDS_WORKERS
     run:
         import parse_compounds
@@ -250,7 +250,7 @@ rule load_mibig_compounds:
         # extract_mibig_compounds rule above) -- not from the GBKs, which no longer carry it.
         **(
             {
-                "results": WORKDIR / "mibig_json" / "results.jsonl",
+                "results": WORKDIR / "mibig_json" / "results.jsonl.gz",
                 "versions": WORKDIR / "mibig_json" / "versions.json",
                 "annotations": WORKDIR / "mibig_json" / "annotations.json",
             }

database/scripts/load_compounds.py

@@ -40,6 +40,7 @@
     primary_sequence_from_result,
 )
 from retromol.model.result import Result
+from retromol.io.json import open_json_file
 from retromol_database.duckdb import RetroMolDuckDB
 from retromol_fingerprint.fingerprint import TOKEN_LINK
 from taxonomy import TaxonomyDB, resolve_phylogeny
@@ -152,7 +153,7 @@ def run(
 
     db = RetroMolDuckDB.open(db_path)
     try:
-        with open(results_path) as fh:
+        with open_json_file(results_path, "rt") as fh:
             with tqdm(desc=f"load_compounds[{source}]", unit="cmpd") as pbar:
                 for line in fh:
                     line = line.strip()

database/scripts/parse_compounds.py

@@ -9,7 +9,6 @@
 """
 
 import argparse
-import json
 import logging
 from pathlib import Path
 from typing import Literal
@@ -18,6 +17,7 @@
 
 from common import load_ruleset, run_retromol_stream_quiet
 from retromol.io.streaming import stream_json_records, stream_sdf_records
+from retromol.io.json import dumps_compact, open_json_file
 
 log = logging.getLogger(__name__)
 
@@ -47,7 +47,7 @@ def run(
     successes = 0
     errors = 0
 
-    with open(output_path, "w", buffering=1) as out:
+    with open_json_file(output_path, "wt") as out:
         with tqdm(desc="parse_compounds", unit="cmpd") as pbar:
             for evt in run_retromol_stream_quiet(
                 ruleset=ruleset,
@@ -59,7 +59,7 @@ def run(
                 if evt.error is not None:
                     errors += 1
                 elif evt.result is not None:
-                    out.write(json.dumps(evt.result) + "\n")
+                    out.write(dumps_compact(evt.result) + "\n")
                     successes += 1
 
                 pbar.update(1)
@@ -78,7 +78,7 @@ def main() -> None:
     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("--output", required=True, help="result JSONL path; use .jsonl.gz for gzip compression")
     ap.add_argument("--rxn-rules", default=None)
     ap.add_argument("--mxn-rules", default=None)
     ap.add_argument("--match-stereochemistry", action="store_true")

src/retromol/io/json.py

@@ -1,20 +1,46 @@
 """This module provides JSON I/O functions for the RetroMol package."""
 
 import json
+import gzip
+from pathlib import Path
 from typing import Any, Generator
 
 import ijson
 
 
-def iter_json(path: str, jsonl: bool = False) -> Generator[Any, None, None]:
+def open_json_file(path: str | Path, mode: str = "rt"):
+    """Open plain or gzip JSON/JSONL using its .gz suffix, with bounded buffers.
+
+    gzip level 6 balances file size and writing speed. Text files use UTF-8.
+    Append mode supports concatenated gzip members, which gzip readers handle
+    transparently. Callers must close the file (prefer a context manager).
+    """
+    if "b" not in mode and "t" not in mode:
+        mode += "t"
+    kwargs = {} if "b" in mode else {"encoding": "utf-8"}
+    if str(path).lower().endswith(".gz"):
+        return gzip.open(path, mode, compresslevel=6, **kwargs)
+    return open(path, mode, **kwargs)
+
+
+def dumps_compact(data: Any) -> str:
+    """JSON without formatting whitespace; preserve all values and Unicode."""
+    return json.dumps(data, separators=(",", ":"), ensure_ascii=False)
+
+
+def iter_json(path: str | Path, jsonl: bool | None = None) -> Generator[Any, None, None]:
     """
-    Stream items from a JSON array or a JSON Lines (JSONL) file.
+    Stream items from a plain or gzip JSON array or JSON Lines (JSONL) file.
 
     :param path: Path to the JSON or JSONL file.
-    :param jsonl: If True, treat the file as JSONL (one JSON object per line). If False, assume a single JSON array.
+    :param jsonl: True for JSONL, False for a JSON array; by default infer JSONL
+        from .jsonl/.ndjson (optionally followed by .gz).
     :yield: Parsed JSON objects.
     """
-    with open(path, "rb") as f:
+    if jsonl is None:
+        name = str(path).lower().removesuffix(".gz")
+        jsonl = name.endswith((".jsonl", ".ndjson"))
+    with open_json_file(path, "rb") as f:
         if jsonl:
             for line in f:
                 line = line.strip()

src/retromol/io/streaming.py

@@ -177,12 +177,12 @@ def stream_sdf_records(sdf_path: str, fast: bool = False) -> Iterator[dict[str,
             continue
 
 
-def stream_json_records(path: str, jsonl: bool = False) -> Iterator[dict[str, Any]]:
+def stream_json_records(path: str, jsonl: bool | None = None) -> Iterator[dict[str, Any]]:
     """
-    Stream JSON or JSONL records as dicts.
+    Stream plain or gzip JSON or JSONL records as dicts.
 
     :param path: Path to JSON or JSONL file.
-    :param jsonl: If True, treat as JSONL (one JSON object per line).
+    :param jsonl: True for JSONL, False for an array, or infer from suffix (default).
     :return: Iterator over record dicts.
     """
     for rec in iter_json(path, jsonl=jsonl):

src/retromol/model/assembly_graph.py

@@ -1,15 +1,17 @@
 """Module contains utilities for defining and working with assembly graphs."""
 
 from dataclasses import dataclass, asdict
-from typing import Any, Iterable, Iterable, Iterator, Generator
+from typing import Any, Iterable, Iterator, Generator, Literal
 
 from rdkit.Chem.rdchem import Mol
 import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
 import networkx as nx
 
 from retromol.model.reaction_graph import MolNode
 from retromol.model.identity import MolIdentity
 from retromol.chem.tagging import get_tags_mol
+from retromol.visualization.assembly_layout import MoleculeLayout
 
 
 @dataclass(frozen=True)
@@ -73,13 +75,14 @@ def build_assembly_graph(
     :return: NetworkX graph representing the assembly graph.
     """
     monomers = list(monomers)
+    root_tags = get_tags_mol(root_mol)
 
     tag_to_monomer: dict[int, str] = {}
     monomer_to_tags: dict[str, set[int]] = {}
 
     # Map root tags -> monomers
     for m in monomers:
-        tags = get_tags_mol(m.mol)
+        tags = get_tags_mol(m.mol) & root_tags
         monomer_to_tags[m.enc] = tags
 
         for t in tags:
@@ -89,6 +92,9 @@ def build_assembly_graph(
 
     # Initialize empty graph
     g = nx.Graph()
+    # Visualization context is kept only in memory. NetworkX copies/filtered
+    # views retain it; to_dict deliberately serializes only nodes and bonds.
+    g.graph["root_mol"] = root_mol
 
     # Add monomer nodes
     for m in monomers:
@@ -98,7 +104,7 @@ def build_assembly_graph(
 
     UNASSIGNED = "unassigned"
     if include_unassigned:
-        g.add_node(UNASSIGNED, molnode=None, tags=set(), identity=None)
+        g.add_node(UNASSIGNED, molnode=None, tags=root_tags - tag_to_monomer.keys(), identity=None)
 
     # Scan root bonds
     for b in root_mol.GetBonds():
@@ -638,24 +644,47 @@ def build(
 
         return ag
 
+    def molecular_layout(self) -> MoleculeLayout:
+        """Lazily create a layout of the original root, independent of monomers."""
+        if "molecule_layout" not in self.g.graph:
+            root = self.g.graph.get("root_mol")
+            if root is None:
+                raise ValueError("Molecular drawing requires the original root molecule; "
+                                 "pass MoleculeLayout.from_mol(root_mol) as layout, or use layout='spring'")
+            self.g.graph["molecule_layout"] = MoleculeLayout.from_mol(root)
+        return self.g.graph["molecule_layout"]
+
+    def monomer_positions(self, layout: MoleculeLayout | None = None) -> dict[str, tuple[float, float]]:
+        """Original-heavy-atom centroids, including nonempty unassigned regions."""
+        layout = self.molecular_layout() if layout is None else layout
+        return {node: position for node, data in self.g.nodes(data=True)
+                if (position := layout.centroid(data.get("tags", set()))) is not None}
+
     def draw(
         self,
         with_labels: bool = True,
         show_unassigned: bool = False,
-        node_size: int = 1600,
+        node_size: int = 900,
         font_size: int = 9,
         edge_with_scale: float = 1.0,
         savepath: str | None = None,
-    ) -> None:
+        *,
+        layout: MoleculeLayout | Literal["molecule", "spring"] = "molecule",
+        ax: Axes | None = None,
+    ) -> Axes:
         """
         Visualize the assembly graph using Matplotlib.
 
         :param with_labels: Whether to show node labels (default: True).
         :param show_unassigned: Whether to show the unassigned node (default: False).
-        :param node_size: Size of the nodes (default: 1600).
+        :param node_size: Size of the nodes (default: 900).
         :param font_size: Font size for labels (default: 9).
         :param edge_with_scale: Scale factor for edge widths (default: 1.0).
         :param savepath: Optional path to save the figure (default: None).
+        :param layout: Original-heavy-atom centroids by default. Pass one shared
+            MoleculeLayout to align alternatives, or 'spring' for a topology view.
+        :param ax: Optional existing axes. These are neither shown nor closed.
+        :return: Axes containing the drawing.
         """
         # Hide unassigned if requested
         if not show_unassigned:
@@ -666,8 +695,18 @@ def draw(
         if g.number_of_nodes() == 0:
             raise ValueError("AssemblyGraph has no nodes to show!")
 
-        # Layout
-        pos = nx.spring_layout(g, seed=42)
+        # Molecular coordinates and bounds depend only on the root, so alternative
+        # projections keep exactly the same orientation and scale.
+        molecular_layout = None
+        if isinstance(layout, MoleculeLayout) or layout == "molecule":
+            molecular_layout = self.molecular_layout() if layout == "molecule" else layout
+            positions = self.monomer_positions(molecular_layout)
+            pos = {node: positions[node] for node in g if node in positions}
+            g = g.subgraph(pos).copy()
+        elif layout == "spring":
+            pos = nx.spring_layout(g, seed=42)
+        else:
+            raise ValueError("Layout must be 'molecule', 'spring', or a MoleculeLayout")
 
         # Node colors
         node_colors = []
@@ -690,19 +729,25 @@ def draw(
         # Edge widths from number of root bonds
         widths = [edge_with_scale * max(1, data.get("n_bonds", 1)) for _, _, data in g.edges(data=True)]
 
-        plt.figure(figsize=(8, 8))
+        own_figure = ax is None
+        if ax is None:
+            figure, ax = plt.subplots(figsize=(8, 8))
+        else:
+            figure = ax.figure
         nx.draw_networkx_nodes(
             g,
             pos,
             node_color=node_colors,
             node_size=node_size,
             edgecolors="black",
+            ax=ax,
         )
         nx.draw_networkx_edges(
             g,
             pos,
             width=widths,
             alpha=0.8,
+            ax=ax,
         )
 
         if with_labels:
@@ -711,12 +756,29 @@ def draw(
                 pos,
                 labels=labels,
                 font_size=font_size,
+                ax=ax,
             )
 
-        plt.axis("off")
-        plt.tight_layout()
+        if molecular_layout is not None:
+            xmin, xmax, ymin, ymax = molecular_layout.bounds
+            ax.set_xlim(xmin, xmax)
+            ax.set_ylim(ymin, ymax)
+            ax.set_aspect("equal", adjustable="box")
+            if not pos:
+                ax.text(.5, .5, "No original heavy atoms to display", transform=ax.transAxes,
+                        ha="center", va="center")
+        ax.axis("off")
+        if own_figure:
+            # Fixed margins: tight bounding boxes based on different labels would
+            # subtly rescale the otherwise aligned assembly drawings.
+            figure.subplots_adjust(left=.04, right=.96, bottom=.04, top=.96)
 
         if savepath is not None:
-            plt.savefig(savepath, dpi=300)
-        else:
+            try:
+                figure.savefig(savepath, dpi=300)
+            finally:
+                if own_figure:
+                    plt.close(figure)
+        elif own_figure:
             plt.show()
+        return ax