rdkit-cli 0.1.0__py3-none-any.whl

rdkit_cli/commands/standardize.py

@@ -0,0 +1,138 @@
+"""Standardize command implementation."""
+
+import sys
+from pathlib import Path
+
+from rdkit_cli.cli import RdkitHelpFormatter, add_common_io_options, add_common_processing_options
+
+
+def register_parser(subparsers):
+    """Register the standardize command and subcommands."""
+    parser = subparsers.add_parser(
+        "standardize",
+        help="Standardize and canonicalize molecules",
+        description="Apply standardization transforms to molecular structures.",
+        formatter_class=RdkitHelpFormatter,
+    )
+
+    add_common_io_options(parser)
+    add_common_processing_options(parser)
+
+    # Standardization options
+    parser.add_argument(
+        "--no-canonicalize",
+        action="store_true",
+        help="Don't canonicalize output SMILES",
+    )
+    parser.add_argument(
+        "--remove-stereo",
+        action="store_true",
+        help="Remove stereochemistry information",
+    )
+    parser.add_argument(
+        "--disconnect-metals",
+        action="store_true",
+        help="Disconnect metal atoms from molecules",
+    )
+    parser.add_argument(
+        "--normalize",
+        action="store_true",
+        help="Apply normalization transforms",
+    )
+    parser.add_argument(
+        "--reionize",
+        action="store_true",
+        help="Standardize ionization states",
+    )
+    parser.add_argument(
+        "--uncharge",
+        action="store_true",
+        help="Neutralize charges",
+    )
+    parser.add_argument(
+        "--fragment-parent",
+        action="store_true",
+        help="Keep only the largest fragment",
+    )
+    parser.add_argument(
+        "--tautomer-parent",
+        action="store_true",
+        help="Canonicalize tautomer form",
+    )
+    parser.add_argument(
+        "--cleanup",
+        action="store_true",
+        help="Apply standard cleanup (normalize + uncharge + fragment-parent)",
+    )
+    parser.add_argument(
+        "--include-original",
+        action="store_true",
+        help="Include original SMILES in output",
+    )
+
+    parser.set_defaults(func=run_standardize)
+
+
+def run_standardize(args) -> int:
+    """Run the standardize command."""
+    # Lazy imports
+    from rdkit_cli.core.standardizer import MoleculeStandardizer
+    from rdkit_cli.io import create_reader, create_writer
+    from rdkit_cli.parallel.batch import process_molecules
+
+    # Handle --cleanup shortcut
+    normalize = args.normalize or args.cleanup
+    uncharge = args.uncharge or args.cleanup
+    fragment_parent = args.fragment_parent or args.cleanup
+
+    # Create standardizer
+    standardizer = MoleculeStandardizer(
+        canonicalize=not args.no_canonicalize,
+        remove_stereo=args.remove_stereo,
+        disconnect_metals=args.disconnect_metals,
+        normalize=normalize,
+        reionize=args.reionize,
+        uncharge=uncharge,
+        fragment_parent=fragment_parent,
+        tautomer_parent=args.tautomer_parent,
+        include_original=args.include_original,
+    )
+
+    # Create reader
+    input_path = Path(args.input)
+    if not input_path.exists():
+        print(f"Error: Input file not found: {input_path}", file=sys.stderr)
+        return 1
+
+    reader = create_reader(
+        input_path,
+        smiles_column=args.smiles_column,
+        name_column=args.name_column,
+        has_header=not args.no_header,
+    )
+
+    # Create writer
+    output_path = Path(args.output)
+    writer = create_writer(
+        output_path,
+        columns=standardizer.get_column_names(),
+    )
+
+    # Process
+    with reader, writer:
+        result = process_molecules(
+            reader=reader,
+            writer=writer,
+            processor=standardizer.standardize,
+            n_workers=args.ncpu,
+            quiet=args.quiet,
+        )
+
+    if not args.quiet:
+        print(
+            f"Processed {result.successful}/{result.total_processed} molecules "
+            f"({result.failed} failed) in {result.elapsed_time:.1f}s",
+            file=sys.stderr,
+        )
+
+    return 0 if result.failed == 0 else 1

rdkit_cli/core/__init__.py

@@ -0,0 +1 @@
+"""Core processing logic for rdkit-cli."""

rdkit_cli/core/conformers.py

@@ -0,0 +1,197 @@
+"""Conformer generation engine."""
+
+from typing import Optional, Any
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, rdDistGeom
+
+from rdkit_cli.io.readers import MoleculeRecord
+
+
+class ConformerGenerator:
+    """Generate 3D conformers for molecules."""
+
+    def __init__(
+        self,
+        num_conformers: int = 10,
+        method: str = "etkdgv3",
+        optimize: bool = True,
+        force_field: str = "mmff",
+        max_iterations: int = 200,
+        random_seed: int = 42,
+    ):
+        """
+        Initialize conformer generator.
+
+        Args:
+            num_conformers: Number of conformers to generate
+            method: Embedding method (etkdgv3, etkdgv2, etdg)
+            optimize: Whether to optimize conformers
+            force_field: Force field for optimization (mmff, uff)
+            max_iterations: Maximum optimization iterations
+            random_seed: Random seed for reproducibility
+        """
+        self.num_conformers = num_conformers
+        self.method = method.lower()
+        self.optimize = optimize
+        self.force_field = force_field.lower()
+        self.max_iterations = max_iterations
+        self.random_seed = random_seed
+
+        # Set up embedding parameters
+        if self.method == "etkdgv3":
+            self.params = rdDistGeom.ETKDGv3()
+        elif self.method == "etkdgv2":
+            self.params = rdDistGeom.ETKDGv2()
+        elif self.method == "etdg":
+            self.params = rdDistGeom.ETDG()
+        else:
+            raise ValueError(f"Unknown method: {method}")
+
+        self.params.randomSeed = random_seed
+        self.params.numThreads = 0  # Use all available threads
+
+    def generate(self, record: MoleculeRecord) -> Optional[dict[str, Any]]:
+        """
+        Generate conformers for a molecule.
+
+        Args:
+            record: MoleculeRecord to process
+
+        Returns:
+            Dictionary with molecule and conformer info, or None if failed
+        """
+        if record.mol is None:
+            return None
+
+        try:
+            # Add hydrogens
+            mol = Chem.AddHs(record.mol)
+
+            # Embed conformers
+            conf_ids = AllChem.EmbedMultipleConfs(
+                mol,
+                numConfs=self.num_conformers,
+                params=self.params,
+            )
+
+            if len(conf_ids) == 0:
+                return None
+
+            # Optimize if requested
+            energies = []
+            if self.optimize:
+                if self.force_field == "mmff":
+                    results = AllChem.MMFFOptimizeMoleculeConfs(
+                        mol,
+                        maxIters=self.max_iterations,
+                        numThreads=0,
+                    )
+                    energies = [r[1] for r in results]
+                elif self.force_field == "uff":
+                    results = AllChem.UFFOptimizeMoleculeConfs(
+                        mol,
+                        maxIters=self.max_iterations,
+                        numThreads=0,
+                    )
+                    energies = [r[1] for r in results]
+
+            # Get lowest energy conformer
+            if energies:
+                best_conf = min(range(len(energies)), key=lambda i: energies[i])
+                best_energy = energies[best_conf]
+            else:
+                best_conf = 0
+                best_energy = None
+
+            result: dict[str, Any] = {
+                "smiles": record.smiles,
+                "mol": mol,
+                "num_conformers": len(conf_ids),
+                "best_conformer": best_conf,
+            }
+
+            if best_energy is not None:
+                result["energy"] = round(best_energy, 2)
+
+            if record.name:
+                result["name"] = record.name
+
+            return result
+
+        except Exception:
+            return None
+
+
+class ConformerOptimizer:
+    """Optimize existing 3D structures."""
+
+    def __init__(
+        self,
+        force_field: str = "mmff",
+        max_iterations: int = 200,
+    ):
+        """
+        Initialize conformer optimizer.
+
+        Args:
+            force_field: Force field (mmff, uff)
+            max_iterations: Maximum iterations
+        """
+        self.force_field = force_field.lower()
+        self.max_iterations = max_iterations
+
+    def optimize(self, record: MoleculeRecord) -> Optional[dict[str, Any]]:
+        """
+        Optimize a molecule's 3D structure.
+
+        Args:
+            record: MoleculeRecord with 3D coordinates
+
+        Returns:
+            Dictionary with optimized molecule, or None if failed
+        """
+        if record.mol is None:
+            return None
+
+        try:
+            mol = Chem.Mol(record.mol)
+
+            # Check if molecule has 3D coordinates
+            if mol.GetNumConformers() == 0:
+                # Try to generate 3D structure
+                mol = Chem.AddHs(mol)
+                AllChem.EmbedMolecule(mol, rdDistGeom.ETKDGv3())
+
+            if mol.GetNumConformers() == 0:
+                return None
+
+            # Optimize
+            if self.force_field == "mmff":
+                result = AllChem.MMFFOptimizeMolecule(mol, maxIters=self.max_iterations)
+                props = AllChem.MMFFGetMoleculeProperties(mol)
+                if props:
+                    ff = AllChem.MMFFGetMoleculeForceField(mol, props)
+                    energy = ff.CalcEnergy() if ff else None
+                else:
+                    energy = None
+            else:
+                result = AllChem.UFFOptimizeMolecule(mol, maxIters=self.max_iterations)
+                ff = AllChem.UFFGetMoleculeForceField(mol)
+                energy = ff.CalcEnergy() if ff else None
+
+            output: dict[str, Any] = {
+                "smiles": Chem.MolToSmiles(Chem.RemoveHs(mol)),
+                "mol": mol,
+            }
+
+            if energy is not None:
+                output["energy"] = round(energy, 2)
+
+            if record.name:
+                output["name"] = record.name
+
+            return output
+
+        except Exception:
+            return None

rdkit_cli/core/depict.py

@@ -0,0 +1,241 @@
+"""Molecular depiction/visualization engine."""
+
+from typing import Optional, Any
+from pathlib import Path
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Draw, rdDepictor
+from rdkit.Chem.Draw import rdMolDraw2D
+
+from rdkit_cli.io.readers import MoleculeRecord
+
+
+class MoleculeDepiction:
+    """Generate 2D depictions of molecules."""
+
+    def __init__(
+        self,
+        width: int = 300,
+        height: int = 300,
+        image_format: str = "svg",
+        add_atom_indices: bool = False,
+        add_stereo_annotation: bool = False,
+        highlight_atoms: Optional[list[int]] = None,
+        highlight_bonds: Optional[list[int]] = None,
+        use_kekulize: bool = True,
+        wedge_bonds: bool = True,
+        add_chiral_hs: bool = True,
+    ):
+        """
+        Initialize molecule depiction.
+
+        Args:
+            width: Image width in pixels
+            height: Image height in pixels
+            image_format: Output format ('svg' or 'png')
+            add_atom_indices: Add atom index labels
+            add_stereo_annotation: Add stereo annotations
+            highlight_atoms: Atom indices to highlight
+            highlight_bonds: Bond indices to highlight
+            use_kekulize: Use Kekule form for drawing
+            wedge_bonds: Draw wedged bonds
+            add_chiral_hs: Add chiral Hs
+        """
+        self.width = width
+        self.height = height
+        self.image_format = image_format.lower()
+        self.add_atom_indices = add_atom_indices
+        self.add_stereo_annotation = add_stereo_annotation
+        self.highlight_atoms = highlight_atoms or []
+        self.highlight_bonds = highlight_bonds or []
+        self.use_kekulize = use_kekulize
+        self.wedge_bonds = wedge_bonds
+        self.add_chiral_hs = add_chiral_hs
+
+    def depict(self, mol: Chem.Mol) -> Optional[str]:
+        """
+        Generate depiction of a molecule.
+
+        Args:
+            mol: RDKit molecule
+
+        Returns:
+            SVG or PNG data as string/bytes
+        """
+        if mol is None:
+            return None
+
+        try:
+            # Prepare molecule
+            mol = Chem.Mol(mol)  # Copy
+            if self.add_chiral_hs:
+                mol = Chem.AddHs(mol)
+                AllChem.EmbedMolecule(mol, AllChem.ETKDGv3())
+                mol = Chem.RemoveHs(mol)
+
+            # Generate 2D coords
+            rdDepictor.Compute2DCoords(mol)
+
+            # Create drawer
+            if self.image_format == "svg":
+                drawer = rdMolDraw2D.MolDraw2DSVG(self.width, self.height)
+            else:
+                drawer = rdMolDraw2D.MolDraw2DCairo(self.width, self.height)
+
+            # Configure options
+            opts = drawer.drawOptions()
+            opts.addAtomIndices = self.add_atom_indices
+            opts.addStereoAnnotation = self.add_stereo_annotation
+
+            # Draw
+            if self.highlight_atoms or self.highlight_bonds:
+                drawer.DrawMolecule(
+                    mol,
+                    highlightAtoms=self.highlight_atoms,
+                    highlightBonds=self.highlight_bonds,
+                )
+            else:
+                drawer.DrawMolecule(mol)
+
+            drawer.FinishDrawing()
+
+            return drawer.GetDrawingText()
+
+        except Exception:
+            return None
+
+    def depict_record(self, record: MoleculeRecord) -> Optional[dict[str, Any]]:
+        """
+        Generate depiction of a molecule record.
+
+        Args:
+            record: MoleculeRecord to process
+
+        Returns:
+            Dictionary with image data or None
+        """
+        if record.mol is None:
+            return None
+
+        image_data = self.depict(record.mol)
+        if image_data is None:
+            return None
+
+        result: dict[str, Any] = {
+            "smiles": record.smiles,
+            "image": image_data,
+        }
+
+        if record.name:
+            result["name"] = record.name
+
+        return result
+
+
+class GridDepiction:
+    """Generate grid of molecule depictions."""
+
+    def __init__(
+        self,
+        mols_per_row: int = 4,
+        mol_width: int = 200,
+        mol_height: int = 200,
+        legends: Optional[list[str]] = None,
+        use_svg: bool = True,
+    ):
+        """
+        Initialize grid depiction.
+
+        Args:
+            mols_per_row: Molecules per row
+            mol_width: Width per molecule
+            mol_height: Height per molecule
+            legends: List of labels for molecules
+            use_svg: Output SVG instead of PNG
+        """
+        self.mols_per_row = mols_per_row
+        self.mol_width = mol_width
+        self.mol_height = mol_height
+        self.legends = legends
+        self.use_svg = use_svg
+
+    def depict(self, mols: list[Chem.Mol]) -> Optional[str]:
+        """
+        Generate grid depiction.
+
+        Args:
+            mols: List of molecules
+
+        Returns:
+            SVG or PNG data
+        """
+        if not mols:
+            return None
+
+        try:
+            # Prepare molecules
+            prepared_mols = []
+            for mol in mols:
+                if mol is not None:
+                    mol = Chem.Mol(mol)
+                    rdDepictor.Compute2DCoords(mol)
+                    prepared_mols.append(mol)
+                else:
+                    prepared_mols.append(None)
+
+            legends = self.legends or [""] * len(prepared_mols)
+
+            if self.use_svg:
+                return Draw.MolsToGridImage(
+                    prepared_mols,
+                    molsPerRow=self.mols_per_row,
+                    subImgSize=(self.mol_width, self.mol_height),
+                    legends=legends[:len(prepared_mols)],
+                    useSVG=True,
+                )
+            else:
+                img = Draw.MolsToGridImage(
+                    prepared_mols,
+                    molsPerRow=self.mols_per_row,
+                    subImgSize=(self.mol_width, self.mol_height),
+                    legends=legends[:len(prepared_mols)],
+                )
+                # Convert to bytes
+                import io
+                buf = io.BytesIO()
+                img.save(buf, format="PNG")
+                return buf.getvalue()
+
+        except Exception:
+            return None
+
+
+def depict_smiles(
+    smiles: str,
+    width: int = 300,
+    height: int = 300,
+    image_format: str = "svg",
+) -> Optional[str]:
+    """
+    Convenience function to depict a SMILES string.
+
+    Args:
+        smiles: SMILES string
+        width: Image width
+        height: Image height
+        image_format: Output format
+
+    Returns:
+        Image data or None
+    """
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return None
+
+    depictor = MoleculeDepiction(
+        width=width,
+        height=height,
+        image_format=image_format,
+    )
+
+    return depictor.depict(mol)