workbench 0.8.162__py3-none-any.whl → 0.8.220__py3-none-any.whl

workbench/utils/chem_utils/salts.py

@@ -0,0 +1,256 @@
+"""Salt feature extraction utilities for molecular analysis"""
+
+import logging
+import pandas as pd
+from typing import List, Dict, Tuple, Optional, Union
+
+# Molecular Descriptor Imports
+from rdkit import Chem
+from rdkit.Chem import Descriptors
+
+# Set up the logger
+log = logging.getLogger("workbench")
+
+
+def _get_salt_feature_columns() -> List[str]:
+    """Internal: Return list of all salt feature column names"""
+    return [
+        "has_salt",
+        "mw_ratio",
+        "salt_to_api_ratio",
+        "has_metal_salt",
+        "has_halide",
+        "ionic_strength_proxy",
+        "has_organic_salt",
+    ]
+
+
+def _classify_salt_types(salt_frags: List[Chem.Mol]) -> Dict[str, int]:
+    """Internal: Classify salt fragments into categories"""
+    features = {
+        "has_organic_salt": 0,
+        "has_metal_salt": 0,
+        "has_halide": 0,
+    }
+
+    for frag in salt_frags:
+        # Get atoms
+        atoms = [atom.GetSymbol() for atom in frag.GetAtoms()]
+
+        # Metal detection
+        metals = ["Na", "K", "Ca", "Mg", "Li", "Zn", "Fe", "Al"]
+        if any(metal in atoms for metal in metals):
+            features["has_metal_salt"] = 1
+
+        # Halide detection
+        halides = ["Cl", "Br", "I", "F"]
+        if any(halide in atoms for halide in halides):
+            features["has_halide"] = 1
+
+        # Organic vs inorganic (simple heuristic: contains C)
+        if "C" in atoms:
+            features["has_organic_salt"] = 1
+
+    return features
+
+
+def extract_advanced_salt_features(
+    mol: Optional[Chem.Mol],
+) -> Tuple[Optional[Dict[str, Union[int, float]]], Optional[Chem.Mol]]:
+    """Extract comprehensive salt-related features from RDKit molecule"""
+    if mol is None:
+        return None, None
+
+    # Get fragments
+    fragments = Chem.GetMolFrags(mol, asMols=True)
+
+    # Identify API (largest organic fragment) vs salt fragments
+    fragment_weights = [(frag, Descriptors.MolWt(frag)) for frag in fragments]
+    fragment_weights.sort(key=lambda x: x[1], reverse=True)
+
+    # Find largest organic fragment as API
+    api_mol = None
+    salt_frags = []
+
+    for frag, mw in fragment_weights:
+        atoms = [atom.GetSymbol() for atom in frag.GetAtoms()]
+        if "C" in atoms and api_mol is None:  # First organic fragment = API
+            api_mol = frag
+        else:
+            salt_frags.append(frag)
+
+    # Fallback: if no organic fragments, use largest
+    if api_mol is None:
+        api_mol = fragment_weights[0][0]
+        salt_frags = [frag for frag, _ in fragment_weights[1:]]
+
+    # Initialize all features with default values
+    features = {col: 0 for col in _get_salt_feature_columns()}
+    features["mw_ratio"] = 1.0  # default for no salt
+
+    # Basic features
+    features.update(
+        {
+            "has_salt": int(len(salt_frags) > 0),
+            "mw_ratio": Descriptors.MolWt(api_mol) / Descriptors.MolWt(mol),
+        }
+    )
+
+    if salt_frags:
+        # Salt characterization
+        total_salt_mw = sum(Descriptors.MolWt(frag) for frag in salt_frags)
+        features.update(
+            {
+                "salt_to_api_ratio": total_salt_mw / Descriptors.MolWt(api_mol),
+                "ionic_strength_proxy": sum(abs(Chem.GetFormalCharge(frag)) for frag in salt_frags),
+            }
+        )
+
+        # Salt type classification
+        features.update(_classify_salt_types(salt_frags))
+
+    return features, api_mol
+
+
+def add_salt_features(df: pd.DataFrame) -> pd.DataFrame:
+    """Add salt features to dataframe with 'molecule' column containing RDKit molecules"""
+    salt_features_list = []
+
+    for idx, row in df.iterrows():
+        mol = row["molecule"]
+        features, clean_mol = extract_advanced_salt_features(mol)
+
+        if features is None:
+            # Handle invalid molecules
+            features = {col: None for col in _get_salt_feature_columns()}
+
+        salt_features_list.append(features)
+
+    # Convert to DataFrame and concatenate
+    salt_df = pd.DataFrame(salt_features_list)
+    return pd.concat([df, salt_df], axis=1)
+
+
+if __name__ == "__main__":
+    print("Running salt feature extraction tests...")
+
+    # Test molecules with various salt forms
+    test_molecules = {
+        "aspirin": "CC(=O)OC1=CC=CC=C1C(=O)O",  # No salt
+        "sodium_acetate": "CC(=O)[O-].[Na+]",  # Simple sodium salt
+        "calcium_acetate": "CC(=O)[O-].[Ca+2].CC(=O)[O-]",  # Calcium salt (divalent)
+        "ammonium_chloride": "[NH4+].[Cl-]",  # Inorganic salt
+        "methylamine_hcl": "CN.Cl",  # Organic salt
+        "benzene": "c1ccccc1",  # No salt
+        "aspirin_sodium": "CC(=O)OC1=CC=CC=C1C(=O)[O-].[Na+]",  # API with sodium salt
+        "complex_salt": "CC(C)(C)c1ccc(O)cc1.Cl.Cl.[Zn+2]",  # Complex with metal
+    }
+
+    # Test 1: Basic salt feature extraction
+    print("\n1. Testing basic salt feature extraction...")
+
+    salt_test_df = pd.DataFrame({"smiles": list(test_molecules.values()), "name": list(test_molecules.keys())})
+
+    salt_test_df["molecule"] = salt_test_df["smiles"].apply(Chem.MolFromSmiles)
+    salt_result = add_salt_features(salt_test_df)
+
+    print("   Salt feature results:")
+    print(f"   {'Molecule':20} has_salt  metal  halide  organic  mw_ratio")
+    print("   " + "-" * 65)
+    for _, row in salt_result.iterrows():
+        print(
+            f"   {row['name']:20} {row['has_salt']:^8} {row['has_metal_salt']:^6} "
+            f"{row['has_halide']:^7} {row['has_organic_salt']:^8} {row['mw_ratio']:>8.3f}"
+        )
+
+    # Test 2: Detailed feature extraction for specific cases
+    print("\n2. Testing detailed salt feature extraction...")
+
+    for name, smiles in [
+        ("sodium_acetate", test_molecules["sodium_acetate"]),
+        ("calcium_acetate", test_molecules["calcium_acetate"]),
+        ("aspirin", test_molecules["aspirin"]),
+    ]:
+        mol = Chem.MolFromSmiles(smiles)
+        features, api_mol = extract_advanced_salt_features(mol)
+
+        print(f"\n   {name}:")
+        print(f"     Total MW: {Descriptors.MolWt(mol):.2f}")
+        if api_mol:
+            print(f"     API MW: {Descriptors.MolWt(api_mol):.2f}")
+        if features:
+            print(f"     Salt/API ratio: {features['salt_to_api_ratio']:.3f}")
+            print(f"     Ionic strength proxy: {features['ionic_strength_proxy']}")
+
+    # Test 3: Edge cases
+    print("\n3. Testing edge cases...")
+
+    # Empty molecule (None)
+    empty_features, empty_api = extract_advanced_salt_features(None)
+    print(f"   None molecule: features={empty_features}, api={empty_api}")
+
+    # Single fragment (no salt)
+    benzene_mol = Chem.MolFromSmiles("c1ccccc1")
+    benzene_features, benzene_api = extract_advanced_salt_features(benzene_mol)
+    print(
+        f"   Single fragment (benzene): has_salt={benzene_features['has_salt']}, "
+        f"mw_ratio={benzene_features['mw_ratio']:.3f}"
+    )
+
+    # Multiple organic fragments
+    multi_org = Chem.MolFromSmiles("c1ccccc1.CC(=O)O")
+    multi_features, multi_api = extract_advanced_salt_features(multi_org)
+    print(f"   Multiple organic fragments: has_salt={multi_features['has_salt']}")
+
+    # Test 4: Salt type classification
+    print("\n4. Testing salt type classification...")
+
+    test_salts = [
+        ("NaCl", "[Na+].[Cl-]", "Sodium chloride"),
+        ("KBr", "[K+].[Br-]", "Potassium bromide"),
+        ("CaCl2", "[Ca+2].[Cl-].[Cl-]", "Calcium chloride"),
+        ("Organic salt", "CC(=O)[O-].CN", "Acetate with methylamine"),
+    ]
+
+    for name, smiles, description in test_salts:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol:
+            features, _ = extract_advanced_salt_features(mol)
+            print(f"   {name:15} ({description})")
+            print(
+                f"     Metal: {features['has_metal_salt']}, "
+                f"Halide: {features['has_halide']}, "
+                f"Organic: {features['has_organic_salt']}"
+            )
+
+    # Test 5: DataFrame integration
+    print("\n5. Testing DataFrame integration...")
+
+    # Create a mixed DataFrame
+    mixed_df = pd.DataFrame(
+        {
+            "smiles": [
+                test_molecules["aspirin"],
+                test_molecules["sodium_acetate"],
+                test_molecules["calcium_acetate"],
+            ],
+            "name": ["aspirin", "sodium_acetate", "calcium_acetate"],
+            "existing_col": [1, 2, 3],  # Test that existing columns are preserved
+        }
+    )
+
+    mixed_df["molecule"] = mixed_df["smiles"].apply(Chem.MolFromSmiles)
+    result_df = add_salt_features(mixed_df)
+
+    # Check that all columns are present
+    expected_cols = ["smiles", "name", "existing_col", "molecule"] + _get_salt_feature_columns()
+    missing_cols = [col for col in expected_cols if col not in result_df.columns]
+
+    if missing_cols:
+        print(f"   ✗ Missing columns: {missing_cols}")
+    else:
+        print("   ✓ All expected columns present")
+        print(f"   ✓ Original columns preserved: 'existing_col' in result = {('existing_col' in result_df.columns)}")
+        print(f"   ✓ Salt features added: {len(_get_salt_feature_columns())} new columns")
+
+    print("\n✅ All salt feature tests completed!")

workbench/utils/chem_utils/sdf.py

@@ -0,0 +1,292 @@
+"""SDF File utilities for molecular data in Workbench"""
+
+import logging
+import pandas as pd
+from typing import List, Optional
+from rdkit import Chem
+from rdkit.Chem import AllChem, SDWriter
+
+# Set up the logger
+log = logging.getLogger("workbench")
+
+
+def df_to_sdf_file(
+    df: pd.DataFrame,
+    output_file: str,
+    smiles_col: str = "smiles",
+    id_col: Optional[str] = None,
+    include_cols: Optional[List[str]] = None,
+    skip_invalid: bool = True,
+    generate_3d: bool = True,
+):
+    """
+    Convert DataFrame with SMILES to SDF file.
+
+    Args:
+        df: DataFrame containing SMILES and other data
+        output_file: Path to output SDF file
+        smiles_col: Column name containing SMILES strings
+        id_col: Column to use as molecule ID/name
+        include_cols: Specific columns to include as properties (default: all except smiles and molecule columns)
+        skip_invalid: Skip invalid SMILES instead of raising error
+        generate_3d: Generate 3D coordinates and optimize geometry
+    """
+    written_count = 0
+
+    with SDWriter(output_file) as writer:
+        writer.SetForceV3000(True)
+        for idx, row in df.iterrows():
+            mol = Chem.MolFromSmiles(row[smiles_col])
+            if mol is None:
+                if not skip_invalid:
+                    raise ValueError(f"Invalid SMILES at row {idx}: {row[smiles_col]}")
+                continue
+
+            # Generate 3D coordinates
+            if generate_3d:
+                mol = Chem.AddHs(mol)
+
+                # Try progressively more aggressive embedding strategies
+                embed_strategies = [
+                    {"maxAttempts": 1000, "randomSeed": 42},
+                    {"maxAttempts": 1000, "randomSeed": 42, "useRandomCoords": True},
+                    {"maxAttempts": 1000, "randomSeed": 42, "boxSizeMult": 5.0},
+                ]
+
+                embedded = False
+                for strategy in embed_strategies:
+                    if AllChem.EmbedMolecule(mol, **strategy) != -1:
+                        embedded = True
+                        break
+
+                if not embedded:
+                    if not skip_invalid:
+                        raise ValueError(f"Could not generate 3D coords for row {idx}")
+                    continue
+
+                AllChem.MMFFOptimizeMolecule(mol)
+
+            # Set molecule name/ID
+            if id_col and id_col in df.columns:
+                mol.SetProp("_Name", str(row[id_col]))
+
+            # Determine which columns to include
+            if include_cols:
+                cols_to_add = [col for col in include_cols if col in df.columns and col != smiles_col]
+            else:
+                # Auto-exclude common molecule column names and SMILES column
+                mol_col_names = ["mol", "molecule", "rdkit_mol", "Mol"]
+                cols_to_add = [col for col in df.columns if col != smiles_col and col not in mol_col_names]
+
+            # Add properties
+            for col in cols_to_add:
+                mol.SetProp(col, str(row[col]))
+
+            writer.write(mol)
+            written_count += 1
+
+    log.info(f"Wrote {written_count} molecules to SDF: {output_file}")
+    return written_count
+
+
+def sdf_file_to_df(
+    sdf_file: str,
+    include_smiles: bool = True,
+    smiles_col: str = "smiles",
+    id_col: Optional[str] = None,
+    include_props: Optional[List[str]] = None,
+    exclude_props: Optional[List[str]] = None,
+) -> pd.DataFrame:
+    """
+    Convert SDF file to DataFrame.
+
+    Args:
+        sdf_file: Path to input SDF file
+        include_smiles: Add SMILES column to output
+        smiles_col: Name for SMILES column
+        id_col: Column name for molecule ID/name (uses _Name property)
+        include_props: Specific properties to include (default: all)
+        exclude_props: Properties to exclude from output
+
+    Returns:
+        DataFrame with molecules and their properties
+    """
+    data = []
+
+    suppl = Chem.SDMolSupplier(sdf_file)
+    for idx, mol in enumerate(suppl):
+        if mol is None:
+            log.warning(f"Could not parse molecule at index {idx}")
+            continue
+
+        row_data = {}
+
+        # Add SMILES if requested
+        if include_smiles:
+            row_data[smiles_col] = Chem.MolToSmiles(mol)
+
+        # Add molecule name/ID if requested
+        if id_col and mol.HasProp("_Name"):
+            row_data[id_col] = mol.GetProp("_Name")
+
+        # Get all properties
+        prop_names = mol.GetPropNames()
+
+        # Filter properties based on include/exclude lists
+        if include_props:
+            prop_names = [p for p in prop_names if p in include_props]
+        if exclude_props:
+            prop_names = [p for p in prop_names if p not in exclude_props]
+
+        # Add properties to row
+        for prop in prop_names:
+            if prop != "_Name":  # Skip _Name if we already handled it
+                row_data[prop] = mol.GetProp(prop)
+
+        data.append(row_data)
+
+    df = pd.DataFrame(data)
+    log.info(f"Read {len(df)} molecules from SDF: {sdf_file}")
+
+    return df
+
+
+if __name__ == "__main__":
+    import tempfile
+    import os
+
+    print("Running SDF utilities tests...")
+
+    # Create test data
+    test_data = pd.DataFrame(
+        {
+            "smiles": [
+                "CCO",  # Ethanol
+                "c1ccccc1",  # Benzene
+                "CC(=O)O",  # Acetic acid
+                "INVALID_SMILES",  # Invalid
+                "CN1C=NC2=C1C(=O)N(C(=O)N2C)C",  # Caffeine
+            ],
+            "name": ["Ethanol", "Benzene", "Acetic Acid", "Invalid", "Caffeine"],
+            "mol_weight": [46.07, 78.11, 60.05, 0, 194.19],
+            "category": ["alcohol", "aromatic", "acid", "error", "alkaloid"],
+            "mol": ["should_exclude", "should_exclude", "should_exclude", "should_exclude", "should_exclude"],
+        }
+    )
+
+    # Test 1: Basic DataFrame to SDF conversion
+    print("\n1. Testing DataFrame to SDF conversion...")
+    with tempfile.NamedTemporaryFile(suffix=".sdf", delete=False) as tmp:
+        tmp_path = tmp.name
+
+    try:
+        # Test with 3D generation
+        count = df_to_sdf_file(
+            test_data, tmp_path, smiles_col="smiles", id_col="name", skip_invalid=True, generate_3d=True
+        )
+        print(f"   ✓ Wrote {count} molecules with 3D coords (expected 4, skipped 1 invalid)")
+
+        # Test without 3D generation
+        count = df_to_sdf_file(
+            test_data, tmp_path, smiles_col="smiles", id_col="name", skip_invalid=True, generate_3d=False
+        )
+        print(f"   ✓ Wrote {count} molecules without 3D coords")
+
+    except Exception as e:
+        print(f"   ✗ Error writing SDF: {e}")
+
+    # Test 2: SDF to DataFrame conversion
+    print("\n2. Testing SDF to DataFrame conversion...")
+    try:
+        # Read back the SDF
+        df_read = sdf_file_to_df(tmp_path, include_smiles=True, smiles_col="SMILES", id_col="mol_name")
+        print(f"   ✓ Read {len(df_read)} molecules from SDF")
+        print(f"   ✓ Columns: {list(df_read.columns)}")
+
+    except Exception as e:
+        print(f"   ✗ Error reading SDF: {e}")
+
+    # Test 3: Column filtering
+    print("\n3. Testing column inclusion/exclusion...")
+    try:
+        # Test with specific columns only
+        count = df_to_sdf_file(
+            test_data,
+            tmp_path,
+            smiles_col="smiles",
+            include_cols=["name", "mol_weight"],
+            skip_invalid=True,
+            generate_3d=False,
+        )
+
+        df_filtered = sdf_file_to_df(tmp_path)
+        excluded_mol = "mol" not in df_filtered.columns
+        included_weight = any("mol_weight" in str(col) for col in df_filtered.columns)
+
+        print(f"   {'✓' if excluded_mol else '✗'} 'mol' column excluded")
+        print(f"   {'✓' if included_weight else '✗'} 'mol_weight' included")
+
+    except Exception as e:
+        print(f"   ✗ Error with column filtering: {e}")
+
+    # Test 4: Error handling
+    print("\n4. Testing error handling...")
+
+    # Test with skip_invalid=False
+    try:
+        count = df_to_sdf_file(test_data, tmp_path, smiles_col="smiles", skip_invalid=False, generate_3d=False)
+        print("   ✗ Should have raised error for invalid SMILES")
+    except ValueError:
+        print("   ✓ Correctly raised error for invalid SMILES")
+
+    # Test 5: Property filtering on read
+    print("\n5. Testing property filtering on read...")
+    try:
+        # Write full data
+        df_to_sdf_file(test_data, tmp_path, smiles_col="smiles", skip_invalid=True, generate_3d=False)
+
+        # Read with include filter
+        df_include = sdf_file_to_df(tmp_path, include_props=["mol_weight", "category"])
+        print(f"   ✓ Include filter: {list(df_include.columns)}")
+
+        # Read with exclude filter
+        df_exclude = sdf_file_to_df(tmp_path, exclude_props=["category"])
+        has_category = "category" in df_exclude.columns
+        print(f"   {'✗' if has_category else '✓'} Exclude filter: 'category' excluded")
+
+    except Exception as e:
+        print(f"   ✗ Error with property filtering: {e}")
+
+    # Test 6: Edge cases
+    print("\n6. Testing edge cases...")
+
+    # Empty DataFrame
+    empty_df = pd.DataFrame(columns=["smiles", "name"])
+    try:
+        count = df_to_sdf_file(empty_df, tmp_path)
+        print(f"   ✓ Empty DataFrame: wrote {count} molecules")
+    except Exception as e:
+        print(f"   ✗ Empty DataFrame error: {e}")
+
+    # Missing columns
+    bad_df = pd.DataFrame({"not_smiles": ["CCO"]})
+    try:
+        count = df_to_sdf_file(bad_df, tmp_path, smiles_col="smiles")
+        print("   ✗ Should have raised error for missing column")
+    except KeyError:
+        print("   ✓ Correctly raised error for missing SMILES column")
+
+    # Large molecule test (3D generation stress test)
+    large_mol_df = pd.DataFrame({"smiles": ["C" * 50], "name": ["Long Chain"]})  # Very long carbon chain
+    try:
+        count = df_to_sdf_file(large_mol_df, tmp_path, generate_3d=True, skip_invalid=True)
+        print(f"   ✓ Large molecule: wrote {count} molecule(s)")
+    except Exception as e:
+        print(f"   ✗ Large molecule error: {e}")
+
+    # Cleanup
+    if os.path.exists(tmp_path):
+        os.remove(tmp_path)
+        print(f"\n✓ Cleaned up temp file: {tmp_path}")
+
+    print("\n✅ All SDF utilities tests completed!")