workbench 0.8.171__py3-none-any.whl → 0.8.173__py3-none-any.whl

workbench/utils/chem_utils/mol_standardize.py

@@ -0,0 +1,428 @@
+"""
+mol_standardize.py - Molecular Standardization for ADMET Preprocessing
+Following ChEMBL structure standardization pipeline
+
+Purpose:
+    Standardizes chemical structures to ensure consistent molecular representations
+    for ADMET modeling. Handles tautomers, salts, charges, and structural variations
+    that can cause the same compound to be represented differently.
+
+Standardization Pipeline:
+    1. Cleanup
+       - Removes explicit hydrogens
+       - Disconnects metal atoms from organic fragments
+       - Normalizes functional groups (e.g., nitro, sulfoxide representations)
+
+    2. Fragment Parent Selection (optional, controlled by extract_salts parameter)
+       - Identifies and keeps the largest organic fragment
+       - Removes salts, solvents, and counterions
+       - Example: [Na+].CC(=O)[O-] → CC(=O)O (keeps acetate, removes sodium)
+
+    3. Charge Neutralization (optional, controlled by extract_salts parameter)
+       - Neutralizes charges where possible
+       - Only applied when extract_salts=True (following ChEMBL pipeline)
+       - Skipped when extract_salts=False to preserve ionic character
+       - Example: CC(=O)[O-] → CC(=O)O
+
+    4. Tautomer Canonicalization (optional, default=True)
+       - Generates canonical tautomer form for consistency
+       - Example: Oc1ccccn1 → O=c1cccc[nH]1 (2-hydroxypyridine → 2-pyridone)
+
+Output DataFrame Columns:
+    - orig_smiles: Original input SMILES (preserved for traceability)
+    - smiles: Standardized molecule (with or without salts based on extract_salts)
+    - salt: Removed salt/counterion as SMILES (only populated if extract_salts=True)
+
+Salt Handling:
+    Salt forms can dramatically affect properties like solubility.
+    This module offers two modes for handling salts:
+
+    When extract_salts=True (default, ChEMBL standard):
+        - Removes salts/counterions to get parent molecule
+        - Neutralizes charges on the parent
+        - Records removed salts in 'salt' column
+        Input: [Na+].CC(=O)[O-]  →  Parent: CC(=O)O, Salt: [Na+]
+
+    When extract_salts=False (preserve full salt form):
+        - Keeps all fragments including salts/counterions
+        - Preserves ionic charges (no neutralization)
+
+Mixture Detection:
+    The module detects and logs potential mixtures (vs true salt forms):
+    - Multiple large neutral organic fragments indicate a mixture
+    - Mixtures are logged but NOT recorded in the salt column
+    - True salts (small/charged fragments) are properly extracted
+
+    Downstream modeling options:
+    1. Use parent only (standard approach for most ADMET properties)
+    2. Include salt as a categorical or computed feature
+    3. Model parent + salt effects hierarchically
+    4. Use full salt form for properties like solubility/formulation
+
+References:
+    - "ChEMBL Structure Pipeline" (Bento et al., 2020)
+      https://doi.org/10.1186/s13321-020-00456-1
+    - "Standardization and Validation with the RDKit" (Greg Landrum, RSC Open Science 2021)
+      https://github.com/greglandrum/RSC_OpenScience_Standardization_202104/blob/main/Standardization%20and%20Validation%20with%20the%20RDKit.ipynb
+
+Usage:
+    from mol_standardize import standardize
+
+    # Basic usage (removes salts by default, ChEMBL standard)
+    df_std = standardize(df, smiles_column='smiles')
+
+    # Keep salts in the molecule (preserve ionic forms)
+    df_std = standardize(df, extract_salts=False)
+
+    # Without tautomer canonicalization (faster, less aggressive)
+    df_std = standardize(df, canonicalize_tautomer=False)
+"""
+
+import logging
+from typing import Optional, Tuple
+import pandas as pd
+from rdkit import Chem
+from rdkit.Chem import Mol
+from rdkit.Chem.MolStandardize import rdMolStandardize
+from rdkit import RDLogger
+
+log = logging.getLogger("workbench")
+RDLogger.DisableLog("rdApp.warning")
+
+
+class MolStandardizer:
+    """
+    Streamlined molecular standardizer for ADMET preprocessing
+    Uses ChEMBL standardization pipeline with RDKit
+    """
+
+    def __init__(self, canonicalize_tautomer: bool = True, remove_salts: bool = True):
+        """
+        Initialize standardizer with ChEMBL defaults
+
+        Args:
+            canonicalize_tautomer: Whether to canonicalize tautomers (default True)
+            remove_salts: Whether to remove salts/counterions (default True)
+        """
+        self.canonicalize_tautomer = canonicalize_tautomer
+        self.remove_salts = remove_salts
+        self.params = rdMolStandardize.CleanupParameters()
+        self.tautomer_enumerator = rdMolStandardize.TautomerEnumerator(self.params)
+
+    def standardize(self, mol: Mol) -> Tuple[Optional[Mol], Optional[str]]:
+        """
+        Main standardization pipeline for ADMET
+
+        Pipeline:
+        1. Cleanup (remove Hs, disconnect metals, normalize)
+        2. Get largest fragment (optional - only if remove_salts=True)
+        3. Neutralize charges
+        4. Canonicalize tautomer (optional)
+
+        Args:
+            mol: RDKit molecule object
+
+        Returns:
+            Tuple of (standardized molecule or None if failed, salt SMILES or None)
+        """
+        if mol is None:
+            return None, None
+
+        try:
+            # Step 1: Cleanup
+            mol = rdMolStandardize.Cleanup(mol, self.params)
+            if mol is None:
+                return None, None
+
+            salt_smiles = None
+
+            # Step 2: Fragment handling (conditional based on remove_salts)
+            if self.remove_salts:
+                # Get parent molecule and extract salt information
+                parent_mol = rdMolStandardize.FragmentParent(mol, self.params)
+                if parent_mol:
+                    salt_smiles = self._extract_salt(mol, parent_mol)
+                    mol = parent_mol
+                else:
+                    return None, None
+            # If not removing salts, keep the full molecule intact
+
+            # Step 3: Neutralize charges (skip if keeping salts to preserve ionic forms)
+            if self.remove_salts:
+                mol = rdMolStandardize.ChargeParent(mol, self.params, skipStandardize=True)
+                if mol is None:
+                    return None, salt_smiles
+
+            # Step 4: Canonicalize tautomer
+            if self.canonicalize_tautomer:
+                mol = self.tautomer_enumerator.Canonicalize(mol)
+
+            return mol, salt_smiles
+
+        except Exception as e:
+            log.warning(f"Standardization failed: {e}")
+            return None, None
+
+    def _extract_salt(self, orig_mol: Mol, parent_mol: Mol) -> Optional[str]:
+        """
+        Extract salt/counterion by comparing original and parent molecules.
+
+        Detects and handles mixtures vs true salt forms:
+        - True salts: small (<= 6 heavy atoms) or charged fragments
+        - Mixtures: multiple large neutral organic fragments
+
+        Args:
+            orig_mol: Original molecule (before FragmentParent)
+            parent_mol: Parent molecule (after FragmentParent)
+
+        Returns:
+            SMILES string of salt components or None if no salts/mixture detected
+        """
+        try:
+            # Get all fragments from original molecule
+            orig_frags = Chem.GetMolFrags(orig_mol, asMols=True)
+
+            # If only one fragment, no salt
+            if len(orig_frags) <= 1:
+                return None
+
+            # Get canonical SMILES of parent for comparison
+            parent_smiles = Chem.MolToSmiles(parent_mol, canonical=True)
+
+            # Separate fragments into salts vs potential mixture components
+            salt_frags = []
+            mixture_frags = []
+
+            for frag in orig_frags:
+                frag_smiles = Chem.MolToSmiles(frag, canonical=True)
+
+                # Skip the parent fragment
+                if frag_smiles == parent_smiles:
+                    continue
+
+                # Classify fragment as salt or mixture component
+                num_heavy = frag.GetNumHeavyAtoms()
+                has_charge = any(atom.GetFormalCharge() != 0 for atom in frag.GetAtoms())
+
+                # More nuanced classification
+                if has_charge and num_heavy <= 10:  # Small charged fragment - likely a salt
+                    salt_frags.append(frag_smiles)
+                elif not has_charge and num_heavy <= 6:  # Small neutral - could be solvent/salt
+                    salt_frags.append(frag_smiles)
+                else:
+                    # Large neutral fragment - likely part of a mixture
+                    mixture_frags.append(frag_smiles)
+
+            # Check if this looks like a mixture
+            if mixture_frags:
+                # Log mixture detection
+                total_frags = len(orig_frags)
+                log.warning(
+                    f"Mixture detected: {total_frags} total fragments, "
+                    f"{len(mixture_frags)} large neutral organics. "
+                    f"Removing: {'.'.join(mixture_frags + salt_frags)}"
+                )
+                # Return None for mixtures - don't pollute the salt column
+                return None
+
+            # Return actual salts only
+            return ".".join(salt_frags) if salt_frags else None
+
+        except Exception as e:
+            log.info(f"Salt extraction failed: {e}")
+            return None
+
+
+def standardize(
+    df: pd.DataFrame,
+    canonicalize_tautomer: bool = True,
+    extract_salts: bool = True,
+) -> pd.DataFrame:
+    """
+    Standardize molecules in a DataFrame for ADMET modeling
+
+    Args:
+        df: Input DataFrame with SMILES column
+        canonicalize_tautomer: Whether to canonicalize tautomers (default: True)
+        extract_salts: Whether to remove and extract salts (default: True)
+                      If False, keeps full molecule with salts/counterions intact,
+                      skipping charge neutralization to preserve ionic character
+
+    Returns:
+        DataFrame with:
+        - orig_smiles: Original SMILES (preserved)
+        - smiles: Standardized SMILES (working column for downstream)
+        - salt: Removed salt/counterion SMILES (only if extract_salts=True)
+                None for mixtures or when no true salts present
+    """
+
+    # Check for the smiles column (any capitalization)
+    smiles_column = next((col for col in df.columns if col.lower() == "smiles"), None)
+    if smiles_column is None:
+        raise ValueError("Input DataFrame must have a 'smiles' column")
+
+    # Copy input DataFrame to avoid modifying original
+    result = df.copy()
+
+    # Preserve original SMILES if not already saved
+    if "orig_smiles" not in result.columns:
+        result["orig_smiles"] = result[smiles_column]
+
+    # Initialize standardizer with salt removal control
+    standardizer = MolStandardizer(canonicalize_tautomer=canonicalize_tautomer, remove_salts=extract_salts)
+
+    def process_smiles(smiles: str) -> pd.Series:
+        """
+        Process a single SMILES string through standardization pipeline
+
+        Args:
+            smiles: Input SMILES string
+
+        Returns:
+            Series with standardized SMILES and extracted salt (if applicable)
+        """
+        # Handle missing values
+        if pd.isna(smiles) or smiles == "":
+            log.error("Encountered missing or empty SMILES string")
+            return pd.Series({"smiles": None, "salt": None})
+
+        # Parse molecule
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            log.error(f"Invalid SMILES: {smiles}")
+            return pd.Series({"smiles": None, "salt": None})
+
+        # Full standardization with optional salt removal
+        std_mol, salt_smiles = standardizer.standardize(mol)
+
+        # After standardization, validate the result
+        if std_mol is not None:
+            # Check if molecule is reasonable
+            if std_mol.GetNumAtoms() == 0 or std_mol.GetNumAtoms() > 200:  # Arbitrary limits
+                log.error(f"Unusual molecule size: {std_mol.GetNumAtoms()} atoms")
+
+        if std_mol is None:
+            return pd.Series(
+                {
+                    "smiles": None,
+                    "salt": salt_smiles,  # May have extracted salt even if full standardization failed
+                }
+            )
+
+        # Convert back to SMILES
+        return pd.Series(
+            {"smiles": Chem.MolToSmiles(std_mol, canonical=True), "salt": salt_smiles if extract_salts else None}
+        )
+
+    # Process molecules
+    processed = result[smiles_column].apply(process_smiles)
+
+    # Update the dataframe with processed results
+    for col in ["smiles", "salt"]:
+        result[col] = processed[col]
+
+    return result
+
+
+if __name__ == "__main__":
+    import time
+    from workbench.api import DataSource
+
+    # Test with DataFrame including various salt forms
+    test_data = pd.DataFrame(
+        {
+            "smiles": [
+                # Organic salts
+                "[Na+].CC(=O)[O-]",  # Sodium acetate
+                "CC(=O)O.CCN",  # Acetic acid + ethylamine (acid-base pair)
+                # Tautomers
+                "CC(=O)CC(C)=O",  # Acetylacetone - tautomer
+                "c1ccc(O)nc1",  # 2-hydroxypyridine/2-pyridone - tautomer
+                # Multi-fragment
+                "CCO.CC",  # Ethanol + methane mixture
+                # Simple organics
+                "CC(C)(C)c1ccccc1",  # tert-butylbenzene
+                # Carbonate salts
+                "[Na+].[Na+].[O-]C([O-])=O",  # Sodium carbonate
+                "[Li+].[Li+].[O-]C([O-])=O",  # Lithium carbonate
+                "[K+].[K+].[O-]C([O-])=O",  # Potassium carbonate
+                "[Mg++].[O-]C([O-])=O",  # Magnesium carbonate
+                "[Ca++].[O-]C([O-])=O",  # Calcium carbonate
+                # Drug salts
+                "CC(C)NCC(O)c1ccc(O)c(O)c1.Cl",  # Isoproterenol HCl
+                "CN1CCC[C@H]1c2cccnc2.[Cl-]",  # Nicotine HCl
+                # Tautomer with salt
+                "c1ccc(O)nc1.Cl",  # 2-hydroxypyridine with HCl
+                # Edge cases
+                None,  # Missing value
+                "INVALID",  # Invalid SMILES
+            ],
+            "compound_id": [f"C{i:03d}" for i in range(1, 17)],
+        }
+    )
+
+    # General test
+    standardize(test_data)
+
+    # Remove the last two rows to avoid errors with None and INVALID
+    test_data = test_data.iloc[:-2].reset_index(drop=True)
+
+    # Test WITHOUT salt removal (keeps full molecule)
+    print("\nStandardization KEEPING salts (extract_salts=False):")
+    print("This preserves the full molecule including counterions")
+    result_keep = standardize(test_data, extract_salts=False, canonicalize_tautomer=True)
+    display_cols = ["compound_id", "orig_smiles", "smiles", "salt"]
+    print(result_keep[display_cols].to_string())
+
+    # Test WITH salt removal
+    print("\n" + "=" * 70)
+    print("Standardization REMOVING salts (extract_salts=True):")
+    print("This extracts parent molecule and records salt information")
+    result_remove = standardize(test_data, extract_salts=True, canonicalize_tautomer=True)
+    print(result_remove[display_cols].to_string())
+
+    # Test WITHOUT tautomerization (keeping salts)
+    print("\n" + "=" * 70)
+    print("Standardization KEEPING salts, NO tautomerization:")
+    result_no_taut = standardize(test_data, extract_salts=False, canonicalize_tautomer=False)
+    print(result_no_taut[display_cols].to_string())
+
+    # Show the difference for salt-containing molecules
+    print("\n" + "=" * 70)
+    print("Comparison showing differences:")
+    for idx, row in result_keep.iterrows():
+        keep_smiles = row["smiles"]
+        remove_smiles = result_remove.loc[idx, "smiles"]
+        no_taut_smiles = result_no_taut.loc[idx, "smiles"]
+        salt = result_remove.loc[idx, "salt"]
+
+        # Show differences when they exist
+        if keep_smiles != remove_smiles or keep_smiles != no_taut_smiles:
+            print(f"\n{row['compound_id']} ({row['orig_smiles']}):")
+            if keep_smiles != no_taut_smiles:
+                print(f"  With salt + taut:    {keep_smiles}")
+                print(f"  With salt, no taut:  {no_taut_smiles}")
+            if keep_smiles != remove_smiles:
+                print(f"  Parent only + taut:  {remove_smiles}")
+            if salt:
+                print(f"  Extracted salt:      {salt}")
+
+    # Summary statistics
+    print("\n" + "=" * 70)
+    print("Summary:")
+    print(f"Total molecules: {len(result_remove)}")
+    print(f"Molecules with salts: {result_remove['salt'].notna().sum()}")
+    unique_salts = result_remove["salt"].dropna().unique()
+    print(f"Unique salts found: {unique_salts[:5].tolist()}")
+
+    # Get a real dataset from Workbench and time the standardization
+    ds = DataSource("aqsol_data")
+    df = ds.pull_dataframe()[["id", "smiles"]]
+    start_time = time.time()
+    std_df = standardize(df, extract_salts=True, canonicalize_tautomer=True)
+    end_time = time.time()
+    print(f"\nStandardized {len(std_df)} molecules from Workbench in {end_time - start_time:.2f} seconds")
+    print(std_df.head())
+    print(f"Molecules with salts: {std_df['salt'].notna().sum()}")
+    unique_salts = std_df["salt"].dropna().unique()
+    print(f"Unique salts found: {unique_salts[:5].tolist()}")