rdkit-cli 0.1.0__py3-none-any.whl

rdkit_cli/core/similarity.py

@@ -0,0 +1,206 @@
+"""Molecular similarity computation engine."""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Any
+
+from rdkit import Chem, DataStructs
+from rdkit.Chem import AllChem, rdMolDescriptors
+from rdkit.ML.Cluster import Butina
+
+from rdkit_cli.io.readers import MoleculeRecord
+
+
+class SimilarityMetric(Enum):
+    """Supported similarity metrics."""
+
+    TANIMOTO = "tanimoto"
+    DICE = "dice"
+    COSINE = "cosine"
+    SOKAL = "sokal"
+    RUSSEL = "russel"
+
+
+def get_morgan_fingerprint(mol: Chem.Mol, radius: int = 2, n_bits: int = 2048):
+    """Get Morgan fingerprint for a molecule."""
+    return rdMolDescriptors.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)
+
+
+def compute_similarity(
+    fp1,
+    fp2,
+    metric: SimilarityMetric = SimilarityMetric.TANIMOTO,
+) -> float:
+    """
+    Compute similarity between two fingerprints.
+
+    Args:
+        fp1: First fingerprint
+        fp2: Second fingerprint
+        metric: Similarity metric to use
+
+    Returns:
+        Similarity score (0-1)
+    """
+    if metric == SimilarityMetric.TANIMOTO:
+        return DataStructs.TanimotoSimilarity(fp1, fp2)
+    elif metric == SimilarityMetric.DICE:
+        return DataStructs.DiceSimilarity(fp1, fp2)
+    elif metric == SimilarityMetric.COSINE:
+        return DataStructs.CosineSimilarity(fp1, fp2)
+    elif metric == SimilarityMetric.SOKAL:
+        return DataStructs.SokalSimilarity(fp1, fp2)
+    elif metric == SimilarityMetric.RUSSEL:
+        return DataStructs.RusselSimilarity(fp1, fp2)
+    else:
+        raise ValueError(f"Unknown metric: {metric}")
+
+
+def bulk_tanimoto_similarity(query_fp, fps: list) -> list[float]:
+    """Compute Tanimoto similarity of query against multiple fingerprints."""
+    return list(DataStructs.BulkTanimotoSimilarity(query_fp, fps))
+
+
+class SimilaritySearcher:
+    """Search for similar molecules."""
+
+    def __init__(
+        self,
+        query_smiles: str,
+        threshold: float = 0.7,
+        metric: SimilarityMetric = SimilarityMetric.TANIMOTO,
+        radius: int = 2,
+        n_bits: int = 2048,
+    ):
+        """
+        Initialize similarity searcher.
+
+        Args:
+            query_smiles: Query molecule SMILES
+            threshold: Minimum similarity threshold
+            metric: Similarity metric
+            radius: Morgan fingerprint radius
+            n_bits: Fingerprint bit size
+        """
+        self.threshold = threshold
+        self.metric = metric
+        self.radius = radius
+        self.n_bits = n_bits
+
+        # Generate query fingerprint
+        query_mol = Chem.MolFromSmiles(query_smiles)
+        if query_mol is None:
+            raise ValueError(f"Invalid query SMILES: {query_smiles}")
+
+        self.query_fp = get_morgan_fingerprint(query_mol, radius, n_bits)
+
+    def search(self, record: MoleculeRecord) -> Optional[dict[str, Any]]:
+        """
+        Check if molecule is similar to query.
+
+        Args:
+            record: MoleculeRecord to check
+
+        Returns:
+            Dictionary with similarity score if above threshold, None otherwise
+        """
+        if record.mol is None:
+            return None
+
+        fp = get_morgan_fingerprint(record.mol, self.radius, self.n_bits)
+        similarity = compute_similarity(self.query_fp, fp, self.metric)
+
+        if similarity < self.threshold:
+            return None
+
+        result: dict[str, Any] = {
+            "smiles": record.smiles,
+            "similarity": round(similarity, 4),
+        }
+
+        if record.name:
+            result["name"] = record.name
+
+        return result
+
+
+def compute_similarity_matrix(
+    mols: list[Chem.Mol],
+    metric: SimilarityMetric = SimilarityMetric.TANIMOTO,
+    radius: int = 2,
+    n_bits: int = 2048,
+) -> list[list[float]]:
+    """
+    Compute pairwise similarity matrix.
+
+    Args:
+        mols: List of molecules
+        metric: Similarity metric
+        radius: Morgan fingerprint radius
+        n_bits: Fingerprint bit size
+
+    Returns:
+        Symmetric similarity matrix
+    """
+    # Generate fingerprints
+    fps = [get_morgan_fingerprint(mol, radius, n_bits) for mol in mols if mol is not None]
+    n = len(fps)
+
+    # Compute pairwise similarities
+    matrix = [[0.0] * n for _ in range(n)]
+
+    for i in range(n):
+        matrix[i][i] = 1.0
+        for j in range(i + 1, n):
+            sim = compute_similarity(fps[i], fps[j], metric)
+            matrix[i][j] = sim
+            matrix[j][i] = sim
+
+    return matrix
+
+
+def cluster_molecules(
+    mols: list[Chem.Mol],
+    cutoff: float = 0.3,
+    radius: int = 2,
+    n_bits: int = 2048,
+) -> list[list[int]]:
+    """
+    Cluster molecules using Butina algorithm.
+
+    Args:
+        mols: List of molecules
+        cutoff: Distance cutoff (1 - similarity)
+        radius: Morgan fingerprint radius
+        n_bits: Fingerprint bit size
+
+    Returns:
+        List of clusters (each cluster is a list of molecule indices)
+    """
+    # Generate fingerprints
+    fps = []
+    valid_indices = []
+    for i, mol in enumerate(mols):
+        if mol is not None:
+            fps.append(get_morgan_fingerprint(mol, radius, n_bits))
+            valid_indices.append(i)
+
+    n = len(fps)
+    if n == 0:
+        return []
+
+    # Compute distance matrix (lower triangle)
+    dists = []
+    for i in range(1, n):
+        sims = DataStructs.BulkTanimotoSimilarity(fps[i], fps[:i])
+        dists.extend([1 - s for s in sims])
+
+    # Cluster using Butina
+    clusters = Butina.ClusterData(dists, n, cutoff, isDistData=True)
+
+    # Map back to original indices
+    result = []
+    for cluster in clusters:
+        result.append([valid_indices[i] for i in cluster])
+
+    return result

rdkit_cli/core/standardizer.py

@@ -0,0 +1,141 @@
+"""Molecule standardization engine."""
+
+from typing import Optional, Any
+
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from rdkit.Chem.MolStandardize import rdMolStandardize
+
+from rdkit_cli.io.readers import MoleculeRecord
+
+
+class MoleculeStandardizer:
+    """Standardizer for molecular structures."""
+
+    def __init__(
+        self,
+        canonicalize: bool = True,
+        remove_stereo: bool = False,
+        disconnect_metals: bool = False,
+        normalize: bool = False,
+        reionize: bool = False,
+        uncharge: bool = False,
+        fragment_parent: bool = False,
+        tautomer_parent: bool = False,
+        include_original: bool = False,
+    ):
+        """
+        Initialize standardizer.
+
+        Args:
+            canonicalize: Canonicalize SMILES
+            remove_stereo: Remove stereochemistry information
+            disconnect_metals: Disconnect metal atoms
+            normalize: Apply normalization transforms
+            reionize: Standardize ionization state
+            uncharge: Neutralize charges
+            fragment_parent: Keep only largest fragment
+            tautomer_parent: Canonicalize tautomer
+            include_original: Include original SMILES in output
+        """
+        self.canonicalize = canonicalize
+        self.remove_stereo = remove_stereo
+        self.disconnect_metals = disconnect_metals
+        self.normalize = normalize
+        self.reionize = reionize
+        self.uncharge = uncharge
+        self.fragment_parent = fragment_parent
+        self.tautomer_parent = tautomer_parent
+        self.include_original = include_original
+
+        # Initialize standardizers
+        self._metal_disconnector = rdMolStandardize.MetalDisconnector() if disconnect_metals else None
+        self._normalizer = rdMolStandardize.Normalizer() if normalize else None
+        self._reionizer = rdMolStandardize.Reionizer() if reionize else None
+        self._uncharger = rdMolStandardize.Uncharger() if uncharge else None
+        self._fragment_chooser = rdMolStandardize.LargestFragmentChooser() if fragment_parent else None
+        self._tautomer_canon = rdMolStandardize.TautomerCanonicalizer() if tautomer_parent else None
+
+    def standardize(self, record: MoleculeRecord) -> Optional[dict[str, Any]]:
+        """
+        Standardize a molecule record.
+
+        Args:
+            record: MoleculeRecord to process
+
+        Returns:
+            Dictionary with standardized SMILES or None if failed
+        """
+        if record.mol is None:
+            return None
+
+        try:
+            mol = record.mol
+
+            # Apply transformations in order
+            if self._metal_disconnector:
+                mol = self._metal_disconnector.Disconnect(mol)
+
+            if self._normalizer:
+                mol = self._normalizer.normalize(mol)
+
+            if self._reionizer:
+                mol = self._reionizer.reionize(mol)
+
+            if self._uncharger:
+                mol = self._uncharger.uncharge(mol)
+
+            if self._fragment_chooser:
+                mol = self._fragment_chooser.choose(mol)
+
+            if self._tautomer_canon:
+                mol = self._tautomer_canon.canonicalize(mol)
+
+            if self.remove_stereo:
+                Chem.RemoveStereochemistry(mol)
+
+            # Generate output SMILES
+            if self.canonicalize:
+                output_smiles = Chem.MolToSmiles(mol, canonical=True)
+            else:
+                output_smiles = Chem.MolToSmiles(mol)
+
+            result: dict[str, Any] = {}
+
+            if self.include_original:
+                result["original_smiles"] = record.smiles
+
+            result["smiles"] = output_smiles
+
+            if record.name:
+                result["name"] = record.name
+
+            return result
+
+        except Exception:
+            return None
+
+    def get_column_names(self) -> list[str]:
+        """Get output column names in order."""
+        cols = []
+        if self.include_original:
+            cols.append("original_smiles")
+        cols.append("smiles")
+        cols.append("name")
+        return cols
+
+
+def canonicalize_smiles(smiles: str) -> Optional[str]:
+    """
+    Canonicalize a SMILES string.
+
+    Args:
+        smiles: Input SMILES
+
+    Returns:
+        Canonical SMILES or None if parsing failed
+    """
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return None
+    return Chem.MolToSmiles(mol, canonical=True)

rdkit_cli/io/__init__.py

@@ -0,0 +1,7 @@
+"""I/O handling for multiple file formats."""
+
+from rdkit_cli.io.formats import FileFormat, FormatConfig, detect_format
+from rdkit_cli.io.readers import create_reader
+from rdkit_cli.io.writers import create_writer
+
+__all__ = ["FileFormat", "FormatConfig", "detect_format", "create_reader", "create_writer"]

rdkit_cli/io/formats.py

@@ -0,0 +1,109 @@
+"""File format detection and configuration."""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Optional
+
+
+class FileFormat(Enum):
+    """Supported file formats."""
+
+    CSV = "csv"
+    TSV = "tsv"
+    SMI = "smi"
+    SDF = "sdf"
+    PARQUET = "parquet"
+
+
+@dataclass
+class FormatConfig:
+    """Configuration for file format handling."""
+
+    format: FileFormat
+    has_header: bool = True
+    smiles_column: str = "smiles"
+    name_column: Optional[str] = None
+    delimiter: str = ","
+    extra_columns: list[str] = field(default_factory=list)
+
+    def __post_init__(self):
+        """Set format-specific defaults."""
+        if self.format == FileFormat.TSV:
+            self.delimiter = "\t"
+        elif self.format == FileFormat.SMI:
+            self.has_header = False
+            self.delimiter = " "
+
+
+# File extension to format mapping
+EXTENSION_MAP: dict[str, FileFormat] = {
+    ".csv": FileFormat.CSV,
+    ".tsv": FileFormat.TSV,
+    ".smi": FileFormat.SMI,
+    ".smiles": FileFormat.SMI,
+    ".sdf": FileFormat.SDF,
+    ".mol": FileFormat.SDF,
+    ".parquet": FileFormat.PARQUET,
+    ".pq": FileFormat.PARQUET,
+}
+
+
+def detect_format(path: str | Path) -> FileFormat:
+    """
+    Detect file format from file extension.
+
+    Args:
+        path: Path to the file
+
+    Returns:
+        Detected FileFormat
+
+    Raises:
+        ValueError: If format cannot be detected
+    """
+    path = Path(path)
+    suffix = path.suffix.lower()
+
+    if suffix in EXTENSION_MAP:
+        return EXTENSION_MAP[suffix]
+
+    raise ValueError(
+        f"Cannot detect format for '{path}'. "
+        f"Supported extensions: {', '.join(EXTENSION_MAP.keys())}"
+    )
+
+
+def create_format_config(
+    path: str | Path,
+    format_override: Optional[FileFormat] = None,
+    has_header: Optional[bool] = None,
+    smiles_column: str = "smiles",
+    name_column: Optional[str] = None,
+) -> FormatConfig:
+    """
+    Create a FormatConfig for a file.
+
+    Args:
+        path: Path to the file
+        format_override: Override auto-detected format
+        has_header: Override default header setting
+        smiles_column: Name of the SMILES column
+        name_column: Name of the molecule name column
+
+    Returns:
+        Configured FormatConfig
+    """
+    file_format = format_override or detect_format(path)
+
+    config = FormatConfig(
+        format=file_format,
+        smiles_column=smiles_column,
+        name_column=name_column,
+    )
+
+    # Override header if explicitly specified
+    if has_header is not None:
+        config.has_header = has_header
+
+    return config