retrocast 0__py3-none-any.whl

retrocast/__init__.py

@@ -0,0 +1,56 @@
+"""
+retrocast: A unified toolkit for retrosynthesis benchmark analysis.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from packaging.version import Version
+
+from retrocast.adapters import ADAPTER_MAP, adapt_routes, adapt_single_route, get_adapter
+from retrocast.curation.filtering import deduplicate_routes
+from retrocast.curation.sampling import sample_k_by_length, sample_random_k, sample_top_k
+from retrocast.models.chem import Molecule, ReactionStep, Route, TargetInput
+
+
+def _normalize_version_with_patch(version_str: str) -> str:
+    """Ensure version always has explicit major.minor.micro format (e.g., 0.3.0.dev16 not 0.3.dev16)."""
+    v = Version(version_str)
+    # Reconstruct with explicit patch version
+    base = f"{v.major}.{v.minor}.{v.micro}"
+
+    # Add pre-release, post-release, dev, local parts if present
+    parts = [base]
+    if v.pre:
+        parts.append(f"{v.pre[0]}{v.pre[1]}")
+    if v.post is not None:
+        parts.append(f".post{v.post}")
+    if v.dev is not None:
+        parts.append(f".dev{v.dev}")
+    if v.local:
+        parts.append(f"+{v.local}")
+
+    return "".join(parts)
+
+
+try:
+    __version__ = _normalize_version_with_patch(version("retrocast"))
+except PackageNotFoundError:
+    # Package not installed (running from source without editable install)
+    __version__ = "0.0.0.dev0+unknown"
+__all__ = [
+    # Core schemas
+    "Route",
+    "Molecule",
+    "ReactionStep",
+    "TargetInput",
+    # Adapter functions
+    "adapt_single_route",
+    "adapt_routes",
+    "get_adapter",
+    "ADAPTER_MAP",
+    # Route processing utilities
+    "deduplicate_routes",
+    "sample_top_k",
+    "sample_random_k",
+    "sample_k_by_length",
+]

retrocast/adapters/__init__.py

@@ -0,0 +1,169 @@
+from typing import Any
+
+from retrocast.adapters.aizynth_adapter import AizynthAdapter
+from retrocast.adapters.askcos_adapter import AskcosAdapter
+from retrocast.adapters.base_adapter import BaseAdapter
+from retrocast.adapters.dms_adapter import DMSAdapter
+from retrocast.adapters.dreamretro_adapter import DreamRetroAdapter
+from retrocast.adapters.multistepttl_adapter import TtlRetroAdapter
+from retrocast.adapters.paroutes_adapter import PaRoutesAdapter
+from retrocast.adapters.retrochimera_adapter import RetrochimeraAdapter
+from retrocast.adapters.retrostar_adapter import RetroStarAdapter
+from retrocast.adapters.synllama_adapter import SynLlaMaAdapter
+from retrocast.adapters.synplanner_adapter import SynPlannerAdapter
+from retrocast.adapters.syntheseus_adapter import SyntheseusAdapter
+from retrocast.exceptions import RetroCastException
+from retrocast.models.chem import Route, TargetIdentity
+
+ADAPTER_MAP: dict[str, BaseAdapter] = {
+    "aizynth": AizynthAdapter(),
+    "askcos": AskcosAdapter(),
+    "dms": DMSAdapter(),
+    "dreamretro": DreamRetroAdapter(),
+    "multistepttl": TtlRetroAdapter(),
+    "paroutes": PaRoutesAdapter(),
+    "retrochimera": RetrochimeraAdapter(),
+    "retrostar": RetroStarAdapter(),
+    "synplanner": SynPlannerAdapter(),
+    "syntheseus": SyntheseusAdapter(),
+    "synllama": SynLlaMaAdapter(),
+}
+
+# Adapters that expect target-centric data format (dict with metadata + nested routes)
+# vs route-centric format (list of route objects)
+TARGET_CENTRIC_ADAPTERS = {"askcos", "retrochimera", "paroutes"}
+
+
+def get_adapter(adapter_name: str) -> BaseAdapter:
+    """
+    Retrieves an adapter instance from the `ADAPTER_MAP`.
+    """
+    adapter = ADAPTER_MAP.get(adapter_name)
+    if adapter is None:
+        raise RetroCastException(
+            f"unknown adapter '{adapter_name}'. Check `retrocast.adapters.ADAPTER_MAP` for available adapters."
+        )
+    return adapter
+
+
+def adapt_single_route(
+    raw_route: Any,
+    target: TargetIdentity,
+    adapter_name: str,
+) -> Route | None:
+    """
+    Adapt a single raw route to the unified Route format.
+
+    This is a convenience function for users who want to adapt individual routes
+    programmatically without the full batch processing pipeline. It intelligently
+    handles both route-centric and target-centric adapter formats.
+
+    Args:
+        raw_route: A single route or target data in the model's native format.
+            - For route-centric adapters (DMS, AiZynth, SynPlanner): Pass a single
+              route object/dict from the model's output list.
+            - For target-centric adapters (RetroChimera, ASKCOS): Pass the complete
+              target data dict (containing target metadata and nested routes).
+        target: Target molecule information (id and canonical SMILES).
+        adapter_name: Name of the adapter to use (e.g., "dms", "aizynth", "retrostar").
+            See ADAPTER_MAP.keys() for available adapters.
+
+    Returns:
+        Route object if successful, None if adaptation failed.
+
+    Examples:
+        Route-centric adapter (DMS):
+        >>> from retrocast.adapters import adapt_single_route
+        >>> from retrocast.models.chem import TargetIdentity
+        >>>
+        >>> target = TargetIdentity(id="aspirin", smiles="CC(=O)Oc1ccccc1C(=O)O")
+        >>> raw_dms_route = {"smiles": "CC(=O)Oc1ccccc1C(=O)O", "children": [...]}
+        >>>
+        >>> route = adapt_single_route(raw_dms_route, target, "dms")
+        >>> if route:
+        ...     print(f"Route depth: {route.length}")
+        ...     print(f"Starting materials: {len(route.leaves)}")
+
+        Target-centric adapter (RetroChimera):
+        >>> target = TargetIdentity(id="mol1", smiles="CCO")
+        >>> retrochimera_data = {
+        ...     "smiles": "CCO",
+        ...     "result": {"outputs": [{"routes": [...]}]}
+        ... }
+        >>> route = adapt_single_route(retrochimera_data, target, "retrochimera")
+    """
+    adapter = get_adapter(adapter_name)
+
+    # Determine if this adapter expects target-centric or route-centric format
+    if adapter_name in TARGET_CENTRIC_ADAPTERS:
+        # Target-centric adapters (RetroChimera, ASKCOS) expect a dict directly
+        raw_data = raw_route
+    else:
+        # Route-centric adapters (DMS, AiZynth, etc.) expect a list of routes
+        raw_data = [raw_route] if not isinstance(raw_route, list) else raw_route
+
+    # Get first successful route from the generator
+    for route in adapter.cast(raw_data, target):
+        return route
+
+    return None
+
+
+def adapt_routes(
+    raw_routes: Any,
+    target: TargetIdentity,
+    adapter_name: str,
+    max_routes: int | None = None,
+) -> list[Route]:
+    """
+    Adapt multiple raw routes to the unified Route format.
+
+    Args:
+        raw_routes: Routes in the model's native format (typically a list or dict).
+        target: Target molecule information (id and canonical SMILES).
+        adapter_name: Name of the adapter to use.
+        max_routes: Maximum number of routes to return (None for all successful routes).
+
+    Returns:
+        List of successfully adapted Route objects.
+
+    Example:
+        >>> from retrocast.adapters import adapt_routes
+        >>> from retrocast.models.chem import TargetIdentity
+        >>>
+        >>> target = TargetIdentity(id="ibuprofen", smiles="CC(C)Cc1ccc(cc1)C(C)C(=O)O")
+        >>> raw_routes = [route1, route2, route3, ...]  # Your model's output
+        >>>
+        >>> routes = adapt_routes(raw_routes, target, "aizynth", max_routes=10)
+        >>> print(f"Adapted {len(routes)} routes successfully")
+    """
+    adapter = get_adapter(adapter_name)
+    routes = []
+
+    for i, route in enumerate(adapter.cast(raw_routes, target)):
+        routes.append(route)
+        if max_routes and i + 1 >= max_routes:
+            break
+
+    return routes
+
+
+__all__ = [
+    "adapt_single_route",
+    "adapt_routes",
+    "get_adapter",
+    "ADAPTER_MAP",
+    "TARGET_CENTRIC_ADAPTERS",
+    "BaseAdapter",
+    "AizynthAdapter",
+    "AskcosAdapter",
+    "DMSAdapter",
+    "DreamRetroAdapter",
+    "TtlRetroAdapter",
+    "PaRoutesAdapter",
+    "RetrochimeraAdapter",
+    "RetroStarAdapter",
+    "SynPlannerAdapter",
+    "SyntheseusAdapter",
+    "SynLlaMaAdapter",
+]

retrocast/adapters/aizynth_adapter.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Generator
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field, RootModel, ValidationError
+
+from retrocast.adapters.base_adapter import BaseAdapter
+from retrocast.adapters.common import build_molecule_from_bipartite_node
+from retrocast.exceptions import AdapterLogicError, RetroCastException
+from retrocast.models.chem import Route, TargetIdentity
+
+logger = logging.getLogger(__name__)
+
+# --- pydantic models for input validation ---
+# these models validate the raw aizynthfinder output format before any transformation.
+
+
+class AizynthBaseNode(BaseModel):
+    """a base model for shared fields between node types."""
+
+    smiles: str
+    children: list[AizynthNode] = Field(default_factory=list)
+
+
+class AizynthMoleculeInput(AizynthBaseNode):
+    """represents a 'mol' node in the raw aizynth tree."""
+
+    type: Literal["mol"]
+    in_stock: bool = False
+
+
+class AizynthReactionInput(AizynthBaseNode):
+    """represents a 'reaction' node in the raw aizynth tree."""
+
+    type: Literal["reaction"]
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+# a discriminated union to handle the bipartite graph structure.
+AizynthNode = Annotated[AizynthMoleculeInput | AizynthReactionInput, Field(discriminator="type")]
+
+
+class AizynthRouteList(RootModel[list[AizynthMoleculeInput]]):
+    """the top-level object for a single target is a list of potential routes."""
+
+    pass
+
+
+class AizynthAdapter(BaseAdapter):
+    """adapter for converting aizynthfinder-style outputs to the benchmarktree schema."""
+
+    def cast(self, raw_target_data: Any, target: TargetIdentity) -> Generator[Route, None, None]:
+        """
+        validates raw aizynthfinder data, transforms it, and yields route objects.
+        """
+        try:
+            validated_routes = AizynthRouteList.model_validate(raw_target_data)
+        except ValidationError as e:
+            logger.warning(f"  - raw data for target '{target.id}' failed aizynth schema validation. error: {e}")
+            return
+
+        for rank, aizynth_tree_root in enumerate(validated_routes.root, start=1):
+            try:
+                route = self._transform(aizynth_tree_root, target, rank)
+                yield route
+            except RetroCastException as e:
+                logger.warning(f"  - route for '{target.id}' failed transformation: {e}")
+                continue
+
+    def _transform(self, aizynth_root: AizynthMoleculeInput, target: TargetIdentity, rank: int) -> Route:
+        """
+        orchestrates the transformation of a single aizynthfinder output tree.
+        raises RetroCastException on failure.
+        """
+        # use the common recursive builder with new schema
+        target_molecule = build_molecule_from_bipartite_node(raw_mol_node=aizynth_root)
+
+        if target_molecule.smiles != target.smiles:
+            msg = (
+                f"mismatched smiles for target {target.id}. "
+                f"expected canonical: {target.smiles}, but adapter produced: {target_molecule.smiles}"
+            )
+            logger.error(msg)
+            raise AdapterLogicError(msg)
+
+        return Route(target=target_molecule, rank=rank, metadata={})

retrocast/adapters/askcos_adapter.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from collections.abc import Generator
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field, ValidationError
+
+from retrocast.adapters.base_adapter import BaseAdapter
+from retrocast.chem import canonicalize_smiles, get_inchi_key
+from retrocast.exceptions import AdapterLogicError, RetroCastException
+from retrocast.models.chem import Molecule, ReactionStep, Route, TargetIdentity
+from retrocast.typing import ReactionSmilesStr, SmilesStr
+
+logger = logging.getLogger(__name__)
+
+# --- pydantic models for input validation ---
+
+
+class AskcosBaseNode(BaseModel):
+    smiles: str
+    id: str
+
+
+class AskcosChemicalNode(AskcosBaseNode):
+    type: Literal["chemical"]
+    terminal: bool
+
+
+class AskcosTemplateSource(BaseModel):
+    """Nested structure for template information."""
+
+    reaction_smarts: str | None = None
+
+
+class AskcosModelMetadata(BaseModel):
+    """Model metadata containing template information."""
+
+    source: dict[str, Any] = Field(default_factory=dict)
+
+    def get_template(self) -> str | None:
+        """Extract reaction_smarts from nested template structure."""
+        template_dict = self.source.get("template", {})
+        return template_dict.get("reaction_smarts") if isinstance(template_dict, dict) else None
+
+
+class AskcosReactionProperties(BaseModel):
+    mapped_smiles: str | None = None
+
+
+class AskcosReactionNode(AskcosBaseNode):
+    type: Literal["reaction"]
+    reaction_properties: AskcosReactionProperties | None = None
+    model_metadata: list[AskcosModelMetadata] = Field(default_factory=list)
+
+
+AskcosNode = Annotated[AskcosChemicalNode | AskcosReactionNode, Field(discriminator="type")]
+
+
+class AskcosPathwayEdge(BaseModel):
+    source: str
+    target: str
+
+
+class AskcosUDS(BaseModel):
+    node_dict: dict[str, AskcosNode]
+    uuid2smiles: dict[str, str]
+    pathways: list[list[AskcosPathwayEdge]]
+
+
+class AskcosResults(BaseModel):
+    uds: AskcosUDS
+
+
+class AskcosOutput(BaseModel):
+    results: AskcosResults
+
+
+class AskcosAdapter(BaseAdapter):
+    """adapter for converting askcos outputs to the benchmarktree schema."""
+
+    def __init__(self, use_full_graph: bool = False):
+        """
+        initializes the adapter.
+
+        args:
+            use_full_graph: if true, attempts to extract all possible routes
+                from the full search graph instead of using the pre-computed
+                pathways. defaults to false.
+        """
+        self.use_full_graph = use_full_graph
+
+    def cast(self, raw_target_data: Any, target: TargetIdentity) -> Generator[Route, None, None]:
+        """validates raw askcos data, transforms its pathways, and yields route objects."""
+        if self.use_full_graph:
+            raise NotImplementedError("extracting routes from the full askcos search graph is not yet implemented.")
+
+        try:
+            validated_output = AskcosOutput.model_validate(raw_target_data)
+        except ValidationError as e:
+            logger.warning(f"  - raw data for target '{target.id}' failed askcos schema validation. error: {e}")
+            return
+
+        uds = validated_output.results.uds
+
+        # Extract metadata from stats if available
+        stats = raw_target_data.get("results", {}).get("stats", {})
+        metadata = {
+            "total_iterations": stats.get("total_iterations"),
+            "total_chemicals": stats.get("total_chemicals"),
+            "total_reactions": stats.get("total_reactions"),
+            "total_templates": stats.get("total_templates"),
+            "total_paths": stats.get("total_paths"),
+        }
+
+        for i, pathway_edges in enumerate(uds.pathways):
+            try:
+                route = self._transform_pathway(
+                    pathway_edges=pathway_edges,
+                    uuid2smiles=uds.uuid2smiles,
+                    node_dict=uds.node_dict,
+                    target_input=target,
+                    rank=i + 1,
+                    metadata=metadata,
+                )
+                yield route
+            except RetroCastException as e:
+                logger.warning(f"  - pathway {i} for target '{target.id}' failed transformation: {e}")
+                continue
+
+    def _transform_pathway(
+        self,
+        pathway_edges: list[AskcosPathwayEdge],
+        uuid2smiles: dict[str, str],
+        node_dict: dict[str, AskcosNode],
+        target_input: TargetIdentity,
+        rank: int,
+        metadata: dict[str, Any],
+    ) -> Route:
+        """transforms a single askcos pathway (represented by its edges) into a route."""
+        adj_list = defaultdict(list)
+        for edge in pathway_edges:
+            adj_list[edge.source].append(edge.target)
+
+        root_uuid = "00000000-0000-0000-0000-000000000000"
+        if root_uuid not in uuid2smiles:
+            raise AdapterLogicError("root uuid not found in pathway data.")
+
+        target_molecule = self._build_molecule(
+            chem_uuid=root_uuid,
+            path_prefix="retrocast-mol-root",
+            adj_list=adj_list,
+            uuid2smiles=uuid2smiles,
+            node_dict=node_dict,
+        )
+
+        if target_molecule.smiles != target_input.smiles:
+            msg = (
+                f"mismatched smiles for target {target_input.id}. "
+                f"expected canonical: {target_input.smiles}, but adapter produced: {target_molecule.smiles}"
+            )
+            raise AdapterLogicError(msg)
+
+        return Route(target=target_molecule, rank=rank, metadata=metadata)
+
+    def _build_molecule(
+        self,
+        chem_uuid: str,
+        path_prefix: str,
+        adj_list: dict[str, list[str]],
+        uuid2smiles: dict[str, str],
+        node_dict: dict[str, AskcosNode],
+    ) -> Molecule:
+        """recursively builds a canonical molecule from a chemical uuid."""
+        raw_smiles = uuid2smiles.get(chem_uuid)
+        if not raw_smiles:
+            raise AdapterLogicError(f"uuid '{chem_uuid}' not found in uuid2smiles map.")
+
+        node_data = node_dict.get(raw_smiles)
+        if not node_data or not isinstance(node_data, AskcosChemicalNode):
+            raise AdapterLogicError(f"node data for smiles '{raw_smiles}' not found or not a chemical node.")
+
+        canon_smiles = canonicalize_smiles(node_data.smiles)
+        is_leaf = node_data.terminal
+        synthesis_step = None
+
+        if not is_leaf and chem_uuid in adj_list:
+            child_reaction_uuids = adj_list[chem_uuid]
+            if len(child_reaction_uuids) > 1:
+                logger.warning(f"molecule {canon_smiles} has multiple child reactions in pathway; using first one.")
+
+            rxn_uuid = child_reaction_uuids[0]
+            synthesis_step = self._build_reaction_step(
+                rxn_uuid=rxn_uuid,
+                product_smiles=canon_smiles,
+                path_prefix=path_prefix,
+                adj_list=adj_list,
+                uuid2smiles=uuid2smiles,
+                node_dict=node_dict,
+            )
+
+        return Molecule(
+            smiles=canon_smiles,
+            inchikey=get_inchi_key(canon_smiles),
+            synthesis_step=synthesis_step,
+        )
+
+    def _build_reaction_step(
+        self,
+        rxn_uuid: str,
+        product_smiles: SmilesStr,
+        path_prefix: str,
+        adj_list: dict[str, list[str]],
+        uuid2smiles: dict[str, str],
+        node_dict: dict[str, AskcosNode],
+    ) -> ReactionStep:
+        """builds a canonical reaction step from a reaction uuid."""
+        raw_smiles = uuid2smiles.get(rxn_uuid)
+        if not raw_smiles:
+            raise AdapterLogicError(f"uuid '{rxn_uuid}' not found in uuid2smiles map.")
+
+        node_data = node_dict.get(raw_smiles)
+        if not node_data or not isinstance(node_data, AskcosReactionNode):
+            raise AdapterLogicError(f"node data for reaction '{raw_smiles}' not found or not a reaction node.")
+
+        reactants: list[Molecule] = []
+        reactant_smiles_list: list[SmilesStr] = []
+
+        reactant_uuids = adj_list.get(rxn_uuid, [])
+        for i, reactant_uuid in enumerate(reactant_uuids):
+            reactant_molecule = self._build_molecule(
+                chem_uuid=reactant_uuid,
+                path_prefix=f"{path_prefix}-{i}",
+                adj_list=adj_list,
+                uuid2smiles=uuid2smiles,
+                node_dict=node_dict,
+            )
+            reactants.append(reactant_molecule)
+            reactant_smiles_list.append(reactant_molecule.smiles)
+
+        # Extract mapped_smiles from reaction_properties if available
+        mapped_smiles = None
+        if node_data.reaction_properties and node_data.reaction_properties.mapped_smiles:
+            mapped_smiles = ReactionSmilesStr(node_data.reaction_properties.mapped_smiles)
+
+        # Extract template from model_metadata if available
+        template = None
+        if node_data.model_metadata and len(node_data.model_metadata) > 0:
+            template = node_data.model_metadata[0].get_template()
+
+        return ReactionStep(
+            reactants=reactants,
+            mapped_smiles=mapped_smiles,
+            template=template,
+        )

retrocast/adapters/base_adapter.py

@@ -0,0 +1,43 @@
+from abc import ABC, abstractmethod
+from collections.abc import Generator
+from typing import Any
+
+from retrocast.models.chem import Route, TargetIdentity
+
+
+class BaseAdapter(ABC):
+    """
+    Abstract base class for all model output adapters.
+
+    An adapter's role is to transform a model's raw output format into the
+    canonical `Route` schema.
+    """
+
+    @abstractmethod
+    def cast(self, raw_target_data: Any, target: TargetIdentity) -> Generator[Route, None, None]:
+        """
+        Validates, transforms, and yields Routes from raw model data.
+
+        This is the primary method for an adapter. It encapsulates all model-specific
+        logic. It should be a generator that yields successful routes and handles its
+        own exceptions internally by logging and continuing.
+
+        Args:
+            raw_target_data: The raw data blob from a file for a single target.
+                This blob can follow one of two common patterns:
+
+                1.  **Route-Centric**: The data is a list of route objects, where the
+                    root of each route object contains the target SMILES (e.g.,
+                    AiZynthFinder, DMS). `raw_target_data` is typically a `list`.
+
+                2.  **Target-Centric**: The data is a single JSON object that contains
+                    metadata (like a top-level `smiles` key) and a nested list of
+                    routes (e.g., RetroChimera). `raw_target_data` is typically a `dict`.
+
+                The adapter is responsible for handling the specific structure of its model.
+            target: The identity of the target molecule (id and canonical SMILES).
+
+        Yields:
+            Successfully transformed Route objects.
+        """
+        raise NotImplementedError