diff --git a/README.md b/README.md index 87919f4..ab6481e 100644 --- a/README.md +++ b/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. diff --git a/database/Snakemake b/database/Snakemake index d743c9c..b8dcfce 100644 --- a/database/Snakemake +++ b/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", } diff --git a/database/scripts/load_compounds.py b/database/scripts/load_compounds.py index fa16632..0e0d899 100644 --- a/database/scripts/load_compounds.py +++ b/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() diff --git a/database/scripts/parse_compounds.py b/database/scripts/parse_compounds.py index 6d4b33a..0386e5b 100644 --- a/database/scripts/parse_compounds.py +++ b/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") diff --git a/src/retromol/io/json.py b/src/retromol/io/json.py index 78283c3..19f9be4 100644 --- a/src/retromol/io/json.py +++ b/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() diff --git a/src/retromol/io/streaming.py b/src/retromol/io/streaming.py index 3c62091..5a1b659 100644 --- a/src/retromol/io/streaming.py +++ b/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): diff --git a/src/retromol/model/assembly_graph.py b/src/retromol/model/assembly_graph.py index fadf612..88ec291 100644 --- a/src/retromol/model/assembly_graph.py +++ b/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 diff --git a/src/retromol/model/assembly_sampling.py b/src/retromol/model/assembly_sampling.py new file mode 100644 index 0000000..d3e6c63 --- /dev/null +++ b/src/retromol/model/assembly_sampling.py @@ -0,0 +1,331 @@ +"""Distinct assembly projections of an explored reaction hypergraph. + +A frontier stops at an identified node or an unresolved leaf; otherwise it +chooses one reaction and includes ALL its products. Different derivations of +the same frontier count once. Atom assignments distinguish repeated monomers. +""" + +import random +from functools import cached_property + +import networkx as nx + +from retromol.chem.tagging import get_tags_mol +from retromol.model.assembly_graph import AssemblyGraph +from retromol.model.reaction_graph import ReactionGraph + + +Frontier = frozenset[str] + + +class AssemblyLimitError(RuntimeError): + """Exact enumeration exceeded its budget; no partial count is returned.""" + + +class AssemblyConstraintError(ValueError): + """No explored assembly satisfies the terminal-unit constraints.""" + + +class AssemblySpace: + """Enumerate once, then count, inspect or sample distinct assemblies.""" + + def __init__( + self, reaction_graph: ReactionGraph, root_enc: str, *, + max_assemblies: int = 10_000, max_combinations: int = 1_000_000, + input_heavy_atoms: int | None = None, + ): + if root_enc not in reaction_graph.nodes: + raise ValueError(f"Root {root_enc!r} is absent from the reaction graph") + if max_assemblies < 1 or max_combinations < 1: + raise ValueError("Enumeration limits must be positive") + self.reaction_graph = reaction_graph + self.root_enc = root_enc + self.max_assemblies = max_assemblies + self.max_combinations = max_combinations + root_heavy_atoms = reaction_graph.nodes[root_enc].mol.GetNumHeavyAtoms() + if input_heavy_atoms is None: + input_heavy_atoms = root_heavy_atoms + if type(input_heavy_atoms) is not int or input_heavy_atoms < root_heavy_atoms: + raise ValueError("Input heavy-atom count must be an integer at least as large as the root's count") + self.input_heavy_atoms = input_heavy_atoms + + @cached_property + def _node_tags(self) -> dict[str, frozenset[int]]: + root_tags = get_tags_mol(self.reaction_graph.nodes[self.root_enc].mol) + return {enc: frozenset(get_tags_mol(node.mol) & root_tags) + for enc, node in self.reaction_graph.nodes.items()} + + @cached_property + def _terminal_tags(self) -> dict[str, frozenset[int]]: + """Terminal regions reachable from this root, stopping at terminal nodes.""" + graph = self.reaction_graph + pending = [self.root_enc] + seen = set() + terminal = {} + while pending: + enc = pending.pop() + if enc in seen: + continue + seen.add(enc) + node = graph.nodes[enc] + if node.is_identified and node.identity is not None and node.identity.terminal: + if self._node_tags[enc]: + terminal[enc] = self._node_tags[enc] + continue + for index in graph.out_edges.get(enc, []): + edge = graph.edges[index] + if enc not in edge.dsts: + pending.extend(edge.dsts) + return terminal + + def preserves_terminals(self, frontier: Frontier) -> bool: + """Whether a selection keeps terminal regions intact across all routes.""" + selected = [self._node_tags[enc] for enc in frontier] + selected_terminals = [self._terminal_tags[enc] for enc in frontier if enc in self._terminal_tags] + for region in set(self._terminal_tags.values()): + if any(region <= tags for tags in selected): + continue + if any(region & tags and tags - region for tags in selected_terminals): + continue + return False + return True + + @cached_property + def frontiers(self) -> tuple[Frontier, ...]: + """Distinct selections that preserve terminal units, at all coverages.""" + return tuple(frontier for frontier in self._raw_frontiers if self.preserves_terminals(frontier)) + + @cached_property + def _raw_frontiers(self) -> tuple[Frontier, ...]: + """Enumerate routes before enforcing terminal boundaries across routes.""" + graph = self.reaction_graph + tags = self._node_tags + + # Memoization only depends on ancestors in the same strongly connected + # component. Ordinary acyclic branches share all their cached results. + dependencies = nx.DiGraph() + dependencies.add_nodes_from(graph.nodes) + for edge in graph.edges: + if edge.src not in edge.dsts: + dependencies.add_edges_from((edge.src, child) for child in edge.dsts) + components = {} + for component in nx.strongly_connected_components(dependencies): + members = frozenset(component) + for enc in members: + components[enc] = members + + memo: dict[tuple[str, frozenset[str]], dict[Frontier, frozenset[int]]] = {} + combinations = 0 + + def check_size(options: dict) -> None: + if len(options) > self.max_assemblies: + raise AssemblyLimitError( + f"Exact assembly enumeration exceeded max_assemblies={self.max_assemblies} " + "at an intermediate state; increase the limit to count or sample this graph." + ) + + def visit(enc: str, ancestors: frozenset[str]) -> dict[Frontier, frozenset[int]]: + nonlocal combinations + if enc in ancestors: + return {} # no cyclic derivations + key = (enc, ancestors & components[enc]) + if key in memo: + return memo[key] + if not tags[enc]: + return {frozenset(): frozenset()} # reaction-introduced byproduct + if len(ancestors) >= 400: + raise AssemblyLimitError("Exact assembly enumeration exceeded the 400-step depth limit") + node = graph.nodes[enc] + stop = {frozenset([enc]): tags[enc]} + options = dict(stop) if node.is_identified else {} + terminal = node.is_identified and node.identity is not None and node.identity.terminal + if not terminal: + for index in graph.out_edges.get(enc, []): + edge = graph.edges[index] + if enc in edge.dsts or not edge.dsts: + continue + partial: dict[Frontier, frozenset[int]] = {frozenset(): frozenset()} + for child in edge.dsts: + child_options = visit(child, ancestors | {enc}) + joined = {} + for frontier, occupied in partial.items(): + for child_frontier, child_tags in child_options.items(): + combinations += 1 + if combinations > self.max_combinations: + raise AssemblyLimitError( + "Exact assembly enumeration exceeded " + f"max_combinations={self.max_combinations}; increase the limit." + ) + if not occupied.isdisjoint(child_tags): + continue + joined[frontier | child_frontier] = occupied | child_tags + check_size(joined) + partial = joined + if not partial: + break + options.update(partial) + check_size(options) + # A fragment with no finite, compatible expansion remains unresolved. + # This also preserves original atoms at dead ends of cyclic branches. + if not options: + options = stop + memo[key] = options + return options + + options = visit(self.root_enc, frozenset()) + return tuple(sorted(options, key=lambda frontier: tuple(sorted(frontier)))) + + @cached_property + def _coverages(self) -> tuple[float, ...]: + """Score frontiers without constructing all their assembly graphs.""" + graph = self.reaction_graph + root_heavy_tags = { + atom.GetIsotope() for atom in graph.nodes[self.root_enc].mol.GetAtoms() + if atom.GetAtomicNum() > 1 and atom.GetIsotope() != 0 + } + identified_tags = { + enc: get_tags_mol(node.mol) & root_heavy_tags + for enc, node in graph.nodes.items() if node.is_identified + } + return tuple( + len(set().union(*(identified_tags[enc] for enc in frontier if enc in identified_tags))) + / self.input_heavy_atoms if self.input_heavy_atoms else 0.0 + for frontier in self.frontiers + ) + + @cached_property + def max_coverage(self) -> float: + """Highest attainable coverage among all compatible assembly views.""" + return max(self._coverages, default=0.0) + + def eligible_frontiers(self, min_coverage: float | None = None) -> tuple[Frontier, ...]: + """Keep maximum-coverage views by default, or views meeting a minimum. + + An explicit minimum is an inclusive fraction from 0 to 1. Zero includes + all views; an unattainable minimum returns no views. Filtering happens + after exact enumeration, so the enumeration budgets still apply. + """ + if min_coverage is not None and not 0 <= min_coverage <= 1: + raise ValueError("Minimum coverage must be a finite fraction between 0 and 1") + if min_coverage == 0: + return self.frontiers + threshold = self.max_coverage if min_coverage is None else min_coverage + return tuple(frontier for frontier, coverage in zip(self.frontiers, self._coverages) + if coverage >= threshold) + + def count(self, *, min_coverage: float | None = 0.0) -> int: + """Exact count: all views by default; None counts only maximum coverage.""" + return len(self.eligible_frontiers(min_coverage)) + + def build(self, frontier: Frontier) -> AssemblyGraph: + """Build a selected frontier using bonds of the stored root molecule.""" + if frontier not in self.frontiers: + raise ValueError("Selection is not a valid assembly frontier") + graph = self.reaction_graph + return AssemblyGraph.build( + root_mol=graph.nodes[self.root_enc].mol, + monomers=[graph.nodes[enc] for enc in sorted(frontier)], + include_unassigned=True, + ) + + def reaction_graph_for(self, frontier: Frontier) -> ReactionGraph: + """Recover one explored route for a valid frontier, without inventing edges. + + Used when the cost-based default route violates terminal boundaries. + Every branch must yield exactly its assigned frontier nodes; cyclic + derivations and overlapping product assignments are rejected. + """ + if frontier not in self.frontiers: + raise ValueError("Selection is not a valid assembly frontier") + graph = self.reaction_graph + tags = self._node_tags + required = { + enc: frozenset(target for target in frontier if tags[target] & node_tags) + for enc, node_tags in tags.items() + } + memo = {} + attempts = 0 + + def visit(enc: str, ancestors: frozenset[str]) -> frozenset[int] | None: + nonlocal attempts + if enc in ancestors: + return None + if len(ancestors) >= 400: + raise AssemblyLimitError("Assembly route recovery exceeded the 400-step depth limit") + key = (enc, ancestors) + if key in memo: + return memo[key] + wanted = required[enc] + if enc in frontier or not tags[enc]: + return frozenset() + if enc in self._terminal_tags or any(not tags[target] <= tags[enc] for target in wanted): + return None + for index in graph.out_edges.get(enc, []): + attempts += 1 + if attempts > self.max_combinations: + raise AssemblyLimitError("Assembly route recovery exceeded max_combinations=" + f"{self.max_combinations}; increase the limit.") + edge = graph.edges[index] + if enc in edge.dsts or not edge.dsts: + continue + assigned = frozenset() + occupied = frozenset() + edges = frozenset([index]) + for child in edge.dsts: + if occupied & tags[child] or assigned & required[child]: + break + child_edges = visit(child, ancestors | {enc}) + if child_edges is None: + break + occupied |= tags[child] + assigned |= required[child] + edges |= child_edges + else: + if assigned == wanted: + memo[key] = edges + return edges + memo[key] = None + return None + + indices = visit(self.root_enc, frozenset()) + if indices is None: + raise AssemblyConstraintError("No explored reaction route yields the selected assembly") + selected = ReactionGraph() + nodes = {self.root_enc} + for index in sorted(indices): + edge = graph.edges[index] + nodes.add(edge.src) + nodes.update(edge.dsts) + selected.out_edges.setdefault(edge.src, []).append(len(selected.edges)) + selected.edges.append(edge) + for enc in sorted(nodes): + selected.nodes[enc] = graph.nodes[enc] + selected.out_edges.setdefault(enc, []) + return selected + + def sample( + self, n: int = 1, *, seed: int | None = None, min_coverage: float | None = None, + ) -> list[AssemblyGraph]: + """Uniform sampling among eligible views, without replacement.""" + if type(n) is not int or n < 0: + raise ValueError("Sample size must be a nonnegative integer") + if n == 0: + return [] + eligible = self.eligible_frontiers(min_coverage) + chosen = random.Random(seed).sample(eligible, min(n, len(eligible))) + return [self.build(frontier) for frontier in chosen] + + +def count_unique_assemblies( + reaction_graph: ReactionGraph, root_enc: str, *, min_coverage: float | None = 0.0, **limits, +) -> int: + """Count all views, or qualifying views with the given coverage policy.""" + return AssemblySpace(reaction_graph, root_enc, **limits).count(min_coverage=min_coverage) + + +def sample_assemblies( + reaction_graph: ReactionGraph, root_enc: str, n: int = 1, *, seed: int | None = None, + min_coverage: float | None = None, **limits, +) -> list[AssemblyGraph]: + """Sample distinct assemblies, defaulting to the highest coverage only.""" + return AssemblySpace(reaction_graph, root_enc, **limits).sample(n, seed=seed, min_coverage=min_coverage) diff --git a/src/retromol/model/reaction_graph.py b/src/retromol/model/reaction_graph.py index df69206..542f505 100644 --- a/src/retromol/model/reaction_graph.py +++ b/src/retromol/model/reaction_graph.py @@ -1,7 +1,9 @@ """Data structures for representing reaction application graphs.""" import logging -from dataclasses import dataclass, field +import json +from collections import defaultdict, deque +from dataclasses import dataclass, field, replace from typing import Any, Iterable, Literal from retromol.chem.mol import Mol, encode_mol, mol_to_smiles, smiles_to_mol @@ -32,6 +34,7 @@ class MolNode: smiles: str identity: MolIdentity | None = None identified: bool | None = None # None=unknown, False=checked-no, True=checked-yes + pruned: bool = False # uninformative descendants were collapsed here @property def is_checked(self) -> bool: @@ -86,6 +89,7 @@ def to_dict(self) -> dict[str, Any]: "smiles": self.smiles, "identity": self.identity.to_dict() if self.identity else None, "identified": self.identified, + "pruned": self.pruned, } @classmethod @@ -104,6 +108,7 @@ def from_dict(cls, data: dict[str, Any]) -> "MolNode": smiles=data["smiles"], identity=identity, identified=data["identified"], + pruned=data.get("pruned", False), ) return node @@ -157,13 +162,13 @@ class RxnEdge: """ Directed hyper-edge parent -> children, labeled by ReactionStep. - :var src: int: Encoding of source molecule node. - :var dsts: Tuple[int, ...]: encodings of child molecule nodes. + :var src: str: Encoding of source molecule node. + :var dsts: Tuple[str, ...]: encodings of child molecule nodes. :var step: ReactionStep: details of the reaction application. """ - src: int - dsts: tuple[int, ...] + src: str + dsts: tuple[str, ...] step: ReactionStep def to_dict(self) -> dict[str, Any]: @@ -224,7 +229,7 @@ def __str__(self) -> str: """ return f"ReactionGraph(num_nodes={len(self.nodes)}, num_edges={len(self.edges)})" - def add_node(self, mol: Mol) -> int: + def add_node(self, mol: Mol) -> str: """ Add a molecule node to the graph if not already present. @@ -276,30 +281,124 @@ def get_leaf_nodes(self, identified_only: bool = True) -> list[MolNode]: return leaves - - def to_dict(self) -> dict[str, Any]: + def prune_unidentified(self, root_enc: str) -> "ReactionGraph": + """Collapse wholly unidentified subtrees to unresolved boundary nodes. + + Keep every explored alternative leading to an identification, together + with all its sibling products. Identified intermediates remain selectable. + Direct self-dependent reactions cannot form finite decompositions and + are discarded. Longer cycles are handled when selecting assemblies. """ - Serialize the ReactionGraph to a dictionary. + if root_enc not in self.nodes: + raise ValueError(f"Root {root_enc!r} is absent from the reaction graph") + incoming: dict[str, set[str]] = defaultdict(set) + outgoing: dict[str, list[RxnEdge]] = defaultdict(list) + for edge in self.edges: + if edge.src in edge.dsts: + continue + outgoing[edge.src].append(edge) + for child in edge.dsts: + incoming[child].add(edge.src) + + productive = set(self.identified_nodes) + queue = deque(productive) + while queue: + for parent in incoming[queue.popleft()]: + if parent not in productive: + productive.add(parent) + queue.append(parent) + + graph = ReactionGraph() + queue = deque([root_enc]) + while queue: + enc = queue.popleft() + if enc in graph.nodes: + continue + node = self.nodes[enc] + # A wholly unproductive reaction alternative adds no information. + # Retain unresolved products only when they are siblings required + # by an informative reaction, rather than creating cheap dead-end + # alternatives by truncating those reactions to unknown leaves. + edges = [edge for edge in outgoing[enc] if any(d in productive for d in edge.dsts)] + collapsed = bool(self.out_edges.get(enc)) and not edges + graph.nodes[enc] = replace(node, pruned=True) if collapsed else node + graph.out_edges[enc] = [] + for edge in edges: + graph.out_edges[enc].append(len(graph.edges)) + graph.edges.append(edge) + queue.extend(edge.dsts) + return graph - :return: Dictionary representation of the ReactionGraph. + def to_dict(self) -> dict[str, Any]: + """Compact graph: indexed nodes/edges and interned identities/steps. + + Atom-tagged SMILES preserve structure, stereo, and original-atom mapping. + Encodings, untagged SMILES and adjacency are reconstructed on load. + Node indices follow sorted encodings, also used by Result.root. """ - return { - "nodes": {enc: node.to_dict() for enc, node in self.nodes.items()}, - "edges": [edge.to_dict() for edge in self.edges], - "out_edges": {enc: indices for enc, indices in self.out_edges.items()}, - } + indices = {enc: i for i, enc in enumerate(sorted(self.nodes))} + identities: list[dict[str, Any]] = [] + identity_indices: dict[str, int] = {} + nodes = [] + for enc in indices: + node = self.nodes[enc] + record: dict[str, Any] = {"smiles": mol_to_smiles(node.mol, include_tags=True)} + if node.is_identified: + if node.identity is None: + raise ValueError(f"Identified node {enc!r} has no identity") + identity = node.identity.matched_rule.to_dict() + key = json.dumps(identity, sort_keys=True, separators=(",", ":")) + if key not in identity_indices: + identity_indices[key] = len(identities) + identities.append(identity) + record["identity"] = identity_indices[key] + elif not node.is_checked: + record["checked"] = False + if node.pruned: + record["pruned"] = True + nodes.append(record) + + steps: list[dict[str, Any]] = [] + step_indices: dict[ReactionStep, int] = {} + edges = [] + for edge in self.edges: + if edge.step not in step_indices: + step_indices[edge.step] = len(steps) + steps.append(edge.step.to_dict()) + edges.append({"src": indices[edge.src], "dsts": [indices[d] for d in edge.dsts], + "step": step_indices[edge.step]}) + return {"nodes": nodes, "identities": identities, "steps": steps, "edges": edges} @classmethod def from_dict(cls, data: dict[str, Any]) -> "ReactionGraph": - """ - Deserialize a ReactionGraph from a dictionary. - - :param data: Dictionary representation of the ReactionGraph. - :return: ReactionGraph object. - """ - reaction_graph = cls( - nodes={enc: MolNode.from_dict(node_data) for enc, node_data in data["nodes"].items()}, - edges=[RxnEdge.from_dict(edge_data) for edge_data in data["edges"]], - out_edges={enc: indices for enc, indices in data["out_edges"].items()}, - ) - return reaction_graph + """Load the compact graph format; rebuild encodings and adjacency.""" + identities = [MolIdentity(MatchingRule.from_dict(d)) for d in data["identities"]] + steps = [ReactionStep.from_dict(d) for d in data["steps"]] + graph = cls() + encodings = [] + for record in data["nodes"]: + mol = smiles_to_mol(record["smiles"]) + enc = encode_mol(mol) + if enc in graph.nodes: + raise ValueError("Duplicate molecule in reaction graph nodes") + identity = cls._at(identities, record["identity"]) if "identity" in record else None + graph.nodes[enc] = MolNode( + enc, mol, mol_to_smiles(mol), identity, + True if identity else (False if record.get("checked", True) else None), + pruned=record.get("pruned", False), + ) + graph.out_edges[enc] = [] + encodings.append(enc) + for record in data["edges"]: + edge = RxnEdge(cls._at(encodings, record["src"]), + tuple(cls._at(encodings, d) for d in record["dsts"]), + cls._at(steps, record["step"])) + graph.out_edges[edge.src].append(len(graph.edges)) + graph.edges.append(edge) + return graph + + @staticmethod + def _at(items: list, index: int): + if type(index) is not int or not 0 <= index < len(items): + raise ValueError(f"Invalid graph reference: {index!r}") + return items[index] diff --git a/src/retromol/model/readout.py b/src/retromol/model/readout.py index ac5ed46..1c64a31 100644 --- a/src/retromol/model/readout.py +++ b/src/retromol/model/readout.py @@ -97,7 +97,10 @@ def from_reaction_graph( include_unassigned: bool = True, ) -> "LinearReadout": """ - Create a LinearReadout from a Result object. + Create a LinearReadout from one selected reaction route. + + An explored graph can contain mutually exclusive, overlapping leaves. + Select a route first, or use from_assembly_graph with a sampled assembly. :param root_enc: Encoding of the root molecule. :param reaction_graph: ReactionGraph object. @@ -116,7 +119,12 @@ def from_reaction_graph( # Create assembly graph of monomers; first collect nodes to include collected = g.get_leaf_nodes(identified_only=identified_only) a = AssemblyGraph.build(root_mol=root, monomers=collected, include_unassigned=include_unassigned) + return cls.from_assembly_graph(a) + @classmethod + def from_assembly_graph(cls, assembly_graph: AssemblyGraph) -> "LinearReadout": + """Extract sequences from a particular assembly projection on demand.""" + a = assembly_graph # Break bonds between monomers that are not backbone-related bonds (i.e., keep C-C and C-N bonds only) f = a.filtered_by_root_bond_elements(allow_pairs={frozenset(("C", "C")), frozenset(("C", "N"))}, drop_isolated=False) diff --git a/src/retromol/model/result.py b/src/retromol/model/result.py index 0f96b6a..75cfcb5 100644 --- a/src/retromol/model/result.py +++ b/src/retromol/model/result.py @@ -1,85 +1,168 @@ -"""Module defining the Result data class.""" +"""Compact parsing results, with assembly projections computed on demand.""" from dataclasses import dataclass +from functools import cached_property from typing import Any +from retromol.chem.mol import Mol, encode_mol, mol_to_inchikey, smiles_to_mol +from retromol.chem.stereo import capture_double_bond_stereo +from retromol.chem.tagging import get_tags_mol, remove_tags +from retromol.model.assembly_graph import AssemblyGraph +from retromol.model.assembly_sampling import AssemblyConstraintError, AssemblySpace from retromol.model.submission import Submission -from retromol.model.reaction_graph import ReactionGraph +from retromol.model.reaction_graph import MolNode, ReactionGraph from retromol.model.readout import LinearReadout -from retromol.chem.mol import smiles_to_mol -from retromol.chem.tagging import get_tags_mol @dataclass(frozen=True) class Result: - """ - Represents a RetroMol parsing result. + """An explored reaction graph and its input/provenance metadata. - :var submission: Submission: The original submission associated with this result. - :var reaction_graph: ReactionGraph: The reaction graph generated from retrosynthetic analysis. - :var linear_readout: LinearReadout: The linear readout representation of the reaction graph. + Only these source data are serialized. Assemblies, coverage and sequences + are projections requested by the consumer. The default projection uses a + minimum-cost, fine-grained policy, weighting collapsed unresolved regions + by their original heavy atoms. """ submission: Submission reaction_graph: ReactionGraph - linear_readout: LinearReadout + reaction_rules_hash: str | None = None + matching_rules_hash: str | None = None + match_stereochemistry: bool = False + root_enc: str | None = None + + def __post_init__(self) -> None: + root = self.root_enc if self.root_enc is not None else encode_mol(self.submission.mol) + if root not in self.reaction_graph.nodes: + raise ValueError(f"Root {root!r} is absent from the reaction graph") + object.__setattr__(self, "root_enc", root) def __str__(self) -> str: - """ - String representation of the Result. - - :return: String representation of the Result. - """ - return f"Result(submission={self.submission}, reaction_graph={self.reaction_graph}, linear_readout={self.linear_readout})" - - def calculate_coverage(self) -> float: - """ - Calculate the fraction of input heavy atoms in identified assembly monomers. + return f"Result(submission={self.submission}, reaction_graph={self.reaction_graph})" + + @property + def root(self) -> MolNode: + return self.reaction_graph.nodes[self.root_enc] + + @property + def props(self) -> dict[str, Any] | None: + return self.submission.props + + @cached_property + def selected_reaction_graph(self) -> ReactionGraph: + """One fine-grained route; selection never alters the stored graph.""" + from retromol.pipelines.parsing import extract_min_edge_synthesis_subgraph + + selected = extract_min_edge_synthesis_subgraph( + self.reaction_graph, self.root_enc, + edge_base_cost=0.25, nonterminal_leaf_penalty=100.0, + weight_pruned_atoms=True, + ).graph + space = self.assembly_space() + frontier = frozenset(node.enc for node in selected.get_leaf_nodes(identified_only=False)) + if space.preserves_terminals(frontier): + return selected + # A cheap route can cut a terminal unit before recognizing it. Choose + # a valid, maximally covered alternative, preferring finer resolution, + # and retain its actual reaction route for downstream reconstruction. + candidates = space.eligible_frontiers() + if not candidates: + raise AssemblyConstraintError("No explored assembly preserves the identified terminal units") + return space.reaction_graph_for(max(candidates, key=len)) + + @cached_property + def assembly_graph(self) -> AssemblyGraph: + """Default assembly, calculated only when requested.""" + return AssemblyGraph.build( + root_mol=self.root.mol, + monomers=self.selected_reaction_graph.get_leaf_nodes(identified_only=False), + include_unassigned=True, + ) + + @cached_property + def linear_readout(self) -> LinearReadout: + """Default sequence view; neither it nor its assembly is saved.""" + return LinearReadout.from_assembly_graph(self.assembly_graph) + + def assembly_space(self, **limits) -> AssemblySpace: + """Create a reusable enumeration for counting, selecting and sampling.""" + return AssemblySpace(self.reaction_graph, self.root_enc, input_heavy_atoms=self.input_heavy_atoms, **limits) + + def count_unique_assemblies(self, *, min_coverage: float | None = 0.0, **limits) -> int: + """Count all views; pass None for maximum coverage, or a minimum fraction.""" + return self.assembly_space(**limits).count(min_coverage=min_coverage) - :return: Coverage score as a float. + def sample_assemblies( + self, n: int = 1, *, seed: int | None = None, min_coverage: float | None = None, **limits, + ) -> list[AssemblyGraph]: + """Sample maximum-coverage views, or explicitly allow a lower minimum.""" + return self.assembly_space(**limits).sample(n, seed=seed, min_coverage=min_coverage) + + @cached_property + def input_heavy_atoms(self) -> int: + """All submitted heavy atoms, including fragments removed in preparation.""" + return smiles_to_mol(self.submission.smiles).GetNumHeavyAtoms() + + def calculate_coverage(self, assembly_graph: AssemblyGraph | None = None) -> float: + """Identified assembly heavy atoms / ALL submitted heavy atoms. + + Supply an assembly to measure a particular sampled resolution. Without + one, use the default projection, never the union of alternative routes. """ - input_heavy_atoms = smiles_to_mol(self.submission.smiles).GetNumHeavyAtoms() + input_heavy_atoms = self.input_heavy_atoms if not input_heavy_atoms: return 0.0 - root_heavy_tags = { - atom.GetIsotope() - for atom in self.submission.mol.GetAtoms() + atom.GetIsotope() for atom in self.root.mol.GetAtoms() if atom.GetAtomicNum() > 1 and atom.GetIsotope() != 0 } + assembly = self.assembly_graph if assembly_graph is None else assembly_graph identified_tags: set[int] = set() - for node in self.linear_readout.assembly_graph.monomer_nodes(): + for node in assembly.monomer_nodes(): if node.is_identified: identified_tags.update(get_tags_mol(node.mol)) - - return len(identified_tags.intersection(root_heavy_tags)) / input_heavy_atoms + return len(identified_tags & root_heavy_tags) / input_heavy_atoms def to_dict(self) -> dict[str, Any]: - """ - Serialize the Result to a dictionary. - - :return: Dictionary representation of the Result. - """ + """Serialize source data only, using the compact version 2 schema.""" return { - "submission": self.submission.to_dict(), + "schema_version": 2, + "root": { + "node": sorted(self.reaction_graph.nodes).index(self.root_enc), + "smiles": self.submission.smiles, + "name": self.submission.name, + "keep_stereo": self.submission.keep_stereo, + "neutralize": self.submission.neutralize, + "canonicalize_tautomer": self.submission.canonicalize_tautomer, + "match_stereochemistry": self.match_stereochemistry, + }, + "props": self.props, + "reaction_rules_hash": self.reaction_rules_hash, + "matching_rules_hash": self.matching_rules_hash, "reaction_graph": self.reaction_graph.to_dict(), - "linear_readout": self.linear_readout.to_dict(), } - + @classmethod def from_dict(cls, data: dict[str, Any]) -> "Result": - """ - Deserialize a Result from a dictionary. - - :param data: Dictionary representation of the Result. - :return: Result object. - """ - submission = Submission.from_dict(data["submission"]) - reaction_graph = ReactionGraph.from_dict(data["reaction_graph"]) - linear_readout = LinearReadout.from_dict(data["linear_readout"]) - + """Load version 2. Older result formats are intentionally unsupported.""" + if data.get("schema_version") != 2: + raise ValueError("Unsupported result schema; reparse the input to produce schema_version 2") + graph = ReactionGraph.from_dict(data["reaction_graph"]) + root = data["root"] + enc = ReactionGraph._at(list(graph.nodes), root["node"]) + submission = Submission( + smiles=root["smiles"], name=root["name"], props=data["props"], + keep_stereo=root["keep_stereo"], neutralize=root["neutralize"], + canonicalize_tautomer=root["canonicalize_tautomer"], + ) + # The saved root is authoritative for atom numbering and chemistry. + mol = Mol(graph.nodes[enc].mol) + object.__setattr__(submission, "mol", mol) + object.__setattr__(submission, "inchikey", mol_to_inchikey(remove_tags(mol, in_place=False))) + object.__setattr__(submission, "stereo_registry", capture_double_bond_stereo(mol)) return cls( - submission=submission, - reaction_graph=reaction_graph, - linear_readout=linear_readout, + submission, graph, + reaction_rules_hash=data["reaction_rules_hash"], + matching_rules_hash=data["matching_rules_hash"], + match_stereochemistry=root["match_stereochemistry"], root_enc=enc, ) diff --git a/src/retromol/model/rules.py b/src/retromol/model/rules.py index 67c82f3..183af92 100644 --- a/src/retromol/model/rules.py +++ b/src/retromol/model/rules.py @@ -3,8 +3,10 @@ import logging import itertools import hashlib +import json from collections import Counter from dataclasses import dataclass, field +from functools import cached_property from importlib.resources import files from pathlib import Path from typing import Any @@ -64,6 +66,14 @@ def id(self) -> str: """ return hashlib.sha256(self.smarts.encode("utf-8")).hexdigest() + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "smarts": self.smarts, + "props": self.props, + "allowed_in_bulk": self.allowed_in_bulk, + } + @classmethod def from_dict(cls, data: dict[str, Any]) -> "ReactionRule": """ @@ -468,6 +478,27 @@ class RuleSet: reaction_rules: list[ReactionRule] matching_rules: list[MatchingRule] + @cached_property + def reaction_rules_hash(self) -> str: + """SHA-256 of ordered, canonical rule definitions (not YAML formatting). + + Treat a RuleSet and its nested lists/props as immutable once used. + """ + return self._hash_rules(self.reaction_rules) + + @cached_property + def matching_rules_hash(self) -> str: + """SHA-256 of ordered matching definitions, including names and flags.""" + return self._hash_rules(self.matching_rules) + + @staticmethod + def _hash_rules(rules: list) -> str: + content = json.dumps( + [rule.to_dict() for rule in rules], sort_keys=True, + separators=(",", ":"), ensure_ascii=False, allow_nan=False, + ) + return hashlib.sha256(content.encode("utf-8")).hexdigest() + def __str__(self) -> str: """ String representation of the RuleSet. @@ -478,18 +509,7 @@ def __str__(self) -> str: def find_structural_matches(self, mol: Mol, match_stereochemistry: bool = False) -> list["MatchingRule"]: """ - Find every matching rule whose motif has the same structure as `mol`. - - Unlike a name lookup, this doesn't require knowing which specific monomer a - molecule is -- useful for reconciling a molecule from outside the ruleset - (e.g. a substrate-specificity model's predicted SMILES) with whatever - monomer(s) in this ruleset it is chemically identical to. Matching with - `match_stereochemistry=False` (the default) means two rules that only differ - by stereochemistry (e.g. "A2^R" and "A2^S") both match a molecule with - undetermined or different stereochemistry, and multiple rules can match at - once (e.g. distinct rules that describe the same free-molecule structure but - differ in which atom continues the polymer chain, such as "aspartic acid" vs. - an isoAsp-linked variant sharing the same 2D graph). + Find every matching rule whose motif has the same structure as mol. :param mol: Molecule to match against every rule in this ruleset. :param match_stereochemistry: Whether to consider stereochemistry in matching. @@ -524,14 +544,6 @@ def load_from_files( reaction_rules = [ReactionRule.from_dict(d) for d in reaction_rules_data] matching_rules = [MatchingRule.from_dict(d) for d in matching_rules_data] - # `match_mol` (see chem/matching.py) is greedy: it returns the first rule that - # matches and stops looking. Sort stereochemistry-specific rules (e.g. - # "alanine^L") ahead of their achiral fallback (e.g. "alanine") so a specific - # rule is always tried before its fallback, regardless of the order they - # happen to appear in mxn.yml -- an achiral rule's query has no chirality tags - # to violate, so with `useChirality=True` it still matches a stereo-defined - # molecule and would otherwise win by appearing first. Stable sort: relative - # order within each group (stereo / non-stereo) is preserved from mxn.yml. matching_rules.sort(key=lambda rule: not rule.stereochemistry) return RuleSet(match_stereochemistry, reaction_rules, matching_rules) @@ -544,4 +556,4 @@ def load_default(cls, match_stereochemistry: bool = False) -> "RuleSet": :param match_stereochemistry: Whether to consider stereochemistry in matching rules. :return: The loaded RuleSet instance. """ - return cls.load_from_files(match_stereochemistry=match_stereochemistry) \ No newline at end of file + return cls.load_from_files(match_stereochemistry=match_stereochemistry) diff --git a/src/retromol/pipelines/parsing.py b/src/retromol/pipelines/parsing.py index 8688505..1988e5e 100644 --- a/src/retromol/pipelines/parsing.py +++ b/src/retromol/pipelines/parsing.py @@ -11,7 +11,6 @@ from retromol.model.rules import RuleSet, index_uncontested, apply_uncontested from retromol.model.result import Result from retromol.model.reaction_graph import ReactionGraph, ReactionStep, RxnEdge -from retromol.model.readout import LinearReadout from retromol.model.synthesis import SynthesisExtractResult from retromol.chem.mol import Mol, encode_mol, mol_to_smiles from retromol.chem.tagging import get_tags_mol @@ -125,6 +124,7 @@ def extract_min_edge_synthesis_subgraph( edge_base_cost: float = 1.0, nonterminal_leaf_penalty: float = 0.25, unsolved_leaf_penalty: float = 5.0, + weight_pruned_atoms: bool = False, ) -> SynthesisExtractResult: """ Extract a minimum-edge synthesis subgraph from a retrosynthesis ReactionGraph. @@ -145,6 +145,8 @@ def extract_min_edge_synthesis_subgraph( :param edge_base_cost: Base cost per reaction edge :param nonterminal_leaf_penalty: Penalty for identified leaves that are non-terminal :param unsolved_leaf_penalty: Penalty for unsolved leaves (i.e., "give up" cost) + :param weight_pruned_atoms: Charge collapsed unresolved regions per original + heavy atom, so pruning does not make a large unknown region cheap :return: The extracted synthesis subgraph and status :raises ValueError: If root_enc is not in the graph """ @@ -161,6 +163,16 @@ def extract_min_edge_synthesis_subgraph( out_edges[e.src].append(ei) kind_rank = {k: i for i, k in enumerate(prefer_kind)} + root_heavy_tags = { + atom.GetIsotope() for atom in g.nodes[root_enc].mol.GetAtoms() + if atom.GetAtomicNum() > 1 and atom.GetIsotope() != 0 + } + + def unresolved_cost(enc: str) -> float: + node = g.nodes[enc] + if not weight_pruned_atoms or not node.pruned: + return unsolved_leaf_penalty + return unsolved_leaf_penalty * len(get_tags_mol(node.mol) & root_heavy_tags) def edge_cost(e: RxnEdge) -> float: # Primary objective: fewer edges @@ -196,7 +208,7 @@ def solve_cost(enc: str) -> float: # Identified nonterminal with no edges is still fine as 0, else unsolved penalty if n and n.is_identified: return 0.0 - return unsolved_leaf_penalty + return unresolved_cost(enc) # Soft leaves: identified + terminal=False leaf_cost = inf @@ -240,7 +252,7 @@ def solve_cost(enc: str) -> float: if best == inf and allow_cycle_frontier: # All continuations return to an ancestor. Preserve this molecule as # an unresolved leaf rather than discarding its entire parent route. - best = unsolved_leaf_penalty + best = unresolved_cost(enc) memo_cost[enc] = best memo_choice[enc] = best_ei @@ -319,26 +331,16 @@ def run_retromol(submission: Submission, rules: RuleSet) -> Result: g = process_mol(submission, rules) log.debug(f"Retrosynthesis graph has {len(g.nodes)} ({len(g.identified_nodes)} identified) nodes and {len(g.edges)} edges") - # Extract minimum-edge synthesis subgraph root = encode_mol(submission.mol) - r = extract_min_edge_synthesis_subgraph( - g, - root_enc=root, - edge_base_cost=0.25, # low base cost encourages longer syntheses - nonterminal_leaf_penalty=100.0, # high penalty forces expansion of non-terminal leaves (e.g., fatty acids) - ) - log.debug(f"Extracted synthesis subgraph has {len(r.graph.nodes)} ({len(r.graph.identified_nodes)} identified) nodes and {len(r.graph.edges)} edges") - - if not r.solved: - log.debug("Retrosynthesis extraction retained unresolved frontier fragments") - - # Calculate the linear readouts for the synthesis graph - linear_readout = LinearReadout.from_reaction_graph(root, r.graph) - + # Retain all informative explored alternatives. Readout policy belongs to + # consumers, and is not applied while parsing or serializing a result. return Result( submission=submission, - reaction_graph=r.graph, - linear_readout=linear_readout, + reaction_graph=g.prune_unidentified(root), + root_enc=root, + reaction_rules_hash=rules.reaction_rules_hash, + matching_rules_hash=rules.matching_rules_hash, + match_stereochemistry=rules.match_stereochemistry, ) diff --git a/src/retromol/visualization/assembly_layout.py b/src/retromol/visualization/assembly_layout.py new file mode 100644 index 0000000..27a9430 --- /dev/null +++ b/src/retromol/visualization/assembly_layout.py @@ -0,0 +1,60 @@ +"""Shared RDKit coordinates for comparing assembly projections of one molecule.""" + +from dataclasses import dataclass, field +from typing import Iterable + +from rdkit import Chem +from rdkit.Chem import rdDepictor +from rdkit.Chem.rdchem import Mol + + +@dataclass(frozen=True) +class MoleculeLayout: + """A 2D depiction and its original-heavy-atom coordinates, indexed by tag.""" + + atom_positions: dict[int, tuple[float, float]] + bounds: tuple[float, float, float, float] + molecule: Mol = field(repr=False, compare=False) + + @classmethod + def from_mol(cls, root_mol: Mol) -> "MoleculeLayout": + # Restore tag order so serialization/reloading or RDKit atom renumbering + # cannot rotate/rearrange equivalent projections of the same tagged root. + order = sorted(range(root_mol.GetNumAtoms()), + key=lambda i: (root_mol.GetAtomWithIdx(i).GetIsotope() or float("inf"), i)) + molecule = Chem.RenumberAtoms(root_mol, order) if order else Mol(root_mol) + tags = [atom.GetIsotope() for atom in molecule.GetAtoms()] + # Isotopes are RetroMol's atom identifiers, not labels for the depiction. + for atom in molecule.GetAtoms(): + atom.SetIsotope(0) + atom.SetAtomMapNum(0) + positions = {} + heavy_positions = [] + if molecule.GetNumAtoms(): + rdDepictor.Compute2DCoords(molecule, canonOrient=True, clearConfs=True, forceRDKit=True) + conformer = molecule.GetConformer() + for index, atom in enumerate(molecule.GetAtoms()): + if atom.GetAtomicNum() <= 1: + continue + point = conformer.GetAtomPosition(index) + xy = (float(point.x), float(point.y)) + heavy_positions.append(xy) + tag = tags[index] + if tag: + if tag in positions: + raise ValueError(f"Root contains duplicate atom tag {tag}") + positions[tag] = xy + if heavy_positions: + xs, ys = zip(*heavy_positions) + padding = max(2.0, .12 * max(max(xs) - min(xs), max(ys) - min(ys))) + bounds = (min(xs) - padding, max(xs) + padding, min(ys) - padding, max(ys) + padding) + else: + bounds = (-2.0, 2.0, -2.0, 2.0) + return cls(positions, bounds, molecule) + + def centroid(self, tags: Iterable[int]) -> tuple[float, float] | None: + """Centroid of unique original heavy atoms, excluding introduced atoms/H.""" + points = [self.atom_positions[tag] for tag in sorted(set(tags)) if tag in self.atom_positions] + if not points: + return None + return (sum(x for x, _ in points) / len(points), sum(y for _, y in points) / len(points)) diff --git a/src/retromol_cli/cli.py b/src/retromol_cli/cli.py index c494ac5..32df9dd 100644 --- a/src/retromol_cli/cli.py +++ b/src/retromol_cli/cli.py @@ -4,7 +4,9 @@ import json import logging import os +import re from collections import Counter +from contextlib import ExitStack from datetime import datetime from typing import Any @@ -14,11 +16,13 @@ from retromol.utils.logging import setup_logging, add_file_handler from retromol.model.rules import RuleSet from retromol.model.result import Result +from retromol.model.assembly_sampling import AssemblyConstraintError, AssemblyLimitError from retromol.model.submission import Submission from retromol.pipelines.parsing import run_retromol_with_timeout from retromol.io.streaming import run_retromol_stream, stream_sdf_records, stream_table_rows, stream_json_records -from retromol.chem.mol import encode_mol +from retromol.io.json import dumps_compact, open_json_file from retromol.visualization.reaction_graph import visualize_reaction_graph +from retromol.visualization.assembly_layout import MoleculeLayout log = logging.getLogger(__name__) @@ -43,6 +47,28 @@ def add_rule_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--mxn-rules", type=str, default=None, help="path to matching rules YAML (default: None, default ruleset)") +def nonnegative_int(value: str) -> int: + """Validate counts at argument parsing time, before running chemistry.""" + try: + count = int(value) + except ValueError: + raise argparse.ArgumentTypeError("must be a nonnegative integer") from None + if count < 0: + raise argparse.ArgumentTypeError("must be a nonnegative integer") + return count + + +def coverage_fraction(value: str) -> float: + """Validate an inclusive coverage threshold before running chemistry.""" + try: + coverage = float(value) + except ValueError: + raise argparse.ArgumentTypeError("must be a finite fraction between 0 and 1") from None + if not 0 <= coverage <= 1: + raise argparse.ArgumentTypeError("must be a finite fraction between 0 and 1") + return coverage + + def cli() -> argparse.Namespace: """ Parse command line arguments. @@ -64,6 +90,24 @@ def cli() -> argparse.Namespace: # For 'single' mode user should just give a SMILES as input single_parser.add_argument("-s", "--smiles", type=str, help="SMILES string of the compound to process") + single_parser.add_argument( + "--max-assembly-graphs", type=nonnegative_int, default=1, metavar="N", + help="maximum assembly drawings: 0 skips, 1 prefers the default assembly if eligible, " + "larger values sample distinct eligible assemblies (default: 1)", + ) + single_parser.add_argument( + "--min-coverage", type=coverage_fraction, default=None, metavar="FRACTION", + help="minimum identified heavy-atom coverage from 0 to 1; " + "default: only assemblies with the highest attainable coverage; 0 allows all views", + ) + single_parser.add_argument( + "--seed", type=int, default=42, + help="assembly sampling seed (default: 42)", + ) + single_parser.add_argument( + "--compression", choices=["none", "gzip"], default="none", + help="result JSON compression (default: none; gzip writes result.json.gz)", + ) # For 'batch' mode user should provide a path to an SDF file input_group = batch_parser.add_mutually_exclusive_group(required=True) @@ -75,7 +119,11 @@ def cli() -> argparse.Namespace: batch_parser.add_argument("--pool-chunksize", type=int, default=50, help="chunksize hint for imap_unordered (default: 50)") batch_parser.add_argument("--maxtasksperchild", type=int, default=2000, help="recycle worker after N tasks (default: 2000)") batch_parser.add_argument("--results", choices=["files", "jsonl"], default="jsonl", help="write each result to a file or append to JSONL (default: jsonl)") - batch_parser.add_argument("--jsonl-path", type=str, default=None, help="path to results jsonl (default: /results.jsonl)") + batch_parser.add_argument("--jsonl-path", type=str, default=None, help="path to results JSONL (.gz appended for gzip; default: /results.jsonl.gz)") + batch_parser.add_argument( + "--compression", choices=["none", "gzip"], default="gzip", + help="result compression (default: gzip; use none for plain JSONL/JSON)", + ) batch_parser.add_argument("--no-tqdm", action="store_true", help="disable progress bars for lowest overhead") batch_parser.add_argument("--rdkit-fast", action="store_true", help="use fast SDF parse (sanitize=False, removeHs=True); we'll sanitize only when needed") @@ -91,20 +139,29 @@ def cli() -> argparse.Namespace: add_rule_args(single_parser) add_rule_args(batch_parser) - return parser.parse_args() + args = parser.parse_args() + if (args.mode == "batch" and args.jsonl_path and args.jsonl_path.lower().endswith(".gz") + and args.compression == "none"): + parser.error("--jsonl-path ending in .gz requires --compression gzip") + return args -def _open_jsonl(outdir: str, jsonl_path: str | None) -> tuple[Any, str]: +def _open_jsonl(outdir: str, jsonl_path: str | None, compression: str = "gzip") -> tuple[Any, str]: """ - Open a JSONL file for appending results. + Open a plain or gzip JSONL stream for appending results. :param outdir: str: output directory :param jsonl_path: str | None: path to JSONL file, or None to use default + :param compression: gzip (default) or none; append .gz when needed :return: tuple[file handle, path]: opened file handle and the path used """ path = jsonl_path or os.path.join(outdir, "results.jsonl") - os.makedirs(os.path.dirname(path), exist_ok=True) - return open(path, "a", buffering=1), path # line-buffered + if compression == "gzip" and not path.lower().endswith(".gz"): + path += ".gz" + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + # One continuous stream per run, not one gzip member per compound. This + # lets compression reuse repeated structures/metadata across records. + return open_json_file(path, "at"), path def main() -> None: @@ -153,29 +210,62 @@ def main() -> None: # Write out result to file and then read back in again for visualization (test I/O) result_dict = result.to_dict() - with open(os.path.join(args.outdir, "result.json"), "w") as f: - json.dump(result_dict, f, indent=4) + result_path = os.path.join(args.outdir, "result.json" + (".gz" if args.compression == "gzip" else "")) + with open_json_file(result_path, "wt") as f: + f.write(dumps_compact(result_dict)) - with open(os.path.join(args.outdir, "result.json"), "r") as f: + with open_json_file(result_path, "rt") as f: result_data = json.load(f) result2 = Result.from_dict(result_data) - # Report the percentage of input heavy atoms in identified assembly monomers - coverage = result2.calculate_coverage() - log.info(f"coverage: {coverage:.2%}") - - # Get linear readout; draw assembly graph - linear_readout = result2.linear_readout - out_assembly_graph_fig = os.path.join(args.outdir, "assembly_graph.png") - linear_readout.assembly_graph.draw(show_unassigned=True, savepath=out_assembly_graph_fig) - log.info(f"linear readout: {linear_readout}") + if args.max_assembly_graphs == 0: + assemblies = [] + log.info("Skipping assembly drawings (--max-assembly-graphs 0)") + else: + space = result2.assembly_space() + try: + threshold = args.min_coverage + default_coverage = result2.calculate_coverage() if args.max_assembly_graphs == 1 else None + if threshold is None: + # A fully covered default already proves the maximum and + # avoids enumeration for the usual single-drawing case. + threshold = 1.0 if default_coverage == 1.0 else space.max_coverage + if default_coverage is not None and default_coverage >= threshold: + assemblies = [result2.assembly_graph] + else: + assemblies = space.sample(args.max_assembly_graphs, seed=args.seed, min_coverage=threshold) + log.info( + "Drawing %d of %d eligible assemblies (%d total; minimum coverage %.2f%%; seed=%d)", + len(assemblies), space.count(min_coverage=threshold), space.count(), 100 * threshold, args.seed, + ) + if not assemblies: + if not space.frontiers: + log.warning("No explored assemblies preserve the identified terminal units; skipping drawings.") + else: + log.warning("No assemblies meet minimum coverage %.2f%%; highest attainable coverage is %.2f%%.", + 100 * threshold, 100 * space.max_coverage) + except (AssemblyLimitError, AssemblyConstraintError) as exc: + log.error("Cannot sample assembly drawings: %s. The reaction graph is saved in %s.", exc, result_path) + raise SystemExit(1) from None + + # Rerunning into the same directory must not leave earlier, now-invalid + # projections visible beside the new drawings. Only our generated names + # are replaced; user-named images in the directory are left alone. + for entry in os.scandir(args.outdir): + if entry.is_file() and re.fullmatch(r"assembly_graph(?:_[0-9]{3,})?\.png", entry.name): + os.remove(entry.path) + + drawing_layout = MoleculeLayout.from_mol(result2.root.mol) if assemblies else None + for index, assembly in enumerate(assemblies, start=1): + filename = "assembly_graph.png" if args.max_assembly_graphs == 1 else f"assembly_graph_{index:03d}.png" + assembly.draw(show_unassigned=True, savepath=os.path.join(args.outdir, filename), layout=drawing_layout) + log.info("%s coverage: %.2f%%", filename, 100 * result2.calculate_coverage(assembly)) # Visualize reaction graph - root = encode_mol(result2.submission.mol) visualize_reaction_graph( result2.reaction_graph, html_path=os.path.join(args.outdir, "reaction_graph.html"), - root_enc=root + root_enc=result2.root_enc ) result_counts["successes"] += 1 @@ -194,58 +284,55 @@ def main() -> None: else: source_iter = stream_json_records(args.json) - # Progress bars: outer ~batches, inner = molecules processed - pbar_outer = tqdm(desc="Batches", unit="batch", disable=args.no_tqdm) - pbar_inner = tqdm(desc="Processed", unit="mol", disable=args.no_tqdm) - - # Results saved into JSONL format to limit file operations - jsonl_fh = None - jsonl_path = None - if args.results == "jsonl": - jsonl_fh, jsonl_path = _open_jsonl(args.outdir, args.jsonl_path) - log.info(f"Appending results to JSONL file at: {jsonl_path}") - result_counts = Counter() - processed_in_current_batch = 0 - - for evt in run_retromol_stream( - ruleset=ruleset, - row_iter=source_iter, - smiles_col=smiles_col, - workers=args.workers, - batch_size=args.batch_size, - pool_chunksize=args.pool_chunksize, - maxtasksperchild=args.maxtasksperchild, - ): - # evt has: result (dict or None) and error (str or None) - if evt.error is not None: - log.error(evt.error) - result_counts["errors"] += 1 - elif evt.result is not None: - # Result is already serialized as dict - jsonl_fh.write(json.dumps(evt.result) + "\n") - result_counts["successes"] += 1 - else: - log.error("received empty result with no error message") - result_counts["failures"] += 1 - - # Progress - pbar_inner.update(1) - processed_in_current_batch += 1 - if processed_in_current_batch >= args.batch_size: + # Close streams even if parsing/writing is interrupted, so completed + # records and the gzip footer are flushed to disk. + with ExitStack() as resources: + pbar_outer = resources.enter_context(tqdm(desc="Batches", unit="batch", disable=args.no_tqdm)) + pbar_inner = resources.enter_context(tqdm(desc="Processed", unit="mol", disable=args.no_tqdm)) + jsonl_fh = None + if args.results == "jsonl": + jsonl_fh, jsonl_path = _open_jsonl(args.outdir, args.jsonl_path, args.compression) + resources.enter_context(jsonl_fh) + log.info(f"Appending results to JSONL file at: {jsonl_path}") + + for record_index, evt in enumerate(run_retromol_stream( + ruleset=ruleset, + row_iter=source_iter, + smiles_col=smiles_col, + workers=args.workers, + batch_size=args.batch_size, + pool_chunksize=args.pool_chunksize, + maxtasksperchild=args.maxtasksperchild, + ), start=1): + if evt.error is not None: + log.error(evt.error) + result_counts["errors"] += 1 + elif evt.result is not None: + serialized = dumps_compact(evt.result) + if jsonl_fh is not None: + jsonl_fh.write(serialized + "\n") + else: + # Numbered files preserve repeated input IDs. Exclusive + # creation avoids silently replacing a previous run. + suffix = ".json.gz" if args.compression == "gzip" else ".json" + path = os.path.join(args.outdir, f"result_{record_index:09d}{suffix}") + with open_json_file(path, "xt") as f: + f.write(serialized) + result_counts["successes"] += 1 + else: + log.error("received empty result with no error message") + result_counts["failures"] += 1 + + pbar_inner.update(1) + processed_in_current_batch += 1 + if processed_in_current_batch >= args.batch_size: + pbar_outer.update(1) + processed_in_current_batch = 0 + + if processed_in_current_batch > 0: pbar_outer.update(1) - processed_in_current_batch = 0 - - # If there was a final partial batch, tick the outer bar once more - if processed_in_current_batch > 0: - pbar_outer.update(1) - - pbar_inner.close() - pbar_outer.close() - - if jsonl_fh: - jsonl_fh.close() log.info(f"Streaming complete. Summary: {dict(result_counts)}") diff --git a/src/retromol_synthesis/reconstruction.py b/src/retromol_synthesis/reconstruction.py index c913788..b6db9e0 100644 --- a/src/retromol_synthesis/reconstruction.py +++ b/src/retromol_synthesis/reconstruction.py @@ -9,7 +9,7 @@ from rdkit.Chem.rdchem import Atom, BondStereo, BondType from rdkit.Chem.rdmolops import AssignStereochemistry, SetDoubleBondNeighborDirections -from retromol.chem.mol import encode_mol, smiles_to_mol, smarts_to_mol, mol_to_smiles, sanitize_mol +from retromol.chem.mol import smiles_to_mol, smarts_to_mol, mol_to_smiles, sanitize_mol from retromol.chem.tagging import get_tags_mol from retromol.chem.reaction import smarts_to_reaction from retromol.chem.stereo import BondStereoRecord @@ -486,8 +486,9 @@ def reconstruct_linear_readout(result: Result) -> list[Reconstruction]: :return: The reconstructed candidates, one per eligible path. Empty if no path in the readout consists of (mostly) eligible primary-sequence building blocks. """ - root_enc = encode_mol(result.submission.mol) - readout = LinearReadout.from_reaction_graph(root_enc, reaction_graph=result.reaction_graph, identified_only=True) + readout = LinearReadout.from_reaction_graph( + result.root_enc, reaction_graph=result.selected_reaction_graph, identified_only=True, + ) tagged_input_smiles = mol_to_smiles(result.submission.mol, include_tags=True) reconstructions: list[Reconstruction] = [] diff --git a/tests/helpers_retromol.py b/tests/helpers_retromol.py index e6bb2cb..a97168c 100644 --- a/tests/helpers_retromol.py +++ b/tests/helpers_retromol.py @@ -3,7 +3,6 @@ from retromol.model.rules import RuleSet from retromol.model.result import Result from retromol.model.submission import Submission -from retromol.model.reaction_graph import MolNode from retromol.pipelines.parsing import run_retromol_with_timeout @@ -71,7 +70,7 @@ def assert_result(result: Result, expected_coverage: float, expected_monomers: l coverage: float = result.calculate_coverage() assert compare_floats(coverage, expected_coverage), f"expected coverage {expected_coverage}, got {coverage}" - ident_nodes: MolNode = result.reaction_graph.identified_nodes.values() + ident_nodes = result.selected_reaction_graph.identified_nodes.values() assert all(n.is_identified for n in ident_nodes), "not all identified nodes are marked as identified" found_monomers: list[str] = [n.identity.name for n in ident_nodes] diff --git a/tests/test_assembly_layout.py b/tests/test_assembly_layout.py new file mode 100644 index 0000000..f6ee6cb --- /dev/null +++ b/tests/test_assembly_layout.py @@ -0,0 +1,129 @@ +"""Molecular centroids keep alternative assembly drawings spatially comparable.""" + +import json + +import matplotlib.pyplot as plt +import pytest +from rdkit import Chem + +from retromol.model.assembly_graph import AssemblyGraph +from retromol.model.result import Result +from retromol.model.submission import Submission +from retromol.pipelines.parsing import run_retromol +from retromol.visualization.assembly_layout import MoleculeLayout + +from .test_assembly_sampling import layered_result +from .test_result import make_result + + +def test_positions_average_original_heavy_atoms_and_place_real_unassigned_regions(): + result = make_result('[2H]CCO', [('[1H][2CH2][3CH2][99CH2]O', True)]) + assembly = result.assembly_graph + layout = MoleculeLayout.from_mol(result.root.mol) + assert set(layout.atom_positions) == {2, 3, 4} + positions = assembly.monomer_positions(layout) + expected = tuple((layout.atom_positions[2][axis] + layout.atom_positions[3][axis]) / 2 for axis in [0, 1]) + assert positions[assembly.monomer_ids()[0]] == pytest.approx(expected) + assert positions[assembly.unassigned] == layout.atom_positions[4] + assert layout.centroid([1, 99, 0]) is None + assert layout.centroid([2, 2, 3, 99]) == pytest.approx(expected) + + +def test_alternatives_keep_common_nodes_and_plot_bounds_at_identical_coordinates(): + result, left, right, singles = layered_result() + space = result.assembly_space() + coarse = space.build(frozenset([left, right])) + fine = space.build(frozenset([left, *singles[2:]])) + layout = MoleculeLayout.from_mol(result.root.mol) + assert coarse.monomer_positions(layout)[left] == fine.monomer_positions(layout)[left] + assert coarse.unassigned not in coarse.monomer_positions(layout) + fig, axes = plt.subplots(1, 2) + try: + coarse.draw(layout=layout, ax=axes[0], show_unassigned=True) + fine.draw(layout=layout, ax=axes[1], show_unassigned=True) + for ax in axes: + assert ax.get_xlim() == layout.bounds[:2] + assert ax.get_ylim() == layout.bounds[2:] + assert ax.get_aspect() == 1 + assert len(axes[0].collections[0].get_offsets()) == 2 + assert len(axes[1].collections[0].get_offsets()) == 3 + finally: + plt.close(fig) + + +def test_depiction_is_lazy_and_never_changes_source_or_saved_payload(monkeypatch): + result, *_ = layered_result() + before = json.loads(json.dumps(result.to_dict())) + conformers = result.root.mol.GetNumConformers() + with monkeypatch.context() as context: + context.setattr(MoleculeLayout, 'from_mol', lambda *args: pytest.fail('Unexpected depiction during parsing/readout')) + assemblies = result.sample_assemblies(10) + _ = result.linear_readout + assert json.loads(json.dumps(result.to_dict())) == before + layout = assemblies[0].molecular_layout() + assert assemblies[0].molecular_layout() is layout + assert result.root.mol.GetNumConformers() == conformers + assert all(a.GetIsotope() == 0 for a in layout.molecule.GetAtoms()) + assert json.loads(json.dumps(result.to_dict())) == before + + +def test_filtered_graphs_keep_layout_context(): + result, *_ = layered_result() + assembly = result.assembly_graph + expected = assembly.monomer_positions() + filtered = assembly.drop_unassigned().filtered_by_root_bond_elements() + assert filtered.monomer_positions() == expected + for component in filtered.connected_components(): + assert all(position == expected[node] for node, position in component.monomer_positions().items()) + + +def test_detached_assembly_accepts_explicit_layout_and_spring_remains_available(): + result, *_ = layered_result() + assembly = AssemblyGraph.from_dict(result.assembly_graph.to_dict()) + with pytest.raises(ValueError, match='original root molecule'): + assembly.monomer_positions() + layout = MoleculeLayout.from_mol(result.root.mol) + assert assembly.monomer_positions(layout) == result.assembly_graph.monomer_positions(layout) + fig, ax = plt.subplots() + try: + assembly.draw(layout='spring', ax=ax) + finally: + plt.close(fig) + + +def test_zero_heavy_atom_depiction_does_not_invent_a_node_position(): + result = make_result('[2H][2H]', []) + fig, ax = plt.subplots() + try: + result.assembly_graph.draw(ax=ax, show_unassigned=True) + assert any(text.get_text() == 'No original heavy atoms to display' for text in ax.texts) + assert result.assembly_graph.monomer_positions() == {} + finally: + plt.close(fig) + + +def test_daptomycin_layout_survives_loading_and_atom_renumbering(ruleset, tmp_path): + from .data.integration_retromol import CASES + + smiles = next(smiles for name, smiles, *_ in CASES if name == 'daptomycin') + result = run_retromol(Submission(smiles), ruleset) + loaded = Result.from_dict(json.loads(json.dumps(result.to_dict()))) + reordered = Chem.RenumberAtoms(result.root.mol, list(reversed(range(result.root.mol.GetNumAtoms())))) + layout = MoleculeLayout.from_mol(result.root.mol) + for molecule in [loaded.root.mol, reordered]: + other = MoleculeLayout.from_mol(molecule) + assert other.bounds == pytest.approx(layout.bounds) + assert other.atom_positions.keys() == layout.atom_positions.keys() + for tag, position in layout.atom_positions.items(): + assert other.atom_positions[tag] == pytest.approx(position) + assemblies = loaded.sample_assemblies(10, seed=42) + assert len(assemblies) == 5 + positions = [assembly.monomer_positions(layout) for assembly in assemblies] + common = set.intersection(*(set(position) for position in positions)) + assert len(common) == 13 # all thirteen amino acids stay fixed + assert all(position[node] == positions[0][node] for position in positions for node in common) + image = tmp_path / 'daptomycin.png' + figures = plt.get_fignums() + assemblies[0].draw(layout=layout, savepath=str(image), show_unassigned=True) + assert image.read_bytes().startswith(b'\x89PNG\r\n\x1a\n') + assert plt.get_fignums() == figures diff --git a/tests/test_assembly_sampling.py b/tests/test_assembly_sampling.py new file mode 100644 index 0000000..c4d2cec --- /dev/null +++ b/tests/test_assembly_sampling.py @@ -0,0 +1,523 @@ +"""Assembly alternatives, compact persistence, pruning and provenance.""" + +import json +from dataclasses import replace + +import pytest +import yaml + +from retromol.chem.mol import smiles_to_mol +from retromol.chem.tagging import get_tags_mol +from retromol.model.assembly_sampling import AssemblyConstraintError, AssemblyLimitError, AssemblySpace, count_unique_assemblies, sample_assemblies +from retromol.model.reaction_graph import ReactionGraph, ReactionStep +from retromol.model.readout import LinearReadout +from retromol.model.result import Result +from retromol.model.rules import MatchingRule, ReactionRule, RuleSet +from retromol.model.submission import Submission +from retromol.pipelines import parsing + + +STEP = ReactionStep('contested', ('split',), ('rule-1',)) + + +def add(graph, smiles, *, known=False, terminal=False): + enc = graph.add_node(smiles_to_mol(smiles)) + if known: + graph.nodes[enc].identify([MatchingRule( + name=graph.nodes[enc].smiles, smiles=graph.nodes[enc].smiles, + props={}, pseudonyms=[], terminal=terminal, + )]) + else: + graph.nodes[enc].identify([]) + return enc + + +def edge(graph, source, *children, step=STEP): + graph.add_edge(source, [graph.nodes[c].mol for c in children], step) + + +def layered_result(): + submission = Submission('CCCC') + graph = ReactionGraph() + root = graph.add_node(submission.mol) + left = add(graph, '[1CH3][2CH3]', known=True) + right = add(graph, '[3CH3][4CH3]', known=True) + singles = [add(graph, f'[{tag}CH4]', known=True, terminal=True) for tag in range(1, 5)] + edge(graph, root, left, right) + edge(graph, left, *singles[:2]) + edge(graph, right, *singles[2:]) + return Result(submission, graph), left, right, singles + + +def signature(assembly): + nodes = tuple(sorted(assembly.monomer_ids())) + bonds = sorted(tuple(sorted((link.a1_tag, link.a2_tag))) + for _, _, links in assembly.edges_with_bonds() for link in links) + return nodes, bonds + + +def test_independent_regions_allow_mixed_resolution_and_rebuild_root_bonds(): + result, left, right, singles = layered_result() + space = result.assembly_space() + assert space.count() == 4 + assert set(space.frontiers) == { + frozenset([left, right]), frozenset([*singles[:2], right]), + frozenset([left, *singles[2:]]), frozenset(singles), + } + coarse = space.build(frozenset([left, right])) + assert signature(coarse)[1] == [(2, 3)] + fine = space.build(frozenset(singles)) + assert signature(fine)[1] == [(1, 2), (2, 3), (3, 4)] + for assembly in space.sample(10, seed=17): + assert result.calculate_coverage(assembly) == 1 + assert assembly.g.nodes['unassigned']['tags'] == set() + tags = [get_tags_mol(n.mol) for n in assembly.monomer_nodes()] + assert sum(map(len, tags)) == len(set.union(*tags)) == 4 + + +def test_duplicate_routes_and_product_permutations_count_once(): + result, left, right, singles = layered_result() + edge(result.reaction_graph, result.root_enc, right, left, + step=ReactionStep('contested', ('other route',), ('rule-2',))) + edge(result.reaction_graph, result.root_enc, *singles) + assert result.count_unique_assemblies() == 4 + assert len({signature(a)[0] for a in result.sample_assemblies(100, seed=42)}) == 4 + + +def terminal_bypass_result(): + result, left, right, singles = layered_result() + graph = result.reaction_graph + node = graph.nodes[left] + graph.nodes[left] = replace(node, identity=replace( + node.identity, matched_rule=replace(node.identity.matched_rule, terminal=True), + )) + # This shortcut decomposes the terminal left unit before recognizing it. + edge(graph, result.root_enc, *singles) + return result, left, right, singles + + +def test_terminal_boundaries_apply_across_routes_before_coverage_filtering(): + result, left, right, singles = terminal_bypass_result() + payload = json.loads(json.dumps(result.to_dict())) + space = result.assembly_space() + assert set(space.frontiers) == {frozenset([left, right]), frozenset([left, *singles[2:]])} + assert space.count() == space.count(min_coverage=None) == 2 + for minimum in [None, 0, .5, 1]: + samples = space.sample(10, seed=42, min_coverage=minimum) + assert len(samples) == 2 + assert all(left in assembly.monomer_ids() for assembly in samples) + assert all(result.calculate_coverage(assembly) == 1 for assembly in samples) + with pytest.raises(ValueError, match='not a valid assembly frontier'): + space.build(frozenset(singles)) + loaded = Result.from_dict(payload) + assert [signature(a) for a in loaded.sample_assemblies(10, seed=42)] == [ + signature(a) for a in result.sample_assemblies(10, seed=42) + ] + assert json.loads(json.dumps(result.to_dict())) == payload + + +def test_default_assembly_and_readout_replace_a_cheaper_terminal_bypass(): + result, left, right, singles = terminal_bypass_result() + graph = result.reaction_graph + old_default = parsing.extract_min_edge_synthesis_subgraph( + graph, result.root_enc, edge_base_cost=.25, nonterminal_leaf_penalty=100, + ).graph + assert {n.enc for n in old_default.get_leaf_nodes()} == set(singles) + expected = {left, *singles[2:]} + assert set(result.assembly_graph.monomer_ids()) == expected + assert set(result.linear_readout.assembly_graph.monomer_ids()) == expected + assert {n.enc for n in result.selected_reaction_graph.get_leaf_nodes()} == expected + assert result.calculate_coverage() == 1 + assert all(e in graph.edges for e in result.selected_reaction_graph.edges) + + +def test_nonterminal_setting_explicitly_allows_finer_views_again(): + result, left, _, singles = terminal_bypass_result() + graph = result.reaction_graph + node = graph.nodes[left] + graph.nodes[left] = replace(node, identity=replace( + node.identity, matched_rule=replace(node.identity.matched_rule, terminal=False), + )) + assert result.count_unique_assemblies() == 4 + assert frozenset(singles) in result.assembly_space().frontiers + + +def test_crossing_terminal_assignments_compete_but_subdivisions_are_excluded(): + submission = Submission('CCC') + graph = ReactionGraph() + root = graph.add_node(submission.mol) + ab = add(graph, '[1CH3][2CH3]', known=True, terminal=True) + bc = add(graph, '[2CH3][3CH3]', known=True, terminal=True) + a, b, c = [add(graph, f'[{tag}CH4]', known=True, terminal=True) for tag in range(1, 4)] + edge(graph, root, ab, c) + edge(graph, root, a, bc) + edge(graph, root, a, b, c) + result = Result(submission, graph) + space = result.assembly_space() + assert set(space.frontiers) == {frozenset([ab, c]), frozenset([a, bc])} + assert space.count() == 2 + assert not space.preserves_terminals(frozenset([a, b, c])) + assert space.preserves_terminals(frozenset(result.assembly_graph.monomer_ids())) + + +def test_disconnected_terminal_match_does_not_constrain_another_root(): + result, *_ = layered_result() + add(result.reaction_graph, '[1CH3][2CH2][3CH3]', known=True, terminal=True) + assert result.count_unique_assemblies() == 4 + + +def test_terminal_unit_cannot_be_hidden_by_losing_some_of_its_original_atoms(): + result, left, right, singles = terminal_bypass_result() + # Dropping the second carbon is not an acceptable way to avoid the boundary. + edge(result.reaction_graph, result.root_enc, singles[0], right) + assert result.count_unique_assemblies() == 2 + assert all(left in a.monomer_ids() for a in result.sample_assemblies(10, min_coverage=0)) + + +def test_terminal_boundaries_ignore_reaction_added_atoms(): + submission = Submission('CCC') + graph = ReactionGraph() + root = graph.add_node(submission.mol) + known = add(graph, '[1CH3][2CH2][99OH]', known=True, terminal=True) + singles = [add(graph, f'[{tag}CH4]', known=True, terminal=True) for tag in range(1, 4)] + edge(graph, root, known, singles[2]) + edge(graph, root, *singles) + result = Result(submission, graph) + assert result.count_unique_assemblies() == 1 + assert set(result.sample_assemblies()[0].monomer_ids()) == {known, singles[2]} + + +def test_incompatible_terminal_constraints_do_not_invent_a_route_or_relax_protection(): + result, left, right, singles = layered_result() + graph = result.reaction_graph + for enc in [left, right]: + node = graph.nodes[enc] + graph.nodes[enc] = replace(node, identity=replace( + node.identity, matched_rule=replace(node.identity.matched_rule, terminal=True), + )) + # Each route preserves one terminal unit and splits the other. The graph + # never explored a route that keeps both intact, so we must not invent one. + graph.edges.clear() + graph.out_edges = {enc: [] for enc in graph.nodes} + edge(graph, result.root_enc, left, *singles[2:]) + edge(graph, result.root_enc, *singles[:2], right) + assert result.count_unique_assemblies() == 0 + assert result.sample_assemblies(10, min_coverage=0) == [] + with pytest.raises(AssemblyConstraintError, match='preserves the identified terminal units'): + _ = result.assembly_graph + + +def test_equal_names_at_different_atom_positions_remain_distinct(): + submission = Submission('CCC') + graph = ReactionGraph() + root = graph.add_node(submission.mol) + a = add(graph, '[1CH4]', known=True) + bc = add(graph, '[2CH3][3CH3]', known=True) + ab = add(graph, '[1CH3][2CH3]', known=True) + c = add(graph, '[3CH4]', known=True) + edge(graph, root, a, bc) + edge(graph, root, ab, c) + space = AssemblySpace(graph, root) + assert space.count() == 2 + assert len({signature(x)[0] for x in space.sample(10, seed=1)}) == 2 + + +def test_overlapping_products_are_not_compatible_assemblies(): + result, left, right, singles = layered_result() + edge(result.reaction_graph, result.root_enc, left, singles[0], right) + assert result.count_unique_assemblies() == 4 + with pytest.raises(ValueError, match='not a valid'): + result.assembly_space().build(frozenset([left, singles[0], right])) + + +def test_reaction_introduced_byproducts_do_not_create_extra_assemblies(): + result, left, right, _ = layered_result() + water = add(result.reaction_graph, 'O', known=True) + edge(result.reaction_graph, result.root_enc, left, right, water) + assert result.count_unique_assemblies() == 4 + assert all(water not in f for f in result.assembly_space().frontiers) + + +def test_seeded_sampling_is_reproducible_after_json_roundtrip(): + result, *_ = layered_result() + loaded = Result.from_dict(json.loads(json.dumps(result.to_dict()))) + expected = [signature(a) for a in result.sample_assemblies(3, seed=13)] + assert [signature(a) for a in loaded.sample_assemblies(3, seed=13)] == expected + assert count_unique_assemblies(loaded.reaction_graph, loaded.root_enc) == 4 + assert len(sample_assemblies(loaded.reaction_graph, loaded.root_enc, 10, seed=13)) == 4 + assert result.sample_assemblies(0) == [] + with pytest.raises(ValueError, match='nonnegative integer'): + result.sample_assemblies(-1) + + +def partially_identified_result(): + result, _, _, singles = layered_result() + graph = result.reaction_graph + graph.nodes[singles[0]] = replace(graph.nodes[singles[0]], identified=False, identity=None) + return result + + +def test_default_sampling_uses_global_maximum_not_default_projection_coverage(): + result = partially_identified_result() + # The cost-based fine-grained projection is not necessarily the best-covered. + assert result.calculate_coverage() == .75 + space = result.assembly_space() + assert space.count() == 4 + assert space.max_coverage == 1 + assert space.count(min_coverage=None) == 2 + assert result.count_unique_assemblies(min_coverage=None) == 2 + best = space.sample(10, seed=42) + assert len(best) == len({signature(a)[0] for a in best}) == 2 + assert all(result.calculate_coverage(a) == 1 for a in best) + for minimum, expected_count in [(0, 4), (.75, 4), (.76, 2), (1, 2)]: + samples = space.sample(10, seed=42, min_coverage=minimum) + assert len(samples) == space.count(min_coverage=minimum) == expected_count + assert all(result.calculate_coverage(a) >= minimum for a in samples) + + +def test_maximum_coverage_below_one_includes_all_input_atoms_and_roundtrips(): + result = partially_identified_result() + result = replace(result, submission=replace(result.submission, smiles='CCCC.[Na+]')) + payload = json.loads(json.dumps(result.to_dict())) + loaded = Result.from_dict(payload) + space = loaded.assembly_space() + assert space.max_coverage == .8 # largest fragment has 4 of the 5 input heavy atoms + assert len(space.sample(10, seed=1)) == 2 + assert all(loaded.calculate_coverage(a) == .8 for a in space.sample(10, seed=1)) + assert space.count(min_coverage=.81) == 0 + assert space.sample(10, min_coverage=.81) == [] # never backfill with lower coverage + assert [signature(a) for a in loaded.sample_assemblies(1, seed=17)] == [ + signature(a) for a in result.sample_assemblies(1, seed=17) + ] + assert json.loads(json.dumps(loaded.to_dict())) == payload + assert count_unique_assemblies(loaded.reaction_graph, loaded.root_enc, + input_heavy_atoms=5, min_coverage=.81) == 0 + assert sample_assemblies(loaded.reaction_graph, loaded.root_enc, 10, + input_heavy_atoms=5, min_coverage=.81) == [] + + +@pytest.mark.parametrize('minimum', [-.1, 1.1, float('nan'), float('inf'), -float('inf')]) +def test_sampling_rejects_invalid_coverage_thresholds(minimum): + result, *_ = layered_result() + with pytest.raises(ValueError, match='finite fraction between 0 and 1'): + result.sample_assemblies(10, min_coverage=minimum) + with pytest.raises(ValueError, match='finite fraction between 0 and 1'): + result.count_unique_assemblies(min_coverage=minimum) + + +def test_zero_coverage_still_has_a_best_assembly(): + submission = Submission('CCO') + graph = ReactionGraph() + graph.add_node(submission.mol) + result = Result(submission, graph) + space = result.assembly_space() + assert space.max_coverage == 0 + assert space.count(min_coverage=None) == 1 + assert result.calculate_coverage(space.sample()[0]) == 0 + assert space.sample(min_coverage=.01) == [] + + +def test_count_and_sample_raise_instead_of_silently_truncating(): + result, *_ = layered_result() + with pytest.raises(AssemblyLimitError, match='max_assemblies=2'): + result.count_unique_assemblies(max_assemblies=2) + with pytest.raises(AssemblyLimitError, match='max_assemblies=2'): + result.sample_assemblies(1, seed=1, max_assemblies=2) + with pytest.raises(AssemblyLimitError, match='max_combinations=1'): + result.count_unique_assemblies(max_combinations=1) + with pytest.raises(ValueError, match='positive'): + result.assembly_space(max_assemblies=0) + + +def test_pruning_preserves_unknown_boundary_and_identified_siblings(): + submission = Submission('CCCO') + graph = ReactionGraph() + root = graph.add_node(submission.mol) + known = add(graph, '[1CH4]', known=True) + unknown = add(graph, '[2CH3][3CH2][4OH]') + tail = add(graph, '[2CH3][3CH]=[4O]') + leaves = [add(graph, '[2CH4]'), add(graph, '[3CH2]=[4O]')] + edge(graph, root, known, unknown) + edge(graph, unknown, tail) + edge(graph, tail, *leaves) + pruned = graph.prune_unidentified(root) + assert set(pruned.nodes) == {root, known, unknown} + assert pruned.nodes[unknown].pruned + assert not graph.nodes[unknown].pruned # source graph was not mutated + result = Result(submission, pruned) + assert result.count_unique_assemblies() == 1 + assembly = result.sample_assemblies(seed=3)[0] + assert set.union(*(get_tags_mol(n.mol) for n in assembly.monomer_nodes())) == {1, 2, 3, 4} + assert result.calculate_coverage(assembly) == .25 + restored = Result.from_dict(result.to_dict()) + assert restored.reaction_graph.nodes[unknown].pruned + + +def test_wholly_unidentified_graph_collapses_to_root(): + submission = Submission('CCO') + graph = ReactionGraph() + root = graph.add_node(submission.mol) + child = add(graph, '[1CH3][2CH]=[3O]') + edge(graph, root, child) + edge(graph, child, root) + result = Result(submission, graph.prune_unidentified(root)) + assert set(result.reaction_graph.nodes) == {root} + assert result.reaction_graph.edges == [] + assert result.root.pruned + assert result.count_unique_assemblies() == 1 + assert result.calculate_coverage() == 0 + + +def test_pruning_drops_wholly_unproductive_alternative_instead_of_making_it_cheaper(): + result, *_ = layered_result() + graph = result.reaction_graph + unknown = add(graph, '[1CH2]=[2CH][3CH2][4CH3]') + deeper = add(graph, '[1CH2]=[2CH][3CH]=[4CH2]') + edge(graph, result.root_enc, unknown) + edge(graph, unknown, deeper) + pruned = graph.prune_unidentified(result.root_enc) + assert unknown not in pruned.nodes and deeper not in pruned.nodes + assert len(pruned.out_edges[result.root_enc]) == 1 + assert Result(result.submission, pruned).count_unique_assemblies() == 4 + + +def test_identified_parent_drops_wholly_unknown_descendants(): + graph = ReactionGraph() + root = add(graph, '[1CH4]', known=True) + child = add(graph, '[1CH3]O') + edge(graph, root, child) + pruned = graph.prune_unidentified(root) + assert set(pruned.nodes) == {root} + assert pruned.nodes[root].is_identified and pruned.nodes[root].pruned + assert AssemblySpace(pruned, root).frontiers == (frozenset([root]),) + + +def test_collapsed_unknown_region_does_not_become_a_cheap_default_shortcut(): + result, _, _, singles = layered_result() + graph = result.reaction_graph + unknown = add(graph, '[2CH3][3CH2][4CH3]') + graph.nodes[unknown] = replace(graph.nodes[unknown], pruned=True) + edge(graph, result.root_enc, singles[0], unknown) + costs = dict(edge_base_cost=4, nonterminal_leaf_penalty=100) + plain = parsing.extract_min_edge_synthesis_subgraph(graph, result.root_enc, **costs) + weighted = parsing.extract_min_edge_synthesis_subgraph( + graph, result.root_enc, weight_pruned_atoms=True, **costs, + ) + assert not plain.solved + assert weighted.solved + assert {n.enc for n in weighted.graph.get_leaf_nodes()} == set(singles) + + +def test_productive_cycles_terminate_and_keep_identified_exit(): + graph = ReactionGraph() + root = add(graph, '[1CH4]') + intermediate = add(graph, '[1CH3]O') + identified = add(graph, '[1CH3]N', known=True, terminal=True) + edge(graph, root, intermediate) + edge(graph, intermediate, root) + edge(graph, intermediate, identified) + pruned = graph.prune_unidentified(root) + assert len(pruned.edges) == 3 + assert AssemblySpace(pruned, root).frontiers == (frozenset([identified]),) + edge(graph, root, root, identified) + assert len(graph.prune_unidentified(root).edges) == 3 + + +def test_compact_schema_interns_records_and_keeps_input_options_and_props(): + result, *_ = layered_result() + result = replace(result, submission=Submission( + 'CCCC.[Na+]', name='example', props={'source': 1, 'root': 'user property'}, + keep_stereo=False, neutralize=False, canonicalize_tautomer=True, + ), reaction_rules_hash='a' * 64, matching_rules_hash='b' * 64, match_stereochemistry=True) + # Derive views first; these caches must still never enter the saved result. + _ = result.linear_readout + data = json.loads(json.dumps(result.to_dict())) + assert set(data) == {'schema_version', 'root', 'props', 'reaction_graph', + 'reaction_rules_hash', 'matching_rules_hash'} + g = data['reaction_graph'] + assert len(g['identities']) == 2 # ethane and methane, each shared across occurrences + assert len(g['steps']) == 1 + assert 'out_edges' not in g + assert all('enc' not in n and 'tagged_smiles' not in n for n in g['nodes']) + loaded = Result.from_dict(data) + assert loaded.root_enc == result.root_enc + assert loaded.props == {'source': 1, 'root': 'user property'} + assert loaded.submission.name == 'example' + assert not loaded.submission.keep_stereo and not loaded.submission.neutralize + assert loaded.submission.canonicalize_tautomer and loaded.match_stereochemistry + assert loaded.reaction_rules_hash == 'a' * 64 + assert loaded.matching_rules_hash == 'b' * 64 + assert loaded.calculate_coverage() == .8 + assert loaded.count_unique_assemblies() == 4 + assert 'linear_readout' not in loaded.__dict__ + assert json.loads(json.dumps(loaded.to_dict())) == data + + +@pytest.mark.parametrize('index', [-1, 1_000_000, True]) +def test_invalid_serialized_node_references_are_rejected(index): + result, *_ = layered_result() + data = result.to_dict() + data['root']['node'] = index + with pytest.raises(ValueError, match='Invalid graph reference'): + Result.from_dict(data) + + +def test_old_result_format_is_explicitly_rejected(): + with pytest.raises(ValueError, match='reparse'): + Result.from_dict({'submission': {}, 'reaction_graph': {}, 'linear_readout': {}}) + + +def test_parse_keeps_competing_routes_and_does_not_compute_readouts(monkeypatch): + result, left, right, singles = layered_result() + edge(result.reaction_graph, result.root_enc, *singles) + rules = RuleSet(False, [], []) + monkeypatch.setattr(parsing, 'process_mol', lambda *args: result.reaction_graph) + + def unexpected_selection(*args, **kwargs): + raise AssertionError('Parsing and saving must not choose an assembly') + + monkeypatch.setattr(parsing, 'extract_min_edge_synthesis_subgraph', unexpected_selection) + monkeypatch.setattr(LinearReadout, 'from_assembly_graph', unexpected_selection) + parsed = parsing.run_retromol(result.submission, rules) + assert len(parsed.reaction_graph.out_edges[parsed.root_enc]) == 2 + assert parsed.count_unique_assemblies() == 4 + assert parsed.to_dict()['reaction_rules_hash'] == rules.reaction_rules_hash + assert parsed.matching_rules_hash == rules.matching_rules_hash + + +def test_decanoic_lipopeptide_has_five_resolutions(ruleset): + result = parsing.run_retromol(Submission('CCCCCCCCCC(=O)NCC(=O)NC(C)C(=O)O'), ruleset) + space = result.assembly_space() + assert space.count() == 5 + assemblies = space.sample(10, seed=42) + assert sorted(len(a.monomer_ids()) for a in assemblies) == [3, 4, 5, 6, 7] + fatty_names = {n.identity.name for a in assemblies for n in a.monomer_nodes() + if 'acid' in n.identity.name} + assert fatty_names == {'decanoic acid', 'octanoic acid', 'hexanoic acid', 'butanoic acid', 'acetic acid'} + assert all(result.calculate_coverage(a) == 1 for a in assemblies) + sequences = [LinearReadout.from_assembly_graph(a).primary_sequence() for a in assemblies] + assert sorted(map(len, sequences)) == [3, 4, 5, 6, 7] + + +def test_rule_hashes_include_definition_metadata_and_order(tmp_path): + reaction = ReactionRule('cut', '[C:1]-[C:2]>>[C:1].[C:2]', {'a': 1, 'b': 2}) + matching = MatchingRule('methane', 'C', {}, [], terminal=False) + rules = RuleSet(False, [reaction], [matching]) + changed_rxn = RuleSet(False, [replace(reaction, allowed_in_bulk=True)], [matching]) + changed_match = RuleSet(False, [reaction], [replace(matching, terminal=True)]) + assert rules.reaction_rules_hash != changed_rxn.reaction_rules_hash + assert rules.matching_rules_hash == changed_rxn.matching_rules_hash + assert rules.matching_rules_hash != changed_match.matching_rules_hash + assert rules.reaction_rules_hash == changed_match.reaction_rules_hash + assert rules.reaction_rules_hash == RuleSet(False, [replace(reaction, props={'b': 2, 'a': 1})], []).reaction_rules_hash + a = replace(matching, name='first') + b = replace(matching, name='second') + assert RuleSet(False, [], [a, b]).matching_rules_hash != RuleSet(False, [], [b, a]).matching_rules_hash + rxn_path, mxn_path = tmp_path / 'rxn.yml', tmp_path / 'mxn.yml' + rxn_path.write_text(yaml.safe_dump([reaction.to_dict()])) + mxn_path.write_text(yaml.safe_dump([matching.to_dict()])) + first = RuleSet.load_from_files(rxn_path, mxn_path) + rxn_path.write_text('# different formatting\n' + yaml.safe_dump([reaction.to_dict()], default_flow_style=True)) + second = RuleSet.load_from_files(rxn_path, mxn_path) + assert first.reaction_rules_hash == second.reaction_rules_hash == rules.reaction_rules_hash + assert first.matching_rules_hash == second.matching_rules_hash == rules.matching_rules_hash diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..64507a6 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,174 @@ +"""Single-mode drawing limits use distinct, reproducible assembly projections.""" + +import json +from pathlib import Path + +import pytest + +from retromol.model.assembly_graph import AssemblyGraph +from retromol.model.assembly_sampling import AssemblyLimitError, AssemblySpace +from retromol.model.result import Result +from retromol.visualization.assembly_layout import MoleculeLayout +from retromol_cli import cli as cli_module + + +SMILES = 'CCCCCCCCCC(=O)NCC(=O)NC(C)C(=O)O' + + +@pytest.fixture +def single_run(monkeypatch, tmp_path): + # Exercise parsing, serialization, selection and coverage; only the expensive + # rendering and global logging configuration are replaced. + drawn = [] + visualized = [] + layouts = [] + monkeypatch.setattr(cli_module, 'setup_logging', lambda **kwargs: None) + monkeypatch.setattr(cli_module, 'add_file_handler', lambda *args, **kwargs: None) + + def record_drawing(assembly, **kwargs): + drawn.append((Path(kwargs['savepath']).name, tuple(sorted(assembly.monomer_ids())))) + assert kwargs['show_unassigned'] + layout = kwargs['layout'] + assert isinstance(layout, MoleculeLayout) + if layouts: + assert layout is layouts[0] + layouts.append(layout) + + monkeypatch.setattr(AssemblyGraph, 'draw', record_drawing) + monkeypatch.setattr(cli_module, 'visualize_reaction_graph', lambda *args, **kwargs: visualized.append(kwargs)) + + def run(*options, smiles=SMILES): + drawn.clear() + visualized.clear() + layouts.clear() + output = tmp_path / 'output' + monkeypatch.setattr('sys.argv', ['retromol', '-o', str(output), 'single', '-s', smiles, *options]) + cli_module.main() + payload = json.loads((output / 'result.json').read_text()) + assert 'linear_readout' not in payload + assert 'assembly_graph' not in payload + assert len(visualized) == 1 + return Result.from_dict(payload), list(drawn) + + return run + + +def test_single_default_keeps_default_assembly_filename_and_selection(single_run): + result, drawn = single_run() + assert drawn == [('assembly_graph.png', tuple(sorted(result.assembly_graph.monomer_ids())))] + + +@pytest.mark.parametrize('maximum,expected_count', [(0, 0), (1, 1), (3, 3), (10, 5)]) +def test_single_respects_drawing_limit_and_available_assemblies(single_run, maximum, expected_count): + result, drawn = single_run('--max-assembly-graphs', str(maximum)) + assert len(drawn) == expected_count + assert len({nodes for _, nodes in drawn}) == expected_count + if maximum > 1: + assert [name for name, _ in drawn] == [f'assembly_graph_{index:03d}.png' for index in range(1, expected_count + 1)] + expected = [tuple(sorted(a.monomer_ids())) for a in result.sample_assemblies(maximum, seed=42)] + assert [nodes for _, nodes in drawn] == expected + + +def test_single_passes_seed_to_sampling(single_run): + result, first = single_run('--max-assembly-graphs', '3', '--seed', '17') + _, repeated = single_run('--max-assembly-graphs', '3', '--seed', '17') + assert first == repeated + assert [nodes for _, nodes in first] == [tuple(sorted(a.monomer_ids())) for a in result.sample_assemblies(3, seed=17)] + + +def test_single_rerun_removes_stale_generated_drawings(single_run, tmp_path): + output = tmp_path / 'output' + output.mkdir() + stale = [output / 'assembly_graph.png', *(output / f'assembly_graph_{i:03d}.png' for i in range(6, 11))] + for path in stale: + path.write_bytes(b'old drawing') + user_image = output / 'assembly_graph_notes.png' + user_image.write_bytes(b'user image') + _, drawn = single_run('--max-assembly-graphs', '10') + assert len(drawn) == 5 + assert not any(path.exists() for path in stale) + assert user_image.read_bytes() == b'user image' + + +@pytest.mark.parametrize('maximum', ['-1', '1.5', 'invalid']) +def test_single_rejects_invalid_limit_before_creating_output(monkeypatch, tmp_path, capsys, maximum): + output = tmp_path / 'output' + monkeypatch.setattr('sys.argv', ['retromol', '-o', str(output), 'single', '-s', SMILES, '--max-assembly-graphs', maximum]) + with pytest.raises(SystemExit) as error: + cli_module.main() + assert error.value.code == 2 + assert 'must be a nonnegative integer' in capsys.readouterr().err + assert not output.exists() + + +def test_single_enumeration_limit_fails_clearly_and_preserves_result(single_run, monkeypatch, tmp_path, caplog): + def over_budget(*args, **kwargs): + raise AssemblyLimitError('enumeration budget exceeded') + + monkeypatch.setattr(AssemblySpace, 'sample', over_budget) + with pytest.raises(SystemExit) as error: + single_run('--max-assembly-graphs', '10') + assert error.value.code == 1 + assert 'enumeration budget exceeded' in caplog.text + assert 'reaction graph is saved' in caplog.text + assert Result.from_dict(json.loads((tmp_path / 'output' / 'result.json').read_text())).root + + +def test_daptomycin_samples_only_full_coverage_and_keeps_all_resolutions(single_run, caplog): + from .data.integration_retromol import CASES + + smiles = next(smiles for name, smiles, *_ in CASES if name == 'daptomycin') + with caplog.at_level('INFO'): + result, drawn = single_run('--max-assembly-graphs', '20', smiles=smiles) + space = result.assembly_space() + assert space.count() == 5 + assert space.count(min_coverage=None) == len(drawn) == 5 + assert 'Drawing 5 of 5 eligible assemblies (5 total' in caplog.text + assemblies = [space.build(frozenset(nodes)) for _, nodes in drawn] + assert all(result.calculate_coverage(a) == 1 for a in assemblies) + fatty_acids = {'decanoic acid', 'octanoic acid', 'hexanoic acid', 'butanoic acid', 'acetic acid'} + assert fatty_acids <= {node.identity.name for a in assemblies for node in a.monomer_nodes()} + assert len({nodes for _, nodes in drawn}) == 5 + threonine = next(n.enc for n in result.reaction_graph.identified_nodes.values() + if n.identity.name == 'threonine') + for assembly in space.sample(20, seed=42, min_coverage=0): + assert threonine in assembly.monomer_ids() + assert 'B12' not in {n.identity.name for n in assembly.monomer_nodes()} + assert result.calculate_coverage(assembly) == 1 + # The source graph remains agnostic, including the bypass's identified B12. + assert any(n.identity.name == 'B12' for n in result.reaction_graph.identified_nodes.values()) + + +def test_single_default_also_uses_maximum_if_default_projection_is_lower(single_run, monkeypatch): + from .test_assembly_sampling import partially_identified_result + + parsed = partially_identified_result() + monkeypatch.setattr(cli_module, 'run_retromol_with_timeout', lambda *args: parsed) + result, drawn = single_run() + assert result.calculate_coverage() == .75 + assert len(drawn) == 1 + assert result.calculate_coverage(result.assembly_space().build(frozenset(drawn[0][1]))) == 1 + _, lower = single_run('--min-coverage', '.75') + assert lower == [('assembly_graph.png', tuple(sorted(result.assembly_graph.monomer_ids())))] + _, all_views = single_run('--max-assembly-graphs', '10', '--min-coverage', '.75') + assert len(all_views) == 4 + + +@pytest.mark.parametrize('maximum', [1, 10]) +def test_single_unattainable_minimum_draws_nothing_and_keeps_result(single_run, caplog, maximum): + result, drawn = single_run('--max-assembly-graphs', str(maximum), '--min-coverage', '1', + smiles=SMILES + '.[Na+]') + assert not drawn + assert result.assembly_space().max_coverage < 1 + assert 'No assemblies meet minimum coverage 100.00%' in caplog.text + + +@pytest.mark.parametrize('minimum', ['-0.1', '1.1', 'NaN', 'inf', 'invalid']) +def test_single_rejects_invalid_coverage_before_creating_output(monkeypatch, tmp_path, capsys, minimum): + output = tmp_path / 'output' + monkeypatch.setattr('sys.argv', ['retromol', '-o', str(output), 'single', '-s', SMILES, '--min-coverage', minimum]) + with pytest.raises(SystemExit) as error: + cli_module.main() + assert error.value.code == 2 + assert 'must be a finite fraction between 0 and 1' in capsys.readouterr().err + assert not output.exists() diff --git a/tests/test_json_io.py b/tests/test_json_io.py new file mode 100644 index 0000000..59a3387 --- /dev/null +++ b/tests/test_json_io.py @@ -0,0 +1,164 @@ +"""Lossless, streaming result storage in plain JSON and gzip JSONL.""" + +import gzip +import json +from pathlib import Path + +import pytest + +from retromol.io.json import dumps_compact, iter_json, open_json_file +from retromol.io.streaming import ResultEvent, stream_json_records +from retromol.model.result import Result +from retromol.model.submission import Submission +from retromol.pipelines.parsing import run_retromol +from retromol_cli import cli as cli_module + + +RECORDS = [ + {'id': 'same', 'note': 'Ångström, α/β; keep spaces inside strings', 'nested': {'flag': False, 'missing': None}}, + {'id': 'same', 'note': 'line one\nline two', 'numbers': [0, -3, 1.25]}, +] + + +@pytest.mark.parametrize('suffix', ['.jsonl', '.jsonl.gz', '.ndjson', '.ndjson.gz', '.json', '.json.gz']) +def test_streaming_plain_and_gzip_preserve_all_json_values(tmp_path, suffix): + path = tmp_path / ('records' + suffix) + with open_json_file(path, 'wt') as handle: + if suffix.startswith('.jsonl') or suffix.startswith('.ndjson'): + for record in RECORDS: + handle.write(dumps_compact(record) + '\n') + else: + handle.write(dumps_compact(RECORDS)) + assert list(iter_json(path)) == RECORDS + assert list(stream_json_records(str(path))) == RECORDS + if suffix.endswith('.gz'): + assert path.read_bytes().startswith(b'\x1f\x8b') + + +def test_compact_json_only_removes_formatting_whitespace(): + text = dumps_compact(RECORDS) + assert json.loads(text) == RECORDS + assert 'Ångström' in text + assert len(text.encode()) < len(json.dumps(RECORDS).encode()) + + +def test_appending_gzip_runs_reads_all_members(tmp_path): + for record in RECORDS: + handle, path = cli_module._open_jsonl(str(tmp_path), None) + with handle: + handle.write(dumps_compact(record) + '\n') + assert path == str(tmp_path / 'results.jsonl.gz') + assert list(iter_json(path)) == RECORDS + + +def test_custom_jsonl_filename_supports_current_directory(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + handle, path = cli_module._open_jsonl(str(tmp_path), 'custom.jsonl') + with handle: + handle.write(dumps_compact(RECORDS[0]) + '\n') + assert path == 'custom.jsonl.gz' + assert list(iter_json(path)) == RECORDS[:1] + + +@pytest.fixture +def batch_run(monkeypatch, tmp_path, ruleset): + result = run_retromol(Submission('CCO.[Na+]', props=RECORDS[0]), ruleset) + payload = result.to_dict() + input_path = tmp_path / 'input.jsonl.gz' + with open_json_file(input_path, 'wt') as handle: + handle.write(dumps_compact({'smiles': 'CCO.[Na+]', **RECORDS[0]}) + '\n') + output = tmp_path / 'output' + monkeypatch.setattr(cli_module, 'setup_logging', lambda **kwargs: None) + monkeypatch.setattr(cli_module, 'add_file_handler', lambda *args, **kwargs: None) + + def events(**kwargs): + rows = list(kwargs['row_iter']) + assert rows == [{'smiles': 'CCO.[Na+]', **RECORDS[0]}] + yield ResultEvent(payload, None) + yield ResultEvent(None, 'intentional failed input') + yield ResultEvent(payload, None) # repeated IDs must not overwrite results + + monkeypatch.setattr(cli_module, 'run_retromol_stream', events) + + def run(*options): + monkeypatch.setattr('sys.argv', ['retromol', '-o', str(output), 'batch', '--json', str(input_path), '--no-tqdm', *options]) + cli_module.main() + return output, payload + + return run + + +@pytest.mark.parametrize('compression,suffix', [('gzip', '.jsonl.gz'), ('none', '.jsonl')]) +def test_batch_writes_lossless_results_with_selected_compression(batch_run, compression, suffix): + output, payload = batch_run('--compression', compression) + path = output / ('results' + suffix) + records = list(iter_json(path)) + assert records == [payload, payload] + for record in records: + loaded = Result.from_dict(record) + assert loaded.props == RECORDS[0] + assert loaded.calculate_coverage() == .75 + assert len(loaded.reaction_rules_hash) == len(loaded.matching_rules_hash) == 64 + with open_json_file(path) as handle: + assert handle.readline().rstrip('\n') == dumps_compact(payload) + + +def test_batch_compresses_by_default(batch_run): + output, payload = batch_run() + assert not (output / 'results.jsonl').exists() + assert list(iter_json(output / 'results.jsonl.gz')) == [payload, payload] + + +@pytest.mark.parametrize('error', [RuntimeError, KeyboardInterrupt]) +def test_interrupted_batch_closes_gzip_and_preserves_completed_records(batch_run, monkeypatch, error): + original = cli_module.run_retromol_stream + + def interrupted(**kwargs): + yield next(original(**kwargs)) + raise error('interrupted run') + + monkeypatch.setattr(cli_module, 'run_retromol_stream', interrupted) + with pytest.raises(error): + batch_run() + path = Path(cli_module.cli().outdir) / 'results.jsonl.gz' + records = list(iter_json(path)) + assert len(records) == 1 + assert Result.from_dict(records[0]).props == RECORDS[0] + + +@pytest.mark.parametrize('compression,suffix', [('gzip', '.json.gz'), ('none', '.json')]) +def test_separate_result_files_preserve_repeated_ids(batch_run, compression, suffix): + output, payload = batch_run('--results', 'files', '--compression', compression) + paths = sorted(output.glob('result_*')) + assert [p.name for p in paths] == [f'result_000000001{suffix}', f'result_000000003{suffix}'] + before = [p.read_bytes() for p in paths] + for path in paths: + with open_json_file(path) as handle: + assert json.load(handle) == payload + with pytest.raises(FileExistsError): + batch_run('--results', 'files', '--compression', compression) + assert [p.read_bytes() for p in paths] == before + + +def test_plain_output_rejects_misleading_gzip_extension(tmp_path, monkeypatch, capsys): + output = tmp_path / 'output' + monkeypatch.setattr('sys.argv', ['retromol', '-o', str(output), 'batch', '--json', 'input.jsonl', + '--compression', 'none', '--jsonl-path', 'results.jsonl.gz']) + with pytest.raises(SystemExit) as error: + cli_module.main() + assert error.value.code == 2 + assert 'requires --compression gzip' in capsys.readouterr().err + assert not output.exists() + + +def test_single_can_save_compressed_result(tmp_path, monkeypatch): + monkeypatch.setattr(cli_module, 'setup_logging', lambda **kwargs: None) + monkeypatch.setattr(cli_module, 'add_file_handler', lambda *args, **kwargs: None) + monkeypatch.setattr(cli_module, 'visualize_reaction_graph', lambda *args, **kwargs: None) + monkeypatch.setattr('sys.argv', ['retromol', '-o', str(tmp_path), 'single', '-s', 'CCO.[Na+]', + '--compression', 'gzip', '--max-assembly-graphs', '0']) + cli_module.main() + assert not (tmp_path / 'result.json').exists() + with gzip.open(tmp_path / 'result.json.gz', 'rt') as handle: + result = Result.from_dict(json.load(handle)) + assert result.calculate_coverage() == .75 diff --git a/tests/test_parsing.py b/tests/test_parsing.py index ad60bd9..3694e6e 100644 --- a/tests/test_parsing.py +++ b/tests/test_parsing.py @@ -5,7 +5,6 @@ from retromol.chem.mol import encode_mol, smiles_to_mol from retromol.model.reaction_graph import ReactionGraph, ReactionStep -from retromol.model.readout import LinearReadout from retromol.model.result import Result from retromol.model.rules import MatchingRule, RuleSet from retromol.model.submission import Submission @@ -57,7 +56,7 @@ def test_cyclic_fragment_does_not_discard_identified_siblings() -> None: root = encode_mol(submission.mol) extracted = extract_min_edge_synthesis_subgraph(graph, root) - result = Result(submission, extracted.graph, LinearReadout.from_reaction_graph(root, extracted.graph)) + result = Result(submission, extracted.graph) assert root in extracted.graph.nodes assert not extracted.solved @@ -109,7 +108,7 @@ def test_default_rules_preserve_partial_results_for_sulfated_compounds(smiles: s result = run_retromol_with_timeout(Submission(smiles), ruleset) assert encode_mol(result.submission.mol) in result.reaction_graph.nodes - assert_acyclic(result.reaction_graph) + assert_acyclic(result.selected_reaction_graph) assert 0 < result.calculate_coverage() < 1 assert any(not node.is_identified for node in result.linear_readout.assembly_graph.monomer_nodes()) restored = Result.from_dict(result.to_dict()) diff --git a/tests/test_result.py b/tests/test_result.py index a376d2b..bd2b4db 100644 --- a/tests/test_result.py +++ b/tests/test_result.py @@ -32,13 +32,7 @@ def make_result(smiles: str, fragments: list[tuple[str, bool]]) -> Result: node = graph.nodes[enc] rules = [MatchingRule(name=node.smiles, smiles=tagged_smiles, props={}, pseudonyms=[])] if identified else [] node.identify(rules) - assembly = AssemblyGraph.build( - root_mol=submission.mol, - monomers=[graph.nodes[enc] for enc in children], - include_unassigned=True, - ) - # No extracted paths: identified assembly nodes must count independently of sequences. - return Result(submission, graph, LinearReadout(assembly_graph=assembly, paths=[])) + return Result(submission, graph) def test_coverage_ignores_identified_intermediates_and_unknown_leaves() -> None: @@ -49,10 +43,13 @@ def test_coverage_ignores_identified_intermediates_and_unknown_leaves() -> None: assert Result.from_dict(result.to_dict()).calculate_coverage() == pytest.approx(0.5) -def test_coverage_counts_assembly_monomers_without_linear_paths() -> None: +def test_coverage_counts_assembly_monomers_without_linear_paths(monkeypatch) -> None: result = make_result("CCOC", [("[1CH3][2CH3]", True), ("[3OH][4CH3]", True)]) - assert result.linear_readout.paths == [] + def unexpected_readout(*args, **kwargs): + raise AssertionError("Coverage must not request sequences") + + monkeypatch.setattr(LinearReadout, "from_assembly_graph", unexpected_readout) assert result.calculate_coverage() == 1.0 @@ -68,6 +65,8 @@ def test_coverage_includes_all_submitted_fragments_in_denominator() -> None: assert result.submission.mol.GetNumHeavyAtoms() == 3 assert result.calculate_coverage() == pytest.approx(3 / 4) assert Result.from_dict(result.to_dict()).calculate_coverage() == pytest.approx(3 / 4) + assert result.assembly_space().max_coverage == pytest.approx(3 / 4) + assert result.sample_assemblies(min_coverage=1) == [] def test_coverage_ignores_hydrogens_and_reaction_introduced_atoms() -> None: @@ -78,11 +77,14 @@ def test_coverage_ignores_hydrogens_and_reaction_introduced_atoms() -> None: assert result.submission.mol.GetNumAtoms() == 4 assert result.calculate_coverage() == pytest.approx(2 / 3) + space = result.assembly_space() + assert {result.calculate_coverage(a) for a in space.sample(10, min_coverage=0)} == {2 / 3, 1} + assert space.count(min_coverage=.7) == 1 def test_coverage_does_not_drop_untagged_input_atoms_from_denominator() -> None: result = make_result("CCO", [("[1CH3][2CH2]O", True)]) - result.submission.mol.GetAtomWithIdx(2).SetIsotope(0) + result.root.mol.GetAtomWithIdx(2).SetIsotope(0) assert result.calculate_coverage() == pytest.approx(2 / 3) @@ -92,18 +94,20 @@ def test_coverage_with_no_heavy_atoms_is_zero(smiles: str) -> None: result = make_result(smiles, []) assert result.calculate_coverage() == 0.0 + assert result.assembly_space().max_coverage == 0.0 + assert result.calculate_coverage(result.sample_assemblies()[0]) == 0.0 def test_coverage_of_empty_assembly_is_zero() -> None: result = make_result("CCO", []) - assert result.calculate_coverage() == 0.0 + assert result.calculate_coverage(AssemblyGraph.build(result.root.mol, [])) == 0.0 def test_coverage_counts_repeated_atom_tags_only_once() -> None: result = make_result("CCO", [("[1CH3][2CH3]", True)]) - graph = result.linear_readout.assembly_graph.g - node = result.linear_readout.assembly_graph.monomer_nodes()[0] + graph = result.assembly_graph.g + node = result.assembly_graph.monomer_nodes()[0] graph.add_node( "overlapping monomer", tags=graph.nodes[node.enc]["tags"], identity=node.identity, molnode=replace(node, enc="overlapping monomer"),